ATFlash-Kernel-FA4 — FlashAttention-4 with the per-RoPE-wavelength window applied

A ready-to-build copy of FlashAttention at PR #2656 head (4f1fde1, Pack-GQA on SM120) with the ATFlash window ports already applied. This exists so that sm_120 users can build directly instead of applying patches; the patches themselves, the window rule, the acceptance gates and all measurements live in the companion repository ATFlash-Kernel.

This is a standalone copy, not a fork. It has no upstream link, sends nothing upstream, and is not affiliated with or endorsed by the FlashAttention authors. The upstream README is preserved as UPSTREAM_README.md. One deliberate omission from the snapshot: upstream's CI definitions (.github/workflows) are not carried, so this copy neither ships nor triggers upstream's CI; everything else is byte-for-byte upstream.

History = provenance

The git history is the audit trail, one meaningful step per commit:

1. upstream snapshot (4f1fde1; workflows omitted as noted above)
2. the prefill window — 4 files, +296/−13, the diff the paper reports
3. a non-semantic trace/constexpr compile fix (+12) — the measured binary
4. the decode split-KV window — cumulative 4 files, +562/−64

git diff <first-commit> HEAD reproduces exactly what was changed and where.

What the window does, in one figure

The FlashAttention online-softmax loop with the insertion marked: the window merely advances the start of the pair reduction per distance band (center), and the out-of-window pair components of K^T and q are never loaded into SRAM at all (right). Everything else -- tiling, running max, normalizer, PV -- is untouched, which is what the bit-identical window-off path checks.

Build (sm_120)

pip install "torch==2.11.*" --index-url https://download.pytorch.org/whl/cu128
pip install nvidia-cutlass-dsl==4.2.0
pip install -e .   # CuTeDSL path: python-only install

The window is opt-in through the interface (atflash_window_zhi); without a window table the code path — and the generated trace — is stock. Verify with the gates in ATFlash-Kernel before trusting any number.

Note on bases: upstream later squash-merged PR #2656 into main as 5835c73 with small review edits in flash_attn/cute/flash_fwd.py. The prefill window (commits 2–3) applies cleanly to either tree; the decode split-KV commit requires this snapshot's base, which is what the measured binaries used.

License

BSD-3-Clause, unchanged from upstream (LICENSE, AUTHORS). The window commits are contributed under the same license.

Citation

Paper: https://arxiv.org/abs/2608.02947 — see ATFlash-Kernel for BibTeX.