NVIDIA CUDA Tile es un modelo de programación basado en GPU que busca la portabilidad para los núcleos Tensor de NVIDIA, maximizando así el rendimiento de la GPU. Una de las ventajas de CUDA Tile es la posibilidad de construir su propio DSL sobre él.
Este artículo detalla el trabajo de NVIDIA para integrar CUDA Tile como un backend para OpenAI Triton, un DSL de Python de código abierto diseñado para escribir kernels de DL para GPUs. OpenAI Triton soporta el cálculo en bloques, una técnica que divide datos y tareas en pequeñas secciones. Triton cuenta con un compilador basado en MLIR que genera PTX, lo que permite a los investigadores sin experiencia en CUDA escribir código eficiente para GPU.
¿Qué son CUDA Tile y CUDA Tile IR?
CUDA Tile amplía el modelo de programación CUDA para proporcionar un soporte de primera clase para la programación en bloques. Introducido en CUDA 13.1, representa un cambio de paradigma en la programación de GPU. En lugar de requerir que los desarrolladores piensen en términos de hilos individuales mediante el modelo SIMT, este modelo basado en bloques permite expresar cálculos a un nivel más alto de abstracción.
Los desarrolladores solo necesitan especificar operaciones sobre bloques de datos (tiles), mientras que el compilador y el sistema en tiempo de ejecución manejan automáticamente la planificación de hilos, el mapeo de hardware y la asignación de recursos. Este diseño reduce la complejidad de la programación y permite optimizaciones más agresivas por parte del compilador.
CUDA Tile IR es una representación intermedia y una infraestructura de compilador basada en MLIR. El desarrollo de CUDA Tile está impulsado por la especificación de CUDA Tile IR, que define la semántica formal, las operaciones y el sistema de tipos para cálculos basados en bloques en GPUs de NVIDIA.
¿Qué es Triton-to-TileIR?
El backend Triton-to-TileIR actúa como un puente para Triton, permitiéndole apuntar a CUDA Tile IR en lugar de PTX. Esto amplía el ecosistema del compilador de Triton, permitiendo a los desarrolladores compilar y ejecutar kernels de GPU escritos en OpenAI Triton al nuevo backend CUDA Tile IR. Este puente conecta lenguajes de programación de alto nivel (Triton) con el modelo de programación de GPU de próxima generación de NVIDIA, ofreciendo un camino fluido para aprovechar las capacidades modernas del hardware sin necesidad de reescribir el código.
A medida que la programación de GPU continúa evolucionando más allá de los modelos SIMT tradicionales hacia abstracciones basadas en bloques, esta integración permite a los desarrolladores beneficiarse de la sintaxis accesible de Python de Triton mientras acceden al soporte nativo de TileIR para los núcleos Tensor y la portabilidad arquitectónica.
Triton-to-TileIR democratiza el acceso a estas nuevas capacidades. Es importante destacar que Triton en sí mismo es un lenguaje de programación basado en bloques, donde los desarrolladores expresan cálculos en términos de secciones de datos (tiles) en lugar de hilos individuales, lo que se alinea conceptualmente con CUDA Tile IR.
Esto proporciona un camino directo para la compilación del backend: en lugar de compilar las abstracciones a nivel de bloques de Triton a código a nivel de hilos SIMT, Triton-to-TileIR preserva la semántica a nivel de bloques y compila directamente a CUDA Tile IR, que entiende nativamente el cálculo a granularidad de bloques.
La comunidad de usuarios existente de Triton puede aprovechar las ventajas de CUDA Tile IR sin necesidad de aprender nuevos lenguajes o reescribir el código existente. Una simple configuración de variable de entorno cambia el pipeline de compilación del backend PTX al backend CUDA Tile IR, desbloqueando así un rendimiento mejorado y compatibilidad arquitectónica a futuro.
Los usuarios de Triton podrán seleccionar qué backend (PTX o CUDA Tile IR) utilizar en función de cada kernel en sus aplicaciones.
Hoja de ruta de desarrollo de Triton-to-TileIR
Como un proyecto incubador dentro de la organización triton-lang, Triton-to-TileIR está en desarrollo activo. El repositorio sirve como un espacio de colaboración para implementar y refinar el backend CUDA Tile IR antes de su posible integración en el compilador principal de Triton.
Varios flujos de trabajo técnicos probablemente componen la hoja de ruta de desarrollo, incluyendo:
- Infraestructura de conversión central: Implementación de patrones de conversión de dialectos de MLIR para mapear las operaciones de Triton a equivalentes de CUDA Tile IR.
- Pruebas y validación: Desarrollo de suites de pruebas exhaustivas para verificar la corrección semántica de la transformación, incluyendo casos extremos en el flujo de control, patrones de acceso a la memoria y precisión numérica.
- Evaluación del rendimiento: Establecimiento de líneas base de rendimiento comparando kernels compilados con TileIR frente a equivalentes compilados con PTX en diversas operaciones (multiplicación de matrices, convoluciones, operaciones elemento a elemento, reducciones, etc.).
- Integración en proyectos de código abierto: Coordinación con la comunidad de código abierto para habilitar un mejor soporte del backend CUDA Tile IR en proyectos de código abierto, como Helion.
Cómo usar Triton-to-TileIR
Triton-to-TileIR actualmente solo admite compilación basada en código fuente. No hay binarios preconstruidos disponibles, lo que requiere construir el proyecto desde el código fuente en su entorno local.
Requisitos previos:
- Versión de CUDA: CUDA 13.1 o superior.
- Arquitectura de GPU: GPUs NVIDIA Blackwell (por ejemplo, GeForce RTX 5080); arquitecturas de GPU anteriores se habilitarán en próximas versiones de CUDA.
Construir desde el código fuente
Cuando se cumplan los requisitos previos, clone y construya el proyecto desde el código fuente:
# Clone el repositorio git clone https://github.com/triton-lang/Triton-to-tile-IR.git cd Triton-to-tile-IR # Construir e instalar # Siga las instrucciones de construcción específicas según el README del proyecto pip install -e .
Tenga en cuenta que los pasos de construcción detallados pueden variar. Consulte el README de Triton-to-TileIR y la documentación de construcción para configuraciones específicas de arquitectura, gestión de dependencias y orientación sobre solución de problemas.
Verificar la compilación de Tile IR
Después de construir, verifique la instalación ejecutando el tutorial de suma de vectores y confirmando que se está utilizando el backend Tile IR:
# Navegar al directorio del tutorial cd python/tutorials # Ejecutar el ejemplo de suma de vectores con Tile IR habilitado export ENABLE_TILE=1 python 01-vector-add.py
Cuando el backend Tile IR está activo, Triton almacena en caché los kernels compilados con extensiones de archivo .tileIR en lugar de los archivos .cubin estándar utilizados por el backend SIMT. Verifique estos archivos en caché:
# Encontrar el directorio de caché de Triton (típicamente en ~/.triton/cache)
Limitaciones de Triton-to-TileIR
Aunque Triton-to-TileIR abre nuevas posibilidades prometedoras, el proyecto se encuentra en una etapa de desarrollo relativamente temprana con varias limitaciones conocidas, incluyendo operaciones no soportadas y problemas temporales de rendimiento.
Operaciones no soportadas
No todas las operaciones soportadas por Triton están implementadas aún en el backend de Tile IR. Para más información sobre operaciones y características no soportadas o soportadas completamente.
A medida que CUDA continúa lanzando nuevas versiones, la compatibilidad del backend de Triton CUDA Tile IR seguirá mejorando.
Degradación del tensor de punteros y rendimiento subóptimo
El patrón de “tensor de punteros” en Triton, donde los tensores se componen de punteros para describir patrones de acceso a la memoria, muestra un rendimiento subóptimo en el backend de Tile IR con CUDA 13.1. Esta es una situación de rendimiento temporal. Para las cargas de trabajo afectadas, puede:
- Retornarse temporalmente al backend SIMT para ciertas operaciones críticas.
- Esperar las próximas optimizaciones en futuras versiones del proyecto.
- Refinar el código para adoptar la API de carga/almacenamiento TMA.
En cuanto al último punto, refinar el código para adoptar la API de carga/almacenamiento TMA: muchos de los tensores cargados en kernels tienen bloques contiguos y formas y desplazamientos bien definidos. Por lo tanto, materializar un tensor de punteros dentro del kernel ya no es necesario. En su lugar, esta información de diseño puede pasarse a la API de carga/almacenamiento TMA, permitiendo al backend Tile IR obtener un mejor rendimiento.
Por ejemplo, un patrón típico de tensor de punteros puede verse de la siguiente manera:
# Antes: estilo tensor-de-punteros offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) offs_k = tl.arange(0, BLOCK_K) a_ptrs = a_ptr + (offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak) b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn) a = tl.load(a_ptrs) b = tl.load(b_ptrs)
Aquí, cada elemento en a_ptrs es un puntero explícito calculado en el kernel, aunque el bloque en sí es contiguo y su diseño puede describirse completamente por (shape, strides, block_shape).
Con TMA, la misma operación puede reescribirse como:
desc_a = tl.make_tensor_descriptor( a, # puntero base shape=(M, K), strides=(stride_am, stride_ak), block_shape=(BLOCK_M, BLOCK_K) # tamaño del bloque ) desc_b = tl.make_tensor_descriptor( b, shape=(K, N), strides=(stride_bk, stride_bn), block_shape=(BLOCK_K, BLOCK_N) ) offs_m = pid_m * BLOCK_M offs_n = pid_n * BLOCK_N a_tile = desc_a.load([offs_m, 0]) # [BLOCK_M, BLOCK_K] b_tile = desc_b.load([0, offs_n]) # [BLOCK_K, BLOCK_N] desc_c.store([offs_m, offs_n], acc) # Almacenamiento respaldado por TMA
Aprende más sobre Triton-to-TileIR
El proyecto Triton-to-TileIR representa un paso significativo en la evolución de la programación de GPU, cerrando la brecha entre la productividad del desarrollador y la eficiencia del hardware. Al permitir que el modelo de programación basado en bloques y accesible de Triton apunte al conjunto de instrucciones virtuales CUDA Tile IR, la integración promete ofrecer rendimiento, portabilidad y preparación para el futuro para los practicantes de machine learning y desarrolladores de GPU.
Para los desarrolladores que ya usan Triton, el backend TileIR ofrecerá un camino para aprovechar arquitecturas de GPU de próxima generación con cambios mínimos en el código. Para el ecosistema más amplio de programación de GPU, esta colaboración demuestra cómo las asociaciones estratégicas entre diseñadores de lenguajes y proveedores de hardware pueden crear beneficios acumulativos, haciendo que las capacidades avanzadas del hardware sean accesibles sin sacrificar las abstracciones de alto nivel que permiten una rápida innovación.
A medida que el proyecto madura y pasa de la incubación a la preparación para producción, será fascinante observar cómo la integración influye tanto en la adopción de Triton como en la trayectoria más amplia de la programación de GPU basada en bloques. La métrica de éxito final será simple: ¿pueden los investigadores con experiencia limitada en GPU escribir código Triton que se ejecute con un rendimiento casi óptimo en GPUs de NVIDIA?
Para aprender más, visita el repositorio triton-lang/Triton-to-tile-IR y las Consejos de optimización del rendimiento para el backend de CUDA Tile IR.
