Training-free sparse attention for long-context LLM decode

Training-free KV-cache routing and sparse attention for long-context decode on frozen pretrained LLMs: a from-scratch Triton sparse-decode kernel, a Blackwell wall-clock replication of ClusterKV-style cluster routing, and eleven rigorously falsified negative results. A lab notebook, not a library or a SOTA claim.

Three load-bearing results (each fully scoped):

• Sink-correct Triton sparse-decode kernel — on gpt-oss-20b, 200/200 token-identical to full attention on one 32K retrieval prompt (C=1024/N=64); fp32 rel-L2 1.5e-7. Correctly folds gpt-oss's learned attention sink into the softmax denominator, which published cluster kernels omit.
• Real wall-clock speedup where attention dominates — Qwen2.5-7B-Instruct-1M (dense, no MoE), single RTX PRO 6000 Blackwell (sm_120), vs HuggingFace flex_attention: 1.92x at 128K, 3.73x at 256K, 5.37x at 512K. On gpt-oss-20b (MoE + sliding-window) the same kernel is roughly a wash end-to-end — the honest mechanism is documented.
• Eleven documented negative results, each built correctly and falsified with a mechanism (FAVOR+, CSAttention reimplementation, HippoRAG-style diffusion, learned page scorers, synth-KV distillation, …).

Scope. Frozen pretrained models, no weight training. Single author, exploratory, not peer-reviewed; numbers not independently reproduced and are specific to the stated model / hardware / context / metric / baseline. The baseline throughout is HuggingFace flex_attention on a nightly Triton 3.6 / CUDA 13 / sm_120 stack — not an optimized production serving stack (vLLM / paged FlashAttention). Two model regimes are studied and never blurred: gpt-oss-20b (head_dim 64, learned attention sinks, MoE + alternating sliding-window layers) for correctness/output-error, and Qwen2.5-7B-Instruct-1M (head_dim 128, GQA, no sinks) for wall-clock and downstream quality.

---

Strongest signals

Conditions are inline. Every headline number is re-derivable from a checked-in data/ artifact except where noted.

1. Sink-correct Triton sparse-decode kernel

scripts/sparse_decode_triton.py (844 lines). Direct gather of per-head-selected KV positions + FlashAttention-style online (numerically stable) softmax. The non-obvious correctness detail: gpt-oss has a learned per-head attention sink — a scalar bias added only to the softmax denominator — and this kernel folds it in as a virtual key with V=0 (sparse_decode_kernel, lines 115-123). Published cluster kernels (ClusterKV, Multipole, Tactic) target sink-free Llama/Mistral/Qwen2 and have no such term; applying them unmodified to gpt-oss is incorrect, and stock SDPA silently drops the sink.

