Resumen Técnico: Atención Global Jerárquica (HGA)

Planteamiento del Problema

El cuello de botella principal para el despliegue de modelos de lenguaje de gran escala (LLMs) preentrenados de contexto largo no es únicamente la complejidad computacional $O(n^2)$ de la atención causal densa, sino el requisito de memoria del caché de Clave/Valor (K/V). En la atención densa, el caché K/V para todos los tokens históricos debe residir en la memoria del acelerador (VRAM). Para modelos grandes cuantizados (por ejemplo, FP8), los pesos del modelo en sí mismos suelen consumir la mayor parte del presupuesto de VRAM, dejando un espacio insuficiente para almacenar el caché K/V para contextos largos (por ejemplo, 32K–64K tokens). Esta limitación impide la inferencia "out-of-the-box" de checkpoints preentrenados en hardware de grado de consumo o de una sola GPU sin reentrenamiento o modificación arquitectónica.

Los métodos de atención dispersa existentes a menudo requieren entrenamiento desde cero, introducen parámetros de calibración o aproximan el resultado de la atención, lo que potencialmente degrada la calidad de las representaciones aprendidas. Existe la necesidad de un mecanismo de sustitución directa (drop-in) que permita a los checkpoints preentrenados existentes utilizar una recuperación dispersa y eficiente en memoria preservando la exactitud del cálculo de la atención final.

Metodología: Atención Global Jerárquica (HGA)

HGA es un parche a nivel de sistemas diseñado como un reemplazo directo para la atención causal densa en transformadores preentrenados. Preserva los parámetros originales del checkpoint ($W_Q, W_K, W_V, W_O$) y no requiere reentrenamiento ni calibración.

Mecanismo Central

HGA desplaza el problema de la recuperación de largo alcance de la materialización densa de token-a-token a un proceso de enrutamiento jerárquico de dos niveles, seguido de una atención exacta a nivel de token:

1. Particionamiento: La secuencia de entrada se divide en fragmentos (chunks) de tamaño fijo (por ejemplo, $C=64$ tokens). Cada fragmento se subdivide además en grupos.
2. Enrutamiento Jerárquico:
   • Nivel 1 (Enrutamiento de Fragmentos): El fragmento actual consulta una tabla compacta de "resúmenes de claves de fragmentos" derivados de fragmentos previos.
   • Nivel 2 (Enrutamiento de Grupos): Dentro de los fragmentos seleccionados, la consulta se enruta a "resúmenes de claves de grupo" específicos para refinar la selección.
   • Visibilidad Determinista: La política de enrutamiento combina la recuperación basada en contenido con ventanas de visibilidad fijas: el fragmento actual, un número configurable de fragmentos "sumidero" (sink) iniciales y los fragmentos locales recientes siempre son visibles.
3. Resúmenes Sensibles a RoPE: Para asegurar la compatibilidad con los Embeddings de Posición Rotatorios (RoPE), HGA construye los resúmenes utilizando una estrategia mixta: los pares de RoPE de alta frecuencia se rotan por token antes del promedio, mientras que los pares de baja frecuencia se promedian en el espacio de la clave bruta y se rotan en la posición representativa del fragmento. Estos resúmenes se computan a partir de las claves proyectadas existentes sin nuevas capas lineales.
4. Atención Exacta de Token: Una vez seleccionados los fragmentos o grupos relevantes, el enrutador extrae los vectores K/V originales a nivel de token desde la memoria RAM del host (o NVMe) hacia la VRAM. El resultado de la atención final se calcula utilizando un softmax estándar sobre este conjunto acotado y exacto de tokens. No se utilizan valores de resumen sintéticos en el resultado final.

Arquitectura de Memoria

HGA introduce una abstracción de almacenamiento K/V de niveles jerárquicos:

• Hot (Caliente): Pesos del modelo, resúmenes de fragmentos y ventanas siempre visibles residen en VRAM.
• Warm (Templada): Un caché LRU acotado en VRAM mantiene los fragmentos enrutados recientemente.
• Cold (Fría): El historial completo de tokens del caché K/V reside en la memoria RAM del host (o NVMe).
  Esto desacopla el consumo de memoria de la GPU de la longitud total del contexto; el uso de VRAM depende principalmente de los pesos del modelo y del conjunto de trabajo enrutado, no de la longitud total de la secuencia.

