CSAttention: Centroid-Scoring Attention for Accelerating LLM Inference

Abstract: Long-context LLMs increasingly rely on extended, reusable prefill prompts for agents and domain Q&A, pushing attention and KV-cache to become the dominant decode-time bottlenecks. While sparse attention reduces computation and transfer costs, it often struggles to maintain accuracy at high sparsity levels due to the inherent distribution shift between Queries and Keys. We propose Centroid-Scoring Attention (CSAttention), a training-free sparse attention method optimized for high-throughput serving of reusable contexts. CSAttention adopts a storage-for-computation strategy tailored to the offline-prefill/online-decode setting: it front-loads computation into a one-time offline prefill phase that can be amortized across multiple queries, while aggressively optimizing per-step decoding latency. Specifically, CSAttention constructs query-centric lookup tables during offline prefill, whose size remains fixed during decoding, and enables online decoding to replace full-context scans with efficient table lookups and GPU-friendly score accumulation. Extensive experiments demonstrate that CSAttention achieves near-identical accuracy to full attention. Under high sparsity (95%) and long-context settings (32K-128K), CSAttention consistently outperforms state-of-the-art sparse attention methods in both model accuracy and inference speed, achieving up to 4.6x inference speedup over the most accurate baseline at a context length of 128K.

• The paper introduces a training-free sparse-attention method that clusters prefill queries across subspaces and aggregates centroid-based key scores, avoiding the query–key distribution shift affecting key-centric indexes.
• CSAttention retains about 95% of tokens while staying within 0.7 LongBench points of full attention across three 7B–8B models, including 52.04 versus 52.41 for Llama-3.1-8B.
• The method reduces decode-time work to bounded search plus sparse attention, achieving up to 17.9× speedup over SparQ at 128K context and 4.24× over dense attention in All-GPU serving.

CSAttention is a training-free sparse attention method designed for high-throughput serving of long, reusable prefill contexts. The paper targets the "write-once, read-many" serving regime common to RAG pipelines and agent workloads, where a one-time offline prefill over a long context is amortized across many online decoding requests. Its central contribution is a query-centric retrieval index that replaces full-context scans at decode time with bounded-capacity table lookups and GPU-friendly score accumulation, sustaining near-full accuracy at 95% sparsity where key-centric baselines degrade substantially.

Motivation: query-key distribution shift

The method is grounded in three empirical observations on Llama-3.1-8B-Instruct. First, accuracy under sparse attention correlates strongly with Top-K recall of dense-attention weights, particularly at very high sparsity, so stable recall is the governing design objective. Second, per-subspace contributions to the $q \cdot k$ similarity are heavily tailed across $m$ subspaces, meaning aggregated evidence over subspaces can recover true high-scoring keys even when each subspace alone is only partially informative. Third, PCA visualization of queries and keys from the same layer and head shows a marked distribution shift between $Q$ and $K$, arising because the two are produced by different projections.

This last observation motivates the paper's core architectural departure. Prior index-based methods (PQCache, ClusterKV) build lookup structures by clustering keys, following a search path $Q \to K\text{-centroid} \to K$. Because queries lie out-of-distribution relative to key-built centroids, recall becomes unstable as sparsity increases. CSAttention instead clusters the prefill queries themselves in each subspace, yielding a $Q \to Q\text{-centroid} \to K$ path in which nearest-centroid assignment occurs in the same space as the incoming query, eliminating the OOD hop.

Method

CSAttention augments, rather than replaces, the standard KV cache. During offline prefill, each head's $d$ dimensions are split into $m$ subspaces (default $m{=}8$); within each subspace, prefill queries are clustered with cosine $k$-means into $m$0 centroids ($m$1 by default). For every centroid, partial dot-products against all keys in that subspace are computed via batched GEMM, and a compressed Top-$m$2 list of (index, score) pairs is stored, with $m$3 tied to the prefill length $m$4.

At decode time, an arriving query selects its nearest centroid per subspace (a batched GEMV), gathers the $m$5 short lists, and performs a branchless reduce-by-key accumulation of partial scores over the union of indices, whose size is bounded by $m$6 and therefore constant with respect to generation length. Keys aligned with the query tend to "collide" across multiple subspaces and rise to the top after aggregation. A recent window of $m$7 positions is unioned into the candidate set before final Top-$m$8 selection, and streaming updates try-insert each newly appended key into the fixed-capacity lists without resizing. Two execution modes are supported: All-GPU, where tables and KV reside in HBM, and CPU$m$9GPU, where tables and KV live in DRAM and CPU-side search overlaps asynchronously with GPU attention, transferring only the selected Top-$Q$0 entries per step.

