Triton 3.7 permite extender el compilador de GPU sin forks

Resumen

La version PyTorch-Triton 3.7 introduce el sistema Triton Plugin Extensions, un framework para cargar de forma dinamica pases de compilacion personalizados, dialectos (incluidas sus operaciones) y extensiones de DSL en el Triton oficial en tiempo de ejecucion, 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 estandar con un rendimiento que iguala o supera a las librerias del fabricante, tanto en NVIDIA H100 como en AMD MI350.

¿Por que hacen falta las extensiones?

Escribir kernels de GPU de alto rendimiento a menudo exige ir mas alla de lo que ofrece el pipeline por defecto del compilador Triton. Los pases de optimizacion personalizados, los intrinsecos especificos del hardware y los patrones especializados de gestion de memoria son esenciales para exprimir las ultimas gotas de rendimiento en cargas de produccion. Hasta ahora, habilitar estas capacidades significaba mantener un fork de Triton, y esos forks tienen costos reales.

Los forks quedan rapidamente atras del proyecto original. Cada actualizacion upstream arriesga conflictos de fusion, APIs rotas y cambios sutiles de comportamiento que exigen una reconciliacion cuidadosa. Los equipos que se quedan fijados a una version 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 nucleo de Triton. Un sistema de plugins que cargue extensiones de forma dinamica en tiempo de ejecucion permitiria a investigadores e ingenieros iterar sobre funciones personalizadas a maxima velocidad, siempre corriendo sobre la ultima version oficial, y publicar resultados sin esperar a que los cambios se integren en el repositorio principal.

El sistema Triton Plugin Extensions

La version PyTorch-Triton 3.7 entrega exactamente eso: un sistema de extensiones por plugins de proposito general que abarca todo el pipeline de compilacion e integrado en el Triton oficial. Los plugins son librerias compartidas (archivos .so) que se descubren y cargan en tiempo de ejecucion 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 el y las extensiones quedan disponibles de inmediato.

Pipeline de compilacion sobrescribible

En el corazon 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 reduccion, desde la representacion intermedia de Triton de alto nivel (TTIR) pasando por TritonGPU IR (TTGIR) hasta LLVM IR y el ensamblador especifico del objetivo (PTX, AMDGCN). Con estos hooks, los plugins pueden:

• Insertar uno o mas pases personalizados en puntos arbitrarios de cualquier etapa.
• Deshabilitar pases especificos dentro de una etapa.
• Reemplazar pases existentes por implementaciones personalizadas (por ejemplo, una estrategia propia de especializacion de warps).
• Sobrescribir etapas completas o el pipeline entero.

Esto esta disponible tanto en el backend de NVIDIA como en el de AMD.

Operaciones, dialectos y reduccion personalizados

La API de plugins esta disenada para complementar a PyBind11 y habilita tres niveles de extensibilidad:

• Pases de transformacion personalizados: pases individuales que se pueden insertar en puntos arbitrarios del pipeline sin un dialecto asociado.
• Dialectos MLIR y pases de conversion personalizados: dialectos compilados por separado y cargados en Triton, con pases de plugin que reescriben patrones estandar de Triton IR hacia operaciones de dialectos propios para una reduccion especializada.
• Operaciones DSL de nivel superior: nueva sintaxis y semantica a nivel de Python que habilitan abstracciones de programacion completamente nuevas sin alterar el propio Triton.

Control por kernel

Los plugins se pueden activar y desactivar de forma dinamica a nivel de kernel. Un hook de compilador fijado en el codigo del kernel activa un pipeline personalizado para todos los kernels llamados despues, hasta que se desactive. No hay limite en cuantos pipelines personalizados se pueden definir, y cada plugin es responsable de implementar su propia estrategia de hashing para la gestion de la cache de kernels, de modo que la recompilacion se dispare solo cuando sea necesario. La libreria 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 gestion explicita de memoria y la segmentacion asincronica de computo y carga. Las TLX dan a los autores de kernels control directo sobre la asignacion de memoria compartida, el movimiento de datos y la planificacion de instrucciones, capacidades criticas para escribir kernels persistentes que saturen el hardware moderno de GPU.

Antes, usar TLX requeria 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 quedaran 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 programacion funciona entre fabricantes de hardware, aunque por debajo se mapea a las funciones especificas de cada uno.

NVIDIA H100 (Hopper): GEMM persistente

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

@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 estandar con el plugin de extension TLX frente a cuBLAS. Como el sistema de plugins produce exactamente el mismo codigo generado que el fork compilado, estos resultados aplican a ambos caminos.

En las formas grandes y limitadas por computo que dominan las cargas de LLM en produccion, 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 libreria del fabricante:

AMD MI350: GEMM segmentado

En las GPU AMD MI350, las TLX usan segmentacion explicita basada en registros con operaciones local_store y local_load. Se aplica el mismo patron de gestion de buffers, pero el camino de movimiento de datos pasa por registros en lugar de unidades asincronicas 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 estandar con el plugin de extension TLX frente a rocBLAS. Al producir el mismo codigo generado, los resultados aplican por igual a ambos caminos.

Triton con TLX entrega entre 12 y 15% mas TFLOPS de forma consistente en todos los tamanos de matriz probados, lo que confirma que el camino cargado por plugin no agrega sobrecarga y supera a la libreria 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 librerias del fabricante.