💬 NLP

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

CompactAttention accelerates chunked prefill in long-context large language models by introducing a Block-Union KV Selection mechanism that converts 2D block-sparse masks into efficient, GQA-aware per-group KV block tables, thereby enabling in-place memory access without explicit compaction while maintaining near-dense accuracy and achieving up to 2.72× speedup.

Original authors: Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

Published 2026-05-19

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jiwon Song, Dongwon Jo, Beomseok Kang, Jae-Joon Kim

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a librarian (the AI) trying to answer a question based on a massive library of books (the context). In the past, if you had a huge library, you'd have to read every single book to find the right answer, which takes forever. To speed this up, researchers invented a "chunked prefill" system: instead of reading the whole library at once, you read it in small batches (chunks), adding notes to a notepad (the KV cache) as you go.

However, a new problem arose: How do you quickly find the right pages in your notepad without reading everything again every time you get a new batch of questions?

This paper introduces CompactAttention, a new way to solve this problem. Here is how it works, using simple analogies:

The Problem with Old Methods

The paper identifies two main ways people tried to solve this, and why they failed:

The "Sparse Kernel" Approach (The Inefficient Scanner):
- The Idea: Imagine you have a map of the library with red dots marking only the important books. You try to skip the white spaces and only look at the red dots.
- The Failure: When you are reading a huge library (long context) but only asking a tiny question (small chunk), this method gets slow. It's like having a scanner that is great for scanning a whole wall of text, but when you only have one sentence, the scanner takes too long to set up and calibrate. The overhead of "skipping" the white space actually makes it slower than just reading everything.
The "Query-Subsampled" Approach (The Lazy Librarian):
- The Idea: Instead of checking every question, you pick just a few random questions from your batch, find the important books for those, and assume those books are important for everyone.
- The Failure: This is risky. If you pick the wrong few questions, you might miss a crucial book that only one specific question needed. Also, once you pick those books, you have to physically carry them from the shelves to a special table before you can read them. This "carrying" (copying data) takes a lot of time and energy.

The Solution: CompactAttention

CompactAttention changes the game by separating finding the books from reading them.

Step 1: The "Union" Strategy (Grouping the Search)

Instead of trying to execute a complex "skip list" (sparse kernel) or guessing based on a few questions, CompactAttention uses a smart grouping trick:

Imagine you have a team of detectives (query heads) working on a case. Each detective has their own list of "suspects" (KV blocks) they think are important.
Instead of letting each detective work alone, CompactAttention says: "Let's combine all the suspects from the whole team into one master list."
It does this in two steps:
1. Q-Block Union: It combines the lists for all the questions in the current batch.
2. Intra-Group Union: It combines the lists for detectives who work together.
The Result: You get a single, minimal "Master List" of suspects that covers everyone's needs. No important book is left behind because if any detective needed it, it's on the list.

Step 2: The "Zero-Copy" Execution (Reading in Place)

This is the magic part.

Old Way: Once you have your Master List, you have to physically move all those books from the shelves to a special table so you can read them quickly. This "moving" takes time.
CompactAttention Way: You don't move the books at all. You just hand the librarian a map (metadata) that says, "Go to Shelf A, Row 3, Book 5; then Shelf B, Row 1, Book 2."
The librarian (the computer kernel) goes directly to those spots on the shelves and reads them. This is called "Zero-Copy Paged Attention." It saves all the time and energy spent on moving data around.

Why It's a Big Deal

The paper tested this on a massive AI model (LLaMA-3.1-8B) with a context of 128,000 words (a very long document).

Accuracy: It was just as smart as reading the whole library (Dense Attention). It didn't miss any crucial details.
Speed: It was up to 2.72 times faster than the standard way of doing things.

The Takeaway

Think of CompactAttention as a smart librarian who stops trying to rearrange the library and instead just uses a perfect, combined index card.

By realizing that the "search" (finding the important blocks) and the "execution" (reading them) should be separate, and by using a "grouping" trick to ensure nothing is missed, they managed to make long-document AI processing significantly faster without losing any intelligence. They proved that the bottleneck wasn't just which books to pick, but how you go about picking them up.

Technical Summary: CompactAttention

Problem Statement
As Large Language Models (LLMs) scale to support context windows of hundreds of thousands to millions of tokens, processing long sequences via a single "one-shot" prefill pass has become impractical due to quadratic compute costs and the risk of stalling decode requests in online serving systems. Consequently, chunked prefill—processing inputs sequentially in fixed-size chunks while attending to accumulated Key-Value (KV) caches—has become the standard serving strategy in frameworks like vLLM and SGLang.

However, existing acceleration methods face significant limitations in this regime:

