Don't read the haystack

FlashAttention-4 came out this year and it's a Blackwell kernel: tile-based execution, TMA, hardware I do not own and will not own until someone at NVIDIA loses a bet. What I own is a 4090. But the interesting action in inference right now isn't on the compute side anyway: DeepSeek's sparse attention picks the top 2048 tokens per query with a "lightning indexer" and skips the rest, their V4 stack compresses the KV cache ~10x on top of that, which points at the same thing: generating tokens is a reading problem, not a compute one. So this project is the consumer-hardware version of that idea, built from scratch: two Triton kernels, a page selector, and a small model to check I didn't lobotomize anything.

Decode attention latency vs context length: FlashAttention-2 grows linearly, pageskip stays flat, 16.4x at 64k. — Decode attention at batch 32. FlashAttention-2 reads the whole KV cache, so it scales with context. Mine reads 2048 tokens no matter what, so it doesn't. At 64k context that's 16.4x, and the gap keeps widening forever, because one of these lines has a slope.

Before that number does something dishonest in your head: this is not a drop-in replacement for exact attention. It's a Quest-style approximate decode path, so 16.4x only means anything if the pages it keeps happen to hold the attention mass the softmax wanted, which is a property of your workload and not a theorem I get to prove once and retire on. Read the speed numbers as kernel results, and the quality numbers further down as a sanity check on one small model, not a claim about model quality in general. Two separate threads, and the fast one does not get to vouch for the slow one.

The other thing to say now rather than let you feel clever later: the selection algorithm isn't mine, it's Quest's (2024). What's new here is the from-scratch Triton, the GEMM decomposition that makes page-scoring nearly free, and the consumer-GPU numbers with their traps flagged. Full accounting of what's borrowed at the bottom.

Drag across context length. This is end-to-end speedup, page selection included, under CUDA graphs.

The figure above, with the fine print. Speedup over FlashAttention-2 isn't free: page selection costs a fixed ~35µs, so sparse attention only wins once the dense read it replaces is bigger than that. At batch 1 (grey) it never does, you're slower everywhere, forever. The win needs volume: batch 8 crosses break-even around 16k, batch 32 (blue) crosses by 4k and reaches 16.4× at 64k. Start at 64k and walk left to watch each curve fall through the 1× line.

The same batch story as a table, including the one row a benchmark can't draw for you:

Scenario	vs FlashAttention-2
batch 1, short context	loses
batch 1, long context	marginal
batch 8–32, long context	wins, up to 16.4×
a real serving engine	unproven

Why generating a token is a reading problem

Prefill, chewing through the prompt, is a pile of big matrix multiplies, which is what GPUs are for; FlashAttention exists to keep the tensor cores fed. Decode is a different animal. You generate one token at a time, and for that single token, attention reads every key and value the model has ever cached, does one multiply-accumulate per value read, and moves on. That's an arithmetic intensity of roughly one: the tensor cores are asleep, and your latency is (bytes of KV cache) ÷ (memory bandwidth), full stop. At batch 32 with 64k contexts, the little 1.5B model I use drags 4.3 GB of KV past the memory controller per generated token. The model weights are 3.5 GB. The cache outweighs the brain.

So the plan is two steps: first prove you can read at the speed of the memory (a dense kernel), then prove reading less doesn't break the model (a sparse one).

Part one: hitting the wall on purpose

The dense kernel is the flash-decoding two-stage shape. One query token means there's no parallelism across sequence positions to exploit, so you make some: stage one splits the KV cache across a grid of thread blocks, each runs online softmax over its slice for all the query heads in a GQA group and writes out partial results (running max, normalizer, unnormalized accumulator, in fp32); stage two merges the partials with the standard log-sum-exp rescale. It's about 120 lines of Triton and it matches an fp32 reference to fp16 tolerance across every shape I throw at it, including contexts that aren't multiples of anything.

Dense kernel bandwidth vs context: spikes above 2 TB/s where KV fits in L2, then settles at 915-945 GB/s against a 1008 GB/s peak. — The dense kernel against the 4090's memory. Right of the cliff it sustains 915-945 GB/s, which is 94% of the card's 1008 GB/s peak and within 3% of FlashAttention-2. Left of the cliff: 2.4 TB/s, which would be remarkable if it were real.

