Primus-Turbo 開發指南

Primus-Turbo 是一個針對 AMD GPU (gfx942 / gfx950) 的高效能訓練算子庫。此技能是開發功能以及驗證其準確性和效能的入口點。

選擇你的任務

目標	從此處開始
端到端新增或修改算子/模組	develop-feature/SKILL.md
驗證數值正確性	verify-accuracy/SKILL.md
測量延遲 / 吞吐量	verify-performance/SKILL.md
驅動自主核心最佳化活動	optimize-handoff/SKILL.md → `kernel-optimize/SKILL.md`
使用執行時期證據分析熱點核心	run_profile/tool-rocprof/SKILL.md

本文件的其餘部分是上述文件所依賴的共享參考（建置、架構、程式碼地圖、後端系統）。請閱讀一次，然後跳到特定任務的文件。（run_profile/ 是一個分析工具的類別資料夾，目前包含單一的 tool-rocprof/ 技能。）

建置與迭代

安裝

# 始終先安裝固定的依賴項（Triton / PyTorch 版本至關重要）。
pip install -r requirements.txt

# 可編輯安裝（推薦用於開發）。
GPU_ARCHS=gfx942 pip install --no-build-isolation -e . -v

GPU_ARCHS：gfx942 (MI300X/MI325X)、gfx950 (MI350X/MI355X)、native (自動偵測) 或 "gfx942;gfx950"。
--no-build-isolation 是必需的，以便建置過程能看見已安裝的 torch / triton。
pip install . 會複製到 site-packages（原始碼編輯無效）；pip install -e . 是可編輯的。
建置過程也會自動安裝固定的 amd-aiter 和 origami（參閱 setup.py）。

何時需要重建

你修改了	需要重建？
Python / Triton (`primus_turbo/**.py`)	否（可編輯安裝會立即生效）
C++ / HIP (`csrc/**`) 或 `bindings_pytorch.cpp` 中的 op schema	是：重新執行可編輯安裝命令

架構特定的原始碼透過檔案名稱後綴進行過濾：*_gfx942.{cu,hip} 和 *_gfx950.{cu,hip} 僅在該架構包含於 GPU_ARCHS 中時才編譯（參閱 setup.py 中的 filter_files_by_arch）。

建置產物（3 層，核心庫與前端解耦）

產物	來源	備註
`libprimus_turbo_kernels.so`	`csrc/kernels/`	所有 HIP/CK/hipBLASLt/turbo 核心，與前端無關
`primus_turbo.pytorch._C`	`csrc/pytorch/`	PyTorch 綁定，鏈接上述 `.so`
`primus_turbo.jax._C`	`csrc/jax/`	JAX 綁定 (`PRIMUS_TURBO_FRAMEWORK=JAX`)

關鍵環境變數

變數	用途	預設值
`GPU_ARCHS`	目標架構	自動偵測
`ROCM_HOME`	ROCm 路徑	`/opt/rocm`
`MAX_JOBS`	並行編譯工作數	`64`
`PRIMUS_TURBO_FRAMEWORK`	`PYTORCH` / `JAX`（以 `;` 分隔）	`PYTORCH`
`PRIMUS_TURBO_LOG_LEVEL`	Logger 層級	`WARNING`

使用 pip show primus_turbo（查看 Editable project location 欄位）或 python -c "import primus_turbo; print(primus_turbo.__file__)"（應指向此原始碼樹而非 site-packages）驗證可編輯安裝是否正常運作；如果不是，請以可編輯模式重新安裝。

架構：分層設計

每個算子都是透過這些層的垂直切片。知道要處理哪一層是快速開發的關鍵。

modules/   nn.Module 包裝器 (Linear, GroupedLinear, Attention, ...)
   │ 調用
ops/       Python API + torch.autograd.Function (forward/backward)        ← 面向使用者
   │ 調用
kernels/   AutoKernelDispatcher + KernelBackend  (選擇一個後端)      ← 多後端算子
   │ 分派到
triton/    Triton 核心 (Python, 無需重建)        csrc/  HIP/CK/hipBLASLt 核心 (需重建)
                                                          │ 透過 TORCH_LIBRARY 在 csrc/pytorch/ 中綁定
                                                          └→ torch.ops.primus_turbo_cpp_extension.*

存在三種算子連接模式（參閱 develop-feature/SKILL.md 以獲取每個模式要複製的標準檔案）：

多後端 (gemm, gemm_fp8, grouped_gemm)：ops → kernels 調度器 → Triton 和/或 csrc。
直接 C++ (rmsnorm)：ops autograd Function → torch.ops.primus_turbo_cpp_extension.*（無調度器）。
直接 Triton (swiglu_with_probs)：ops autograd Function → kernels/ 中由 Triton 支援的輔助程式（無調度器）。

