The Hardware Argument for Position-Aware KV-Cache Compression

On modern GPU architectures, a KV-cache compression scheme that selects contiguous positional blocks via cheap centroid routing will outperform a scheme that selects semantically clustered but spatially scattered tokens at equal compression ratios — not because it captures more relevant information, but because coalesced memory access between HBM and SRAM dominates total decode latency, and no amount of FLOP reduction from clever token selection can overcome the gather/scatter penalty of non-contiguous reads.

The Problem

Large language model inference is bottlenecked by the Key-Value cache. It grows linearly with context length, and at 128K tokens on a 70B-parameter model, the KV-cache alone can exceed 40GB of GPU memory. At this scale, compression is not an optimization — it is a deployment requirement.

Two broad strategies have emerged for compressing the KV-cache during inference:

Similarity-aware compression groups tokens by embedding similarity, then prunes or merges low-attention clusters. The logic is sound: tokens that attend to similar things can be represented together.

Position-aware compression groups tokens by their position in the sequence, compresses contiguous blocks, and routes queries to relevant blocks via lightweight centroid scores. The logic is simpler: nearby tokens already share context, and contiguous memory is fast.

This post argues that the second approach — position-aware, block-contiguous compression — will dominate in practice, not because it is more theoretically elegant, but because it aligns with GPU hardware physics in ways that similarity-aware methods fundamentally cannot.

Three Hypotheses

H1: Attention is temporally local

For autoregressive decoder-only transformers, the attention weight distribution for a given query token concentrates in two predictable regions: tokens near the query's own position (local context), and tokens near position zero (the attention sink). Formally, for a query at position $i$:

$\sum_{j \in \mathcal{N}(i)} A_{i,j} + \sum_{j \in \mathcal{B}_0} A_{i,j} \geq (1 - \epsilon) \cdot \sum_{j=1}^{n} A_{i,j}$

where $\mathcal{N}(i)$ is a local window around position $i$, $\mathcal{B}_0$ is the first block of tokens, and $\epsilon$ is small (< 0.05 in most layers).

This is not a new observation. The attention sink phenomenon is well-documented (Xiao et al., 2023), and local-window dominance is the basis of sliding-window attention in models like Mistral. But it has an underappreciated implication for cache compression: if attention mass is positionally concentrated, then a compression scheme that preserves positional neighborhoods loses almost nothing. You don't need to cluster by semantic similarity if proximity already captures the relevant structure.

H2: Memory access patterns dominate total latency

This is the core of the argument.

Similarity-aware methods select a semantically coherent but spatially scattered subset of KV-cache entries. On paper, they reduce FLOP count. In practice, they require gather operations across non-contiguous GPU memory — and on modern GPUs with streaming multiprocessors optimized for coalesced reads, the gather/scatter overhead can negate the FLOP savings entirely:

$T_{\text{similarity}} = T_{\text{compute}} + T_{\text{gather}} \gg T_{\text{position}} = T_{\text{compute}} + T_{\text{contiguous}}$

FlashAttention (Dao et al., 2022) demonstrated that the dominant cost of attention is not arithmetic — it is moving data between HBM and SRAM. That same principle applies here. A compression method that reduces FLOPs by 4x but introduces scattered memory access may deliver less than 2x wall-clock speedup. A method that achieves only 2x FLOP reduction but keeps all accesses contiguous may deliver 4-6x wall-clock speedup.

The prediction is falsifiable: at equal compression ratios, position-aware block selection should outperform similarity-aware cluster selection in wall-clock decode latency, and the gap should widen with sequence length.

H3: Early layers resist compression

Not all transformer layers behave the same way. Early layers (roughly 0–4) tend to compute broad, relatively uniform attention distributions. Later layers produce sparse, peaked distributions where most attention mass concentrates on a few tokens.

This suggests a natural compression boundary. Layers with high attention entropy — where $H_\ell = -\sum_j A_{\ell,j} \log A_{\ell,j}$ is large — should remain uncompressed with full dense attention. Layers with low entropy are safe to compress aggressively.

