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` 中的算子模式	是的：重新运行可编辑安装命令

特定于架构的源文件通过文件名后缀进行过滤：*_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`	日志记录器级别	`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（前向/后向）        ← 面向用户
   │ 调用
kernels/   AutoKernelDispatcher + KernelBackend（选择后端）      ← 多后端算子
   │ 分派到
triton/    Triton 内核（Python，无需重建）        csrc/  HIP/CK/hipBLASLt 内核（需要重建）
                                                          │ 通过 csrc/pytorch/ 中的 TORCH_LIBRARY 绑定
                                                          └→ 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.py`	`quantization/`	`quantization/`	`quantization/`	`ops/test_quantization.py`	`accuracy/eval_sf_accuracy.py`
激活函数（swiglu/geglu）	`activation.py`	`activation/`	`activation/`	—	`ops/test_activation.py`	—
归一化（rmsnorm）	`normalization.py`	—（直接 C++）	—	`normalization/`	`ops/test_normalization.py`	—
MoE 排列/路由/分派	`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 中，模式和 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 分派/合并	需要安装 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"   # 过滤算子和后端
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、后端/自动调优）
benchmark/README.md、benchmark/accuracy/README.md — DeepEP 基准测试、跨平台精度
CONTRIBUTING.md — 分支命名和提交约定
内核优化循环：kernel-optimize/SKILL.md（通过 optimize-handoff/SKILL.md 驱动）

primus-turbo-develop