Triton 3.7 permite extender el compilador de GPU sin forks

Resumen

La versión PyTorch-Triton 3.7 introduce el sistema Triton Plugin Extensions, un framework para cargar de forma dinámica pases de compilación personalizados, dialectos (incluidas sus operaciones) y extensiones de DSL en el Triton oficial en tiempo de ejecución, sin necesidad de mantener un fork ni recompilar. Como primer gran consumidor de este sistema, las Triton Language Extensions (TLX) de Meta quedan habilitadas por defecto, lo que lleva kernels GEMM persistentes y control fino del hardware al Triton estándar con un rendimiento que iguala o supera a las librerías del fabricante, tanto en NVIDIA H100 como en AMD MI350.

¿Por qué hacen falta las extensiones?

Escribir kernels de GPU de alto rendimiento a menudo exige ir más allá de lo que ofrece el pipeline por defecto del compilador Triton. Los pases de optimización personalizados, los intrínsecos específicos del hardware y los patrones especializados de gestión de memoria son esenciales para exprimir las últimas gotas de rendimiento en cargas de producción. Hasta ahora, habilitar estas capacidades significaba mantener un fork de Triton, y esos forks tienen costos reales.

Los forks quedan rápidamente atrás del proyecto original. Cada actualización upstream arriesga conflictos de fusión, APIs rotas y cambios sutiles de comportamiento que exigen una reconciliación cuidadosa. Los equipos que se quedan fijados a una versión bifurcada terminan atrapados en versiones obsoletas, sin poder aprovechar las correcciones de errores, el soporte de nuevo hardware ni las mejoras de la comunidad. La carga de mantenimiento se acumula con el tiempo, y el fork se vuelve un cuello de botella en lugar de un acelerador.

Lo que se necesita es una forma de extender el pipeline del compilador Triton (agregar pases, operaciones e incluso dialectos completos) sin modificar el núcleo de Triton. Un sistema de plugins que cargue extensiones de forma dinámica en tiempo de ejecución permitiría a investigadores e ingenieros iterar sobre funciones personalizadas a máxima velocidad, siempre corriendo sobre la última versión oficial, y publicar resultados sin esperar a que los cambios se integren en el repositorio principal.

El sistema Triton Plugin Extensions

La versión PyTorch-Triton 3.7 entrega exactamente eso: un sistema de extensiones por plugins de propósito general que abarca todo el pipeline de compilación e integrado en el Triton oficial. Los plugins son librerías compartidas (archivos .so) que se descubren y cargan en tiempo de ejecución mediante la variable de entorno TRITON_PLUGIN_PATHS. No hace falta recompilar Triton: basta con instalar el paquete del plugin, apuntar la variable de entorno hacia él y las extensiones quedan disponibles de inmediato.

Pipeline de compilación sobrescribible

En el corazón del sistema hay un conjunto de hooks incrustados en las etapas del compilador de backend (compiler.py) de Triton. Estos hooks dan control fino sobre el pipeline de pases MLIR en cada nivel de reducción, desde la representación intermedia de Triton de alto nivel (TTIR) pasando por TritonGPU IR (TTGIR) hasta LLVM IR y el ensamblador específico del objetivo (PTX, AMDGCN). Con estos hooks, los plugins pueden:

• Insertar uno o más pases personalizados en puntos arbitrarios de cualquier etapa.
• Deshabilitar pases específicos dentro de una etapa.
• Reemplazar pases existentes por implementaciones personalizadas (por ejemplo, una estrategia propia de especialización de warps).
• Sobrescribir etapas completas o el pipeline entero.

Esto está disponible tanto en el backend de NVIDIA como en el de AMD.

Operaciones, dialectos y reducción personalizados

La API de plugins está diseñada para complementar a PyBind11 y habilita tres niveles de extensibilidad:

• Pases de transformación personalizados: pases individuales que se pueden insertar en puntos arbitrarios del pipeline sin un dialecto asociado.
• Dialectos MLIR y pases de conversión personalizados: dialectos compilados por separado y cargados en Triton, con pases de plugin que reescriben patrones estándar de Triton IR hacia operaciones de dialectos propios para una reducción especializada.
• Operaciones DSL de nivel superior: nueva sintaxis y semántica a nivel de Python que habilitan abstracciones de programación completamente nuevas sin alterar el propio Triton.

Control por kernel

Los plugins se pueden activar y desactivar de forma dinámica a nivel de kernel. Un hook de compilador fijado en el código del kernel activa un pipeline personalizado para todos los kernels llamados después, hasta que se desactive. No hay límite en cuántos pipelines personalizados se pueden definir, y cada plugin es responsable de implementar su propia estrategia de hashing para la gestión de la caché de kernels, de modo que la recompilación se dispare solo cuando sea necesario. La librería utlx maneja esto por completo para el usuario.

# Habilitar el plugin TLX es tan simple como fijar una variable de entorno
import os
import sysconfig

dist_packages = sysconfig.get_paths()["purelib"]
libutlx_path = os.path.join(dist_packages, "utlx_plugin", "libutlx.so")
os.environ["TRITON_PLUGIN_PATHS"] = libutlx_path

TLX: las extensiones de lenguaje de Meta, ahora integradas

Las Triton Language Extensions (TLX) son un conjunto de operaciones conscientes del hardware desarrolladas por Meta para la gestión explícita de memoria y la segmentación asincrónica de cómputo y carga. Las TLX dan a los autores de kernels control directo sobre la asignación de memoria compartida, el movimiento de datos y la planificación de instrucciones, capacidades críticas para escribir kernels persistentes que saturen el hardware moderno de GPU.

