primus-turbo-develop

Guide de développement Primus-Turbo

Primus-Turbo est une bibliothèque d'opérateurs d'entraînement haute performance pour GPU AMD (gfx942 / gfx950). Cette compétence est le point d'entrée pour développer une fonctionnalité et valider sa précision et ses performances.

Choisissez votre tâche

Objectif	Commencez ici
Ajouter ou modifier un opérateur / module de bout en bout	develop-feature/SKILL.md
Vérifier la correction numérique	verify-accuracy/SKILL.md
Mesurer la latence / le débit	verify-performance/SKILL.md
Mener une campagne d'optimisation autonome de noyau	optimize-handoff/SKILL.md → kernel-optimize/SKILL.md
Profiler un noyau chaud avec des preuves d'exécution	run_profile/tool-rocprof/SKILL.md

Le reste de ce fichier est la référence partagée (construction, architecture, plan du code, système backend) sur laquelle ces documents s'appuient. Lisez-la une fois, puis passez au fichier spécifique à la tâche. (run_profile/ est un dossier de catégorie pour les outils de profilage et contient actuellement la seule compétence tool-rocprof/.)

Construction et itération

Installation

# Toujours installer d'abord les dépendances épinglées (les versions Triton/PyTorch sont critiques).
pip install -r requirements.txt

# Installation éditable (recommandée pour le développement).
GPU_ARCHS=gfx942 pip install --no-build-isolation -e . -v