About that 2.4 TB/s. A benchmark loop re-reads the same KV cache hundreds of times, and the 4090 has 72 MB of L2. When the cache fits (batch 32 at 2k context is 67 MB), you are benchmarking the L2, not the DRAM, and your kernel reports bandwidth the memory bus physically cannot deliver. The literature calls this a methodology error; I call it the five minutes I thought I was better than the laws of physics. Every microbenchmark you've ever seen with a suspicious bandwidth number is this. The DRAM-bound regime starts when the working set is a few multiples of L2, and there the kernel sits at 94% of peak, right next to FlashAttention-2, because at the memory wall everyone's kernel looks the same. That's what "memory-bound" means: past 90-ish percent, there is nothing left to optimize. The only remaining move is to not read.

One more thing the wall taught me: below about half a gigabyte of KV, none of this matters, because the whole exercise drowns in kernel-launch overhead. In eager PyTorch my select-then-attend step costs ~0.26 ms flat, regardless of context length. That flat cost is the CPU, not the GPU: it queues half a dozen tiny kernels at ~30µs each while the GPU finishes early and waits. Capture the same step into a CUDA graph and it collapses to 0.037-0.134 ms. This is why vLLM captures decode into graphs, and why single-sequence-short-context benchmarks of anything are mostly measuring Python.

Part two: the skipping

The KV cache gets divided into pages of 64 tokens. For each page keep two 128-dim vectors: the elementwise min and max of its keys. Then for a query q, sum_d max(q_d·kmin_d, q_d·kmax_d), picking whichever extreme of the page's range the sign of q makes larger, is a provable upper bound on any attention score inside the page (this is Quest's bound, and one of my tests checks admissibility on every run). Rank pages by the bound, keep the top k plus the attention-sink page and the last few local pages, and hand that page list to a sparse twin of the dense kernel that walks indices instead of a range.

The bound has a cute identity that makes it cheap: max(q·kmin, q·kmax) = q·mid + |q|·half where mid and half are the center and radius of the page's range. So scoring every page in the cache is exactly two batched GEMMs and a top-k, no giant broadcasted intermediate, ~20 microseconds. The summaries are 1/32nd the size of the cache, and in a real serving loop you'd maintain them incrementally as tokens append.

With a budget of 32 pages (2048 tokens, or 3.1% of a 64k cache), the sparse kernel's latency is flat at ~0.05 ms from 1k to 64k context. It doesn't scale with context because it doesn't read the context. Selection plus kernel, all captured in one CUDA graph, beats FlashAttention-2 by 16.4x at batch 32 × 64k, 8.1x at batch 8 × 64k, and loses at batch 1 with short contexts, where attention was never your problem to begin with and you should close the tab.

Part three: does the model still work, though

Fast and wrong is easy. The quality question is whether the pages the bound picks are the pages the softmax actually cares about, so I measured it on real KV caches, captured from DeepSeek-R1-Distill-Qwen-1.5B mid-decode by monkeypatching scaled_dot_product_attention, with 16k tokens of context. The context, for continuity with the last post, is that model's own chain-of-thought transcripts. It is now attending sparsely to its own rationalizations.

Attention mass captured vs selection budget: bound top-k reaches 0.97 at 8192 tokens, sliding window plateaus at 0.81, random trails. — How much of the true attention mass the selected pages contain, averaged over layers 2+. At tight budgets a sliding window is competitive, since reading prose is mostly a local activity. From 2048 tokens up, the bound pulls away: 0.97 vs 0.81 at 8192. The window has a ceiling it can never pass, because some of what the model wants is just not recent.

Two caveats. First, layers 0-1 attend by position, not content, so the bound can't rank their pages (recall 0.004 at layer 0, which is the probe equivalent of a shrug), and those layers stay dense, same as Quest does, for the same reason nobody talks about. Second, on plain prose a recency window is genuinely fine at small budgets. If your model only ever continues text, you don't need any of this machinery.

You need it the moment the answer isn't recent. I planted "the vault access code is 7491" a quarter of the way into 8k tokens of transcript soup, let the model generate the continuation of "the vault access code is" 6k tokens later, and patched each policy into the live model at an equal ~12% budget:

