Hugging Face facilita construir y compartir kernels ROCm | Keryc
Los kernels personalizados son el motor del deep learning rápido: permiten operaciones GPU diseñadas para tu carga de trabajo, desde transformaciones tensoriales hasta multiplicaciones matriciales masivas. ¿El problema? Compilar para la arquitectura correcta, encajar flags, lidiar con errores de compilador y problemas de ABI puede volverse un desastre. Aquí entra Hugging Face con kernel-builder y la librería kernels: reproducibilidad con Nix, soporte para múltiples backends (CUDA, ROCm, Metal, XPU) y una forma ordenada de convertir tu código GPU en un operador nativo de PyTorch.
Qué vas a aprender y por qué importa
En esta guía técnica nos centramos en kernels ROCm. Usamos como ejemplo el kernel GEMM de RadeonFlow_Kernels, ganador del Grand Prize del AMD Developer Challenge 2025, optimizado para la GPU AMD Instinct MI300X y trabajando en formato FP8 e4m3fnuz.
¿Por qué te interesa esto? Porque si haces investigación, desarrollo de aceleradores o infraestructura ML, poder compilar, testear y publicar kernels ROCm reproducibles te ahorra horas —y errores— cuando quieres que otros usen tu trabajo en PyTorch.
Resumen técnico del kernel GEMM (lo esencial)
Tipo: FP8 block-wise GEMM (General Matrix Multiplication) para MI300X.
Formato: e4m3fnuz (FP8 con 4 bits de exponente y 3 de mantisa).
Escala: aplica per-block scaling (a_scale, b_scale) para mantener estabilidad numérica.
Asume shapes precompiladas y layout transpuesto (ajusta el launcher si necesitas más shapes).
Argumentos de la función nativa: a, b, a_scale, b_scale, c con shapes y dtypes específicos (entrada en FP8, salida en bf16, escalas en fp32). Esta convención garantiza throughput alto a costa de precisión reducida, útil en inferencia y ciertas cargas de entrenamiento cuantizadas.
Estructura del proyecto que espera kernel-builder
kernel-builder asume una estructura clara. Un ejemplo mínimo:
build.toml - el manifiesto del build
gemm/ - código fuente HIP/headers (kernel y launcher)
flake.nix - para reproducibilidad con Nix
torch-ext/ - bindings C++ que registran el operador en PyTorch
Ejemplo simplificado de build.toml (secciones clave):
backend = rocm y rocm-archs apuntan a arquitecturas específicas como gfx942 para MI300.
torch en depends indica que el kernel se integra como extensión de PyTorch.
Convenciones de archivos: .h vs .hip
Usa .h para headers con declaraciones, inline o templates que se incluyen.
Usa .hip para implementaciones HIP que necesitan compilarse por separado (launchers, device functions complejas).
Renombrar archivos .cpp a .h o .hip según corresponda ayuda a que kernel-builder identifique y compile correctamente.
El launcher y el binding: cómo se conectan GPU -> PyTorch
El launcher (por ejemplo gemm_launcher.hip) expone una función C run(...) que recibe punteros a memoria GPU y parámetros de shape. Dentro hace dispatch según la forma y llama al kernel más rápido disponible, o aborta si la shape no está soportada.
El binding C++ (torch-ext/torch_binding.cpp) valida tensores, extrae dimensiones y llama a run usando el stream HIP. Finalmente registra la operación en PyTorch con TORCH_LIBRARY_EXPAND para que puedas usar torch.ops o una interfaz Python más amigable.
En Python, torch-ext/gemm/__init__.py suele envolver esa llamada en una función gemm(a, b, a_scale, b_scale, out=None) que crea el tensor de salida en bfloat16 y delega en ops.gemm.
Reproducibilidad con Nix y flake.nix
Para builds reproducibles kernel-builder usa flake.nix. Un flake simple incluye la dependencia del repositorio kernel-builder y genera outputs reproducibles:
Antes de publicar limpia artefactos dev con build2cmake clean build.toml.
Publicar en el Hub de Hugging Face
Pasos esenciales:
hf repo create gemm para crear el repo en tu cuenta o en la org.
Inicializa git, configura LFS para binarios (git lfs track "*.so") y sube sólo los archivos necesarios:
Ejemplo de commit:
git add build/ gemm/ include/ src/utils tests/checker torch-ext flake.nix flake.lock build.toml
git commit -m "feat: Created a compliant gemm kernel"
git push -u origin main
Una vez en el Hub, otros desarrolladores pueden cargar el kernel sin instalarlo localmente usando la librería kernels.
Cargar y usar el kernel desde Python
Esto es lo más cómodo: el kernel se registra y se carga desde el Hub. Ejemplo de uso:
import torch
from kernels import get_kernel
gemm = get_kernel("kernels-community/gemm")
M, N, K = 1024, 1536, 7168
QUANT_SIZE = 128
device = torch.device("cuda")
A_fp32 = torch.randn(M, K, device=device)
B_fp32 = torch.randn(K, N, device=device)
A_fp8 = A_fp32.to(torch.float8_e4m3fnuz)
B_fp8 = B_fp32.to(torch.float8_e4m3fnuz)
A_scale = torch.ones(K // QUANT_SIZE, M, device=device, dtype=torch.float32)
B_scale = torch.ones(K // QUANT_SIZE, N // QUANT_SIZE, device=device, dtype=torch.float32)
C = torch.zeros(M, N, device=device, dtype=torch.bfloat16)
result = gemm.gemm(A_fp8, B_fp8, A_scale, B_scale, C)
Listo: tu kernel ROCm se usa como una librería remota y tu código PyTorch lo invoca como si fuera una extensión local.
Buenas prácticas y consejos técnicos
Mantén flake.lock en el repo para reproducibilidad exacta.
Usa cachix para acelerar builds y evitar que cada desarrollador compile todo desde cero.
Documenta las shapes y layouts soportados; el launcher suele fallar si la input shape no está precompilada.
Incluye tests que validen precisión y rendimiento (por ejemplo, con una carpeta tests/checker).
Si necesitas más shapes, modifica el launcher para añadir dispatch o implementar una fallback menos optimizada.
Controla qué archivos subes: evita artefactos dev innecesarios en el Hub.
Construir kernels ROCm ya no tiene que ser un laberinto de flags y errores crípticos. Con kernel-builder, Nix y la plataforma de Hugging Face puedes pasar de prototipo a un kernel compartible y usable por la comunidad con pasos claros y reproducibles. ¿Listo para que tu kernel corra en MI300X y lo usen miles de personas?