• Correct to the precision floor: 1.5e-7 rel-L2 in fp32 (max abs 2.2e-8), 1.7e-3 rel-L2 in bf16 (the model's own inference dtype), vs a clean PyTorch reference (sparse_decode_reference, lines 590-625; harness lines 808-844).
• Token-identical decode on one 32K retrieval prompt: 200 greedy tokens, C=1024 / N=64, 200/200 token-identical to a full-attention (flex_attention) reference (experiments/2026-05-07-triton-kernel.md; data/cluster_decode_triton_walden_v2_200.json). One prompt/task at one operating point — not a general losslessness guarantee. On a Frankenstein verbatim-quote prompt both this kernel and the reference cluster method diverge from full attention around token 11 (a sparsity-recipe limit, not a kernel bug).
• Per-layer compute: at C=1024 / N=64 it reads ~6.5-7% of K/V; per-layer attention compute is ~32x faster than full (full ~6.7 ms/layer vs Triton 0.209 ms/call median — re-derived from kernel_call_ms in the Walden artifact). A separate microbench on synthetic K/V (experiments/kernel_microbench.py, not persisted to disk) shows a 1-2 order-of-magnitude per-call ceiling in isolation; treat that as illustrative, not a result. The load-bearing number is the verifiable ~32x.

2. Where the speedup is real: per-call win vs wall-clock win

A per-call kernel win only converts to wall-clock when attention dominates per-step time.

• gpt-oss-20b (MoE + half sliding-window layers): roughly a wash end-to-end — ~1.45x at 32K (56.3 ms/step vs 81.5 ms full), bottlenecked by a ~49 ms/step non-attention floor (MoE FFN, sliding-window layers, norms, LM head). Per-step decomposition (kernel ~2.5 ms + index build ~4.7 ms + ~49 ms non-attention) in experiments/2026-05-07-triton-kernel.md.

• Qwen2.5-7B-Instruct-1M (dense, all-full-attention, no MoE): converts to real wall-clock. Single RTX PRO 6000 Blackwell (sm_120), median ms/step decode, 50 greedy tokens, baseline flex_attention:

  context	full (ms/step)	cluster (ms/step)	speedup
  32K	36.6	39.0	0.94x
  64K	58.7	43.4	1.35x
  128K	100.8	52.6	1.92x
  256K	220.7	59.1	3.73x
  512K	499.5	93.0	5.37x
  1M	crashes (Triton compile bug)	~181 (warm)	~5.5x extrapolated

  Straight median of per-step ms from data/qwen7b1m_{32k..512k}_{full,cluster}_blackwell.json. The load-bearing claim is the scaling shape — full attention scales ~linearly in context, cluster sub-linearly — not the absolute factor: the multiplier magnitude partly reflects flex_attention's near-linear scaling on this stack and would shrink against vLLM/paged-FA. The 1M row is cluster-only / extrapolated: flex_attention crashes on a Triton 3.6 / sm_120 compile bug at (Q_LEN=1, KV_LEN=1M), and ~181 ms is the warm step (the first 1M step pays a ~997 ms JIT/alloc cost; n=2, no median possible). Qwen quality was validated qualitatively only at these lengths; token-identical validation is gpt-oss-only.

  The same curve reproduced on an L40S at higher multipliers (1.33x/2.00x/3.41x/5.49x at 32K-256K, experiments/2026-05-07-qwen-long-context.md) because Blackwell makes the full-attention baseline ~2x cheaper per step — the sub-linear-vs-linear shape is preserved; the absolute multiplier shifts with the baseline.

3. Cluster-centroid routing vs Quest (offline output error, gpt-oss-20b, d=64)

experiments/2026-05-06-output-error.md. Metric: per-Q-head relative attention-output L2 ||O_sparse − O_full||₂ / ||O_full||₂, sink correctly in the denominator (experiments/output_error_sweep.py:75-82,227), ~7.4M (layer,head,step) samples across 4 prompts. Cluster-centroid + sink + recency reaches 5.8-8.9% mean output error at ~7% bandwidth, where vanilla Quest needs ~13% bandwidth to match — about half the bandwidth at equal output error.

Two caveats: this is an offline error proxy, not a downstream-accuracy comparison, and the Quest baseline did not receive the same explicit sink+recency injection, so the gap reads optimistic and part of the small-budget advantage is recency mass, not clustering. End-to-end greedy validation (experiments/2026-05-06-decode-validation.md) confirms cluster C=1024 / N=32 (~3.5% bandwidth) is token-identical to full attention (50/50) on Walden retrieval and summarization, where Quest at 1K-2K budget diverges at token 7.

4. Engineering range

• Mini-batch k-means cut 256K clustering (the dominant long-context TTFT cost) from 138.4 s to 22.3 s (6.2x), no measurable quality loss on the tasks tested (experiments/2026-05-08-blackwell-evals-cluster-tuning.md §3).
• Bounded-hierarchical ClusterKV caps cluster size to roughly halve the kernel's padded work-bound at equal recall/F1, via a vectorized stable-order tail-repair that brought 32K clustering from a 28.6 s recursive prototype to ~0.7 s (experiments/2026-05-11-bounded-hierarchical-clusterkv.md). Honest readout: quality effectively tied at n=50 (mean F1 0.2326 vs 0.2355), and the 32K kernel win did not reproduce end-to-end at 256K (dominated by prefill).
• Evidence-gated autoresearch loop (experiments/clusterkv_autoresearch/) drives an LLM coding agent inside disposable sparse-checkout git worktrees, scores each candidate against a high-power offline error proxy, and accepts a patch only if every gate passes AND it beats the prior best. Best run: 18 patches accepted, the rest rejected for no improvement (the gate doing its job), score 114.98 → 122.07 over 80 iterations, zero loop errors. It is not "AI discovered a new attention method": candidates are human-seeded hyperparameter/metric tweaks within a fixed contract, wins are on the proxy, and no loop winner was confirmed on the n=50 F1 benchmark. Its value is the methodology (below).

---

What didn't work (and why) — negative results

For this audience these are first-class. Each is a method built correctly enough that "it didn't work" means "I built it right and it still didn't help," with a mechanism. None refutes the cited paper — each is scoped to this frozen-model / specific-benchmark regime, and several are expected failures.

Idea	What killed it (number)	Mechanism
FAVOR+ / Performer post-hoc linear attention	top-64 recall 0.039→0.046 from m=32→1024 random features (data/favor_plus.csv)	Structural, not statistical: post-hoc random-feature linearization of a model trained for exact softmax is hostile to peaked, sink-bearing attention. Consistent with Performer's own training-time premise (ICLR 2021).
Single-window spatial clustering (gpt-oss)	median head needs a ~12.5K-token window for 95% non-sink mass; only ~13-16% of head-steps fit a 1K window; cross-prompt Spearman +0.86	A 2-parameter (center,width) primitive doesn't describe gpt-oss attention; the misfit is a stable structural property.
Cross-step / cross-layer selection reuse (TidalDecode-style, gpt-oss)	cross-step top-64 Jaccard 0.35 at lag 1; best cross-layer Jaccard 0.18	gpt-oss's alternating full / sliding-window layers diversify adjacent layers by design. Architecture-specific; not a refutation of TidalDecode on dense Llama-class models.
Per-coordinate-extreme candidate selection	dominated at every (recall, bandwidth) point; asymptotes ~0.6 recall@64 (data/coord_extreme.csv)	High-scoring (q,k) pairs are distinguished by many-coordinate alignment, not single-coordinate outliers.
CSAttention reimplementation (Qwen, InfBench En.QA)	pooled n=100 F1 0.251 vs full 0.309 (−1.71σ), below best ClusterKV variant 0.312	Paper-faithful per-KV-head reimplementation (scripts/csattention_decode.py); two real spec deviations found and fixed (+3.6pp), result stayed negative. Attributed to model/benchmark differences; the paper has no public code. (CSAttention's own-setup claims are not disproven.)
HippoRAG-style attention-graph diffusion (multi-hop KV retrieval)	F1 0.295→0.090 across 4 variants, after fixing two real bugs	Two diagnostic-caught bugs (near-uniform R matrix via experiments/inspect_R.py; ~10x norm blowup via experiments/check_diffusion_magnitudes.py) fixed before concluding. Prefill-graph edges encode encoding-time relevance, not query-time relevance. Tests an analogy to HippoRAG (document retrieval), not HippoRAG itself.
Synth-KV distillation (fit compact K',V' to prefill queries)	held-out output error plateaus ~0.71 regardless of m∈{256,1024,4096} (data/synth_kv_proto_v2.json)	Overfits prefill-query distribution; fails under prefill→decode shift.
Learned proposal-MLP page scorer	beats Quest on recall/mass (0.812 vs 0.741 @k=8) but regresses end-to-end F1 (0.198 vs 0.297 plain)	The instructive proxy-vs-task gap: a recall win did not survive integration.
Richer cluster scorers (Bayesian prior, anisotropic LSE, ReSTIR-EMA)	small-n positive, large-n neutral/negative (e.g. prior +15% @n=20, −25% @n=50)	The same prefill→decode distribution-shift signature.
Dream-KV consolidation (offline KV-memory fitting probe)	cross-query generalization fails (mode kv cross_rel 0.821, gain −0.186); 465-trial overnight run, 0 errors	An offline distribution-shift stress test, not a method. Same-session held-out gains exist; cross-query does not generalize.
Pre-RoPE / inverse-RoPE clustering	failed all error gates (err_mean +0.0168, worst-layer +0.0279)	Recorded with a "do not retry unless…" note. Measured post-RoPE K rank; ShadowKV's compressibility claim is for pre-RoPE K — a pre-RoPE capture is an acknowledged unrun experiment, so this does not test ShadowKV's claim.

Unifying hypothesis (a hypothesis, not a theorem). Schemes that push Q-conditional work into prefill fail; schemes that defer it to decode survive.

• Fails: linearization, learned f(Q), priors, diffusion graphs, distilled memories — the decode-time query distribution differs from prefill, and Q has already collapsed the relevant chain into its own direction.
• Survives: methods that precompute Q-agnostic structure and defer the Q-conditional selection to decode (Quest, ClusterKV-style routing).

This restates the well-known worst-case-query argument and is not separately proven (experiments/2026-05-13-synth-kv-and-retrieval-paradigm.md).

Methodology byproduct: n=20 InfBench En.QA F1 reverses sign at n=50 on the same setting (e.g. a Bayesian prior at λ=3.0: 0.342 @n=20 → 0.112 @n=50). This was caught and used to re-architect the autoresearch gate away from the noisy F1 metric toward a ~614K-sample offline output-error proxy with a worst-layer gate (experiments/clusterkv_autoresearch/reports/2026-05-20-output-error-loop.md).

---

Relationship to prior work

This is careful replication + Blackwell engineering + falsification on frozen models, not a new method. Cluster-based KV selection is established prior art; the novel slivers are (a) the gpt-oss learned-sink-correct Triton decode kernel, (b) the sm_120 wall-clock replication with stack gotchas documented, and (c) the breadth and rigor of the negatives.

• ClusterKV (DAC 2025) — k-means on keys + cluster-granularity recallable selection. The method this repo replicates and tunes.
• Quest (ICML 2024) — query-aware per-page min/max selection. The primary baseline.
• TidalDecode (ICLR 2025) — position-persistent cross-layer selection reuse.
• ShadowKV (ICML 2025 Spotlight) — low-rank pre-RoPE K cache + landmarks.
• CSAttention (OpenReview CEpNboUJyw) — query-centric centroid-scoring lookup tables. An ICLR 2026 submission; rejected (four reviewers, all weak-reject), per the OpenReview thread I pulled (experiments/2026-05-08-blackwell-evals-cluster-tuning.md §5). The negative here is in this frozen-Qwen / InfBench harness, not a refutation of the paper's own-setup claims.
• Multipole Attention (NeurIPS 2025), Tactic (arXiv:2502.12216) — closest cluster-based competitors.
• StreamingLLM (ICLR 2024) — positional sink tokens (distinct from gpt-oss's learned per-head sink).
• H2O (NeurIPS 2023), SnapKV (NeurIPS 2024), MInference (NeurIPS 2024 Spotlight) — bracket the eviction / prompt-compression / prefill-sparsity design space.
• Performer / FAVOR+ (ICLR 2021), HippoRAG (NeurIPS 2024) — see the corresponding negatives.

---

Layout

Start with the dated reports in experiments/ — they are the narrative and the strongest signal.

Reports (read in this order):

• 2026-05-07-triton-kernel.md — the kernel: correctness, token-identical decode, the non-attention-floor analysis.
• 2026-05-07-qwen-long-context.md — where the kernel wins end-to-end (dense Qwen; L40S curve, Blackwell table above).
• 2026-05-07-flex-attention-attempt.md — the negative that motivated dropping to Triton (create_block_mask OOM, BlockMask device-asserts, score_mod correct but no speedup).
• 2026-05-06-output-error.md / 2026-05-06-decode-validation.md / 2026-05-06-cluster-quest.md — the cluster-vs-Quest chain (recall → output-error → token-match).
• 2026-05-08-blackwell-evals-cluster-tuning.md — largest report: Blackwell wall-clock, clustering speedups, InfBench/LongBench evals with variance analysis, CSAttention + HippoRAG negatives, replication-vs-novelty accounting.
• 2026-05-11-bounded-hierarchical-clusterkv.md / 2026-05-13-synth-kv-and-retrieval-paradigm.md / 2026-05-14-richer-cluster-scoring.md — bounded variant, unifying framing, richer-scorer sweep.
• 2026-05-13-three-frames-blogpost.md — speculative idea-generation essay (ReSTIR / Mip-NeRF / Hopfield → KV cache). Not validated results.

Reusable code:

• scripts/sparse_decode_triton.py — the kernel (online softmax + sink fold, fused single-launch variant, index-build kernel, PyTorch reference, smoke tests).
• scripts/cluster_decode_triton_llama.py — k-means variants (L2, spherical/cosine, mini-batch, bounded-hierarchical) + the sparse-decode forward patch.
• scripts/cluster_decode_triton.py — gpt-oss integration via GptOssAttention.forward monkey-patch.
• scripts/{quest_decode,csattention_decode,cluster_decode_flex}.py — baseline / comparison reimplementations.
• scripts/{infbench_enqa_eval,longbench_eval,infbench_passkey_eval}.py — checkpointable eval harnesses.
• experiments/clusterkv_autoresearch/ — the evidence-gated method-search harness (scorer, loop, configs, methodology report).

---

Reproducibility & environment

Pinned, uv-locked (pyproject.toml + uv.lock).

• torch 2.10.0+cu128 (PyTorch cu128 index), torchvision/torchaudio matched · transformers 5.8.0 · Triton 3.6 · CUDA 13 · NVIDIA RTX PRO 6000 Blackwell (sm_120); the wall-clock curve was also reproduced on an L40S (48GB).
• Setup: uv sync. Headline numbers each have an exact reproduce command in their report; results are saved as per-step/per-call JSON so the medians above can be re-derived from data/.
• LongBench is not redistributed here — it loads from the 🤗 datasets hub at eval time (scripts/longbench_eval.py); the calibration .pt is regenerable via scripts/calibrate_kmeans.py.

Known issues (documented, reproducible):

• Recurring Triton 3.6 / sm_120 JIT recompile spikes mid-decode (~15-20 s wall-clock per event).
• flex_attention fails to compile at 1M context (Q_LEN=1, KV_LEN=1M int32/int64 ternary error), which blocks the full-attention 1M baseline.
• The kernel sparsifies full-attention layers only (12 of 24 on gpt-oss); a Triton sliding-window variant was correct in fp32 but drifted tokens in bf16 at 64K, so it is left disabled and those layers stay on flex_attention.

Status

Single-author, exploratory, training-free research on frozen pretrained models; not peer-reviewed and the numbers are not independently reproduced (see Scope). No license is attached — the code is provided for reading and reference.