O que Guia de Desenvolvimento Primus-Turbo faz?
Primus-Turbo é uma biblioteca de operadores de treinamento de alto desempenho para GPUs AMD (gfx942 / gfx950). Esta skill é o ponto de entrada para desenvolver um recurso e validar sua precisão e desempenho.
Escolha sua Tarefa
| Objetivo | Comece aqui |
|---|---|
| Adicionar ou alterar um operador / módulo de ponta a ponta | develop-feature/SKILL.md |
| Verificar a correção numérica | verify-accuracy/SKILL.md |
| Medir latência / throughput | verify-performance/SKILL.md |
| Conduzir uma campanha autônoma de otimização de kernel | optimize-handoff/SKILL.md → kernel-optimize/SKILL.md |
| Perfilar um kernel quente com evidências de tempo de execução | run_profile/tool-rocprof/SKILL.md |
O restante deste arquivo é a referência compartilhada (build, arquitetura, mapa de código, sistema de backend) da qual esses documentos dependem. Leia uma vez e depois vá para o arquivo específico da tarefa. (run_profile/ é uma pasta de categoria para ferramentas de perfil e atualmente contém a única skill tool-rocprof/.)
Build e Iteração
Instalação
# Sempre instale as dependências fixadas primeiro (versões do Triton / PyTorch são críticas).
pip install -r requirements.txt
# Instalação editável (recomendada para desenvolvimento).
GPU_ARCHS=gfx942 pip install --no-build-isolation -e . -v
GPU_ARCHS:gfx942(MI300X/MI325X),gfx950(MI350X/MI355X),native(detecção automática) ou"gfx942;gfx950".--no-build-isolationé necessário para que o build veja ostorch/tritonjá instalados.pip install .copia para site-packages (alterações no código-fonte não têm efeito);pip install -e .é editável.- O build também instala automaticamente
amd-aitereorigamifixados (vejasetup.py).
O que precisa de um rebuild
| Você alterou | Precisa de rebuild? |
|---|---|
Python / Triton (primus_turbo/**.py) | Não (instalação editável captura imediatamente) |
C++ / HIP (csrc/**) ou schema de op em bindings_pytorch.cpp | Sim: execute novamente o comando de instalação editável |
Fontes especializadas por arquitetura são filtradas pelo sufixo do nome do arquivo: *_gfx942.{cu,hip} e *_gfx950.{cu,hip} compilam apenas quando essa arquitetura está em GPU_ARCHS (veja setup.py: filter_files_by_arch).
Artefatos de build (3 camadas, biblioteca de kernel desacoplada dos frontends)
| Artefato | Fonte | Notas |
|---|---|---|
libprimus_turbo_kernels.so | csrc/kernels/ | Todos os kernels HIP/CK/hipBLASLt/turbo, independentes de frontend |
primus_turbo.pytorch._C | csrc/pytorch/ | Bindings PyTorch, linka o .so acima |
primus_turbo.jax._C | csrc/jax/ | Bindings JAX (PRIMUS_TURBO_FRAMEWORK=JAX) |
Principais variáveis de ambiente
| Variável | Propósito | Padrão |
|---|---|---|
GPU_ARCHS | Arquitetura(s) alvo | detecção automática |
ROCM_HOME | Caminho do ROCm | /opt/rocm |
MAX_JOBS | Trabalhos de compilação paralelos | 64 |
PRIMUS_TURBO_FRAMEWORK | PYTORCH / JAX (separados por ;) | PYTORCH |
PRIMUS_TURBO_LOG_LEVEL | Nível do logger | WARNING |
Verifique uma instalação editável funcionando com pip show primus_turbo (procure pelo campo Editable project location) ou python -c "import primus_turbo; print(primus_turbo.__file__)" (deve apontar para esta árvore de origem, não para site-packages); reinstale em modo editável se não.
Arquitetura: o design em camadas
Cada operador é uma fatia vertical através dessas camadas. Saber qual camada tocar é o centro do desenvolvimento rápido.
modules/ Wrappers nn.Module (Linear, GroupedLinear, Attention, ...)
│ chama
ops/ API Python + torch.autograd.Function (forward/backward) ← voltado ao usuário
│ chama
kernels/ AutoKernelDispatcher + KernelBackend (seleciona um backend) ← operações multi-backend
│ despacha para
triton/ Kernels Triton (Python, sem rebuild) csrc/ Kernels HIP/CK/hipBLASLt (rebuild)
│ vinculados via TORCH_LIBRARY em csrc/pytorch/
└→ torch.ops.primus_turbo_cpp_extension.*
Existem três padrões de fiação de operadores (veja develop-feature/SKILL.md para o arquivo canônico a ser copiado para cada um):
- Multi-backend (
gemm,gemm_fp8,grouped_gemm):ops→ dispatcherkernels→ Triton e/ou csrc. - C++ direto (
rmsnorm):opsautograd Function →torch.ops.primus_turbo_cpp_extension.*(sem dispatcher). - Triton direto (
swiglu_with_probs):opsautograd Function → um helper baseado em Triton emkernels/(sem dispatcher).
Mapa de Código
As raízes das colunas diferem:
ops/=primus_turbo/pytorch/ops/,kernels/=primus_turbo/pytorch/kernels/,triton/=primus_turbo/triton/(nível superior, não sobpytorch/);csrc/kernels/,tests/,benchmark/são raiz do repositório. A API do usuário fica emops/; a camadakernels/contém o dispatcher ou a implementação do backend (*_impl.py).
| Família de operador | API (ops/) | kernels/ (dispatcher/impl) | Triton (triton/) | C++/HIP (csrc/kernels/) | Testes | 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 |
| Quantização | quantization.py | quantization/ | quantization/ | quantization/ | ops/test_quantization.py | accuracy/eval_sf_accuracy.py |
| Ativação (swiglu/geglu) | activation.py | activation/ | activation/ | — | ops/test_activation.py | — |
| Normalização (rmsnorm) | normalization.py | — (C++ direto) | — | 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 | — |
Encanação dos bindings C++ (compartilhada por todas as operações csrc): declarações em csrc/pytorch/extensions.h, schema + registro CUDA + Meta em csrc/pytorch/bindings_pytorch.cpp.
Sistema de Backend
Duas classes em primus_turbo/pytorch/core/backend.py:
GlobalBackendManager— seleção global por operador × precisão.AutoKernelDispatcher— base por operador com autotune, padrão e fallback.
Prioridade de seleção (alta → baixa): setter de código → variável de ambiente → autotune → padrão no código → fallback (tentar todos can_handle).
BackendType | Usado por | Notas |
|---|---|---|
HIPBLASLT | GEMM (bf16, fp8 tensorwise) | padrão para GEMM densa |
TRITON | GEMM, GroupedGEMM, Attention, ... | ajustável, sem rebuild |
CK | GEMM/GroupedGEMM FP8 (row/block) | Composable Kernel |
TURBO | MXFP8/MXFP4 GEMM, Attention | interno (gfx950) |
AITER | Attention | atenção padrão |
DEEP_EP | MoE dispatch/combine | precisa de instalação DeepEP |
FLYDSL (planejado) | GEMM / Attention (DSL de autoria) | próximo backend AMD tile-DSL; ainda não é um membro de BackendType |
FLYDSLestá no roadmap — uma DSL de tile embutida em Python que reduz via MLIR → ROCDL. Seu conhecimento de ajuste já reside na base de conhecimento de otimização de kernel emknowledge/backend/flydsl/, mas ainda não está registrado emBackendType, entãoGlobalBackendManager.set_*_backend(BackendType.FLYDSL)não está disponível hoje.
from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager
GlobalBackendManager.set_gemm_backend(BackendType.CK) # forçar
GlobalBackendManager.set_auto_tune(True) # ou PRIMUS_TURBO_AUTO_TUNE=1
GlobalBackendManager.reset() # limpar configurações + cache de autotune
Variáveis de ambiente (forma por precisão "fp8:CK,other:TRITON" suportada): PRIMUS_TURBO_GEMM_BACKEND, PRIMUS_TURBO_GROUPED_GEMM_BACKEND, PRIMUS_TURBO_ATTENTION_BACKEND, PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND, PRIMUS_TURBO_AUTO_TUNE.
Referência Rápida de Teste e Bench
# Corretude (uma GPU, paralelo; cada worker xdist fixa uma GPU via conftest.py)
pytest tests/pytorch/ -n 8
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k "blockwise and TRITON" # filtrar op+backend
pytest tests/pytorch/ -n 8 --deterministic-only # suíte de determinismo bit-a-bit
pytest tests/pytorch/ --dist-only # testes multi-GPU
# Desempenho
python benchmark/ops/training/bench_gemm_turbo.py --dtype fp8 --granularity blockwise
python benchmark/ops/training/run_suite.py -d output/ -g gemm_fp8 # suíte em lote
Gates de corretude: bf16/fp16 rtol=atol=1e-2, fp32 1e-4 (tolerâncias allclose de get_tolerances em tests/pytorch/test_utils.py); FP8 SNR ≥ 25 dB (E4M3) / 20 dB (E5M2), FP4 SNR ≥ 10 dB (SNR via compute_snr; limites são fixados nos arquivos de teste/bench, não em get_tolerances); determinismo rtol=atol=0. Detalhes e padrões: verify-accuracy/SKILL.md.
Referências Adicionais
README.md— início rápido, instalação, empacotamentodocs/examples.md— uso da API por operador (GEMM, Attention, GroupedGEMM, FP8/FP4, Backend/AutoTune)benchmark/README.md,benchmark/accuracy/README.md— benchmark DeepEP, precisão entre plataformasCONTRIBUTING.md— convenções de nomeação de branches e commits- Loop de otimização de kernel:
kernel-optimize/SKILL.md(conduzir via optimize-handoff/SKILL.md)