CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

summary

The gist

The gist The authors propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection, which treats 2D block-sparse masks as KV-selection signals rather than direct

In short

CompactAttention proposes a new chunked prefill attention mechanism using Block-Union KV Selection. Instead of traditional sparse kernels, it treats 2D block-sparse masks as selection signals to build GQA-aware KV block tables. This allows selected Key/Value blocks to be accessed in place during paged execution without needing explicit KV compaction, achieving significant speedups while maintaining high accuracy.

Key concepts

Chunked Prefill: This is a serving strategy for long context LLMs where the input sequence is broken into smaller, fixed-size chunks. Each chunk attends to its own tokens and the accumulated Key/Value (KV) cache from all previous chunks. It's used to manage very long contexts efficiently in online serving systems.
Block-Sparse Attention: This method accelerates attention by only computing attention for important blocks of tokens rather than the entire sequence. However, applying it to chunked prefill is inefficient because the query length is limited by chunk size, and repeated searches over accumulated KV caches create overhead.
Block-Union KV Selection: CompactAttention uses this technique to convert 2D block-sparse masks into minimal KV block tables. It employs Q-block union and intra-group union to determine which specific KV blocks are needed. This ensures that selected blocks can be accessed directly in memory without copying them, enabling zero-copy paged execution.
Paged Execution: This is a memory management technique where the KV cache is stored in non-contiguous pages. CompactAttention leverages this by organizing the cache in a KV-head-major layout and using metadata (like kv_indptr) to point to selected blocks, allowing them to be accessed in place without allocating new dense buffers.

Terminology used across episodes

Reinforcement Learning (RL) — explained in 25 episodes
Retrieval-Augmented Generation (RAG) — explained in 20 episodes
Sample Complexity — explained in 20 episodes
Flow Matching — explained in 18 episodes
Federated Learning — explained in 17 episodes
Quantum Fisher Information (QFI) — explained in 17 episodes

This episode discusses

The paper

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection · Read on arXiv

Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon kim

Seoul National University

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 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.79 times attention speedup over dense attention and 1.25 times over FlashPrefill at 128K context length under chunked prefill.

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: We're looking at the paper: "CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection".

Jane: The gist The authors propose CompactAttention, a chunked-prefill attention mechanism based on Block-Union KV Selection,

Tom: First, who's behind it and why it matters.

Paper summary: Tom: So looking at the title CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection, it seems like they’ve found a way to bridge the gap between fast, sequential chunked prefill and accurate attention computation.

Jane: It’s about taking those 2D block-sparse masks and turning them into structured KV tables that let you run dense attention efficiently without needing to copy the data around, which is a big thing for keeping things fast in real serving systems.

Lu: What this paper suggests is that instead of relying on sparse kernels that struggle with chunked prefill, you can use a block-aware selection strategy to keep the computation close to dense attention accuracy while getting significant speedups.

Meng: From a deployment standpoint, if this works reliably across different context lengths and batch sizes—which they tested up to 128K—it means we might be able to deploy these long context models much more affordably because the memory and compute requirements per chunk are optimized.

Tom: It moves the conversation away from just trying to find a single sparse trick for one-shot prefill and shows how you can adapt those concepts to the real-world problem of serving long contexts sequentially.

Jane: It’s about making sure that when we process a massive document or code base in pieces, we don't end up stalling the entire system waiting for slow attention computation across all those accumulated pieces.

Lu: The implication is that by focusing on how to select the right KV blocks using union operations, you can achieve high speedup factors without the accuracy penalty usually associated with skipping parts of the attention map.

Meng: So, for anyone running these models online, CompactAttention suggests a way to keep the latency down while managing that huge KV cache efficiently during chunked prefill operations.

Conclusion: Tom: So we’re wrapping up on CompactAttention, which is about using Block-Union KV Selection to make chunked prefill way faster without sacrificing accuracy, right?

Jane: Yeah, exactly. The title itself says it’s about accelerating chunked prefill by using this specific selection method. It seems they’ve taken the 2D block-sparse masks and turned them into something that lets you run dense attention efficiently.

Lu: The authors are focusing on how to treat those sparse signals as KV selection commands rather than trying to force a sparse kernel execution plan directly, which is where most existing methods get stuck in chunked prefill.

Meng: From my side, the big win here is the zero-copy part. They're not copying the KV data into a new buffer; they're accessing it in place using these block tables. That’s huge for reducing memory bandwidth when we’re handling massive contexts at scale.

Tom: That zero-copy aspect is what I'm really digging into, because explicit compaction overhead just adds more latency that we don't want in a serving system. So, what does this mean for the overall speedup numbers they reported?

Jane: They showed up to two point seven nine times attention speedup over dense attention on LLaMA-three point one-8B-Instruct at a 128K context length when using chunked prefill conditions. That’s a pretty big jump for how fast we can process those long inputs sequentially.

Lu: And they kept the accuracy quite close to what you get with standard dense attention on the RULER benchmark, which is important because you don't want speed at the expense of quality. They also showed similar results on Qwen3-30B-A3B.

Meng: But I gotta ask, how robust is this across different group sizes? Because they mentioned splitting groups for large GQA groups, so I’m curious if that adds too much complexity to the deployment pipeline for production use.

Tom: That splitting part is where the engineering reality comes in. It means we have to manage those smaller execution groups, but it sounds like the core mechanism of using Q-block union and intra-group union is what keeps things stable regardless of group size.

Jane: And it definitely shifts the focus from just finding a single perfect sparse kernel to building this structured KV table approach that works under paged execution constraints for chunked prefill. It’s a new way to handle long context serving.

More episodes

← Home