CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

Resumen Técnico: CompactAttention

Enunciado del Problema
A medida que los Modelos de Lenguaje Grandes (LLM) escalan para admitir ventanas de contexto de cientos de miles a millones de tokens, el procesamiento de secuencias largas mediante una única pasada de prellenado "de una sola vez" se ha vuelto poco práctico debido a los costos computacionales cuadráticos y al riesgo de detener las solicitudes de decodificación en sistemas de servicio en línea. En consecuencia, el prellenado fragmentado (procesar las entradas secuencialmente en fragmentos de tamaño fijo mientras se atiende a las cachés de claves-valor (KV) acumuladas) se ha convertido en la estrategia de servicio estándar en marcos como vLLM y SGLang.

Sin embargo, los métodos de aceleración existentes enfrentan limitaciones significativas en este régimen:

1. Ineficiencia de la Atención Esparsa por Bloques: Los métodos esparsos por bloques tradicionales (por ejemplo, XAttention, SeerAttention) dependen de kernels esparsos que requieren longitudes de consulta grandes para amortizar la sobrecarga del acceso irregular a la memoria y la interpretación de máscaras. En el prellenado fragmentado, la longitud de la consulta está limitada por el tamaño del fragmento ($Q \ll KV$), lo que hace que estos kernels pierdan paralelismo y eficiencia, a menudo fallando en superar a la atención densa.
2. Sobrecarga de Búsqueda de Patrones: Los métodos esparsos requieren encontrar patrones esparsos dependientes de la entrada. En el prellenado fragmentado, esta búsqueda debe repetirse sobre la caché KV acumulada en crecimiento para cada fragmento, convirtiendo la sobrecarga acumulada de búsqueda en un cuello de botella crítico.
3. Limitaciones de la Selección Directa de KV: Métodos como QUOKA evitan los kernels esparsos seleccionando un subconjunto de tokens KV basados en un conjunto submuestreado de consultas y realizando atención densa. Sin embargo, este enfoque adolece de dos problemas:
   • Brechas de Cobertura: Las entradas KV específicas de la consulta, importantes para consultas no muestreadas, pueden perderse, lo que lleva a una degradación de la precisión en tareas que requieren acceso a información distribuida.
   • Sobrecarga de Copia: La selección a nivel de token requiere la recolección y copia explícita de las entradas KV en un búfer compacto antes de la ejecución de la atención, introduciendo una sobrecarga que escala con la longitud del contexto y el tamaño del lote.

Metodología: CompactAttention
Los autores proponen CompactAttention, un mecanismo que desacopla la selección de bloques KV a nivel de bloque de la ejecución de la atención. En lugar de utilizar máscaras esparsas 2D como planes de ejecución directos para kernels esparsos, CompactAttention las trata como señales de selección de KV para construir tablas de bloques KV por grupo conscientes de GQA.

El método opera en dos etapas:

1. Etapa de Selección (Construcción de Tabla de Bloques KV de Unión de Bloques):

   • El sistema acepta una máscara esparsa por bloque y por cabeza 2D generada por cualquier método de búsqueda de patrones ligero (por ejemplo, SeerAttention o FlashPrefill).
   • Aplica una unión de bloques Q a través de los bloques de consulta dentro del fragmento actual. Esto es necesario porque la atención paginada densa consume una única lista de bloques KV para un grupo de bloques de consulta ejecutados conjuntamente, en lugar de decisiones individuales para cada tesela $(Q\text{-bloque}, KV\text{-bloque})$.
   • Aplica una unión intra-grupo a través de las cabezas de consulta que comparten un grupo KV (por defecto, los grupos de Atención de Consulta Agrupada - GQA).
   • El resultado es una tabla mínima de bloques KV por grupo ($T_{b,g}$) que preserva todos los bloques KV seleccionados por las máscaras de entrada originales. Esta construcción garantiza la preservación de la cobertura: cualquier bloque KV seleccionado por cualquier bloque o cabeza de consulta en el grupo permanece seleccionado.

2. Etapa de Ejecución (Atención Paginada sin Copia):

   • A diferencia de los métodos de selección a nivel de token, CompactAttention opera a granularidad de bloque.
   • Utiliza una disposición de memoria KV-cabeza-mayor $[B, H_{kv}, L, D]$ en lugar de la disposición estándar mayor a secuencia. Esto permite que cada tripleta $(lote, KV\text{-cabeza}, bloque)$ se trate como una región de memoria contigua.
   • Las tablas de bloques construidas se pasan como metadatos (específicamente kv_indptr y kv_indices) a un kernel estándar de atención densa paginada.
   • El kernel accede a los bloques KV seleccionados in situ desde la caché KV original sin copia explícita ni compactación en un nuevo búfer.

Contribuciones Clave

• Desacoplamiento de Selección y Ejecución: El artículo introduce un diseño que separa la lógica de qué bloques KV seleccionar de cómo se ejecutan, permitiendo la reutilización de métodos ligeros de búsqueda de patrones esparsos mientras se evitan las ineficiencias de los kernels esparsos.
• Construcción de Unión de Bloques: El uso novedoso de uniones de bloques Q e intra-grupo para convertir máscaras esparsas 2D en tablas de bloques ejecutables, conscientes de GQA, que son mínimas pero preservan la cobertura.
• Ejecución Paginada sin Copia: Al aprovechar una disposición de memoria KV-cabeza-mayor y metadatos de atención paginada, el método elimina la sobrecarga de copia explícita de KV inherente a los métodos de selección a nivel de token como QUOKA.

Resultados Experimentales
Evaluado en LLaMA-3.1-8B-Instruct y Qwen3-30B-A3B-Instruct-2507 bajo prellenado fragmentado:

• Aceleración: En GPUs H200 con una longitud de contexto de 128K, CompactAttention logra hasta 2.72× de aceleración en atención y 1.96× de aceleración de extremo a extremo sobre la línea base de atención densa. Esto supera significativamente a las líneas base esparsas por bloques (que a menudo rinden menos que la atención densa en el régimen $Q \ll KV$) y a QUOKA (limitado por la sobrecarga de copia).
• Precisión: CompactAttention mantiene una precisión cercana a la atención densa en los benchmarks RULER y LongBench V2. Supera a QUOKA, particularmente en tareas que requieren acceso a información distribuida (por ejemplo, Multi-key NIAH), donde la submuestreo de consultas de QUOKA conduce a una degradación de la precisión.
• Eficiencia: Los estudios de ablación confirman que la estrategia de ejecución paginada sin copia es más eficiente que tanto la ejecución de kernel esparso como la compactación explícita de KV, incluso al contabilizar la sobrecarga de construcción de metadatos.

Significado
El artículo afirma que el principal cuello de botella para la atención esparcida en el prellenado fragmentado no es meramente qué bloques KV seleccionar, sino cómo se ejecutan los bloques seleccionados. Al hacer que la selección esparcida sea compatible con kernels eficientes de atención densa paginada mediante la construcción de unión de bloques y la ejecución sin copia, CompactAttention ofrece un camino práctico hacia un servicio de LLM de contexto largo más rápido. Aborda las ineficiencias específicas del régimen $Q \ll KV$ que han obstaculizado la adopción de métodos de atención esparcida existentes en sistemas de prellenado fragmentado.