ATFlash — Per-RoPE-Wavelength Attention Windows

Training-free, input-independent attention pruning for RoPE models, shipped as minimal diffs into released kernels: FlashAttention-4 (prefill) and FlashInfer (decode). With the window off, both ports are bit-identical to stock. With the window on, on RTX PRO 6000 (Blackwell sm_120, GDDR7) both ports outpace their stock implementations at every evaluated Llama context length, and end to end the whole-request speedup reaches 1.31× at a 1M-token context (Qwen2.5-7B-Instruct-1M).

Paper: arXiv:2608.02947 — ATFlash: Per-RoPE-Wavelength Attention Windows for Compute/Memory-Efficient LLM Inference. Section references (§) below follow the paper.

The idea

Zoom into one cell of the attention score matrix: it is a d-dimensional dot product, and under RoPE it is computed two dimensions at a time — rotate the pair, then add.

Rotary position embeddings rotate each 2D frequency pair of q and k by nθ_r and mθ_r, so the pair's contribution to the score depends only on the relative angle (n−m)θ_r — position enters the score exactly, pair by pair.

Over the same distance, how far each pair has rotated differs by orders of magnitude — 3.5 full turns at the highest frequency shown, a fraction of a turn at the lowest. Once a pair has completed k turns it can no longer discriminate position at that distance, so ATFlash sets exactly those terms to zero (gray: never computed, never loaded) and keeps the sum over the rest.

On the memory side this is a prefix rule: when a q row and a K^T column are read for one cell, the gray out-of-window prefix is simply never read — the rotation RoPE(z) is fused at load, and only the surviving dimensions enter SRAM.

Seen over the whole causal triangle (panel (b) above; color depth = the number of retained inner-product terms, and the lifted strip is the last query row split by frequency pair), this is a distance window per pair: ATFlash prunes the query–key inner-product terms of pair r beyond a wavelength-proportional distance w_r = k·λ_r:

• no token is dropped — every key stays reachable through the low-frequency pairs, unlike a sliding window;
• the reduction is input-independent, with a closed form logarithmic in the sequence length N — you know the cost before running (37–48% of the terms within native context, deepening with N);
• it is a slice of the query–key contraction axis — the online-softmax recurrences (running max, normalizer, tiling) are untouched, which is why it ports into released kernels as a small diff, and why the out-of-window pair components are never even loaded into SRAM.

Vertical axis: share of QK^T inner-product terms pruned (higher = more compute and K-traffic saved); dotted segments are extrapolated beyond each model's verified range.

Measured speed (RTX PRO 6000, 96 GB GDDR7)

• Decode (FlashInfer 0.6.13 port, ~100 changed lines): latency reductions of 6.0–22.3% vs stock FlashInfer on Llama-3.2-3B at 8K–128K, growing with context; 1.296× over stock at a 160K shape.
• Prefill (FlashAttention-4 port, 4 files, +296/−13 lines): 12.5–22.7% faster than stock FA4 at 8K–128K, and below cuDNN at every context length.
• Qwen2.5-7B-1M prefill: the windowed kernel is the fastest of all four systems (stock FA4 / cuDNN / FlashInfer / ours) at every N from 8K to 512K (1.234× vs stock at 262K, 1.427× at 512K).
• End to end (FA4-port prefill feeding FlashInfer-port decode, one pipeline, window ON/OFF in the same binary): 1.17× / 1.22× / 1.31× whole-request time at 256K / 512K / 1M tokens.
• Bit-identity: window OFF generates the same numerics as stock on both ports (gates/), so the ratio isolates the window itself.

Recovery of the theoretical 2/(1+s) ceiling lines up across the two implementation families (85–96%), so the conversion of predicted reduction into wall-clock time is not implementation-specific.

Output quality is evaluated in the paper (§6): top-1 match 96–98% and KL at the 1e-3-nat level in the native regime, benchmark scores broadly preserved, with one identified exception (retrieval over structureless strings) that the single parameter k recovers — k is a deployment-time dial between attention accuracy and compute.

Repository layout