Contribuciones Clave

1. Compatibilidad Directa (Drop-In): Un mecanismo para reemplazar la atención densa en transformadores GQA/MoE preentrenados sin modificar pesos, normalizaciones o introducir parámetros de calibración.
2. Recuperación de Token Exacto: A diferencia de los métodos que aproximan los resultados de la atención, HGA aproxima solo la etapa de recuperación, asegurando que la atención final se compute exactamente sobre los tokens seleccionados.
3. Política de Enrutamiento de Dos Niveles: Una estrategia que combina sumideros deterministas/ventanas locales con recuperación basada en contenido del contexto medio, utilizando resúmenes de fragmentos y grupos.
4. Sistema de Almacenamiento por Niveles: Una implementación que permite que todo el historial K/V resida en la memoria RAM del host, permitiendo la inferencia de contexto largo en hardware con VRAM limitada (por ejemplo, una única RTX 5090).
5. Desempeño Zero-Shot: Capacidad demostrada para ejecutar modelos preentrenados (Qwen3-30B-A3B) con contextos de 32K–64K sin ajuste fino (fine-tuning).

Resultados Experimentales

El artículo presenta evaluaciones comparativas en un SmallLM de 40M de parámetros y un modelo Qwen3 de 30B parámetros.

• Calidad (Brecha de Pérdida):
  • En un SmallLM de 40M con copia directa de pesos (sin ajuste fino), el modelo enrutado logra una brecha de pérdida de +0.018 nats en comparación con la atención densa a 8K tokens.
  • A través de longitudes de contexto de 4K a 64K, la brecha de pérdida se mantiene dentro de 0.01–0.02 nats de la atención densa utilizando niveles de dispersión de aproximadamente 3% a 12.5%.
  • A 64K tokens, una dispersión del 6.25% mantiene una brecha de ~0.01 nats.
• Precisión de Recuperación:
  • Una prueba de "Aguja en un Pajar" (Needle-in-a-Haystack) en Qwen3-30B-A3B a 64K tokens logró una tasa de éxito del 100% (3/3 profundidades) usando enrutamiento a nivel de grupo con ~1.9% de dispersión, sin ningún ajuste fino.
• Desempeño:
  • En el SmallLM de 40M, una implementación fusionada en Triton logró una aceleración de 2.72× en el entrenamiento y una aceleración de 2.43× en el prefill a 12K tokens en comparación con una línea base densa.
  • En Qwen3-30B, el modelo parcheado se ejecuta de forma interactiva a 32K de contexto en una única RTX 5090 (32GB VRAM), una configuración donde el almacenamiento denso de K/V es imposible.
• Ajuste Fino (Fine-Tuning): El ajuste fino de un modelo de 0.6B utilizando atención enrutada produjo un rendimiento de ~917 tokens/seg (frente a ~570 para el denso) con un costo de enrutamiento insignificante (+0.015 nats) en comparación con la línea base densa.

Significancia y Reivindicaciones

El artículo afirma que HGA proporciona una solución práctica a nivel de sistemas para el cuello de botella de memoria de los LLMs de contexto largo. Su significancia radica en:

• Preservación del Conocimiento Preentrenado: Al evitar el reentrenamiento y utilizar el K/V exacto de los tokens para el resultado final, HGA permite que los checkpoints existentes se beneficien de la atención dispersa sin degradar sus representaciones aprendidas.
• Eficiencia de Hardware: Permite el despliegue de modelos grandes y cuantizados en hardware de consumo para contextos que exceden con creces la capacidad de VRAM requerida para el caché denso.
• Atribución de Errores: Los autores observan que la pequeña brecha de calidad restante (0.01–0.02 nats) no escala rápidamente con la longitud del contexto. Sugieren que esta brecha probablemente está dominada por la extrapolación de la codificación posicional de contexto largo más que por el algoritmo de enrutamiento en sí, lo que implica que HGA es un mecanismo de recuperación robusto que puede optimizarse aún más mejorando las representaciones posicionales.

Los autores posicionan a HGA no como un rediseño de la arquitectura del transformador, sino como una capa de enrutamiento jerárquica compatible que hace que la inferencia de contexto largo sea viable bajo las restricciones de hardware actuales.