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

Goal	Start here
Add or change an operator / module end-to-end	develop-feature/SKILL.md
Verify numerical correctness	verify-accuracy/SKILL.md
Measure latency / throughput	verify-performance/SKILL.md
Drive an autonomous kernel-optimization campaign	optimize-handoff/SKILL.md → `kernel-optimize/SKILL.md`
Profile a hot kernel with runtime evidence	run_profile/tool-rocprof/SKILL.md

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

# Always install pinned deps first (Triton / PyTorch versions are critical).
pip install -r requirements.txt

# Editable install (recommended for development).
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`	`accuracy/eval_sf_accuracy.py`
Activation (swiglu/geglu)	`activation.py`	`activation/`	`activation/`	—	`ops/test_activation.py`	—
Normalization (rmsnorm)	`normalization.py`	— (direct C++)	—	`normalization/`	`ops/test_normalization.py`	—
MoE permute / router / dispatch	`moe/`	`moe/`, `moe_permute`	`moe/`	`moe_permute/`, `deep_ep/`	`ops/test_moe_permute.py`, `test_fused_moe_router.py`	`bench_deepep_intranode.py`
Async-TP (gemm+comm)	`async_tp.py`	`async_tp/`	`async_tp/`	—	`ops/test_fused_*.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)   # force
GlobalBackendManager.set_auto_tune(True)                # or PRIMUS_TURBO_AUTO_TUNE=1
GlobalBackendManager.reset()                            # clear settings + autotune cache

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

# Correctness (single-GPU, parallel; each xdist worker pins one GPU via conftest.py)
pytest tests/pytorch/ -n 8
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k "blockwise and TRITON"   # filter op+backend
pytest tests/pytorch/ -n 8 --deterministic-only                          # bitwise-determinism suite
pytest tests/pytorch/ --dist-only                                        # multi-GPU tests

# Performance
python benchmark/ops/training/bench_gemm_turbo.py --dtype fp8 --granularity blockwise
python benchmark/ops/training/run_suite.py -d output/ -g gemm_fp8                  # batch suite

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, packaging
docs/examples.md — per-operator API usage (GEMM, Attention, GroupedGEMM, FP8/FP4, Backend/AutoTune)
benchmark/README.md, benchmark/accuracy/README.md — DeepEP bench, cross-platform accuracy
CONTRIBUTING.md — branch naming and commit conventions
Kernel optimization loop: kernel-optimize/SKILL.md (drive via optimize-handoff/SKILL.md)

primus-turbo-develop

What is primus-turbo-develop?

Ask in your favorite AI

Documentation