🤖 AI

TileMix: Tile-Centric Mixed-Precision Attention for LLM Inference Acceleration

TileMix es un kernel de precisión mixta centrado en teselas y libre de entrenamiento que acelera la inferencia de LLM de contexto largo mediante el enrutamiento dinámico de teselas de puntuación de atención alineadas con el hardware entre rutas FP16 e INT8 dentro de una operación de atención densa fusionada, recuperando así la precisión perdida por los métodos uniformes de baja precisión mientras mejora el rendimiento.

Autores originales: Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

Publicado 2026-08-19

📖 6 min de lectura🧠 Análisis profundo

CC BY 4.0

Autores originales: Hanzhi Zhang, Qiao Zhang, Qinglei Cao, Heng Fan, Yan Huang, Kewei Sha, Yunhe Feng

Artículo original bajo licencia CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta es una explicación generada por IA del artículo a continuación. No ha sido escrita ni avalada por los autores. Para mayor precisión técnica, consulte el artículo original. Leer descargo de responsabilidad completo

Los modelos de lenguaje extensos se han convertido en los motores detrás de una nueva generación de inteligencia artificial, capaces de resumir libros enteros, responder preguntas complejas de documentos largos y mantener conversaciones que abarcan miles de palabras. Para hacer esto, estos modelos dependen de un mecanismo llamado atención, que les permite ponderar la importancia de cada palabra en una oración frente a cada otra palabra. Cuando el texto es corto, este proceso es rápido. Pero a medida que el contexto crece para incluir capítulos enteros o contratos legales, el costo computacional explota. El modelo debe calcular una puntuación para cada par posible de palabras, creando una enorme cuadrícula de datos que demanda cantidades masivas de memoria y potencia de procesamiento. Este cuello de botella ha dificultado la ejecución eficiente de estas poderosas herramientas en hardware estándar, especialmente cuando se trata de las secuencias largas requeridas para tareas del mundo real como el análisis legal o la revisión de registros médicos.

Los investigadores han intentado resolver esto simplificando las matemáticas. Un enfoque común es reducir la precisión de los números que utiliza la computadora, pasando de cálculos de alta precisión a otros de menor precisión y más rápidos. Esto es como cambiar la medición de ingredientes con una báscula de laboratorio por el uso de una cuchara de cocina; es mucho más rápido, pero si se hace para cada uno de los pasos de una receta compleja, el plato final podría tener un sabor incorrecto. Otro enfoque es saltarse cálculos por completo, ignorando las palabras que parecen poco importantes. Si bien esto ahorra tiempo, conlleva el riesgo de cortar las conexiones mismas que el modelo necesita para entender la historia. El desafío ha sido encontrar una manera de utilizar la velocidad de las matemáticas de baja precisión sin sacrificar la exactitud de las de alta precisión, todo esto manteniendo la capacidad del modelo para ver el panorama completo.

Un equipo de investigadores ha introducido un nuevo método llamado TileMix que aborda este problema tratando el proceso de atención no como un bloque de trabajo único y uniforme, sino como una colección de teselas (tiles) más pequeñas y manejables. Imagine la enorme cuadrícula de puntuaciones de pares de palabras como un gran mosaico. En lugar de pintar todo el mosaico con un solo tipo de pintura, TileMix lo divide en pequeñas secciones cuadradas alineadas con el hardware. Para cada una de estas secciones, el sistema toma una decisión de milisegundos: ¿necesita este grupo específico de pares de palabras el cálculo lento de alta precisión, o puede conformarse con el de baja precisión y más rápido? Esta decisión se basa en un mapa preplanificado que agrupa estas teselas, permitiendo que la computadora cambie entre alta y baja precisión dentro de la misma operación sin detenerse a reorganizar los datos.

