Expand FLASHATTENTION_DISABLE_DROPOUT to not bring in unneeded headers

[janeyx99]

TL;DR: If we expand the disable dropout macro and drop generator support we can enable stable ABI builds for flash, vLLM, and sglang (and other downstream libs). I did quite a bit of recursive GitHub searching which leads me to believe that majority of C++ users are just passing in an null opt for the gen arg anyway. cc @tridao

So we do the two things:

1. expand disable dropout macro to not pull in at::Generator and Philox related headers, which prevent eventual ABI stable support. This does affect the dropout enabled case because we change from using PhiloxState directly to an opaque buffer, but there's no change in functionality.
2. drop custom generator support in C++ (it already wasn't there in python) -- FA2 will now only ever use the default generator. We DO keep around empty argument slots to minimize the churn of this BC break.

Breaking: Custom RNG generator support removed from the C++ APIs

The raw FA2 forward and backward entry points no longer accept a custom torch.Generator; the argument must be None.

Passing a torch.Generator object into the raw extension entry points (fwd / varlen_fwd / bwd / varlen_bwd) now raises a pybind TypeError: ... incompatible function arguments as the argument slot is now Tensor? and only accepts None.

Passing a Tensor will hit: RuntimeError: flash-attn: the RNG generator argument is no longer supported and must be None; dropout (when enabled) uses the default CUDA generator.

Workaround: If your use case was to be able to get deterministic behavior, use torch.manual_seed to control the default generator in Python instead:

import torch
torch.manual_seed(1234)              # seeds the default CPU + all CUDA generators
out = flash_attn_func(q, k, v, dropout_p=0.1, causal=True)

Test Plan

Everything is run on GPU: H100 (sm_90) with torch 2.10.0+cu129. I ended up building with 12.8 cuz 12.9 causes nans.

I added a new test case to enforce the BC break (note below).

Part 1 — confirming nothing regressed with default dropout enabled build

build command: CUDA_HOME=/usr/local/cuda-12.8 FLASH_ATTN_CUDA_ARCHS=90 pip install -e . --no-build-isolation -v

There are dropout kernels in the binary:

(fa2-pt210) ➜  flash-attention git:(fully-yank-dropout) cuobjdump --dump-elf-symbols flash_attn_2_cuda*.so | c++filt | grep 'flash_fwd_kernel<' | grep -cE '>, true,'
grep: warning: GREP_COLOR='1;32' is deprecated; use GREP_COLORS='mt=1;32'
238

Dense forward+backward across head dims 64 & 128, both dtypes, dropout {0.0, 0.17}, all masking modes (causal/local/alibi/deterministic) — confirms the new opaque-philox plumbing didn't break the default path

$ python -m pytest tests/test_flash_attn.py::test_flash_attn_output -k '(113-203-64 or 113-203-128) and mha and not 50.0'
........................................................................ [ 28%]
........................................................................ [ 56%]
........................................................................ [ 84%]
........................................                                 [100%]
256 passed, 84224 deselected in 21.52s

Variable-length forward+backward, same configs

$ python -m pytest tests/test_flash_attn.py::test_flash_attn_varlen_output -k '(113-203-64 or 113-203-128) and mha and not 50.0'
........................................................................ [ 28%]
........................................................................ [ 56%]
........................................................................ [ 84%]
........................................                                 [100%]
256 passed, 92672 deselected in 22.70s

GQA head-grouping path with dropout_p=0.17 at d=128.

$ python -m pytest tests/test_flash_attn.py::test_flash_attn_output -k '113-203-128 and gqa and 0.17 and not 50.0'
................................................................         [100%]
64 passed, 84416 deselected in 20.65s

KV-cache + splitkv entry point (num_splits 1 & 0), no dropout — confirm that the schema change is okay

$ python -m pytest tests/test_flash_attn.py::test_flash_attn_kvcache -k '(64-256-64 or 1-128-64) and gqa and 0.0 and None and not True'
....                                                                     [100%]
4 passed, 304124 deselected in 26.54s

Part 2 — dropout=0 tests pass dropout disabled build

build command: CUDA_HOME=/usr/local/cuda-12.8 FLASH_ATTN_CUDA_ARCHS=90 FLASH_ATTENTION_DISABLE_DROPOUT=TRUE pip install -e . --no-build-isolation -v

The disabled build should drop the dropout (Is_dropout=true) kernels entirely — confirm none remain in the binary (.so is also 271 MB vs 341 MB for the enabled build).

(fa2-pt210) ➜  flash-attention git:(fully-yank-dropout) cuobjdump --dump-elf-symbols flash_attn_2_cuda*.so | c++filt | grep 'flash_fwd_kernel<' | grep -cE '>, true,'
grep: warning: GREP_COLOR='1;32' is deprecated; use GREP_COLORS='mt=1;32'
0

There are no references to Philox and Generator in the binary:

nm -C flash_attn_2_cuda.cpython-312-x86_64-linux-gnu.so | grep -iE "philox|generator"
grep: warning: GREP_COLOR='1;32' is deprecated; use GREP_COLORS='mt=1;32'
(fa2-pt210) ➜  flash-attention git:(fully-yank-dropout) 

Dense forward+backward at dropout=0 (head dims 64 & 128, both dtypes, all masking modes) — confirms attention is still correct after stripping the ATen RNG headers.

$ python -m pytest tests/test_flash_attn.py::test_flash_attn_output -k '(113-203-64 or 113-203-128) and mha and not 50.0 and not 0.17'
........................................................................ [ 56%]
........................................................                 [100%]
128 passed, 84352 deselected in 21.14s

Variable-length forward+backward at dropout=0, same configs.

$ python -m pytest tests/test_flash_attn.py::test_flash_attn_varlen_output -k '(113-203-64 or 113-203-128) and mha and not 50.0 and not 0.17'
........................................................................ [ 56%]
........................................................                 [100%]
128 passed, 92800 deselected in 21.60s

dropout_p > 0 must still be rejected (not silently ignored) in the disabled build.

(fa2-pt210) ➜  flash-attention git:(fully-yank-dropout) python -c "import torch, flash_attn; q=k=v=torch.randn(2,128,4,64,device='cuda',dtype=torch.float16); flash_attn.flash_attn_func(q,k,v,dropout_p=0.17)"
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/janeyx/repos/flash-attention/flash_attn/flash_attn_interface.py", line 1213, in flash_attn_func
    return FlashAttnFunc.apply(
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/autograd/function.py", line 583, in apply
    return super().apply(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/repos/flash-attention/flash_attn/flash_attn_interface.py", line 851, in forward
    out_padded, softmax_lse, S_dmask, rng_state = _wrapped_flash_attn_forward(
                                                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_ops.py", line 1209, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_library/autograd.py", line 112, in autograd_impl
    result = forward_no_grad(*args, Metadata(keyset, keyword_only_args))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_library/autograd.py", line 41, in forward_no_grad
    result = op.redispatch(keyset & _C._after_autograd_keyset, *args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_ops.py", line 826, in redispatch
    return self._handle.redispatch_boxed(keyset, *args, **kwargs)  # type: ignore[return-value]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 347, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1181, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/.conda/envs/fa2-pt210/lib/python3.12/site-packages/torch/_library/custom_ops.py", line 382, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/janeyx/repos/flash-attention/flash_attn/flash_attn_interface.py", line 99, in _flash_attn_forward
    out, softmax_lse, S_dmask, rng_state = flash_attn_gpu.fwd(
                                           ^^^^^^^^^^^^^^^^^^^
RuntimeError: This flash attention build does not support dropout.

[janeyx99]

To explain why this buffer, the PhiloxCudaState struct is 24 bytes, which we round up here for some cushion:

struct PhiloxCudaState {
    union Payload { uint64_t val; int64_t* ptr; };  // 8 bytes
    Payload seed_;                                   // 8
    Payload offset_;                                 // 8
    uint32_t offset_intragraph_ = 0;                 // 4
    bool captured_ = false;                          // 1 (+3 tail padding to keep struct alignment 8)
  };

[drisspg]

nit; maybe add an abridged verison in the doc block

[janeyx99]

Note: this is a placement-new, which is a variation of new where we give new an existing memory address to put the constructed value in. This is why there's no matching delete!

[Harry-Chen]

It will be really good to see it (and #2688 later) merged upstream, since vllm would need it in our stable ABI transition.