Investigadores han presentado recientemente una nueva técnica de compresión de caché KV denominada TriAttention, diseñada para resolver el cuello de botella de memoria que enfrentan los grandes modelos de lenguaje (LLM) en tareas de razonamiento con contextos largos. Esta tecnología mejora drásticamente la eficiencia de la inferencia al aprovechar la estabilidad de los vectores de consulta (Q) y clave (K) en el espacio de pre-rotación (pre-RoPE).
Al procesar tareas de generación de textos extensos, la caché KV suele agotar rápidamente la memoria de la GPU, lo que provoca la interrupción del modelo. Las soluciones de compresión existentes suelen depender de las puntuaciones de atención en el espacio de post-rotación (post-RoPE), pero debido a la naturaleza de la codificación de posición rotacional, estos métodos a menudo tienen dificultades para identificar información clave, lo que genera una inferencia inestable.
La ventaja de la estabilidad en el espacio de pre-rotación
El equipo de investigación descubrió que, en el espacio de pre-rotación, los vectores Q y K muestran una tendencia a concentrarse alrededor de un centro fijo distinto de cero. Este fenómeno de "concentración Q/K" se mantiene altamente consistente a través de diferentes contextos de entrada y posiciones, lo que permite predecir los patrones de atención mediante series trigonométricas. TriAttention se basa precisamente en este hallazgo: calcula la preferencia de distancia mediante series trigonométricas y, al combinarla con la norma vectorial como señal auxiliar, evalúa con precisión el peso de los valores clave.
Los datos de las pruebas demuestran un rendimiento sobresaliente. En el conjunto de datos AIME25, el algoritmo logró un aumento de 2,5 veces en el rendimiento (throughput) y redujo el consumo de memoria en 10,7 veces, todo ello manteniendo la misma precisión de inferencia que el mecanismo de atención total (Full Attention). En comparación, los métodos de compresión de referencia actuales, con la misma eficiencia, apenas alcanzan la mitad de la precisión del mecanismo de atención total.
En casos de despliegue real, esta tecnología ha demostrado un valor práctico excepcional. Al ejecutar un modelo de 32 mil millones de parámetros en una GPU con 24 GB de VRAM, el mecanismo de atención total falló debido a un desbordamiento de memoria (OOM) al intentar completar tareas de instrucciones largas, mientras que TriAttention logró ejecutar la tarea por completo. En la prueba de referencia MATH 500, alcanzó una velocidad de inferencia de 1.405 tokens por segundo, superando con creces los 223 tokens por segundo del mecanismo de atención total.
Este estudio fue realizado por el equipo de Mao y colaboradores, y el artículo correspondiente ya ha sido publicado como preimpresión. Al introducir un mecanismo de ponderación adaptativa, TriAttention puede ajustar automáticamente los pesos del algoritmo según la concentración de los cabezales de atención, reduciendo eficazmente los costes computacionales del procesamiento de textos largos en grandes modelos sin sacrificar la calidad de la inferencia.