El núcleo de esta innovación es un sistema de enrutamiento que actúa como un controlador de tráfico para el procesador de la computadora. Empaqueta estas decisiones en un código compacto, esencialmente una cadena de bits que le dice al procesador qué camino tomar para cada tesela. Si una tesela está marcada para alta precisión, el procesador utiliza sus unidades matemáticas más precisas. Si está marcada para baja precisión, cambia a sus unidades más rápidas y eficientes energéticamente. Crucialmente, ambos caminos actualizan un estado de memoria compartido que realiza el seguimiento de los resultados generales, asegurando que la salida final sea consistente. Esto permite que el sistema preserve la conectividad total del modelo —lo que significa que nunca se ignoran las interacciones entre palabras— mientras se aprovechan los beneficios de velocidad de las matemáticas de baja precisión para las partes del cálculo que pueden tolerarlo.

En sus experimentos, los investigadores probaron este método en varios modelos de lenguaje extensos populares, incluyendo LLaMA, Qwen y Vicuna, utilizando secuencias que van desde las 16,000 hasta las 64,000 palabras. Compararon su enfoque con el método estándar de alta precisión y una versión que utilizaba matemáticas de baja precisión para todo. Los resultados mostraron que usar matemáticas de baja precisión para todo el cálculo a menudo provocaba una caída significativa en la precisión, causando que el modelo perdiera detalles o fallara en tareas de recuperación. Sin embargo, TileMix fue capaz de recuperar la mayor parte de esa calidad perdida. Al enrutar cuidadosamente solo grupos específicos de teselas hacia la ruta de alta precisión, el sistema mantuvo niveles de precisión muy cercanos a la línea base completa de alta precisión, logrando simultáneamente velocidades de procesamiento mucho más rápidas.

El estudio también exploró diferentes patrones para decidir qué teselas deberían recibir el tratamiento de alta precisión. Encontraron que la disposición importaba; algunos patrones, que mantenían los cálculos de alta precisión en regiones espaciales específicas de la cuadrícula de palabras, funcionaban mejor que otros dependiendo de la tarea. Por ejemplo, ciertos diseños que preservaban la alta precisión para las interacciones locales de palabras funcionaban mejor en tareas de recuperación de hechos, mientras que otros eran más robustos para la respuesta a preguntas generales. Los investigadores demostraron que este método puede aplicarse sin reentrenar los modelos, lo que significa que puede desplegarse inmediatamente en sistemas existentes. También mostraron que el enfoque funciona bien con lotes de longitud variable y diferentes arquitecturas de modelos, lo que sugiere que es una herramienta flexible para mejorar la eficiencia.

En última instancia, TileMix ofrece un equilibrio controlable entre velocidad y precisión. Sugiere que el futuro del procesamiento de contexto largo eficiente no requiere elegir entre ser rápido o ser inteligente. En cambio, al mezclar inteligentemente ambos, los sistemas pueden procesar documentos largos a una velocidad que se acerca a los límites teóricos del hardware, sin las severas penalizaciones de precisión que conlleva el uso de matemáticas de baja precisión en todas partes. Este enfoque proporciona un camino práctico hacia el despliegue de modelos de lenguaje extensos en escenarios donde tanto la velocidad como la precisión son críticas, como el análisis en tiempo real de conjuntos de datos masivos o herramientas interactivas que necesitan comprender contextos largos y complejos instantáneamente.

Resumen Técnico: TileMix: Atención de Precisión Mixta Centrada en Tiles para la Aceleración de la Inferencia de LLM

Planteamiento del Problema

La inferencia de contexto largo en Modelos de Lenguaje de Gran Escala (LLMs) enfrenta un cuello de botella significativo durante la fase de prefill debido a la complejidad computacional cuadrática ( $O(L^2)$ ) de la atención densa de auto-atención. Las estrategias de aceleración existentes generalmente se dividen en tres categorías, cada una con limitaciones:

Cuantización de Baja Precisión: Los métodos como la cuantización INT8 mejoran la eficiencia aritmética y de memoria, pero típicamente aplican una ruta de precisión uniforme en toda la matriz de atención. Esto a menudo deja el enrutamiento de precisión espacial fuera del bucle de atención fusionada, y la ejecución uniforme de baja precisión puede degradar la calidad en contextos largos.
Métodos Basados en la Esparcidad: Los enfoques que seleccionan interacciones de tokens activos (por ejemplo, ventanas deslizantes, patrones esparcidos) reducen la computación pero alteran la conectividad densa del grafo de atención, lo que puede perder interacciones críticas de tokens.
Atención Fusionada Consciente de IO: Los kernels como FlashAttention particionan la atención en tiles alineados con el hardware para optimizar el movimiento de datos y fusionar operaciones (cómputo de puntuación, softmax en línea, agregación de valores). Sin embargo, estos kernels tradicionalmente mantienen una ruta de cómputo de puntuación uniforme, perdiendo la oportunidad de enrutar regiones espaciales específicas a diferentes niveles de precisión dentro de una misma ejecución.

