Support s_aux in GPT-OSS
Author: LoserCheems
Pull request: no
Labels: name: Feature request

### Feature request

Flash Attention currently can’t consume the `s_aux`/learnable bias head that many architectures (e.g. GPT-OSS) rely on—see the open discussion in [Dao-AILab/flash-attention#1797](https://github.com/Dao-AILab/flash-attention/issues/1797). We’d like Transformers to expose an optional backend hook for [Flash-DMATTN](https://github.com/SmallDoges/flash-dmattn), which already supports the same Flash-Attention-style API plus learnable bias tensors out of the box. Integrating Flash-DMATTN as a selectable backend would unblock models that need `s_aux` today while keeping the existing Flash Attention pathways unchanged for users who don’t opt in.

### Motivation

Right now, Flash Attention is the default fast path for many HF models, but the missing `s_aux` handling prevents us from enabling learnable-bias features (attention sinks, trainable mask bias, flash-decoding helpers) in production. Downstream projects either fall back to slower SDPA or maintain private forks. Flash-DMATTN already ships kernels with full bias gradients and has parity with Flash Attention on the rest of the API surface, so wiring it in would immediately restore feature coverage without waiting on upstream CUDA updates. This request is prompted by the stalled upstream issue above and by our own deployments that need learnable bias with long-context inference.

### Your contribution

We can provide a PR that:
- Adds a new optional backend selector (similar to `flash_attn_2` / `flash_attn_3`) that dispatches to Flash-DMATTN when installed.
- Introduces lightweight capability checks so models that request learnable bias automatically choose the compatible backend.
- Supplies documentation snippets and equivalence benchmarks (matching the Flash-DMATTN suite) to ease review.

Comments:
cc @arthurzucker @cyrilvallez for attention

Hey! we are happy to do so via the `kernels` package! TBH we currently use https://huggingface.co/kernels-community/vllm-flash-attn3/blob/main/build/torch26-cxx11-cu124-x86_64-linux/vllm_flash_attn3/flash_attn_interface.py which supports `s_aux` 

As far as I can see, Flash-DMATTN also supports the backwards pass? That indeed would be really cool especially for training. And the bias seems very flexible so that it could be extended to T5, mpnet etc?

We would like to have this as kernel tho as Arthur pointed to. Maybe cc @MekkCyber for kernels

Edit: We likely need to be careful with the naming tho, not that we have 3 different kwargs for og fa, vllm fa, and Flash-DMATTN 

@LoserCheems Hi, may I ask something? Is `s_aux` equal with `attn_bias` in flash sparse attention?

@fahadh4ilyas Hi, `s_aux` is **not** the same as `attn_bias`, though they can have somewhat related effects.

- **`attn_bias`** is an additive term on the attention logits with shape like `(batch, heads, seqlen_q, seqlen_k)`. It changes scores **per (q, k) position** (e.g., ALiBi / relative bias / custom bias patterns).
- **`s_aux`** is a **per-head scalar** (shape `(heads,)`) that is treated as an extra “sink” logit appended to each query row in the softmax. Its probability is computed (it competes in the normalization) but the sink column is **dropped from the output**, so it mainly reduces the probability mass available to real keys.

Extra note (based on our experiments): while learnable attention bias isn’t the “standard” way to implement attention sinks, we found that making `attn_bias` learnable can also **partially mitigate attention-sink behavior** in practice (by shifting the logits distribution and discouraging overly-peaky / sink-dominated patterns). That said, the mechanism is different: learnable bias redistributes mass *among real keys*, whereas `s_aux` explicitly provides an extra competing sink term.