CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

arXiv:2605.16839 · cs.CL · Submitted 2026-05-16 · Read on arXiv

Listen

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.

Seoul National University

cs.CL

Submitted: 2026-05-16

Updated: 2026-10-19

Comments: NeurIPS 2026 Spotlight

Code: https://github.com/jiwonsong-dev/CompactAttention

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 87/100

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

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

Accuracy — 1× in this paper · explained in 1 episode
Attention Mechanism — 1× in this paper · explained in 3 episodes
Benchmark — 1× in this paper · explained in 2 episodes
KV cache — 1× in this paper · explained in 4 episodes
Sparse Attention — 1× in this paper · explained in 5 episodes

Summary

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 sparse-kernel execution plans and converts them into GQA-aware per-group KV block tables through Q-block union and intra-group union to enable zero-copy paged execution.

Motivation

Chunked prefill has become a widely adopted serving strategy for long-context large language models, but efficient attention computation in this regime remains challenging As large language models (LLMs) are increasingly used for long-horizon reasoning, document understanding, code analysis, and agentic workloads, their supported context windows have grown rapidly Processing such long contexts in a single prefill pass is increasingly impractical First, full-sequence attention incurs quadratic compute cost with respect to context length, making one-shot prefill expensive at long contexts Second, in online serving systems where prefill and decode requests are batched together, a long prefill pass can stall decode requests, making it difficult to satisfy time-between-token (TBT) service-level objectives (SLOs) Chunked prefill [6, 7], now adopted in major serving frameworks such as vLLM [8] and SGLang [9], addresses these issues by processing long inputs sequentially in fixed-size chunks, each attending to both its own KVs and the accumulated KV cache from previous chunks The dominant approach to accelerating long-context prefill is block-sparse attention Since FlashAttention [10, 11, 12] operates on blocks of tokens, block-sparse methods [13, 14, 15, 16, 17] first estimate which attention blocks are important and then compute only the selected subset of the attention map However, applying these methods directly to chunked prefill exposes two key limitations First, sparse execution becomes inefficient when Q ≪ KV: block-sparse kernels have too few query blocks to expose sufficient parallelism and amortize irregular access overheads, so the achieved speedup falls far below what the nominal sparsity would suggest Second, sparse pattern search must be repeated at every chunk over the accumulated KV cache, making cumulative search overhead a first-order concern and leaving only lightweight pattern-search mechanisms practical An alternative is to perform dense attention over a selected subset of KV entries, avoiding block-sparse kernel overhead entirely

How it works

CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans It converts these signals 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 The selection stage involves two union operations: first, applying Q-block union across query blocks within the current chunk, and second, applying intra-group union across query heads that are executed together The resulting per-group page table is Tb,g = 1 = 2 = 3 where Tb,g is the minimal table that preserves this coverage The execution stage passes these block tables to a paged attention kernel, which accesses the selected KV blocks in place without copying them into a separate buffer CompactAttention executes the selected KV blocks using a paged dense-attention backend while avoiding explicit K/V compaction

KV Selection Details

The selection stage accepts any block-sparse pattern search method that produces a 2D per-head block mask with sufficiently low per-chunk overhead The first union operation, Q-block union across query blocks, produces a single 1D KV block mask per query head The second operation, intra-group union across query heads that share a KV block table produces the per-group page table Tb,g = 1 = 2 where Tb,g is the minimal table that preserves this coverage For models with large GQA groups, CompactAttention therefore splits each KV group into smaller execution groups and applies intra-group union independently within each group

Execution Strategy

CompactAttention executes the selected KV blocks using a paged dense-attention backend while avoiding explicit K/V compaction The accumulated KV cache is stored in a KV-head-major layout [B, Hkv, L, D], where each (batch, KV head, block) triple corresponds to a contiguous [block size, D] memory region This layout turns selected KV blocks into metadata-addressable pages CompactAttention constructs a ragged page list independently for each (batch, group) row, passing only metadata—kv indptr and kv indices—to the paged attention kernel while reusing the original K/V payloads in place This zero-copy design avoids explicit compaction into a newly allocated dense buffer, whose memory bandwidth overhead grows with context length, batch size, and the number of selected KV blocks

Experimental Results

CompactAttention maintains accuracy close to dense attention while delivering up to 2.72× speedup at 128K context length on H200 In Figure 5, CompactAttention-FP reaches 1.64× end-to-end speedup over the dense-attention baseline on Qwen3-30B-A3B CompactAttention maintains accuracy comparable to dense attention and their corresponding block-sparse counterparts across all difficulty levels and context lengths

Conclusion

CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans Across RULER and LongBench V2, CompactAttention maintains accuracy close to dense attention while improving attention latency under chunked prefill, reaching up to 2.

Improvements for AI systems

Improved chunked prefill efficiency through Block-Union KV Selection: CompactAttention treat[s] 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans, converting them into GQA-aware per-group KV block tables through Q-block union and intra-group union. This construction enables selected KV blocks to be accessed in place without explicit KV compaction, achieving up to 2.72× attention speedup at 128K context length under chunked prefill on LLaMA-3.1-8B.
Reduced KV cache overhead via zero-copy execution: CompactAttention avoids the copy overhead of token-level KV selection methods by executing selected blocks using a paged dense attention backend, storing the accumulated KV cache in a [KV-head major layout] where each (batch, KV head, block) triple corresponds to a contiguous [block size, D] memory region.
Optimized execution for high GQA models: For models with high Group Query Attention ratios (e.g., 8:1), the system employs sub-KV-group union by partitioning KV groups into subgroups of four query heads, which substantially preserves sparsity compared to full KV-group union across context lengths.
Preserved accuracy while increasing speedup: CompactAttention maintains high performance, as it shows that block-union KV table construction and zero-copy paged execution directly address the execution bottlenecks of existing chunked-prefill attention methods while ensuring results are close to dense attention on the RULER benchmark.
Improved batch scalability for serving: The system exhibits better scaling in batched environments, as metadata overhead of CompactAttention-FP grows much more slowly than copy overhead, leading to lower latency at larger batch sizes compared to explicit KV copying methods.

Abstract

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.