This section details on the original issue you should resolve

<issue_title>[BUG REPORT] INF occurs in backward phrase of the first training step</issue_title>
<issue_description>Describe the bug
Using code version v1.1.9, the training configuration is set to seq_len=4096 and window=2048. During the first training step, the backward pass consistently fails with an INF (infinity) error. Preliminary investigation has pinpointed the location in the code where the error occurs https://github.com/SmallDoges/flash-dmattn/blob/main/flash_dmattn/integrations/modeling_flash_dynamic_mask_attention_utils.py#L91. The root cause appears to be that some numerical values are outside the representable range of the bf16 data type.

To Reproduce
Steps to reproduce the behavior:

1. Import flash_dmattn
2. Run the following code:

# Paste your code here

3. See error

Expected behavior
No INF values.

Environment Information
PyTorch: 2.6.0+cu124
CUDA: 12.4
GPU: NVIDIA A800-SXM4-80GB

Additional context

• OS: Ubuntu 20.04
• Python version: 3.12.9
• Flash-DMA version: 1.1.9
• Compute Capability: 8.0

Error traceback

RuntimeError: Rank 0, node job-7db94950-80fb-47db-bacf-a9a63edec186-master-0, device 0, iteration 1: Unexpected result nan (message='found NaN in local grad norm for bucket SmallDoges/flash-dmattn#0 in backward pass 

Debugging Information

Comments on the Issue (you are @copilot in this section)

Following is an analysis done using GPT-5-CodeX regarding the safety of writing bias gradients in the kernel. Next, I will thoroughly check whether the current codebase guarantees safety under all circumstances. Once confirmed, we should analyze whether the caller passes unsafe parameters to the kernel, which might cause NAN issues.

1. Definitions and Goal

We consider one backward iteration over a (query block, key block).
Let (before masking)

raw_ij = (q_i · k_j) * scale + bias_ij

Mask and causal logic zero out (logically remove) invalid positions by setting those score entries to -INF (or a very large negative sentinel) before the softmax. In the backward kernel we reconstruct local softmax probabilities P for the block and compute dS = ∂L/∂raw_ij (which also equals ∂L/∂bias_ij when bias exists). We need to show: at the point labeled “Write dS to dBias” there is no -inf in dS, so writing it directly to global memory is safe.

2. Where -inf Appears (and Disappears)

1. Mask application (apply_mask):
   Masked or causal‑disallowed positions in scores are set to -INF (or a large negative). This is still in fp32 accumulator space (tensor scores built from acc_s).
2. Softmax reconstruction (scale_apply_exp2</*scale_max=*/false>):
   Immediately after masking, the code exponentiates those scores into probabilities.
   exp(-inf) = 0 ⇒ all masked / OOB positions become exactly 0 probability.
   From now on scores holds probabilities P (not raw logits), all finite and non‑negative; masked entries are 0, never -inf.
3. No subsequent step reintroduces -inf: All later operations are linear / multiplicative combinations of finite fp32 values (probabilities P, partial reductions, dP_sum, etc.).

3. Construction of dS in the Kernel

Relevant excerpt (simplified):

Tensor dS = make_tensor(acc_dp.data(), scores.layout());
auto pointwise_mult = [](float p, float dp, float d_row) {
    return p * (p >= 0 ? dp - d_row : d_row);
};
for (mi)
  for (ni) {
     float scaled_ds = pointwise_mult(scores(mi,ni), dS(mi,ni), dP_sum(mi));
     if constexpr (Is_softcap) scaled_ds *= dtanh(mi,ni);
     dS(mi,ni) = scaled_ds;
  }

Interpretation:

• Before the loop, dS(mi,ni) temporarily names dp (an element of ∂L/∂P).

• After the loop, each dS(mi,ni) becomes:

  dS(mi,ni) = P(mi,ni) * (dp(mi,ni) - dP_sum(mi))    (ignoring softcap factor)
  

• For masked entries: P(mi,ni) = 0, so dS(mi,ni) = 0 regardless of the other finite terms.

• All operands (P, dp, dP_sum, optional dtanh) are finite fp32 numbers under normal (non-overflow) regimes; no path produces -inf.

Even in edge cases:

• If dp or dP_sum overflowed to +inf, the multiplication by P=0 for masked entries still yields 0 (IEEE 754 defines 0 * inf = NaN in some cases only when exact 0 and inf; but here P is a finite 0.0f value produced by exp(-large) and the code path historically does not produce a signaling NaN; empirically these kernels rely on having avoided huge mismatched infinities by earlier masking). For active entries, overflow would create +inf or NaN, not -inf. Thus -inf cannot be produced by this formula.
• Negative infinity cannot arise from finite subtraction (dp - dP_sum) unless one term is -inf already; those are not...