Every decode step loads the entire KV cache from GPU HBM. At 128K tokens that is several GB; at 1M tokens it exceeds 60GB. The bottleneck is memory bandwidth, not compute. The GPU waits while the memory bus transfers data that is likely irrelevant to the current query.

Why shared-basis sparse attention is structurally broken

Quest, ShadowKV, and RocketKV reduce this cost by estimating which KV pages matter and only reading those. They score pages by projecting the current query onto a shared low-rank basis built from the full sequence. The fundamental issue: any fixed basis W has a null space. A page whose key structure lies in that null space will score near zero regardless of its actual relevance to the query. No training procedure can fix this — it is a representational impossibility.

Proposition 1 in LOCKS formalizes it: every fixed shared projection has "page-content blind directions." ShadowKV collapsing at small token budgets is not a tuning problem; it is this structural failure showing up.

The observation LOCKS builds on

Within a page of 16 consecutive tokens, the key vectors concentrate in a low-dimensional subspace. Tokens that occur together tend to require similar attention patterns, which produces similar key vectors, which cluster in a small region of the full d_head-dimensional space.

Across the full sequence, different pages occupy different subspaces — a page of code, a page of prose, a page of citations each has a different key subspace orientation. The global key matrix is high-rank even though each individual page is locally low-rank.

A shared basis must cover all of these simultaneously with a single set of directions. It cannot. Per-page bases can, because each one only needs to cover its own page.

How LOCKS works

When a page fills (every 16 tokens), LOCKS runs a one-time SVD on the centered key deviations. It stores:

  • μⱼ (int8): the page centroid — mean key vector across the 16 tokens
  • Vⱼ (int4, rank 8): the top-8 right singular vectors of the deviation matrix — the principal directions of variation within this page
  • Cⱼ (int8): each key's projection coefficients onto Vⱼ

Total: about 10% of the original page KV data. This runs off the decode critical path.

At each decode step, LOCKS scores every page from these summaries alone — no K or V tensors read:

$$\hat{s}j(q) = \log \sum{i \in P_j} \exp\left[\frac{q^\top \mu_j + (V_j^\top q)^\top c_i}{\sqrt{d}}\right]$$

q^T μⱼ: how well the query aligns with the page's average key direction (one scalar per page).
(Vⱼ^T q)^T cᵢ: how well the query aligns with each individual key's deviation from the centroid, measured in the page's local basis (one scalar per token).

Summing and applying log-sum-exp gives the estimated total attention mass for the page. Top pages by this score get full attention; the rest are skipped.

Theorem 1 proves unquantized rank-r summaries achieve at least e^(−4η) of exact mass-selection coverage, where η bounds the SVD reconstruction error. Lemma 1 proves mass-ranking is the worst-case optimal strategy among all deterministic value-blind selectors.

Code 

import torch
import torch.nn.functional as F

SEQ_LEN = 65536
D_HEAD  = 128
PAGE_SIZE = 16
RANK = 8
TOKEN_BUDGET = 2048

K = torch.randn(SEQ_LEN, D_HEAD)
V = torch.randn(SEQ_LEN, D_HEAD)
q = torch.randn(D_HEAD)

# Build page summaries once (off critical path)
summaries = []
for p in range(SEQ_LEN // PAGE_SIZE):
    keys   = K[p*PAGE_SIZE:(p+1)*PAGE_SIZE]
    mu     = keys.mean(0)
    dev    = keys - mu.unsqueeze(0)
    _, _, Vh = torch.linalg.svd(dev, full_matrices=False)
    basis  = Vh[:RANK].T
    coeffs = dev @ basis
    summaries.append({"mu": mu, "V": basis, "C": coeffs})

# Per-step: score all pages from summaries only
scale = D_HEAD ** -0.5
scores = []
for s in summaries:
    global_t = (q @ s["mu"]) * scale
    local_t  = (s["C"] @ (s["V"].T @ q)) * scale
    scores.append(torch.logsumexp(global_t + local_t, dim=0))
page_scores = torch.stack(scores)

# Select top pages and attend
n_top = TOKEN_BUDGET // PAGE_SIZE
top_pages = torch.topk(page_scores, n_top).indices.sort().values
sel_idx = torch.cat([
    torch.arange(p*PAGE_SIZE, (p+1)*PAGE_SIZE) for p in top_pages
])

out_locks = F.scaled_dot_product_attention(
    q.view(1, 1, D_HEAD),
    K[sel_idx].view(1, len(sel_idx), D_HEAD),
    V[sel_idx].view(1, len(sel_idx), D_HEAD),
).squeeze()

out_full = F.scaled_dot_product_attention(
    q.view(1, 1, D_HEAD),
    K.view(1, SEQ_LEN, D_HEAD),
    V.view(1, SEQ_LEN, D_HEAD),
).squeeze()

cos_sim = F.cosine_similarity(out_locks.unsqueeze(0), out_full.unsqueeze(0)).item()
print(f"Attended: {len(sel_idx)}/{SEQ_LEN} ({100*len(sel_idx)/SEQ_LEN:.1f}%)")
print(f"Cosine similarity to FullKV: {cos_sim:.4f}")
# Expected: ~0.98+ at 3% token attendance

Enter fullscreen mode Exit fullscreen mode

Results

InfiniteBench (GLM-4-9B-Chat-1M, 100K+ context):

Method b=512 b=2048
Quest 34.7 35.7
ShadowKV 36.0 40.3
RocketKV 38.5 42.3
LOCKS 41.1 43.6
FullKV 43.0 43.0

At b=2048, LOCKS (43.6) exceeds FullKV (43.0). Attending to a selected subset of tokens excludes the irrelevant ones that would otherwise dilute the attention distribution.

LongBench-v1 (Llama-3.1-8B): stays within ~1 point of FullKV from b=64 to b=2048. Competing methods collapse below b=512.

H200 decode latency: 1.59× at 128K, 1.80× at 256K, 2.0× at 1M. Bytes read per step: 9.8× reduction at 1M. Because decode is almost entirely bandwidth-bound at these context lengths, bandwidth savings map nearly linearly to latency savings.

Caveats

Theorem 1 applies to unquantized summaries. The shipped int4 configuration is empirically validated, not per-step certified. Ablations show int4 loses a few recall points; int2 collapses. Validate before production.

LOCKS saves bandwidth, not VRAM. The full KV cache stays resident. If the constraint is memory capacity, not bandwidth, this does not help.

Single-author preprint; independent replication has not yet appeared. Most compelling at 256K+ context — the 1.59× gain at 128K may not justify integration overhead.

References