El desafío central abordado es cómo integrar la precisión heterogénea (FP16 e INT8) en un único kernel de atención fusionada sin romper la conectividad densa de los tokens ni incurrir en una sobrecarga significativa de metadatos, asegurando al mismo tiempo que los diferentes caminos aritméticos converjan correctamente hacia un estado de softmax en línea compartido.

Metodología: TileMix

TileMix introduce un kernel de enrutamiento de precisión centrado en tiles que trata la precisión numérica como una decisión espacial sobre grupos de tiles de puntuación dentro de un cómputo de atención densa fusionada.

Mecanismos Principales

Enrutamiento de Precisión por Grupos de Tiles:
- La matriz de atención ( $L_q \times L_k$ ) se particiona en tiles de cómputo alineados con el hardware de tamaño $BLOCK_M \times BLOCK_N$ .
- En lugar de asignar la precisión por tile, TileMix agrupa tiles de clave adyacentes a lo largo de la dimensión de la clave. Un único bit de enrutamiento gobierna un grupo de $g$ tiles de clave (donde $g$ es un factor de agrupación).
- Esto permite que el kernel despache cada grupo de tiles ya sea a una ruta FP16 (multiplicación de matrices estándar) o a una ruta INT8 (usando Tensor Core MMA de INT8 con acumulación en INT32) dentro del mismo bucle de streaming.
Enrutamiento por Máscara de Bits Compacta:
- Las decisiones de enrutamiento se empaquetan en enteros de 64 bits (máscaras de bits), donde cada bit corresponde a un grupo de tiles de clave.
- Esto permite una búsqueda de tiempo constante dentro del bucle interno utilizando operaciones de desplazamiento y máscara ( $R_{hk,m,g_j} = ((b_{hk,m} \gg g_j) \& 1)$ ).
- La sobrecarga de metadatos es $O(H_k T_m)$ , donde $H_k$ es el número de cabezas KV y $T_m$ es el número de tiles de consulta, manteniéndose compacta incluso a medida que el contexto crece.
Ejecución Heterogénea de Estado Compartido:
- Ambas rutas, FP16 e INT8, actualizan un estado de softmax en línea compartido (máximo por fila, normalizador y acumulador de salida no normalizado).
- La ruta INT8 computa $Q_8 K_8^\top$ , aplica escalas de bloque y reescala el resultado para entrar en un dominio de punto flotante común antes de actualizar el estado compartido.
- Este diseño preserva el grafo de atención densa (todas las interacciones legales de tokens se computan) mientras varía la precisión aritmética basada en políticas de enrutamiento espacial.
Políticas y Layouts:
- Las políticas de enrutamiento se definen mediante plantillas estructuradas estáticas y libres de datos (ej. Band, Global, Row-Random, Aligned Sparse, BigBird, Sparse Transformer).
- Estas plantillas determinan qué regiones espaciales de la matriz de atención se enrutan a FP16 (alta precisión) y cuáles a INT8 (baja precisión) basándose en una relación de cobertura de INT8 objetivo (ej. 25%, 50%, 75%).
- El sistema soporta Grouped-Query Attention (GQA), batching de longitud variable y cachés KV en INT8.

Contribuciones Clave

