[NVIDIA][CuTe,Fwd/Bwd,sm120] Attention dropout (#2436), forward + backward

[johnnynunez]

Summary

Implements attention dropout (#2436) for the CuTe-DSL forward AND backward on SM120 (Blackwell GeForce / RTX PRO 6000 / DGX Spark), self-contained on upstream main.

This PR is complete and standalone: it also carries the two SM120 prerequisites the dropout backward needs to run at all (so it works on a clean main checkout without depending on any other open PR). The use_tma_O fix is also submitted separately as #2655; reviewers can take it from either place.

Commits

1. SM120 use_tma_O crash fix (FlashAttentionForwardSm120: use_tma_O incorrectly True on SM_121 → AttributeError on tma_atom_O=None #2649) — SM120 forward crashed in tma_partition (TMA O-store atom is null on SM120). Restrict TMA O-store to sm_90..sm_119.
2. SM120 backward enablement — two bugs prevented the SM120 backward from running:
   • dQ_single_wg was referenced in the bwd compile key but never defined in the SM120 branch (UnboundLocalError).
   • __call__ used compute_softmax_scale_log2(), which returns softmax_scale=None when score_mod is None; but the kernel arg is typed Float32 and the backward rescales dK by softmax_scale, so launch crashed (None to Float). Compute softmax_scale_log2 inline (fold LOG2_E) and keep softmax_scale a real Float32, matching the SM100 backward.
   • Also make utils.atomic_add_fp32 robust to the nvvm.atomicrmw binding signature across nvidia-cutlass-dsl CUDA-toolkit builds (cu12.x res=/op=/ptr=/a= vs cu13 positional).
3. Dropout (DropOut Cute DSL #2436), forward + backward:
   • philox.py: Philox4x32 counter-based RNG + apply_dropout(), keyed by global (batch, head, q_idx, kv_idx) so forward and the backward recompute draw the identical keep-mask regardless of fragment layout (backward S is transposed under SdP_swapAB; handled).
   • Forward: inverted dropout (zero dropped, scale kept by 1/(1-p)) on P after online_softmax. One RNG seed per call, reused in backward via ctx.
   • Backward: regenerate the same mask on recomputed P; zeroing dropped P zeroes the dV/dS contributions, matching the forward.

Testing (RTX PRO 6000 Blackwell, sm_120, torch 2.12.0+cu130, cutlass-dsl 4.5.2)

• dropout_p=0 matches PyTorch SDPA (max err <= 9e-3 bf16) with finite grads — no-op default, no regression.
• Backward now runs and matches a PyTorch fp32 reference: |dq/dk/dv - ref| <= 1.3e-2 (hdim 64/96/128, causal & non-causal).
• Dropout changes the output and keeps grads finite; forward+backward are deterministic for a fixed seed.
• Forward keep-fraction 0.702 vs 0.70 target, with exact inverted scaling 1/S/(1-p) (measured 0.02234 vs 0.02232) — statistically correct, unbiased.

API

flash_attn_func(..., dropout_p=0.0). Seed is managed internally (drawn in forward, reused in backward).

[johnnynunez]

Note: the two SM120 backward-enablement fixes carried here (dQ_single_wg, softmax_scale None→Float) are now also submitted as the standalone PR #2658. This PR can be rebased on top of #2658 once it lands so the diff here is dropout-only.

[thad0ctor]

The dropout feature itself is not in #2634. But #2657 is stacked on the same small SM120 fixes from
2655/#2658, so its non-dropout parts overlap/conflict with #2634 if that is adopted by the repo lords

[johnnynunez]

viz @drisspg @v0i0 @tridao
To me it is an important feature to finish adoption for others repositories that uses FA C++ with dropout

[reubenconducts]

Hi @johnnynunez, can you please add tests? Benchmarks would also be helpful for the PR.

[blake-snc]

Validated this on DGX Spark (GB10, sm_121a): dropout_p=0 is bit-exact against the no-dropout baseline for both forward and gradients; same-seed runs are bit-identical and different seeds diverge as expected; dropout backward runs with finite gradients; and the output mean scaling is consistent with the 1/(1-p) rescale (ratio 1.38 at p=0.5). Looks good from the GB10 side. I am closing my #2439 in favor of this, since it covers the same #2436 ask fwd+bwd on current main.