AsyncTLS: Efficient Generative LLM Inference with Asynchronous Two-level Sparse Attention

1. El Problema

La inferencia de modelos de lenguaje grandes (LLM) en contextos largos enfrenta dos desafíos fundamentales:

• Complejidad Computacional Cuadrática: El mecanismo de atención estándar requiere $O(n^2)$ operaciones, lo que se vuelve prohibitivo a medida que la longitud de la secuencia ($n$) crece (cientos de miles de tokens).
• Consumo de Memoria del KV Cache: El almacenamiento de las claves y valores (KV) crece linealmente con la longitud de la secuencia, a menudo excediendo la capacidad de la memoria de alta velocidad (VRAM) de la GPU, obligando a realizar offloading (traslado) a la memoria del CPU, lo que introduce latencia.

Existen soluciones de atención dispersa (sparse attention), pero presentan un dilema:

• Métodos a nivel de token: Ofrecen alta precisión al seleccionar tokens individuales importantes, pero tienen un costo de indexación y selección excesivo que ralentiza la inferencia.
• Métodos a nivel de bloque: Son más eficientes al procesar bloques contiguos, pero sacrifican precisión al incluir tokens irrelevantes dentro de los bloques seleccionados o descartar información crítica fuera de ellos.

2. Metodología: AsyncTLS

Los autores proponen AsyncTLS, un sistema jerárquico que combina lo mejor de ambos mundos mediante una arquitectura de dos niveles y un motor de offloading asíncrono.

A. Atención Dispersa Jerárquica (Dos Niveles)

El sistema utiliza un proceso de filtrado en dos etapas para equilibrar precisión y eficiencia:

1. Filtrado a Nivel de Bloque (Fino): Se divide la memoria KV en bloques. Se calculan puntuaciones de importancia para cada bloque utilizando representaciones comprimidas. Se seleccionan los $K_b$ bloques más relevantes.
   • Innovación: Se reformula el cálculo de puntuación (basado en Quest) para utilizar multiplicaciones de matrices densas (GEMM), optimizando el uso de los núcleos Tensor de las GPUs modernas, especialmente en arquitecturas compartidas como MQA y MLA.
2. Selección a Nivel de Token (Fino): Dentro de los bloques seleccionados, se aplica una selección de tokens individual.
   • Se utiliza un mecanismo de "Doble Dispersidad" (Double Sparsity) para identificar canales importantes y comprimir los vectores de consulta y clave.
   • Se seleccionan los $K_t$ tokens más salientes dentro de los bloques retenidos para el cálculo final de la atención.
   • Resultado: Se reduce drásticamente el espacio de búsqueda para la selección fina, manteniendo la precisión de la atención token-a-token.

B. Motor de Offloading Asíncrono

Para manejar la memoria limitada de la GPU en contextos ultra-largos, AsyncTLS implementa un mecanismo de transferencia de datos optimizado:

• Prefetching Asíncrono: Aprovecha la localidad temporal de los patrones de atención. Los bloques seleccionados en el paso $t-1$ son buenos predictores para el paso $t$.
  • Mientras se calcula la atención fina para el paso actual ($t$) usando los bloques seleccionados en el paso anterior, el sistema calcula simultáneamente los bloques necesarios para el paso $t+1$ y comienza a pre-cargarlos asíncronamente desde la CPU.
• Transferencia Incremental: En lugar de transferir bloques completos en cada paso, el sistema solo transfiere las diferencias (bloques nuevos que entraron y bloques que salieron) entre la selección del paso $t$ y $t-1$. Esto minimiza el ancho de banda utilizado en el bus PCIe.

3. Contribuciones Clave

1. Arquitectura de Atención Dispersa Jerárquica: Propone AsyncTLS, que fusiona el filtrado de bloques (eficiencia) con la selección de tokens (precisión), mitigando el sobrecosto de indexación mediante poda jerárquica.
2. Motor de Offloading Asíncrono: Extiende la atención dispersa a nivel de token al escenario de offloading de KV. Utiliza prefetching asíncrono y transferencia incremental de bloques para superponer el movimiento de memoria con el cómputo, ocultando la latencia.
3. Evaluación Exhaustiva: Demuestra la compatibilidad del método con diversas arquitecturas de atención, incluyendo MHA, GQA (Grouped-Query Attention) y MLA (Multi-head Latent Attention, usada en modelos como DeepSeek), sin necesidad de reentrenamiento (training-free).

4. Resultados Experimentales

El método fue evaluado en modelos de última generación (Qwen3-8B, Qwen3-14B y GLM-4.7-Flash) en contextos de 32k a 128k tokens.

• Precisión: AsyncTLS logra una precisión comparable a la atención completa (Full Attention) en tareas de recuperación en contexto (RULER) y comprensión de contexto largo (LongBench), superando significativamente a los métodos basados solo en bloques (como Quest) y igualando a los métodos basados en tokens (como Double-Sparsity).
• Velocidad de Operador (Kernel):
  • Mejoras de 1.2x a 10.0x en la velocidad de los operadores de atención en comparación con la atención completa.
  • Mejoras de 1.2x a 4.0x frente a métodos de dispersión a nivel de token (DS).
• Rendimiento de Extremo a Extremo (Throughput):
  • Mejoras de 1.3x a 4.7x en el rendimiento general (throughput) en contextos de 48k a 96k tokens.
  • Permite un procesamiento por lotes (batch size) más grande (ej. 6) cuando se habilita el offloading, mientras que la atención completa se limita a lotes de 1 debido a la falta de memoria.

5. Significado e Impacto

AsyncTLS representa un avance significativo en la viabilidad de la inferencia de LLMs en contextos ultra-largos en hardware comercial.

• Rompe la compensación (Trade-off): Resuelve el conflicto histórico entre la precisión de la selección de tokens y la eficiencia de hardware, logrando ambas simultáneamente.
• Eficiencia de Hardware: Al reformular los cálculos para GEMM y optimizar el uso del ancho de banda PCIe mediante transferencias incrementales, hace que la atención dispersa a nivel de token sea práctica para la producción.
• Escalabilidad: Ofrece una solución escalable para la generación de secuencias ultra-largas sin requerir reentrenamiento de los modelos, facilitando su adopción en aplicaciones reales que requieren ventanas de contexto masivas.