程式碼地圖

資料夾根目錄不同：ops/ = primus_turbo/pytorch/ops/，kernels/ = primus_turbo/pytorch/kernels/，triton/ = primus_turbo/triton/（頂層，不在 pytorch/ 下）；csrc/kernels/、tests/、benchmark/ 是倉庫根目錄。使用者 API 位於 ops/ 中；kernels/ 層包含調度器或後端實現 (*_impl.py)。

算子系列	API (`ops/`)	`kernels/` (調度器/實現)	Triton (`triton/`)	C++/HIP (`csrc/kernels/`)	測試	基準測試
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`	— (直接 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++ 綁定管道（所有 csrc 算子共享）：宣告在 csrc/pytorch/extensions.h 中，schema + CUDA + Meta 註冊在 csrc/pytorch/bindings_pytorch.cpp 中。

後端系統

primus_turbo/pytorch/core/backend.py 中的兩個類別：

GlobalBackendManager — 按算子 × 精度進行全域選擇。
AutoKernelDispatcher — 每個算子的基底，包含自動調優、預設值和回退。

選擇優先級（高 → 低）：程式碼設置器 → 環境變數 → 自動調優 → 程式碼內預設值 → 回退（嘗試所有 can_handle）。

`BackendType`	用於	備註
`HIPBLASLT`	GEMM (bf16, fp8 tensorwise)	密集 GEMM 的預設值
`TRITON`	GEMM, GroupedGEMM, Attention, ...	可調優，無需重建
`CK`	GEMM/GroupedGEMM FP8 (row/block)	Composable Kernel
`TURBO`	MXFP8/MXFP4 GEMM, Attention	內部開發 (gfx950)
`AITER`	Attention	預設的 attention
`DEEP_EP`	MoE dispatch/combine	需要安裝 DeepEP
`FLYDSL` (已規劃)	GEMM / Attention (編寫 DSL)	即將推出的 AMD tile-DSL 後端；尚未成為 `BackendType` 成員

FLYDSL 已在路線圖上 — 一個 Python 嵌入式 tile DSL，透過 MLIR 降級到 ROCDL。其調優知識已存在於 knowledge/backend/flydsl/ 的核心最佳化知識庫中，但尚未在 BackendType 中註冊，因此 GlobalBackendManager.set_*_backend(BackendType.FLYDSL) 目前不可用。

from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager
GlobalBackendManager.set_gemm_backend(BackendType.CK)   # 強制
GlobalBackendManager.set_auto_tune(True)                # 或 PRIMUS_TURBO_AUTO_TUNE=1
GlobalBackendManager.reset()                            # 清除設定 + 自動調優快取

支援的環境變數（支援每個精度的形式 "fp8:CK,other:TRITON"）：PRIMUS_TURBO_GEMM_BACKEND、PRIMUS_TURBO_GROUPED_GEMM_BACKEND、PRIMUS_TURBO_ATTENTION_BACKEND、PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND、PRIMUS_TURBO_AUTO_TUNE。

測試與基準測試快速參考

# 正確性（單 GPU，並行；每個 xdist 工作者透過 conftest.py 固定一個 GPU）
pytest tests/pytorch/ -n 8
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k "blockwise and TRITON"   # 過濾 op+backend
pytest tests/pytorch/ -n 8 --deterministic-only                          # 位元級確定性測試套件
pytest tests/pytorch/ --dist-only                                        # 多 GPU 測試

# 效能
python benchmark/ops/training/bench_gemm_turbo.py --dtype fp8 --granularity blockwise
python benchmark/ops/training/run_suite.py -d output/ -g gemm_fp8                  # 批次套件

正確性閘道：bf16/fp16 rtol=atol=1e-2，fp32 1e-4（來自 tests/pytorch/test_utils.py 中 get_tolerances 的 allclose 容差）；FP8 SNR ≥ 25 dB (E4M3) / 20 dB (E5M2)，FP4 SNR ≥ 10 dB（透過 compute_snr 的 SNR；閾值硬編碼在測試/基準檔案中，不在 get_tolerances 中）；確定性 rtol=atol=0。詳細資訊和模式：verify-accuracy/SKILL.md。

其他參考資料

README.md — 快速入門、安裝、打包
docs/examples.md — 每個算子的 API 用法（GEMM、Attention、GroupedGEMM、FP8/FP4、Backend/AutoTune）
benchmark/README.md、benchmark/accuracy/README.md — DeepEP 基準測試、跨平台準確性
CONTRIBUTING.md — 分支命名和提交規範
核心最佳化循環：kernel-optimize/SKILL.md（透過 optimize-handoff/SKILL.md 驅動）

primus-turbo-develop