Enrutamiento de Precisión Centrado en Tiles: El artículo introduce la precisión por grupo de tiles como una abstracción de ejecución espacial para la atención densa fusionada, permitiendo la asignación fina de FP16/INT8 en todas las interacciones legales de tokens sin esparcir el grafo.
Ejecución de Puntuación Heterogénea de Estado Compartido: Se diseña un kernel fusionado que alinea las rutas de puntuación FP16 e INT8 hacia un dominio de puntuación común, integrándolas a través de una única recurrencia de softmax en línea.
Enrutamiento Nativo del Kernel Compacto y Escalable: Los autores desarrollan un enrutamiento por máscara de bits empaquetada con búsqueda de tiempo constante en el bucle interno y un agrupamiento de precisión escalable que mantiene los tiles de cómputo alineados con el hardware mientras soporta contextos largos.
Inferencia Práctica de Contexto Largo: La implementación soporta GQA, batching de longitud variable y cachés KV en INT8, validada a través de múltiples familias de modelos (LLaMA, Qwen, Vicuna) y benchmarks.

Resultados Experimentales

La evaluación se realizó en GPUs NVIDIA A100 utilizando LLaMA 3.2 3B, Qwen 2/2.5 7B y Vicuna 7B.

Recuperación de la Precisión: La cuantización INT8 uniforme ("One") a menudo degrada significamente el rendimiento de recuperación y de respuesta a preguntas a medida que aumenta la longitud del contexto. TileMix recupera esta calidad perdida al retener la precisión FP16 para grupos de tiles seleccionados. Por ejemplo, en el benchmark LongEval, las configuraciones de enrutamiento mixto (como los layouts SpTrans o BigBird) superaron consistentemente a la INT8 uniforme, acercándose o igualando a menudo a las líneas base de FP16 mientras mantenían una alta cobertura de INT8.
Eficiencia de Throughput: TileMix mejora el throughput de prefill sobre el FlashAttention estándar en FP16. A 4k tokens, configuraciones como SpTrans75 (75% INT8) lograron ~31.8 K tokens/s frente a los ~14.3 K tokens/s de FlashAttention. Las configuraciones de TileMix generalmente superaron o igualaron a las líneas base de INT8 uniforme en throughput de extremo a extremo, demostrando que la sobrecarga de enrutamiento y el manejo de precisión mixta es insignificante comparada con las ganancias de los Tensor Cores de INT8.
Estabilidad Numérica: El análisis de la desviación absoluta media respecto a una referencia fija de FP16 mostró que la desviación aumenta con la cobertura de INT8 pero permanece controlable. La elección del layout de enrutamiento (distribución espacial de FP16) impacta significativamente la precisión, mostrando ciertos layouts (como SpTrans) robustez a través de diferentes modelos y tareas.
Protección de Heavy-Hitters: Se demostró que los layouts de enrutamiento estático como SpTrans protegen naturalmente las interacciones de atención de alta importancia (heavy hitters) de ser enrutadas a INT8, enviando solo una pequeña fracción de las interacciones de gran masa a baja precisión a pesar de tener una alta cobertura nominal de INT8.

Significado y Reivindicaciones

El artículo afirma que TileMix establece el enrutamiento de precisión por grupo de tiles de puntuación como una abstracción nativa del kernel para la atención de contexto largo. Su principal significancia radica en:

Desacoplar la Conectividad de la Precisión: Permite preservar la conectividad densa de los tokens (a diferencia de la atención esparcida) mientras permite la ejecución de precisión heterogénea (a diferencia de la cuantización uniforme).
Frontera de Precisión-Eficiencia Controlable: Proporciona un mecanismo práctico para intercambiar precisión y eficiencia ajustando la relación de cobertura de INT8 y el layout de enrutamiento, ofreciendo una "frontera de precisión-eficiencia controlable" a través de diversas familias de modelos.
Despliegue Sin Reentrenamiento: El método opera sin necesidad de reentrenamiento, apoyándose en plantillas espaciales estáticas para guiar el enrutamiento de precisión.
Optimización a Nivel de Kernel: Al integrar el enrutamiento de precisión directamente en el kernel de atención fusionada con búsqueda de tiempo constante, evita las penalizaciones de latencia asociadas con la selección de precisión a nivel de invocación o el programador externo.

Los autores concluyen que el control espacial de la precisión es una dimensión práctica y efectiva para optimizar los kernels de atención densa, recuperando la calidad de contexto largo perdida bajo métodos de baja precisión uniforme, mientras acelera significativamente la inferencia.