Block-Sparse Attention Inefficiency: Traditional block-sparse methods (e.g., XAttention, SeerAttention) rely on sparse kernels that require large query lengths to amortize the overhead of irregular memory access and mask interpretation. In chunked prefill, the query length is limited by the chunk size ( $Q \ll KV$ ), causing these kernels to lose parallelism and efficiency, often failing to outperform dense attention.
Pattern Search Overhead: Sparse methods require finding input-dependent sparse patterns. In chunked prefill, this search must be repeated over the growing accumulated KV cache for every chunk, making cumulative search overhead a critical bottleneck.
Limitations of Direct KV Selection: Methods like QUOKA avoid sparse kernels by selecting a subset of KV tokens based on a subsampled set of queries and performing dense attention. However, this approach suffers from two issues:
- Coverage Gaps: Query-specific KV entries important to non-sampled queries may be missed, leading to accuracy degradation on tasks requiring distributed information access.
- Copy Overhead: Token-level selection requires explicit gathering and copying of KV entries into a compact buffer before attention execution, introducing overhead that scales with context length and batch size.

Methodology: CompactAttention
The authors propose CompactAttention, a mechanism that decouples block-level KV selection from attention execution. Instead of using 2D block-sparse masks as direct execution plans for sparse kernels, CompactAttention treats them as KV-selection signals to construct GQA-aware per-group KV block tables.

The method operates in two stages:

Selection Stage (Block-Union KV Table Construction):
- The system accepts a 2D per-head block-sparse mask generated by any lightweight pattern search method (e.g., SeerAttention or FlashPrefill).
- It applies a Q-block union across query blocks within the current chunk. This is necessary because dense paged attention consumes a single KV block list for a group of query blocks executed together, rather than individual decisions for each $(Q\text{-block}, KV\text{-block})$ tile.
- It applies an intra-group union across query heads that share a KV group (defaulting to Grouped-Query Attention groups).
- The result is a minimal per-group KV block table ( $T_{b,g}$ ) that preserves all KV blocks selected by the original input masks. This construction ensures coverage preservation: any KV block selected by any query block or head in the group remains selected.
Execution Stage (Zero-Copy Paged Attention):
- Unlike token-level selection methods, CompactAttention operates at block granularity.
- It utilizes a KV-head-major memory layout $[B, H_{kv}, L, D]$ rather than the standard sequence-major layout. This allows each $(batch, KV\text{-head}, block)$ triple to be treated as a contiguous memory region.
- The constructed block tables are passed as metadata (specifically kv_indptr and kv_indices) to a standard paged dense attention kernel.
- The kernel accesses the selected KV blocks in-place from the original KV cache without explicit copying or compaction into a new buffer.

Key Contributions

Decoupling Selection and Execution: The paper introduces a design that separates the logic of which KV blocks to select from how they are executed, allowing the reuse of lightweight sparse pattern search methods while avoiding the inefficiencies of sparse kernels.
Block-Union Construction: The novel use of Q-block and intra-group unions to convert 2D sparse masks into executable, GQA-aware block tables that are minimal yet coverage-preserving.
Zero-Copy Paged Execution: By leveraging a KV-head-major layout and paged attention metadata, the method eliminates the explicit KV copy overhead inherent in token-level selection methods like QUOKA.

Experimental Results
Evaluated on LLaMA-3.1-8B-Instruct and Qwen3-30B-A3B-Instruct-2507 under chunked prefill:

Speedup: On H200 GPUs with a 128K context length, CompactAttention achieves up to 2.72× attention speedup and 1.96× end-to-end speedup over the dense attention baseline. This significantly outperforms block-sparse baselines (which often underperform dense attention in the $Q \ll KV$ regime) and QUOKA (which is limited by copy overhead).
Accuracy: CompactAttention maintains accuracy close to dense attention on the RULER and LongBench V2 benchmarks. It outperforms QUOKA, particularly on tasks requiring distributed information access (e.g., Multi-key NIAH), where QUOKA's query subsampling leads to accuracy degradation.
Efficiency: Ablation studies confirm that the zero-copy paged execution strategy is more efficient than both sparse-kernel execution and explicit KV compaction, even when accounting for metadata construction overhead.

Significance
The paper claims that the primary bottleneck for sparse attention in chunked prefill is not merely which KV blocks to select, but how the selected blocks are executed. By making sparse selection compatible with efficient dense paged attention kernels through block-union construction and zero-copy execution, CompactAttention offers a practical path toward faster long-context LLM serving. It addresses the specific inefficiencies of the $Q \ll KV$ regime that have hindered the adoption of existing sparse attention methods in chunked prefill systems.