PaperlayerView exact revision on arXiv

Source-linked AI summary

CompactAttention: Accelerating Chunked Prefill with Block-Union KV Selection

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

arXiv:2605.16839v1cs.CL

TL;DR

Chunked prefill makes long-context attention difficult because short query chunks reduce block-sparse efficiency and repeated pattern search raises overhead, while query-subsampled token selection can miss important KV entries. CompactAttention decouples block-level KV selection from execution, unions selected blocks into GQA-aware tables, and accesses them through zero-copy paged attention. It maintains accuracy close to dense attention and reaches up to 2.72× attention speedup at 128K context length on H200.

  • Problem

    Chunked prefill needs efficient attention despite short query lengths, repeated pattern-search costs, and coverage and copy-overhead limitations in query-subsampled token selection.

  • Method

    CompactAttention converts block-sparse selection masks into GQA-aware per-group KV block tables and executes selected blocks in place with paged attention.

  • Results

    Up to 2.72× attention speedup at 128K context length on H200 is achieved while maintaining accuracy close to dense attention across reported evaluations.

  • Takeaways & Limitations

    Separating KV-block selection from execution lets CompactAttention use block-level coverage with dense paged execution under chunked prefill.

  • Takeaways & Limitations

    CompactAttention inherits missed blocks from its underlying pattern search and may trade sparsity for coverage, with effectiveness varying by architecture and serving configuration.

Abstract

from arXiv · show

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.

1 Introduction

Long-context one-shot prefill is costly, while chunked prefill improves serving practicality but makes efficient attention computation challenging. CompactAttention addresses these chunked-prefill bottlenecks by separating block-level KV selection from execution and achieves a strong accuracy–speedup trade-off.

  • Long-context one-shot prefill incurs quadratic attention cost and can stall batched decode requests, threatening time-between-token service objectives.
  • Chunked prefill processes long inputs sequentially, with each chunk attending to its own KVs and the accumulated KV cache.
  • Block-sparse attention becomes inefficient when Q ≪ KV, while repeated pattern search over the accumulated cache adds substantial overhead across chunks.
  • QUOKA avoids sparse-kernel overhead through query-subsampled dense attention over selected KV entries, but can miss critical entries and requires explicit KV gathering.
  • CompactAttention converts block-sparse masks into GQA-aware per-group KV block tables for in-place paged execution, avoiding sparse-kernel inefficiency and KV-copy overhead.
  • Up to 2.72× attention speedup at 128K context length accompanies accuracy close to dense attention on RULER, yielding the best accuracy–speedup trade-off among baselines.

2 Motivation

Chunked prefill exposes efficiency and coverage weaknesses in existing sparse and query-subsampled attention methods. CompactAttention responds by using block-level selection while executing selected KV blocks through paged attention.

  • Query-Subsampled Direct KV Selection: Mean-attention ranking emphasizes globally important KV positions, whereas max-attention ranking exposes query-specific positions important to only a few queries.
  • Block-sparse methods exploit selected attention tiles, but their one-shot-prefill benefits do not transfer directly to chunked prefill.
  • Kernel Inefficiency: At 90% sparsity and KV length 64K, block-sparse kernels approach ideal 10× speedup for one-shot Q = 64K but fall substantially short for chunked Q = 1024.
  • Pattern Search Overhead: Repeated online pattern search over the accumulated KV cache makes chunked prefill sensitive to selector cost, favoring lightweight methods such as SeerAttention and FlashPrefill.
  • Query-Subsampled Direct KV Selection: QUOKA can miss query-specific KV entries when their corresponding queries are not sampled, degrading on distributed-information tasks while block-sparse methods remain near dense attention.
  • CompactAttention: CompactAttention covers all query blocks, selects KV blocks at block granularity, and enables direct access without explicit compaction by decoupling selection from execution.

3 CompactAttention

CompactAttention decouples block-level KV selection from attention execution, lowering per-head sparse masks into GQA-aware block tables for paged attention. It preserves selected-block coverage while enabling zero-copy execution over the original KV cache.

  • Overview: CompactAttention accepts lightweight block-sparse importance estimates and separates KV selection from attention execution.It converts 2D per-head block masks into KV block tables instead of using them directly as sparse-kernel plans.
  • KV Selection: Q-block union combines selections across query blocks because paged attention uses one KV block list for jointly executed queries.The resulting mask is then prepared for group-wise execution.
  • KV Selection: Intra-group union combines query-head selections into one KV block table per execution group, which defaults to a KV group.For large GQA groups, CompactAttention can split groups into smaller execution groups; the implementation uses four query heads per subgroup.
  • KV Selection: The resulting table is minimal under shared-table constraints while preserving every KV block selected by any query block or head in the group.Blocks absent from the input mask can be excluded without violating coverage preservation.
  • Trade-offs: Union reduces sparsity relative to the original 2D mask, but more aggressive initial pattern search can compensate while preserving accuracy.The paper reports that dense paged execution still achieves higher attention speedup than corresponding block-sparse baselines.
  • Execution: Zero-copy paged execution accesses selected blocks through metadata over the original KV cache rather than explicit K/V compaction.KV-head-major layout exposes KV-head blocks as pages, allowing group-dependent tables while reusing K/V payloads in place.

4 Experiments

