Compiler-Grounded Hierarchical Diagnosis for LLM-Based Triton Kernel Optimization

Resumen Técnico: Diagnóstico Jerárquico Basado en el Compilador para la Optimización de Kernels de Triton mediante LLM

1. Planteamiento del Problema

Si bien los Modelos de Lenguaje de Gran Escala (LLM) han avanzado en la generación y optimización automatizada de kernels, los enfoques existentes dependen predominantemente de "señales superficiales", tales como errores de compilación, métricas de tiempo de ejecución y datos de perfilado. Estas señales indican que un kernel es lento, pero rara vez explican por qué el compilador de backend no logra realizar una optimización rentable. Esto es particularmente agudo en aceleradores emergentes como las NPU (p. ej., Ascend), donde la pila del compilador está evolucionando, es específica para el objetivo y opaca para los desarrolladores.

El desafío central identificado es una brecha semántica entre el código fuente a nivel de operador y el comportamiento de rendimiento del backend. Los fallos de rendimiento suelen emerger solo después de que el código pasa por las etapas de representación intermedia (IR), decisiones de programación (scheduling) y pasadas de compilación específicas del backend. En consecuencia, los agentes de LLM, típicamente preentrenados en patrones centrados en GPU, se comportan como optimizadores de ensayo y error: proponen ediciones de código fuente plausibles, pero carecen de la evidencia para mapear los síntomas de tiempo de ejecución con causas específicas a nivel de compilador.

2. Metodología

El artículo propone un marco de optimización jerárquico basado en el compilador que trata la optimización de kernels como un problema de diagnóstico progresivo entre capas. En lugar de tratar el fundamento del compilador (compiler grounding) como una característica siempre activa, el sistema escala desde evidencia ligera hacia un análisis profundo del compilador solo cuando es necesario.

El Flujo de Trabajo Jerárquico:
El sistema opera a través de un pipeline de bucle cerrado organizado en cuatro niveles de escalada:

• L1: Triaje de Patrones (Pattern Triage): El sistema primero verifica hipótesis respaldadas por patrones fuertes utilizando un índice curado de patrones de optimización, heurísticas y reglas de expertos (incluyendo heurísticas internas destiladas de ejecuciones previas). Si una reescritura de bajo riesgo está justificada por la evidencia existente, el proceso se detiene aquí.
• L2: Diagnóstico por Perfilado (Profiling Diagnosis): Si los patrones fallan, el sistema recopila evidencia de tiempo de ejecución (latencia, rendimiento, contadores de hardware) para localizar cuellos de botella (p. ej., movimiento de memoria, actividad escalar, sincronización). Esto identifica el síntoma, pero no necesariamente el mecanismo.
• L3: Atribución de IR (IR Attribution): Cuando el perfilado es insuficiente, el sistema inspecciona las Representaciones Intermedias (IR) para traducir los síntomas de tiempo de ejecución en diagnósticos estructurales (p. ej., cargas de paso no unitario, falta de fusión, fallos de vectorización inducidos por el diseño/layout). Reconstruye la trayectoria de reducción (lowering) del operador para identificar las transformaciones de etapa a etapa.
• L4: Escalada de Compilador-Fuente (Compiler-Source Escalation): Como último recurso, el sistema invoca el análisis fundamentado en el compilador para explicar por qué un patrón diagnosticado bloquea la optimización del backend. Consulta el conocimiento del compilador (comportamiento de las pasadas, reglas de reducción, restricciones) para recuperar las condiciones a nivel de código fuente que una reescritura exitosa debe satisfacer.

Arquitectura del Sistema y Gestión de Estado:

• Diseño Centrado en Agentes: Implementado como un entorno centrado en agentes (integrado en el repositorio de habilidades de CannBot) en lugar de un servicio de compilación especializado. Utiliza un armazón de ejecución ligero para gestionar la validación, la corrección y la recolección de evidencia, mientras que las "habilidades" específicas del dominio manejan la lógica de optimización.
• Rondas con Estado (Stateful Rounds): El proceso de optimización se organiza en rondas explícitas y auditables. Cada ronda registra su hipótesis, fuentes de evidencia y resultado de validación. Este diseño con estado permite que el sistema sea reanudable y asegura que los análisis más profundos (IR/Compilador) se adjunten como evidencia reutilizable en lugar de observaciones efímeras.
• Síntesis Retrospectiva: Más allá del bucle en línea, el sistema realiza una pasada de síntesis transversal entre operadores. Mapea los historiales de optimización validados a familias de patrones existentes o crea nuevos patrones reutilizables, permitiendo que la capa "más barata" (Triaje de Patrones) mejore con el tiempo.