The hypothesis predicts that this boundary can be identified empirically via per-layer entropy profiling and that it generalizes across model families. Early layers are doing something fundamentally different — likely building broad contextual representations — and compressing them causes cascading errors through the residual stream that cannot be recovered downstream.

The Algorithm (Sketch)

The position-aware approach works as follows:

1. Partition the KV-cache into fixed-size contiguous blocks of $B$ tokens (e.g., $B = 64$).

2. Compute centroids for each block: $c_b = \frac{1}{B} \sum_{j \in b} K_j$

3. Route queries to the most relevant blocks via a cheap dot-product score against centroids: $s_b = q^T \cdot c_b$, then select the Top-$M$ scoring blocks. This costs $\frac{1}{B}$-th the FLOPs of full attention.

4. Load selected blocks contiguously into SRAM and compute exact scaled dot-product attention over them.

5. Always preserve Block 0 (attention sink) and the latest block (local context). These are never pruned regardless of centroid scores.

The key property: every memory access during the attention computation is to a contiguous region. There are no gathers, no index lookups, no scattered reads. The GPU's memory subsystem operates at peak efficiency.

Where This Breaks Down

Intellectual honesty demands acknowledging the failure modes.

Position-aware compression assumes that relevance correlates with proximity. This assumption breaks for retrieval-heavy tasks where the answer is buried thousands of tokens earlier, multi-document QA where relevant passages are spread across the context, code with distant function definitions or imports referenced much later, and any task that requires genuine long-range dependency rather than local context plus an attention sink.

Similarity-aware methods exist precisely because these cases are real. By clustering semantically related tokens regardless of position, they can in principle preserve access to scattered but relevant information that block-based methods would prune.

The empirical question is how often these cases dominate in practice. If 95% of attention mass is captured by local context and the sink, then optimizing for the 5% at the cost of 3-4x throughput is a bad trade. But for applications that live in that 5% — RAG pipelines, legal document analysis, multi-hop reasoning — position-aware compression may not be sufficient.

Proposed Validation

Four experiments to test these hypotheses:

Attention locality profiling. Measure attention mass in local windows vs. distant tokens across layers in Llama-3.1-8B (or comparable open model) at 8K, 32K, and 128K context lengths. This directly tests H1 and provides the data for per-layer entropy analysis (H3).

Memory access benchmark. Compare wall-clock latency of block-contiguous selection vs. scatter-gather selection at equal compression ratios on A100 and H100 GPUs. Isolate the memory access variable by keeping FLOP count identical. This directly tests H2.

Entropy-based layer boundary. Compute per-layer attention entropy across diverse prompts. Determine whether a stable compression boundary exists and whether it generalizes across model sizes and families.

Needle-in-a-haystack stress test. Evaluate retrieval accuracy at increasing compression ratios. Position-aware methods preserve token neighborhoods — the question is whether this is sufficient when the needle's location is unpredictable, or whether similarity-aware methods recover needles more reliably despite their latency cost.

The Bet

If these hypotheses hold, position-aware block compression should achieve comparable or better perplexity to similarity-aware methods, substantially higher decode throughput (due to hardware alignment rather than FLOP reduction), and robust retrieval accuracy up to moderate compression ratios (50-60%), with degradation at extreme ratios where long-range dependencies become critical.

The core claim, stated plainly: on modern GPU architectures, contiguous memory access is not a performance detail — it is the dominant variable in real-world inference latency. Any compression scheme that ignores this in favor of theoretical FLOP savings will underperform a simpler scheme that respects it.

This is ultimately an empirical question, and I intend to run the experiments. But the direction of the bet seems clear: in the contest between algorithmic sophistication and hardware physics, hardware physics tends to win.

References: Xiao et al. (2023), "Efficient Streaming Language Models with Attention Sinks"; Dao et al. (2022), "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"; Jiang et al. (2023), Mistral 7B sliding window attention.