Resumen Técnico: GQLA – Atención Latente de Consulta por Grupos para la Decodificación de Modelos de Lenguaje Grande Adaptativa al Hardware

1. Declaración del Problema

La decodificación moderna de Modelos de Lenguaje Grande (LLM) autoregresivos está fundamentalmente limitada por el tráfico de la memoria caché de claves y valores (KV), ya que cada token generado requiere leer todo el historial de claves y valores almacenados en caché desde la memoria fuera del chip. Aunque arquitecturas recientes como la Atención Latente Multicabeza (MLA), utilizada en DeepSeek-V2/V3, logran una reducción del estado del arte en la memoria caché KV al comprimir conjuntamente claves y valores en un latente de rango bajo, sufren tres desventajas de hardware acopladas al desplegarse en hardware de inferencia diverso:

1. Acoplamiento de Hardware: Los pesos entrenados de MLA exponen únicamente una única ruta de decodificación (una forma "absorbida" de Atención de Consulta Múltiple o MQA). Esta ruta está perfectamente afinada para la relación de cómputo-ancho de banda de GPUs de gama alta como la NVIDIA H100 (cresta $\approx$ 295 FLOPs/byte). Sin embargo, en GPUs restringidas a la exportación o de consumo como la H20 (que conserva el ancho de banda pero reduce el cómputo en $\sim$7x, bajando la cresta a $\approx$ 37 FLOPs/byte), esta ruta se vuelve severamente limitada por el cómputo, resultando en un rendimiento deficiente.
2. Pérdida de Paralelismo Tensorial (TP): La forma MQA absorbida canaliza todas las cabezas de consulta a través de un único latente KV compartido. Esta estructura impide el paralelismo tensorial a lo largo del eje de cabezas, ya que cada dispositivo debe replicar todo el latente, limitando la escalabilidad.
3. Cero Ganancia por Predicción de Múltiples Tokens (MTP): MTP, que genera múltiples tokens por paso para aumentar la intensidad aritmética, empuja la decodificación ya limitada por cómputo en H20 más allá de su cresta de eficiencia, produciendo cero ganancia en el rendimiento. Por el contrario, en H100, MTP supera la cresta, reduciendo la eficiencia de la ganancia.

2. Metodología: Atención Latente de Consulta por Grupos (GQLA)

Los autores proponen Atención Latente de Consulta por Grupos (GQLA), una modificación arquitectónica mínima de MLA que preserva la compresión latente conjunta de rango bajo mientras expone dos rutas de decodificación algebraicamente equivalentes sobre los mismos parámetros entrenados.

Arquitectura

GQLA indexa las matrices de proyección ascendente ($W_{UK}, W_{UV}$) por $g$ grupos en lugar de replicarlas a través de las $h_q$ cabezas de consulta. Este diseño habilita dos rutas de ejecución distintas:

1. Ruta MQA-Absorbida (Amigable con la Memoria):

   • Mecanismo: Idéntica a MLA. Las proyecciones ascendentes se absorben en las proyecciones de consulta y salida. Todas las $h_q$ cabezas atienden directamente a la memoria caché latente compacta ($cKV$).
   • Caché: Almacena únicamente el latente ($r_{kv}$) y una clave RoPE compartida ($d^R_h$).
   • Hardware Objetivo: Optimizado para hardware de clase H100 donde la intensidad aritmética ($\approx$ 242 FLOPs/byte para $s_q=1$) se sitúa justo por debajo de la cresta memoria-cómputo.

2. Ruta GQA (Amigable con el Cómputo):

   • Mecanismo: El latente se expande en $g$ grupos de claves/valores. La decodificación ejecuta Atención de Consulta Agrupada (GQA) estándar sin expansión latente por paso.
   • Caché: Almacena claves y valores expandidos por grupo ($2gd_h$) más la clave RoPE compartida.
   • Hardware Objetivo: Optimizado para hardware de clase H20. Al ajustar la cantidad de grupos ($g$) y el número de tokens de consulta por paso ($s_q$), la intensidad aritmética puede fijarse a la cresta de H20 ($\approx$ 37 FLOPs/byte).
   • Escalabilidad: Soporta hasta paralelismo tensorial de 8 vías sin redundancia a lo largo del eje de grupos.

Selección de Ruta

El tiempo de ejecución selecciona la ruta basándose en el hardware objetivo sin reentrenamiento ni kernels personalizados:

• H100: Utiliza la ruta MQA-absorbida ($s_q=1$).
• H20: Utiliza la ruta GQA (ej. $g=8, s_q=2$ o $g=4, s_q=1$).
• MTP: En H20, la ruta GQA con $s_q=2$ recupera ganancias de rendimiento cercanas a las ideales, mientras que MLA no obtiene ninguna ganancia.

TransGQLA: Conversión desde Modelos Preentrenados