Antes, usar TLX requería compilar desde el fork experimental de Triton de Meta. Con el sistema de extensiones por plugins, las TLX se distribuyen ahora como un paquete de Python independiente (utlx) que funciona con el Triton oficial sin modificar. A partir de PyTorch-Triton 3.7, las TLX quedarán habilitadas por defecto en todas las versiones de Triton en adelante.

Multiplataforma: NVIDIA H100 y AMD MI350

Una de las ventajas clave de las TLX es que el mismo modelo de programación funciona entre fabricantes de hardware, aunque por debajo se mapea a las funciones específicas de cada uno.

NVIDIA H100 (Hopper): GEMM persistente

En las GPU Hopper, las TLX se mapean a las cargas asincrónicas nativas TMA (Tensor Memory Accelerator) y a las instrucciones WGMMA (Warp Group Matrix Multiply-Accumulate). El kernel GEMM persistente usa segmentación por software multietapa con grupos de commit y wait asincrónicos:

@triton.jit
def matmul_kernel_pipelined_hopper(
    a_ptr, b_ptr, c_ptr, M, N, K,
    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,
    NUM_STAGES: tl.constexpr,
):
    # ... indexado de tiles ...

    # Reserva de buffers de memoria compartida multietapa
    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tlx.dtype_of(a_ptr), NUM_STAGES)
    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tlx.dtype_of(b_ptr), NUM_STAGES)

    # Prologo del pipeline de prefetch
    for i in tl.range(0, NUM_STAGES - 1, loop_unroll_factor=NUM_STAGES - 1):
        a = tlx.local_view(buffers_A, i)
        b = tlx.local_view(buffers_B, i)
        token_a = tlx.async_load(a_ptrs, a, mask=...)
        token_b = tlx.async_load(b_ptrs, b, mask=...)
        tlx.async_load_commit_group([token_a, token_b])

    # Bucle principal en K con computo y movimiento de datos superpuestos
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), num_stages=0):
        buf = k % NUM_STAGES
        tlx.async_load_wait_group(NUM_STAGES - 2)
        acc = tlx.async_dot(
            tlx.local_view(buffers_A, buf),
            tlx.local_view(buffers_B, buf),
            acc
        )
        # Prefetch de la siguiente etapa ...

    acc = tlx.async_dot_wait(0, acc)
    # Guardar resultados ...

NVIDIA H100 (Hopper): resultados de rendimiento

La tabla siguiente muestra el rendimiento GEMM en FP16 sobre NVIDIA H100, comparando el Triton estándar con el plugin de extensión TLX frente a cuBLAS. Como el sistema de plugins produce exactamente el mismo código generado que el fork compilado, estos resultados aplican a ambos caminos.

En las formas grandes y limitadas por cómputo que dominan las cargas de LLM en producción, Triton con TLX iguala a cuBLAS en el GEMM cuadrado y lo supera en las formas anchas y grandes, lo que confirma que el camino cargado por plugin no agrega sobrecarga y a la vez alcanza o supera a la librería del fabricante:

AMD MI350: GEMM segmentado

En las GPU AMD MI350, las TLX usan segmentación explícita basada en registros con operaciones local_store y local_load. Se aplica el mismo patrón de gestión de buffers, pero el camino de movimiento de datos pasa por registros en lugar de unidades asincrónicas de hardware:

@triton.jit
def matmul_kernel_pipelined_mi300(
    a_ptr, b_ptr, c_ptr, M, N, K, ...
):
    # ... indexado de tiles ...

    # Reserva de buffers de memoria compartida
    buffers_A = tlx.local_alloc((BLOCK_SIZE_M, BLOCK_SIZE_K), tlx.dtype_of(a_ptr), NUM_STAGES - 1)
    buffers_B = tlx.local_alloc((BLOCK_SIZE_K, BLOCK_SIZE_N), tlx.dtype_of(b_ptr), NUM_STAGES - 1)

    # Prologo: cargar a memoria compartida via registros
    for i in tl.range(0, NUM_STAGES - 1, loop_unroll_factor=NUM_STAGES - 1):
        a_reg = tl.load(a_ptrs, mask=...)
        b_reg = tl.load(b_ptrs, mask=...)
        tlx.local_store(tlx.local_view(buffers_A, i), a_reg)
        tlx.local_store(tlx.local_view(buffers_B, i), b_reg)

    # Bucle principal: operaciones de computo y memoria superpuestas
    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for k in tl.range(NUM_STAGES - 1, K_ITERS, num_stages=0):
        # Cargar el siguiente tile en registros
        a_reg = tl.load(a_ptrs, mask=...)
        b_reg = tl.load(b_ptrs, mask=...)

        # Computo sobre los datos previamente preparados
        a_prev = tlx.local_load(tlx.local_view(buffers_A, buf))
        b_prev = tlx.local_load(tlx.local_view(buffers_B, buf))
        acc = tl.dot(a_prev, b_prev, acc)

        # Guardar los nuevos datos en memoria compartida para la siguiente iteracion
        tlx.local_store(tlx.local_view(buffers_A, ...), a_reg)
        tlx.local_store(tlx.local_view(buffers_B, ...), b_reg)

AMD MI350: resultados de rendimiento

La tabla siguiente muestra el rendimiento GEMM en FP16 sobre AMD MI350, comparando el Triton estándar con el plugin de extensión TLX frente a rocBLAS. Al producir el mismo código generado, los resultados aplican por igual a ambos caminos.

Triton con TLX entrega entre 12 y 15% más TFLOPS de forma consistente en todos los tamaños de matriz probados, lo que confirma que el camino cargado por plugin no agrega sobrecarga y supera a la librería del fabricante:

En conjunto, el sistema de extensiones por plugins convierte lo que antes era un fork de mantenimiento costoso en un paquete que se instala y se activa con una variable de entorno, sin renunciar al rendimiento de las librerías del fabricante.