| name | primus-turbo-develop |
| description | Develop, debug, and validate Primus-Turbo operators and modules on AMD GPUs. Covers the layered architecture (ops / kernels-dispatcher / Triton / HIP-CK csrc / modules), how to add or change a feature end-to-end, accuracy verification (SNR, tolerances, reference implementations), performance benchmarking, the backend dispatch system, and build/test/bench commands. Use for any Primus-Turbo development task (GEMM, Attention, GroupedGEMM, MoE, quantization, normalization, activation) and for accuracy or performance validation. |
Primus-Turbo Development Guide
Primus-Turbo is a high-performance training-operator library for AMD GPUs (gfx942 / gfx950). This skill is the entry point for developing a feature and validating its accuracy and performance.
Pick Your Task
The rest of this file is the shared reference (build, architecture, code map, backend system) those documents rely on. Read it once, then jump to the task-specific file. (run_profile/ is a category folder for profiling tools and currently holds the single tool-rocprof/ skill.)
Build & Iterate
Install
pip install -r requirements.txt
GPU_ARCHS=gfx942 pip install --no-build-isolation -e . -v
GPU_ARCHS: gfx942 (MI300X/MI325X), gfx950 (MI350X/MI355X), native (auto-detect), or "gfx942;gfx950".--no-build-isolation is required so the build sees already-installed torch / triton.pip install . copies into site-packages (source edits have no effect); pip install -e . is editable.- The build also auto-installs pinned
amd-aiter and origami (see setup.py).
What needs a rebuild
| You changed | Rebuild needed? |
|---|
Python / Triton (primus_turbo/**.py) | No (editable install picks it up immediately) |
C++ / HIP (csrc/**) or op schema in bindings_pytorch.cpp | Yes: re-run the editable install command |
Arch-specialized sources are filtered by filename suffix: *_gfx942.{cu,hip} and *_gfx950.{cu,hip} compile only when that arch is in GPU_ARCHS (see setup.py: filter_files_by_arch).
Build artifacts (3 layers, kernel lib decoupled from frontends)
| Artifact | Source | Notes |
|---|
libprimus_turbo_kernels.so | csrc/kernels/ | All HIP/CK/hipBLASLt/turbo kernels, frontend-agnostic |
primus_turbo.pytorch._C | csrc/pytorch/ | PyTorch bindings, links the .so above |
primus_turbo.jax._C | csrc/jax/ | JAX bindings (PRIMUS_TURBO_FRAMEWORK=JAX) |
Key environment variables
| Variable | Purpose | Default |
|---|
GPU_ARCHS | Target arch(s) | auto-detect |
ROCM_HOME | ROCm path | /opt/rocm |
MAX_JOBS | Parallel compile jobs | 64 |
PRIMUS_TURBO_FRAMEWORK | PYTORCH / JAX (;-separated) | PYTORCH |
PRIMUS_TURBO_LOG_LEVEL | Logger level | WARNING |
Verify a working editable install with pip show primus_turbo (look for the Editable project location field) or python -c "import primus_turbo; print(primus_turbo.__file__)" (it should point into this source tree, not site-packages); reinstall in editable mode if not.
Architecture: the layered design
Every operator is a vertical slice through these layers. Knowing which layer to touch is the core of fast development.
modules/ nn.Module wrappers (Linear, GroupedLinear, Attention, ...)
│ calls
ops/ Python API + torch.autograd.Function (forward/backward) ← user-facing
│ calls
kernels/ AutoKernelDispatcher + KernelBackend (selects a backend) ← multi-backend ops
│ dispatches to
triton/ Triton kernels (Python, no rebuild) csrc/ HIP/CK/hipBLASLt kernels (rebuild)
│ bound via TORCH_LIBRARY in csrc/pytorch/
└→ torch.ops.primus_turbo_cpp_extension.*
Three op wiring patterns exist (see develop-feature/SKILL.md for the canonical file to copy for each):
- Multi-backend (
gemm, gemm_fp8, grouped_gemm): ops → kernels dispatcher → Triton and/or csrc.
- Direct C++ (
rmsnorm): ops autograd Function → torch.ops.primus_turbo_cpp_extension.* (no dispatcher).
- Direct Triton (
swiglu_with_probs): ops autograd Function → a Triton-backed helper in kernels/ (no dispatcher).
Code Map
Column roots differ: ops/ = primus_turbo/pytorch/ops/, kernels/ = primus_turbo/pytorch/kernels/, triton/ = primus_turbo/triton/ (top level, not under pytorch/); csrc/kernels/, tests/, benchmark/ are repo-root. The user API lives in ops/; the kernels/ layer holds the dispatcher or the backend impl (*_impl.py).
| Operator family | API (ops/) | kernels/ (dispatcher/impl) | Triton (triton/) | C++/HIP (csrc/kernels/) | Tests | Bench |
|---|
| GEMM (bf16/fp16/fp32) | gemm.py | gemm/gemm_impl.py | gemm/gemm_kernel.py | gemm/ (hipBLASLt) | ops/test_gemm.py | bench_gemm_turbo.py |
| GEMM FP8 | gemm_fp8.py | gemm/gemm_fp8_impl.py | gemm/gemm_fp8_kernel.py | gemm/ck, gemm/turbo | ops/test_gemm_fp8.py | bench_gemm_turbo.py --dtype fp8 |
| GEMM FP4 | gemm_fp4.py | gemm/gemm_fp4_impl.py | — | gemm/ (hipBLASLt fp4) | ops/test_gemm_fp4.py | bench_gemm_turbo.py --dtype fp4 |
| Grouped GEMM | grouped_gemm.py, grouped_gemm_fp8.py | grouped_gemm/ | grouped_gemm/ | grouped_gemm/ck | ops/test_grouped_gemm*.py | bench_grouped_gemm_turbo.py |
| Attention | attention/ | attention/ (aiter/triton/turbo) | attention/attention_kernel.py | attention/turbo | ops/test_attention*.py | bench_attention_turbo.py |
| Quantization | quantization.py | quantization/ | quantization/ | quantization/ | ops/test_quantization.py |
C++ binding plumbing (shared by all csrc ops): declarations in csrc/pytorch/extensions.h, schema + CUDA + Meta registration in csrc/pytorch/bindings_pytorch.cpp.
Backend System
Two classes in primus_turbo/pytorch/core/backend.py:
GlobalBackendManager — global selection by operator × precision.AutoKernelDispatcher — per-operator base with autotune, default, and fallback.
Selection priority (high → low): code setter → env var → autotune → in-code default → fallback (try all can_handle).
BackendType | Used by | Notes |
|---|
HIPBLASLT | GEMM (bf16, fp8 tensorwise) | default for dense GEMM |
TRITON | GEMM, GroupedGEMM, Attention, ... | tunable, no rebuild |
CK | GEMM/GroupedGEMM FP8 (row/block) | Composable Kernel |
TURBO | MXFP8/MXFP4 GEMM, Attention | in-house (gfx950) |
AITER | Attention | default attention |
DEEP_EP | MoE dispatch/combine | needs DeepEP install |
FLYDSL (planned) | GEMM / Attention (authoring DSL) | upcoming AMD tile-DSL backend; not yet a BackendType member |
FLYDSL is on the roadmap — a Python-embedded tile DSL lowering through MLIR → ROCDL. Its tuning knowledge already lives in the kernel-optimize knowledge base at knowledge/backend/flydsl/, but it is not yet registered in BackendType, so GlobalBackendManager.set_*_backend(BackendType.FLYDSL) is not available today.
from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager
GlobalBackendManager.set_gemm_backend(BackendType.CK)
GlobalBackendManager.set_auto_tune(True)
GlobalBackendManager.reset()
Env vars (per-precision form "fp8:CK,other:TRITON" supported): PRIMUS_TURBO_GEMM_BACKEND, PRIMUS_TURBO_GROUPED_GEMM_BACKEND, PRIMUS_TURBO_ATTENTION_BACKEND, PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND, PRIMUS_TURBO_AUTO_TUNE.
Test & Bench Quick Reference
pytest tests/pytorch/ -n 8
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k "blockwise and TRITON"
pytest tests/pytorch/ -n 8 --deterministic-only
pytest tests/pytorch/ --dist-only
python benchmark/ops/training/bench_gemm_turbo.py --dtype fp8 --granularity blockwise
python benchmark/ops/training/run_suite.py -d output/ -g gemm_fp8
Correctness gates: bf16/fp16 rtol=atol=1e-2, fp32 1e-4 (allclose tolerances from get_tolerances in tests/pytorch/test_utils.py); FP8 SNR ≥ 25 dB (E4M3) / 20 dB (E5M2), FP4 SNR ≥ 10 dB (SNR via compute_snr; thresholds are hardcoded in the test/bench files, not in get_tolerances); determinism rtol=atol=0. Details and patterns: verify-accuracy/SKILL.md.
Additional References
README.md — quick start, install, packagingdocs/examples.md — per-operator API usage (GEMM, Attention, GroupedGEMM, FP8/FP4, Backend/AutoTune)benchmark/README.md, benchmark/accuracy/README.md — DeepEP bench, cross-platform accuracyCONTRIBUTING.md — branch naming and commit conventions- Kernel optimization loop:
kernel-optimize/SKILL.md (drive via optimize-handoff/SKILL.md)
