[Cute Sm100] support dropout

[endurehero]

FA4 SM100 Dropout Support

With the help of CodeAgent, this MR adds a dropout implementation
on FA4 SM100 whose per-element mask is aligned with FA2.

---

1. Summary

• New p_dropout, rng_state, return_dropout_mask arguments on the
  FA4 flash_attn_func (forward + backward, SM100 / SM110 only).
• The per-element dropout mask is bit-identical to FA2 when both
  are fed the same Philox (seed, offset) — verified at the cell level
  by the new pytest suite. Forward output and dq / dk / dv errors are
  bounded by FA2's own bf16 error against an fp32 reference.
• End-to-end check on an internal MoE model trained on GB200: the
  per-step training-loss difference between an FA4-dropout run and an
  FA2-dropout run under the same seed schedule stays in the
  1e-3 – 1e-4 band throughout the run, and downstream evaluation
  metrics match within run-to-run noise. See §3.
• Performance vs cuDNN's SDPA dropout backend on GB200: FA4 dropout
  costs ~20–40 % over no-drop FA4, but with dropout enabled it
  dominates cuDNN on backward (up to ~1.7× on hdim = 64 long
  sequences) and is competitive with cuDNN on forward. See §2.

---

2. Performance

Benchmarks captured on a single NVIDIA GB200 with cuDNN 9.22.0 (which was released
in May 2026), dtype=bfloat16, nheads = 16, total_seqlen ≈ 32 k (batch chosen
so batch * seqlen ≈ 32 k). Each cell is averaged over triton.testing.do_bench(warmup=5, rep=10)
with a 1 s GPU-idle pause between cells to suppress DVFS / thermal bleed between runs.

Forward

Backward

Highlights:

• No-drop — cuDNN's fused SDPA wins or ties FA4 on most configs at
  hdim = 64; FA4 catches up at hdim = 128 long sequences. I am not sure
  about this conclusion, please check it together.
• Dropout = 0.125, backward — FA4 leads cuDNN on every config, by
  ~1.4 – 1.7× on hdim = 64 (cuDNN's bwd dropout path appears to fall
  off a fast tile here) and by ~1.1 – 1.2× on hdim = 128.
• Dropout = 0.125, forward — FA4 and cuDNN trade wins: tied on
  hdim = 64 non-causal, FA4 leads on causal long sequences for both
  head dims.
• The FA4 with-drop / no-drop TFLOPS ratio is roughly constant across
  seq lengths (fwd ≈ 75–85 %, bwd ≈ 60–70 %), so dropout-injection
  overhead is amortised inside the tile loop, not seqlen-dependent.
• We chose dropout_rate=0.125 because cudnn requires the input
  dropout_rate to be a multiple of 1/16. However, fa4 dropout does
  not have this limitation.

Reproduction

# Re-run the full benchmark on a fresh GPU (~8 min on GB200).
PYTHONPATH=. python tests/cute/benchmark_dropout.py
# Knobs: --headdim 64,128 --seqlen 1k,4k,8k,16k,32k --causal both \
#        --p-dropout 0.125 --warmup 5 --rep 10

---

3. Correctness

Validated at three levels.

3.1 Bit-identical dropout mask vs FA2 (unit test)

tests/cute/test_flash_attn_dropout.py seeds both FA4 and FA2 with the
same Philox (seed, offset) and asserts that the per-cell dropout
decision in the valid (non-causal-masked) region matches exactly,
then bounds FA4's forward and dq / dk / dv error by FA2's own bf16
error against an fp32 reference. Coverage: causal ∈ {F, T},
p_dropout ∈ {0.1, 0.25}, d ∈ {64, 128}, three (batch, seqlen,
nheads) shapes. A second test confirms the empirical kept-fraction is
1 − p to within ±2 %.

PYTHONPATH=. pytest -q tests/cute/test_flash_attn_dropout.py

3.2 End-to-end training-loss agreement (MoE on GB200)

An internal MoE model was trained on GB200 with FA4 dropout, and the
same schedule was replayed with FA2 dropout under an identical seed
schedule. The per-step training-loss difference stays in the
1e-3 – 1e-4 band throughout the run.

3.3 Downstream evaluation parity

Downstream eval metrics on the same MoE checkpoints match between the
FA4-dropout and FA2-dropout runs to within standard run-to-run noise.

---

4. Implementation Notes

The implementation reuses FA2's Philox layout so the masks line up
cell-for-cell, but generates and consumes the mask inline with the
softmax tile loop — there is no global dropout-mask tensor on the
critical path.

4.1 Forward (flash_fwd_sm100.py + softmax.py + philox.py)

