primus-turbo-develop

Guía de Desarrollo de Primus-Turbo

Primus-Turbo es una biblioteca de operadores de entrenamiento de alto rendimiento para GPUs AMD (gfx942 / gfx950). Esta habilidad es el punto de entrada para desarrollar una funcionalidad y validar su precisión y rendimiento.

Elige tu Tarea

Objetivo	Comienza aquí
Añadir o modificar un operador/módulo de principio a fin	develop-feature/SKILL.md
Verificar la corrección numérica	verify-accuracy/SKILL.md
Medir latencia / rendimiento	verify-performance/SKILL.md
Llevar a cabo una campaña autónoma de optimización de kernels	optimize-handoff/SKILL.md → kernel-optimize/SKILL.md
Perfilar un kernel intensivo con evidencia de ejecución	run_profile/tool-rocprof/SKILL.md

El resto de este archivo es la referencia compartida (compilación, arquitectura, mapa de código, sistema de backend) en la que se apoyan esos documentos. Léelo una vez y luego salta al archivo específico de la tarea. (run_profile/ es una carpeta de categoría para herramientas de perfilado y actualmente contiene la única habilidad tool-rocprof/).

Compilación e Iteración

Instalación

# Instala siempre las dependencias fijadas primero (las versiones de Triton/PyTorch son críticas).
pip install -r requirements.txt

# Instalación editable (recomendada para desarrollo).
GPU_ARCHS=gfx942 pip install --no-build-isolation -e . -v

• GPU_ARCHS: gfx942 (MI300X/MI325X), gfx950 (MI350X/MI355X), native (detección automática), o "gfx942;gfx950".
• --no-build-isolation es necesario para que la compilación vea los torch / triton ya instalados.
• pip install . copia en site-packages (los cambios en el código fuente no tienen efecto); pip install -e . es editable.
• La compilación también instala automáticamente las versiones fijadas de amd-aiter y origami (consulta setup.py).

Qué necesita recompilación

