Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

About

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.72$\times$ attention speedup at 128K context length under chunked prefill.

Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim• 2026

Related benchmarks

TaskDatasetResultRank
Latency MeasurementLLaMA-8B-Instruct Chunked Prefill 3.1 (inference)
Attention Latency (ms)531.5
49
Long-context UnderstandingRULER 32k
Accuracy92.41
38
Long-context UnderstandingRULER 64k
Accuracy88.43
37
Long-context UnderstandingRULER 128k
Accuracy87.03
27
Long-context UnderstandingRULER--
27
Prefill Latency MeasurementLLaMA 8B-Instruct 64K context length 3.1
Attention Latency8.95e+3
7
Prefill Latency MeasurementLLaMA 8B Instruct 128K context length 3.1
Attention Latency (ms)2.41e+4
7
Prefill Latency MeasurementLLaMA 8B Instruct 16K context length 3.1
Attention Prefill Latency1.39e+3
7
Attention LatencyQwen3-30B-A3B-Instruct-2507 64K context window
Speedup Factor1.07
4
Attention LatencyQwen3-30B-A3B-Instruct-2507 128K context window
Speedup Factor1.56
4
Showing 10 of 12 rows

Other info

GitHub

Follow for update