Full attention: "7491." Fine.
Page selection: "7491." Through sparse attention reading 12% of the cache.
Sliding window: "123456." It cannot see the needle, so it makes a number up, with total confidence, in the correct format. There's no error anywhere, just a plausible lie.

Drag the code backward through the context and flip the policy. A sliding window only sees the most recent slice, so once the code falls outside it the model stops recovering it and makes one up in the right format. Page selection keeps finding the page the code lives in, wherever you drag it. Full attention always finds it, at the cost this whole post is about. Shrink the budget to watch the window get hungrier.

The detail I didn't expect: the selection only picked the needle's page in 15.6% of (layer, step) attention calls, and that was enough, because retrieval runs through a minority of heads and layers, and the bound finds the page exactly where it's needed. A sliding window misses it in 100% of them, hence the confident nonsense.

The stuff that broke

Three bugs that cost me real time. The chunked prefill kept dying with Floating point exception (core dumped), no traceback, no mercy, exit code 136, which turned out to be two stacked problems: transformers materializes logits for every position of every chunk against a 152k vocabulary (gigabytes, in fp32, for positions I never look at; logits_to_keep=1 deletes it), and my patched attention was forwarding transformers' attention mask, sized for the full 16k cache, into a call reading a gathered 1k subset, which some kernel expressed as an integer division by zero rather than anything a human could read. And the first version of the quality harness sliced K by query-head count when the new transformers passes it un-expanded with enable_gqa, which at least failed with an actual error message, like a gentleman.

What this is and isn't

It isn't a serving engine: no paged allocator, no scheduler, no incremental summary maintenance, and the quality numbers are one small model. The selection algorithm is Quest's (2024); the two-stage decode structure is Dao et al.'s flash-decoding; the trained version of the whole idea swaps the training-free bound for a learned indexer, which is DeepSeek's DSA lineage. What's mine is the from-scratch Triton implementation, the GEMM decomposition of the bound, the consumer-GPU numbers with their methodology traps flagged, and the patched-into-a-live-model needle harness. Code, kernels, benchmarks, and the needle harness are at github.com/haregali/pageskip.

FlashAttention made attention stop wasting memory bandwidth; sparse attention makes it stop using memory bandwidth. I got to the wall, confirmed it was a wall, and went around it the way the prior work already had. The fastest read is the one you skip.

Benchmark methodology & versions

Because "16.4x" with no rig behind it is trust-me-bro with extra steps.

GPU	RTX 4090, default clocks (not pinned)
Driver / CUDA	570.211.01 / 12.8
PyTorch / Triton	2.11.0+cu128 / 3.6.0
Baseline	torch `scaled_dot_product_attention`, FlashAttention-2 backend
Timing	CUDA events, median of 100–200 iters after 5–10 warmups
CUDA graphs	whole select+sparse step captured as one graph, selection included
Latency KV	synthetic (randn), fp16
Quality / needle KV	real, captured from DeepSeek-R1-Distill-Qwen-1.5B mid-decode

All of it reruns from the repo: python -m pageskip.bench (latency), python -m pageskip.quality (real-KV recall + needle), pytest -q (12 correctness tests vs fp32 references).

References

T. Dao, D. Y. Fu, S. Ermon, A. Rudra, C. Ré. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS, 2022. arXiv:2205.14135
T. Dao, D. Haziza, F. Massa, G. Sizov. Flash-Decoding for Long-Context Inference. 2023. pytorch.org/blog/flash-decoding
J. Tang, Y. Zhao, K. Zhu, G. Xiao, B. Kasikci, S. Han. Quest: Query-Aware Sparsity for Efficient Long-Context LLM Inference. ICML, 2024. arXiv:2406.10774. The page-selection bound implemented here.
J. Yuan et al. Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention. 2025. arXiv:2502.11089
DeepSeek-AI. DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models. 2025. arXiv:2512.02556. DSA and the lightning indexer, the trained version of this idea, extended in DeepSeek-V4 (2026).
Z. Ye et al. FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. 2025. arXiv:2501.01005
NVIDIA. CUDA Graphs, why serving engines capture decode steps. developer.nvidia.com/blog/cuda-graphs