Has cambiado	¿Recompilación necesaria?
Python / Triton (primus_turbo/**.py)	No (la instalación editable lo recoge inmediatamente)
C++ / HIP (csrc/**) o esquema de op en bindings_pytorch.cpp	Sí: vuelve a ejecutar el comando de instalación editable

Los fuentes especializados por arquitectura se filtran mediante el sufijo del nombre de archivo: *_gfx942.{cu,hip} y *_gfx950.{cu,hip} se compilan solo cuando esa arquitectura está en GPU_ARCHS (consulta setup.py: filter_files_by_arch).

Artefactos de compilación (3 capas, biblioteca de kernels desacoplada de los frontends)

Artefacto	Fuente	Notas
libprimus_turbo_kernels.so	csrc/kernels/	Todos los kernels HIP/CK/hipBLASLt/turbo, independientes del frontend
primus_turbo.pytorch._C	csrc/pytorch/	Enlaces de PyTorch, enlaza el .so anterior
primus_turbo.jax._C	csrc/jax/	Enlaces de JAX (PRIMUS_TURBO_FRAMEWORK=JAX)

Variables de entorno clave

Variable	Propósito	Valor por defecto
GPU_ARCHS	Arquitectura(s) objetivo	detección automática
ROCM_HOME	Ruta de ROCm	/opt/rocm
MAX_JOBS	Trabajos de compilación en paralelo	64
PRIMUS_TURBO_FRAMEWORK	PYTORCH / JAX (separados por ;)	PYTORCH
PRIMUS_TURBO_LOG_LEVEL	Nivel del registrador	WARNING

Verifica una instalación editable funcional con pip show primus_turbo (busca el campo Editable project location) o python -c "import primus_turbo; print(primus_turbo.__file__)" (debe apuntar a este árbol fuente, no a site-packages); reinstala en modo editable si no es así.

Arquitectura: el diseño en capas

Cada operador es un corte vertical a través de estas capas. Saber qué capa tocar es el núcleo del desarrollo rápido.

modules/   Envoltorios nn.Module (Linear, GroupedLinear, Attention, ...)
   │ llama
ops/       API de Python + torch.autograd.Function (forward/backward)     ← orientado al usuario
   │ llama
kernels/   AutoKernelDispatcher + KernelBackend (selecciona un backend)   ← ops multi-backend
   │ despacha a
triton/    Kernels Triton (Python, sin recompilación)        csrc/  Kernels HIP/CK/hipBLASLt (requieren recompilación)
                                                                             │ enlazados mediante TORCH_LIBRARY en csrc/pytorch/
                                                                             └→ torch.ops.primus_turbo_cpp_extension.*

Existen tres patrones de cableado de operaciones (consulta develop-feature/SKILL.md para el archivo canónico a copiar para cada uno):

• Multi-backend (gemm, gemm_fp8, grouped_gemm): ops → despachador de kernels → Triton y/o csrc.
• C++ directo (rmsnorm): Función autograd de ops → torch.ops.primus_turbo_cpp_extension.* (sin despachador).
• Triton directo (swiglu_with_probs): Función autograd de ops → un ayudante basado en Triton en kernels/ (sin despachador).

Mapa de Código

Las raíces de las columnas difieren: ops/ = primus_turbo/pytorch/ops/, kernels/ = primus_turbo/pytorch/kernels/, triton/ = primus_turbo/triton/ (nivel superior, no bajo pytorch/); csrc/kernels/, tests/, benchmark/ están en la raíz del repositorio. La API de usuario vive en ops/; la capa kernels/ contiene el despachador o la implementación del backend (*_impl.py).

Familia de operadores	API (ops/)	kernels/ (despachador/impl)	Triton (triton/)	C++/HIP (csrc/kernels/)	Pruebas	Benchmark
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
Cuantización	quantization.py	quantization/	quantization/	quantization/	ops/test_quantization.py	accuracy/eval_sf_accuracy.py
Activación (swiglu/geglu)	activation.py	activation/	activation/	—	ops/test_activation.py	—
Normalización (rmsnorm)	normalization.py	— (C++ directo)	—	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	—

Fontanería de enlace C++ (compartida por todas las ops de csrc): declaraciones en csrc/pytorch/extensions.h, esquema + registro CUDA + Meta en csrc/pytorch/bindings_pytorch.cpp.

Sistema de Backend

Dos clases en primus_turbo/pytorch/core/backend.py:

• GlobalBackendManager — selección global por operador × precisión.
• AutoKernelDispatcher — base por operador con autoajuste, valor por defecto y respaldo.

Prioridad de selección (alta → baja): setter de código → variable de entorno → autoajuste → valor por defecto en código → respaldo (probar todos los can_handle).

BackendType	Usado por	Notas
HIPBLASLT	GEMM (bf16, fp8 tensorwise)	valor por defecto para GEMM denso
TRITON	GEMM, GroupedGEMM, Attention, ...	ajustable, sin recompilación
CK	GEMM/GroupedGEMM FP8 (row/block)	Composable Kernel
TURBO	MXFP8/MXFP4 GEMM, Attention	propio (gfx950)
AITER	Attention	valor por defecto para attention
DEEP_EP	Despacho/combinación MoE	necesita instalación de DeepEP
FLYDSL (planeado)	GEMM / Attention (DSL de autoría)	próximo backend AMD tile-DSL; aún no es un miembro de BackendType

FLYDSL está en la hoja de ruta: un DSL de tiles incrustado en Python que baja a través de MLIR → ROCDL. Su conocimiento de ajuste ya vive en la base de conocimiento de optimización de kernels en knowledge/backend/flydsl/, pero aún no está registrado en BackendType, por lo que GlobalBackendManager.set_*_backend(BackendType.FLYDSL) no está disponible hoy.

from primus_turbo.pytorch.core.backend import BackendType, GlobalBackendManager
GlobalBackendManager.set_gemm_backend(BackendType.CK)    # forzar
GlobalBackendManager.set_auto_tune(True)                 # o PRIMUS_TURBO_AUTO_TUNE=1
GlobalBackendManager.reset()                             # limpiar configuraciones y caché de autoajuste

Variables de entorno (se admite la forma por precisión "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.

Referencia Rápida de Pruebas y Benchmarks

# Corrección (una GPU, paralelo; cada trabajador de xdist fija una GPU mediante conftest.py)
pytest tests/pytorch/ -n 8
pytest tests/pytorch/ops/test_gemm_fp8.py -v -k "blockwise and TRITON"    # filtrar operación+backend
pytest tests/pytorch/ -n 8 --deterministic-only                             # suite de determinismo bit a bit
pytest tests/pytorch/ --dist-only                                           # pruebas multi-GPU

# Rendimiento
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 por lotes

Puertas de corrección: bf16/fp16 rtol=atol=1e-2, fp32 1e-4 (tolerancias allclose de get_tolerances en tests/pytorch/test_utils.py); SNR FP8 ≥ 25 dB (E4M3) / 20 dB (E5M2), SNR FP4 ≥ 10 dB (SNR mediante compute_snr; los umbrales están codificados en los archivos de prueba/benchmark, no en get_tolerances); determinismo rtol=atol=0. Detalles y patrones: verify-accuracy/SKILL.md.

Referencias Adicionales

• README.md — inicio rápido, instalación, empaquetado
• docs/examples.md — uso de la API por operador (GEMM, Attention, GroupedGEMM, FP8/FP4, Backend/AutoTune)
• benchmark/README.md, benchmark/accuracy/README.md — benchmark de DeepEP, precisión multiplataforma
• CONTRIBUTING.md — convenciones de nombres de ramas y commits
• Bucle de optimización de kernels: kernel-optimize/SKILL.md (impulsar mediante optimize-handoff/SKILL.md)