• GPU_ARCHS : gfx942 (MI300X/MI325X), gfx950 (MI350X/MI355X), native (détection auto), ou "gfx942;gfx950".
• --no-build-isolation est requis pour que la construction voie les torch / triton déjà installés.
• pip install . copie dans site-packages (les modifications source n'ont aucun effet) ; pip install -e . est éditable.
• La construction installe aussi automatiquement amd-aiter et origami épinglés (voir setup.py).

Ce qui nécessite une reconstruction

Vous avez modifié	Reconstruction nécessaire ?
Python / Triton (primus_turbo/**.py)	Non (l'installation éditable le prend en compte immédiatement)
C++ / HIP (csrc/**) ou schéma d'op dans bindings_pytorch.cpp	Oui : réexécutez la commande d'installation éditable

Les sources spécialisées par architecture sont filtrées par suffixe de nom de fichier : *_gfx942.{cu,hip} et *_gfx950.{cu,hip} ne sont compilées que lorsque cette architecture est dans GPU_ARCHS (voir setup.py: filter_files_by_arch).

Artefacts de construction (3 couches, bibliothèque de noyaux découplée des interfaces)

Artefact	Source	Notes
libprimus_turbo_kernels.so	csrc/kernels/	Tous les noyaux HIP/CK/hipBLASLt/turbo, indépendants de l'interface
primus_turbo.pytorch._C	csrc/pytorch/	Liaisons PyTorch, lie le .so ci-dessus
primus_turbo.jax._C	csrc/jax/	Liaisons JAX (PRIMUS_TURBO_FRAMEWORK=JAX)

Variables d'environnement clés

Variable	Objectif	Défaut
GPU_ARCHS	Architecture(s) cible(s)	détection auto
ROCM_HOME	Chemin ROCm	/opt/rocm
MAX_JOBS	Tâches de compilation parallèles	64
PRIMUS_TURBO_FRAMEWORK	PYTORCH / JAX (séparés par ;)	PYTORCH
PRIMUS_TURBO_LOG_LEVEL	Niveau du logger	WARNING

Vérifiez une installation éditable fonctionnelle avec pip show primus_turbo (cherchez le champ Editable project location) ou python -c "import primus_turbo; print(primus_turbo.__file__)" (il doit pointer vers cette arborescence source, pas site-packages) ; réinstallez en mode éditable si ce n'est pas le cas.

Architecture : la conception en couches

Chaque opérateur est une tranche verticale à travers ces couches. Savoir quelle couche toucher est le cœur du développement rapide.

modules/   Wrappers nn.Module (Linear, GroupedLinear, Attention, ...)
   │ appelle
ops/       API Python + torch.autograd.Function (forward/backward)     ← face à l'utilisateur
   │ appelle
kernels/   AutoKernelDispatcher + KernelBackend (sélectionne un backend)      ← opérations multi-backend
   │ distribue à
triton/    Noyaux Triton (Python, pas de reconstruction)        csrc/  Noyaux HIP/CK/hipBLASLt (reconstruction)
                                                                  │ lié via TORCH_LIBRARY dans csrc/pytorch/
                                                                  └→ torch.ops.primus_turbo_cpp_extension.*

Trois motifs de câblage d'opérations existent (voir develop-feature/SKILL.md pour le fichier canonique à copier pour chacun) :

• Multi-backend (gemm, gemm_fp8, grouped_gemm) : ops → répartiteur kernels → Triton et/ou csrc.
• C++ direct (rmsnorm) : fonction autograd ops → torch.ops.primus_turbo_cpp_extension.* (pas de répartiteur).
• Triton direct (swiglu_with_probs) : fonction autograd ops → un assistant basé sur Triton dans kernels/ (pas de répartiteur).

Plan du code

Les racines des colonnes diffèrent : ops/ = primus_turbo/pytorch/ops/, kernels/ = primus_turbo/pytorch/kernels/, triton/ = primus_turbo/triton/ (niveau supérieur, pas sous pytorch/) ; csrc/kernels/, tests/, benchmark/ sont à la racine du dépôt. L'API utilisateur vit dans ops/ ; la couche kernels/ contient le répartiteur ou l'implémentation du backend (*_impl.py).

Famille d'opérateurs	API (ops/)	kernels/ (répartiteur/impl)	Triton (triton/)	C++/HIP (csrc/kernels/)	Tests	Bancs d'essai
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
Quantification	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	—
Normalisation (rmsnorm)	normalization.py	— (C++ direct)	—	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	—

Plomberie de liaison C++ (partagée par toutes les opérations csrc) : déclarations dans csrc/pytorch/extensions.h, schéma + enregistrement CUDA + Meta dans csrc/pytorch/bindings_pytorch.cpp.

Système de backend

Deux classes dans primus_turbo/pytorch/core/backend.py :

• GlobalBackendManager — sélection globale par opérateur × précision.
• AutoKernelDispatcher — base par opérateur avec autotune, défaut et repli.

Priorité de sélection (haute → basse) : setter de code → variable d'environnement → autotune → défaut dans le code → repli (essayer tous les can_handle).

BackendType	Utilisé par	Notes
HIPBLASLT	GEMM (bf16, fp8 tensorwise)	défaut pour GEMM dense
TRITON	GEMM, GroupedGEMM, Attention, ...	réglable, pas de reconstruction
CK	GEMM/GroupedGEMM FP8 (row/block)	Composable Kernel
TURBO	MXFP8/MXFP4 GEMM, Attention	interne (gfx950)
AITER	Attention	attention par défaut
DEEP_EP	Dispatching/combinatoire MoE	nécessite l'installation de DeepEP
FLYDSL (prévu)	GEMM / Attention (DSL d'écriture)	futur backend DSL de tuiles AMD ; pas encore un membre de BackendType

FLYDSL est sur la feuille de route — un DSL de tuiles intégré en Python se réduisant via MLIR → ROCDL. Ses connaissances de réglage résident déjà dans la base de connaissances d'optimisation de noyau à knowledge/backend/flydsl/, mais il n'est pas encore enregistré dans BackendType, donc GlobalBackendManager.set_*_backend(BackendType.FLYDSL) n'est pas disponible aujourd'hui.

from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager
GlobalBackendManager.set_gemm_backend(BackendType.CK)   # forcer
GlobalBackendManager.set_auto_tune(True)                # ou PRIMUS_TURBO_AUTO_TUNE=1
GlobalBackendManager.reset()                            # effacer les réglages + le cache d'autotune

Variables d'environnement (forme par précision "fp8:CK,other:TRITON" prise en charge) : PRIMUS_TURBO_GEMM_BACKEND, PRIMUS_TURBO_GROUPED_GEMM_BACKEND, PRIMUS_TURBO_ATTENTION_BACKEND, PRIMUS_TURBO_MOE_DISPATCH_COMBINE_BACKEND, PRIMUS_TURBO_AUTO_TUNE.

Référence rapide des tests et des bancs d'essai

# Correction (GPU unique, parallèle ; chaque travailleur xdist épingle un GPU via conftest.py)
pytest tests/pytorch/ -n 8
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k "blockwise and TRITON"   # filtrer op+backend
pytest tests/pytorch/ -n 8 --deterministic-only                          # suite de déterministe binaire
pytest tests/pytorch/ --dist-only                                        # tests multi-GPU

# 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                  # suite par lots

Seuils de correction : bf16/fp16 rtol=atol=1e-2, fp32 1e-4 (tolérances allclose de get_tolerances dans tests/pytorch/test_utils.py) ; SNR FP8 ≥ 25 dB (E4M3) / 20 dB (E5M2), SNR FP4 ≥ 10 dB (SNR via compute_snr ; les seuils sont codés en dur dans les fichiers de test/bench, pas dans get_tolerances) ; déterminisme rtol=atol=0. Détails et motifs : verify-accuracy/SKILL.md.

Références supplémentaires

• README.md — démarrage rapide, installation, empaquetage
• docs/examples.md — utilisation de l'API par opérateur (GEMM, Attention, GroupedGEMM, FP8/FP4, Backend/AutoTune)
• benchmark/README.md, benchmark/accuracy/README.md — banc d'essai DeepEP, précision multiplateforme
• CONTRIBUTING.md — conventions de nommage de branches et de commits
• Boucle d'optimisation de noyau : kernel-optimize/SKILL.md (pilotée via optimize-handoff/SKILL.md)