• Philox RNG. philox.py exposes a 7-round Philox-4×32 with a
  mul.wide.u32 lowering and a lop3.b32-fused 3-input XOR fast path
  (philox_rounds_with_keys). The per-round key chain is precomputed
  once per tile; the inner softmax step only redoes the counter mix.
• Threshold-based mask generation. Inside
  softmax_step.apply_dropout_rP (softmax.py), each 32-bit Philox
  word is split into 4 bytes and each byte is compared against an
  8-bit threshold derived from p_dropout
  (f32_to_dropout_threshold_byte). Keep ⇔ byte ≥ threshold. Skips
  any float conversion and yields one branch-free predicate per
  element.
• In-place scale + zero. Kept elements are scaled by
  rp = 1 / (1 − p) and dropped ones are zeroed in registers via
  apply_dropout_pair, which fuses the keep-test and * rp into a
  single selp + fma pair on the rP fragment, so dropout does not
  change the MMA tile layout.
• Optional uint8 mask export. When return_dropout_mask=True, the
  same Philox bits are packed into a (b, h, sq, sk) uint8 tensor via
  store_dropout_mask_u8 for testing / debugging; the backward never
  reads it.

4.2 Backward (flash_bwd_sm100.py)

The forward does not store the mask. The backward replays the
Philox stream (same seed / offset carried via rng_state) and
regenerates the mask on the fly per tile. Dropout is then folded into
the existing dS = P · (dP · scale − dPsum) chain.

To stay inside the SM100 register budget once dropout is on, several
dropout-specific knobs are auto-enabled in
FlashAttentionBackwardSm100.__init__:

Knob	Effect
merge_exp_dropout_ds	Fuses P = exp(S − LSE) → drop and dS = P · dP_eff into one per-stage loop; halves the live tSrS_t2r footprint (64 → 32 F32 regs) and keeps drop_bits stage-local.
packed_dropout_fma	Uses mul_packed_f32x2 / fma_packed_f32x2 for pdrop = S · scale and dP_eff = dP · scale − dPsum; drops the FP32 op count of the dropout consumer warpgroup from ~70 M to ~1 M.
direct_sdS_xchg_write	Writes the exchange-stage dS straight into SMEM instead of an intermediate ~16-reg fragment (2-CTA tiles, hdim ≠ 192).
skip_drop_bits_pack	At qhead_per_kvhead ≥ 2, re-walks the SMEM mask in the dS apply step instead of carrying packed bits across phases; avoids the persistent _cached_mask_word spill at GQA = 1.

These are all internal and not user-visible.

4.3 Public API

out, lse, rng_state, dropout_mask = flash_attn_func(
    q, k, v,
    causal=...,
    p_dropout=0.1,            # NEW
    rng_state=None,           # NEW — pass to reproduce or pin the seed
    return_dropout_mask=True, # NEW — uint8 (b, h, sq, sk), test/debug only
    return_lse=True,
)

• p_dropout is rounded to 6 decimals before becoming part of the
  kernel cache key, so nearby probabilities share the same compiled
  binary.
• rng_state is a length-2 int64 tensor [seed, offset]. When it is
  None and p_dropout > 0, FA4 samples it from the default CUDA
  generator and returns it to the caller so the backward (or an
  external replay) can reuse the same Philox stream.
• Dropout is currently asserted off for qv (MLA), split_kv,
  block_sparsity and the dedicated hdim-256 kernel.

[endurehero]

hi @tridao, this is a function implemented according to my actual production requirements on GB200. I would be very happy if you could give me any suggestions

[reubenconducts]

Hi @endurehero, I have a few organization-type comments:

• dropout_rate doesn't need to be a compile-time parameter. Instead, just use the boolean use_dropout (I prefer that nomenclature to is_dropout) to __init__ and pass dropout_rate to __call__, in both forward and backward
• please limit comments to what is necessary, such as notes about experimental results (e.g. "Sweeping showed split=64 (50%) yields ~4% speedup over split=96 (75%)" in flash_fwd_sm100.py). Note the comment concision currently in the codebase. (More generally, please try to keep diff minimal unless necessary)
• The index arithmetic in mask.py is unclear; it can be reformulated in layout algebra and avoid any magic numbers. Additionally, try to deduplicate shared code.
• Remove any dead code paths, such as philox_precompute_round_keys and mRngState, or otherwise utilize them
• I'd suggest splitting the dropout helpers that are currently in softmax.py out into a separate dropout.py module.
• Could you rebase on main and re-benchmark?

[endurehero]

@reubenconducts Thank you very much for your suggestion. I think it is very reasonable. I will modify it today