The complexity analysis makes the amortization explicit: per-step cost reduces from dense $Q$1 to a constant search term $Q$2 plus sparse attention $Q$3 with keep ratio $Q$4. The authors are careful to note that CSAttention does not eliminate the linear KV storage cost; it targets the decode-time compute and transfer bottleneck, and the relative memory overhead of the fixed-size tables shrinks asymptotically toward zero as $Q$5 grows beyond $Q$6.

Accuracy results

On LongBench across Llama-3.1-8B, Qwen3-8B, and Mistral-7B-Instruct-v0.3, CSAttention at ~5% token retention stays within 0.7 points of full attention on macro average — 52.04 vs. 52.41 (Llama), 52.25 vs. 52.30 (Qwen3), and exactly 49.92 vs. 49.92 (Mistral). Baselines trail considerably: PQCache reaches 49.79/50.59/45.49, while H$Q$7O, SparQ, and MagicPig fall further behind, with pronounced failures on multi-document summarization and cross-lingual QA tasks. On LongBench v2 with Llama-3.1-8B, CSAttention scores 31.2 overall, slightly exceeding even the dense baseline (31.0) and improving Hard (29.3 vs. 28.3) and Long (32.4 vs. 30.6) buckets; all sparse baselines drop between 1.2 and 4.8 points. Schedule ablations show that infrequent searching preserves accuracy: keeping 15% or 20% tokens but searching only every 4 or 8 steps remains within 0.61 points of Full, exploiting locality in consecutive tokens' attention patterns.

Efficiency results

In the CPU$Q$8GPU mode, speedups grow with context length because per-step work scales with fixed table sizes rather than total history. Against PQCache, CSAttention achieves 2.95× at 8K rising to 8.26× at 128K; against SparQ, up to 17.9× at 128K; against MagicPig, up to 7.85×. Against H$Q$9O gains are modest (up to 1.33×), reflecting H$K$0O's cheaper but less accurate retention scheme. In All-GPU mode, CSAttention beats full attention itself by 4.24× at 128K. Step-level latency statistics show tight tails (P99 ≈ 1.01 normalized), with a stable per-step composition of Attention : Search : Update ≈ 1.0 : 0.3 : 0.1, supporting the claim of predictable decode-time behavior.

Robustness under chain-of-thought decoding

A stated concern is whether query drift during long CoT generation erodes centroid-scoring recall. The paper evaluates LongBench v2 with CoT decoding (up to 2048 generated tokens) on Llama-70B and Qwen3-32B, using two safeguards: a recent-window candidate union and a centroid backoff that merges top-$K$1 nearest centroids when cosine similarity is low. Results remain near-baseline across difficulty buckets (e.g., 36.4 vs. 36.5 overall for Llama-70B), suggesting that the diversity of the long system prompt's query distribution keeps routing stable. Notably, this robustness claim rests on the assumption that the prefill prompt itself provides a sufficiently diverse query distribution — a condition specific to the reusable-prefill setting the paper targets.

Limitations and open questions

Several constraints bound the applicability of these results. The method presupposes an offline prefill phase whose indexing cost can be amortized; for single-shot requests with unique contexts, the one-time build overhead would not be recoverable, and the paper does not evaluate that regime. The KV cache's linear storage cost is explicitly left unaddressed, so CSAttention complements rather than substitutes for KV offloading or compression. Robustness under CoT depends on prefill query diversity, and behavior under adversarially narrow or drifting query distributions remains untested. Finally, uniform subspace weights are used throughout; learned or confidence-based weighting is mentioned as possible but not explored, leaving open whether non-uniform weighting could further improve recall at fixed table budgets.

Conclusion

CSAttention demonstrates that reorienting sparse-attention retrieval around query-space clustering, combined with subspace partitioning and bounded-capacity centroid-score tables, yields near-lossless accuracy at 95% sparsity and substantial decode speedups that widen with context length. Its strongest empirical claims — parity with full attention on LongBench and LongBench v2, and up to 17.9× speedup over SparQ at 128K — hold specifically in the offline-prefill / online-decode serving model, which the paper's design and evaluation consistently and transparently assume.