Para evitar el preentrenamiento desde cero, los autores introducen TransGQLA, una extensión del pipeline TransMLA. Convierte un punto de control (checkpoint) GQA preentrenado en un modelo GQLA mediante una única modificación dirigida:

• Fusión de Cabezas: En lugar de fusionar cabezas agrupadas en un único latente y replicar proyecciones ascendentes a través de todas las cabezas de consulta (como en TransMLA), TransGQLA mantiene las proyecciones ascendentes indexadas por grupo.
• Resultado: El modelo retiene estructuralmente la ruta de ejecución GQA mientras habilita la ruta MQA-absorbida mediante la operación de absorción estándar. Esto preserva el paralelismo tensorial a lo largo del eje de grupos.

GQLA Disperso

El diseño se extiende a la atención dispersa de grano fino. Debido a que la configuración canónica de GQLA ($h_q/g = 16$) coincide con la dimensión de la baldosa MMA del Tensor-Core ($m=16$), el GQLA disperso puede preservar la ruta GQA en hardware de clase H20. Por el contrario, el MLA disperso está estructuralmente bloqueado en la ruta MQA absorbida debido a las restricciones de baldosas, heredando los mismos problemas de sobrecarga de cómputo.

3. Contribuciones Clave

1. Identificación de Desventajas de MLA: El artículo identifica que el diseño de ruta única de MLA lo acopla a las relaciones de clase H100, sacrifica el paralelismo tensorial en el eje de cabezas y produce cero ganancias de MTP en GPUs de consumo.
2. Arquitectura GQLA: Una variante mínima de MLA que expone dos rutas algebraicamente equivalentes (MQA-absorbida y GQA) sobre los mismos pesos. Una única configuración ($h_q=128, g=8$) con una cabeza MTP fija las líneas de techo tanto de H100 como de H20 simultáneamente.
3. Pipeline TransGQLA: Un método de conversión que transforma puntos de control GQA preentrenados en modelos GQLA con cambios arquitectónicos mínimos, reteniendo el paralelismo tensorial y permitiendo una transición suave a la forma latente.
4. Análisis de Línea de Techo: Una verificación teórica que muestra que los pesos GQLA pueden lograr un rendimiento cercano al máximo tanto en H100 como en H20 cambiando de ruta, mientras que MLA es ineficiente en H20.

4. Resultados

Los autores evaluaron TransGQLA en el modelo LLaMA-3-8B:

• Preservación de Capacidad: Convertir los pesos GQA a la forma latente GQLA (sin entrenamiento adicional) resultó en una pérdida de solo $\sim$9.7 puntos promedio en benchmarks de razonamiento de sentido común (MMLU, ARC, etc.) en comparación con el modelo preentrenado original de 15T tokens.
• Compresión de Caché: En la ruta MQA-absorbida, la memoria caché KV por token se comprimió al 28.125% de la línea base GQA original (una reducción del 71.875%).
• Proyecciones de Rendimiento:
  • H100: La ruta MQA-absorbida alcanza 354K tokens/seg.
  • H20: La ruta GQA (con $g=8, s_q=2$) alcanza 221K tokens/seg, una mejora de 3.4x sobre MLA en el mismo dispositivo (que se queda estancado en 65K tokens/seg debido a estar limitado por cómputo).
• Preentrenamiento Continuo: Los autores proyectan que el preentrenamiento continuo (siguiendo la trayectoria TransMLA) recuperará la capacidad perdida hasta dentro de 0.5 puntos del modelo original después de 30B tokens, manteniendo la compresión de la memoria caché KV.

5. Significado y Afirmaciones

El artículo afirma que GQLA ofrece un principio de diseño práctico y agnóstico al hardware para la inferencia de LLM. Al exponer múltiples rutas de decodificación algebraicamente equivalentes sobre un único conjunto de pesos entrenados, GQLA permite que un modelo se adapte a diferentes restricciones de hardware (rico en cómputo vs. rico en memoria) en el momento del despliegue sin reentrenamiento ni kernels personalizados.

Los autores posicionan a GQLA como una solución al problema del "acoplamiento de hardware", permitiendo un despliegue eficiente tanto en aceleradores insignia (H100) como en hardware de consumo/restringido a la exportación (H20). Enfatizan que este enfoque soporta la Predicción de Múltiples Tokens y el Paralelismo Tensorial en hardware donde MLA falla en hacerlo eficientemente.

Limitaciones Reconocidas:

• Las afirmaciones sobre hardware se basan en modelado de Línea de Techo; las pruebas de rendimiento a nivel de kernel completo en hardware H20/H100 real están pendientes.
• Los resultados de preentrenamiento continuo para TransGQLA están actualmente en progreso y aún no están finalizados en el artículo.
• La evaluación está actualmente limitada al esqueleto LLaMA-3-8B y a benchmarks de razonamiento de sentido común.
• El GQLA disperso se describe estructuralmente pero aún no ha sido comparado con métodos de atención dispersa existentes.