CompactAttention: Aceleración del prellenado fragmentado mediante selección de KV por unión de bloques

Resumen

El prellenado por fragmentos se ha convertido en una estrategia de servicio ampliamente adoptada para modelos de lenguaje grandes de contexto largo, pero el cómputo eficiente de la atención en este régimen sigue siendo un desafío. Los métodos existentes de atención dispersa están diseñados principalmente para prellenado de una sola vez y no se traducen eficientemente al prellenado por fragmentos: los kernels de bloques dispersos pierden eficiencia cuando la longitud de la consulta está limitada por el tamaño del fragmento, mientras que la búsqueda de patrones de grano fino se vuelve costosa cuando se repite sobre la caché de KV acumulada en cada fragmento. QUOKA, un método reciente que apunta directamente al prellenado por fragmentos, evita la sobrecarga de kernels dispersos pero se basa en la selección de KV a nivel de token con submuestreo de consultas, lo cual puede omitir entradas KV específicas de la consulta e introducir una sobrecarga explícita de copia de KV. Para abordar estas limitaciones, proponemos CompactAttention, un mecanismo de atención de prellenado por fragmentos basado en la Selección de KV por Unión de Bloques. CompactAttention trata las máscaras dispersas de bloques 2D como señales de selección de KV en lugar de planes directos de ejecución de kernels dispersos, y las convierte en tablas de bloques KV por grupo conscientes de GQA mediante la unión de bloques Q y la unión intragrupo. Esta construcción produce las tablas de bloques mínimas que preservan todos los bloques KV seleccionados por las máscaras de entrada bajo restricciones de ejecución paginada, permitiendo que los bloques KV seleccionados sean accedidos in situ sin compactación explícita de KV. En LLaMA-3.1-8B-Instruct, CompactAttention mantiene una precisión cercana a la atención densa en el benchmark RULER, al tiempo que proporciona una aceleración de atención de hasta 2.72x en una longitud de contexto de 128K bajo prellenado por fragmentos.

English

Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging. Existing sparse attention methods are primarily designed for one-shot prefill and do not translate efficiently to chunked prefill: block-sparse kernels lose efficiency when the query length is limited by the chunk size, while fine-grained pattern search becomes costly when repeated over the accumulated KV cache at every chunk. QUOKA, a recent method that directly targets chunked prefill, avoids sparse-kernel overhead but relies on query-subsampled, token-level KV selection, which can miss query-specific KV entries and introduce explicit KV-copy overhead. To address these limitations, we propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection. CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction produces the minimal block tables that preserve all KV blocks selected by the input masks under paged execution constraints, enabling selected KV blocks to be accessed in place without explicit KV compaction. On LLaMA-3.1-8B-Instruct, CompactAttention maintains accuracy close to dense attention on the RULER benchmark while delivering up to 2.72times attention speedup at 128K context length under chunked prefill.