CompactAttention is evaluated on two GQA-based long-context models against dense, block-sparse, and chunked-prefill KV-selection baselines. It maintains accuracy close to dense attention while improving long-context speed through zero-copy paged execution.

  • Experimental Setup: Experiments use LLaMA-3.1-8B-Instruct and Qwen3-30B-A3B-Instruct-2507 on RULER and LongBench V2.The models support 128K and 256K tokens, respectively, and both use GQA.
  • Experimental Setup: CompactAttention is compared with dense attention, three block-sparse methods, and QUOKA using paged attention with per-group KV block tables.CompactAttention-SA reuses SeerAttention, while CompactAttention-FP applies FlashPrefill thresholding without model-specific adaptation.
  • Speedup: QUOKA provides limited long-context speedup, while XAttention and SeerAttention can be slower than dense attention; FlashPrefill is the strongest block-sparse baseline.Gather-and-pack, repeated pattern search, and inefficient execution in the Q ≪ KV regime explain these outcomes.
  • Speedup: 2.72× attention speedup and 1.96× end-to-end speedup are achieved by CompactAttention-FP on H200 at 128K context.Both CompactAttention variants show increasing speedup as context length grows.
  • Accuracy: CompactAttention variants remain close to dense accuracy on RULER and LongBench V2, whereas QUOKA degrades more noticeably on tasks requiring broad information access.Block-union preserves the selected blocks and therefore largely preserves the quality of the underlying block-level selectors.
  • Ablation Studies: At comparable or higher accuracy, CompactAttention-FP achieves higher speedup than FlashPrefill in the high-accuracy operating region.At the selected RULER 128K operating point, union reduces initial 89.8% sparsity to 70.2%, comparable to FlashPrefill’s 69.8%.
  • Ablation Studies: CompactAttention-FP has the lowest latency in the matched-mask ablation, outperforming sparse-kernel execution and explicit KV copying despite metadata overhead.The in-place paged path accesses the original KV cache through metadata rather than materializing selected blocks into a contiguous buffer.

5 Limitations

CompactAttention cannot recover KV blocks omitted by its underlying block-sparse selector, and its sparsity-versus-coverage trade-off may vary across deployment settings.

  • Limitations: CompactAttention inherits selector errors because block-union cannot recover KV blocks missed by the input mask.Union also retains any block selected by another query block or head within an execution group, trading sparsity for coverage.
  • Limitations: The sparsity trade-off may vary with model architecture, context length, chunk size, sparsity hyperparameters, and execution-group partitioning.The method is most effective when the accumulated KV cache is large enough to amortize pattern-search and metadata-construction overheads.

6 Conclusion

CompactAttention converts block-sparse selection masks into GQA-aware KV tables and executes selected blocks through zero-copy paged attention. Across RULER and LongBench V2, it stays close to dense accuracy while reaching 2.72× attention speedup at 128K on H200.

  • Conclusion: CompactAttention treats 2D block-sparse masks as KV-selection signals rather than direct sparse-kernel execution plans.Q-block union and intra-group union produce per-group KV block tables for paged execution.
  • Conclusion: 2.72× attention speedup is reached over dense attention at 128K context length on H200 while accuracy remains close to dense attention.The reported accuracy trend spans both RULER and LongBench V2.
  • Conclusion: The results identify execution of selected KV blocks as a central bottleneck alongside block selection under chunked prefill.Making sparse selection compatible with dense paged-attention kernels provides a practical path toward faster long-context serving.

B.1 Sub-KV-Group Union

For GQA ratios above 4:1, CompactAttention uses sub-KV-group union to limit sparsity loss while preserving zero-copy paged execution. Fixed groups of four query heads balance sparsity preservation against block-table metadata overhead.

  • For GQA ratios above 4:1, full KV-group union causes excessive sparsity loss, motivating sub-KV-group union.
  • Sub-KV-group union partitions each KV group into execution groups of four query heads.
  • Smaller subgroups increase block-table metadata overhead enough to outweigh their sparsity benefits.
  • Sub-KV-group union substantially preserves sparsity compared with full KV-group union across context lengths while retaining the same zero-copy execution interface.
  • The current chunk remains fully open because sparsifying it would apply causal masking in compacted-position space and break causal semantics.

C.1 Additional Latency Results

Additional latency experiments show that CompactAttention reduces long-context attention and end-to-end latency, with benefits varying by chunk size and model configuration. Its gains persist on a larger MoE model from 64K context onward.

  • The LLaMA-3.1-8B-Instruct latency measurements use chunk sizes 512 and 1024 on RTX PRO 6000 and H200 SXM GPUs, respectively.
  • Long-context CompactAttention reduces attention latency versus block-sparse baselines and dense attention, translating into lower end-to-end latency.At shorter contexts, pattern-search and metadata-construction costs can make sparse attention slower, though attention contributes less to total prefill time.
  • Chunk Size Sensitivity: At 128K context, CompactAttention-FP improves attention latency across all tested chunk sizes.Its relative speedup decreases at chunk size 2048 because larger chunks increase Q-block union and reduce effective sparsity.
  • Chunk Size Sensitivity: As chunk size increases, fewer chunked-prefill iterations reduce total attention latency.
  • Attention Speedup on Qwen3-30B-A3B: From 64K onward, CompactAttention-FP outperforms QUOKA and the corresponding FlashPrefill baseline on Qwen3-30B-A3B-Instruct-2507.This configuration uses sub-KV-group union with subgroup size 4.

C.2 Additional Accuracy Results

CompactAttention preserves accuracy close to dense attention across the reported RULER and LongBench V2 evaluations. This accuracy advantage is especially notable relative to QUOKA's reported degradation across LongBench V2 categories.

  • RULER with Chunk Size 512: CompactAttention variants maintain RULER accuracy close to dense attention across context lengths, with modest degradation at 128K.
  • LongBench V2: CompactAttention-SA and CompactAttention-FP maintain accuracy comparable to dense attention and block-sparse counterparts across LongBench V2 difficulty levels and context lengths.
  • LongBench V2: QUOKA shows consistent degradation across LongBench V2 categories compared with dense attention.
Loading 2605.16839v1…