3. Contribuciones Clave

1. Formulación del Problema: Los autores formulan la optimización de kernels basada en LLM como un problema de diagnóstico progresivo entre capas, identificando la falta de evidencia fundamentada en el compilador como la razón principal por la cual los agentes basados solo en perfilado fallan en optimizaciones específicas del backend.
2. Flujo de Trabajo Jerárquico: Describen un flujo de trabajo que escala desde el triaje de patrones hasta el perfilado, la atribución de IR y la escalada de compilador-fuente, proponiendo reescrituras de nivel de fuente respaldadas por evidencia solo cuando se requiere evidencia más profunda.
3. Evidencia de Evaluación: Presentan evidencia de evaluación estandarizada y sensible a la distribución en un benchmark de Ascend 950 derivado de NPUKernelBench, reportando ganancias agregadas, la mejor ronda y trazas representativas entre capas.

4. Resultados Experimentales

El sistema se implementó para Triton en NPUs Ascend y se evaluó en 37 operadores convertidos con éxito de un benchmark estandarizado de Ascend 950 (derivado de NPUKernelBench). Cada operador fue probado en 5 casos (185 casos totales).

• Ganancias de Rendimiento:
  • Media Geométrica de Aceleración (Speedup): 4.35× (desde el kernel Triton inicial hasta el optimizado).
  • Mediana de Aceleración: 2.73×.
  • Distribución: 22/37 entradas superaron un aumento de velocidad de 2×; 13/37 superaron 5×. Los resultados muestran una distribución de cola larga, con ganancias sustanciales coexistiendo con entradas cercanas al baseline.
• Contabilidad del Proceso:
  • El índice de la mediana de la mejor ronda fue 10, lo que indica que los resultados óptimos a menudo requirieron múltiples iteraciones.
  • Profundidad de la Evidencia: Aunque 33 de las rondas "mejores" seleccionadas se atribuyeron al Triaje de Patrones, el sistema demostró la utilidad de las capas más profundas: 32 trayectorias alcanzaron el Diagnóstico por Perfilado, 3 la Atribución de IR y 4 la Escalada de Compilador-Fuente.
  • Estudio de Caso: Una traza representativa (relleno de tamaño de Hyena FFT) mostró cómo el perfilado identificó un cuello de botella escalar, la atribución de IR identificó costos de sincronización específicos en la reducción (lowering), y el análisis de compilador-fuente explicó la restricción, logrando una robusta aceleración de 2.08×.

5. Significado y Reivindicaciones

El artículo sostiene que el fundamento del compilador (compiler grounding) debe ser selectivo en lugar de estar siempre activo. Al tratar la optimización como un diagnóstico progresivo, el sistema hace que el razonamiento profundo del compilador sea práctico, invocándolo solo cuando una evidencia más ligera no puede producir una reescritura suficientemente fundamentada.

Los autores enfatizan que sus resultados demuestran un margen sustancial en un subconjunto de operadores, pero no reclaman una mejora uniforme en todos los casos. Establecen explícitamente que la evaluación actual no proporciona aislamiento causal (por ejemplo, comparando la escalada progresiva contra estrategias de solo perfilado o de profundidad constante bajo una herramienta de compilación congelada). El artículo se posiciona como un "argumento de sistemas" apoyado por evidencia transparente de benchmarks y procesos, reservando la descomposición causal de la jerarquía para futuros experimentos controlados.

El trabajo destaca que la relación entre el código fuente y el rendimiento en tiempo de ejecución está mediada por transformaciones invisibles del compilador, y que cerrar esta brecha con un enfoque jerárquico basado en la evidencia es esencial para optimizar kernels en aceleradores emergentes y opacos como las NPU.