dir	contents
atflash/	The method itself: the window rule (window_table.py) and a readable torch reference (reference.py). Start here.
kernels/flashinfer_port/	The FlashInfer 0.6.13 decode window as unified diffs, plus the JIT/shadow-copy helper.
window_tables/	The torch-free table generator, and the exact per-model tables used in every run.
gates/	Bit-identity and neutrality gates (window OFF == stock; chunked == single-shot).
microbench/	The prefill microbenchmark and the raw measurements behind the latency figure above.
e2e_pipeline/	The end-to-end 1M-context harness (FA4 prefill + FlashInfer decode inside HF transformers).
results_raw/	Raw run summaries behind the speed tables.

Each directory carries its own README.md with file-by-file roles. SETUP.md walks through building both ports and reproducing the numbers; VERIFICATION.md records what was actually measured, on what hardware, and — just as importantly — what was not.

Everything derived from FlashAttention lives in the companion repository ATFlash-Kernel-FA4: the upstream tree at the pinned base with the window applied commit by commit, plus the three window patches and their reconciliation ledger under patches/. The split is by license — that repository is BSD-3-Clause throughout (upstream's license), this one is purely Apache-2.0. It is a standalone copy, not a fork, and sends nothing upstream.

Using the ports

FA4 prefill. Public FA2–FA4 wheels do not support sm_120; build from the companion repository ATFlash-Kernel-FA4, which is the upstream tree at the pinned base with the window already applied (its patches/ directory holds the raw diffs and the ledger mapping each to its measured number). Requires nvidia-cutlass-dsl==4.2.0 and PyTorch ≥ 2.11. The window is opt-in through the interface; without a window table the code path is stock.

FlashInfer decode. Start from the released FlashInfer 0.6.13 wheel, apply kernels/flashinfer_port/*.diff as a shadow copy, and point PYTHONPATH + FLASHINFER_WORKSPACE_BASE at it (fi_window_helpers.py shows the exact mechanism). The window condition is added to the K-load predicate; window OFF follows the stock code path.

Window tables. window_tables/ generates the per-pair widths from the model config (RoPE base, head dim, rope_scaling) — nothing is learned and nothing is tuned per input. The tables used in the measurements are included verbatim.

Verify first. Run gates/ against your build: window OFF must be bit-identical to stock before any speed number is meaningful. That test is the one we would ask of any new port, including our own.

Porting to something else. The two ports above are worked examples in two unrelated kernel idioms — a CuTeDSL kernel and a CUDA-template kernel — which is the evidence that the insertion is not tied to one codebase. For a third, you need atflash/window_table.py to produce the per-pair widths, the prefix length D_eff per key tile, and atflash/reference.py to check your output against. One figure carries the whole picture — the nesting that makes the port small: attention's two loops over query n and key m, and inside them the pair loop over r, which is where ATFlash starts the reduction at the window boundary w. Everything to the right of the arrow is the port; the memory side shows the gray out-of-window prefix never entering SRAM, with RoPE(z) fused at load.

(The full FlashAttention online-softmax loop with the same insertion marked — including the running max, normalizer and PV accumulation — is in the companion repository ATFlash-Kernel-FA4, next to the code it explains.)

Scope and portability

Speed numbers above are from a single hardware generation (sm_120, GDDR7) — a representative single-GPU inference platform outside the specialization targets of FA3/FA4 — and the paper reports them as such (§8, Limitations). The window itself is architecture-neutral (a static shortening of the QK reduction); ports to other kernels and other hardware generations are welcome.

Absolute paths in the harnesses (/home/USER/...) are placeholders standing in for the machine the measurements ran on, and agentA / agentB are the tags of the two measurement processes, kept because they are part of the run directory names the raw summaries refer to. Point the paths at your own environment to re-run.

License

Apache-2.0 throughout (LICENSE, NOTICE). The only third-party-derived files here are kernels/flashinfer_port/*.diff, from FlashInfer — itself Apache-2.0; its license and NOTICE are carried under third_party/. All FlashAttention-derived material (BSD-3-Clause) lives in ATFlash-Kernel-FA4. Details in THIRD_PARTY_NOTICES.md. Neither upstream project is affiliated with or endorses these modifications.

Citation

@misc{hayashi2026atflash,
  title         = {ATFlash: Per-RoPE-Wavelength Attention Windows for
                   Compute/Memory-Efficient LLM Inference},
  author        = {Hayashi, Shun-ichiro and Mukunoki, Daichi and
                   Hoshino, Tetsuya and Katagiri, Takahiro},
  year          = {2026},
  eprint        = {2608.02947},
  archivePrefix = {arXiv},
  primaryClass  = {cs.LG}
}