[CuTe, SM120] Optimize backward on consumer Blackwell

[guoriyue]

Summary

This adds a current-main SM120 CuTe DSL backward path for consumer Blackwell, validated on RTX 5090, without changing the public flash_attn.cute API.

• dispatch tuned SM120 backward kernels with matching dQ/dK/dV accumulator and postprocess mappings;
• select dense D64 and D128 configurations, including the D64 64x128 register-resident P/dS path and the long-D128 V-in-register path;
• repair the dP GEMM V-register operand, variable-length dK/dV addressing, and vector atomic stores;
• gate R2P and causal full-mask optimizations to layouts where they are valid;
• add real-GPU and FakeTensor coverage for long sequence, varlen, and CUDA-tensor max_seqlen backward paths.

Mainline performance delta

BF16, 32K total tokens, model width 2048. The run order was patch / unmodified main / patch, with five timing blocks per run. Values are main latency / median(two patch runs); greater than 1 means this PR is faster.

D	sequence	non-causal	causal
64	512	1.095x	1.034x
64	2048	1.136x	1.098x
64	8192	1.154x	1.132x
128	512	1.091x	1.048x
128	2048	1.112x	1.080x
128	8192	1.103x	1.100x

All twelve cells improve. The geometric-mean speedup over current main is 1.098x, with a range of 1.034x to 1.154x.

FA2 parity on RTX 5090

The same 32K-token BF16 matrix compares native FA2 with the patched FA4 path. Values are FA2 latency / FA4 latency.

D	sequence	non-causal	causal
64	512	1.033x	1.058x
64	2048	1.016x	1.028x
64	8192	1.009x	1.024x
128	512	1.019x	1.018x
128	2048	1.013x	1.003x
128	8192	0.995x	0.996x

FA4 wins ten of twelve cells and is within 0.5% of FA2 in the remaining two. FP16 D128 S8192 measures 0.994x non-causal and 0.993x causal.

Correctness

Environment: RTX 5090 (SM120), PyTorch 2.11.0+cu128, CuTe DSL 4.6.0.dev0.

• tests/cute/test_flash_attn_fast.py: 196 passed, 48 expected skips on real GPU;
• FakeTensor compilation: 195 passed, 49 expected skips;
• BF16/FP16 dense D64/D128 output and dQ/dK/dV pass the FP32-reference error criterion for causal and non-causal attention;
• random-padding varlen MHA/GQA/MQA gradients, CUDA-tensor max_seqlen cache reuse, and zero-Q dK/dV behavior pass;
• dedicated S8192 tests assert the actual JIT compile-key selections;
• git diff --check, py_compile, and Ruff pass.

Scope and overlap

The import surface remains:

from flash_attn.cute import flash_attn_func

The patch is limited to SM120 backward dispatch/kernel plumbing and tests. It does not add forward TMA, SplitKV, paged KV, D256, FP8, or FA3 support.

#2634 is a broader SM120 effort and contains overlapping backward work. This PR is intentionally a focused current-main extraction with isolated correctness and performance evidence; if #2634 lands first, this draft should be re-evaluated or closed rather than merged redundantly.

Implementation note

The long-D128 AtomLayoutMdQ=2 writer requires a 256-thread dQ postprocess; the default 128-thread bridge produces large dQ errors. The default AtomLayoutMdQ=4 path intentionally retains its correct, faster 128-thread postprocess. SM120 uses the SM80 universal register-to-shared mapping rather than the SM90 stmatrix mapping.