fix(attention): stop handing FlashAttention 4 the -1 window sentinel

[nvegesna-netizen]

Problem

FlashAttention 4 returns an all-zero output for causal attention on SM100, silently — no
exception, no warning, finite values, every element zero. Fixes #3528.

No context parallelism is needed to reproduce it:

import torch, transformer_engine.pytorch as te
b, s, h, d = 2, 1024, 8, 128
q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=torch.bfloat16) for _ in range(3))
dpa = te.DotProductAttention(h, d, qkv_format="bshd", attn_mask_type="causal",
                             attention_dropout=0.0).to(dtype=torch.bfloat16, device="cuda")
print(dpa(q, k, v).abs().max())   # tensor(0., device='cuda:0')

Run it with NVTE_FUSED_ATTN=0 NVTE_UNFUSED_ATTN=0, or cuDNN FusedAttention wins backend
selection for this shape and hides the defect entirely.

Cause

TE normalises an unbounded window side to -1 (check_set_window_size); FA4's sentinel is None.
Until Dao-AILab/flash-attention#2490 a shim widened any pair summing below zero to full attention,
so TE's causal encoding (-1, 0) was always erased and causal=True alone set the mask. #2490
narrowed that to pairs where both bounds are negative, so (-1, 0) now survives as a literal
band of [row + 1, row] — empty. flash_attn/cute/mask.py guards on is not None and uses the
value arithmetically, so the kernel returns zeros and an all -inf LSE.

Measured directly, with TransformerEngine out of the picture:

flash_attn_func(q, k, v, causal=True, window_size=(-1, 0))      # out.abs().max() = 0.0
flash_attn_func(q, k, v, causal=True, window_size=(None, None)) # out.abs().max() = 3.453125

Change

Normalise negative window bounds to None at the FA4 entry points, where the symbols are bound.

Why there rather than at the call sites. The ~20 sites that build these kwargs are shared with
FA2 >= 2.7 and FA3, for which -1 is the correct spelling; rewriting them in place would mean
forking every one on use_flash_attn_4.

Why all four entry points. The non-CP path calls flash_attn_func / flash_attn_varlen_func;
context parallelism calls _flash_attn_fwd / _flash_attn_bwd. Covering only one pair is worse
than covering neither, because run_attention_with_cp.py grades a CP run against a non-CP run of
the same backend: with both sides returning zeros the comparison agrees and passes, so correcting
one side alone makes the fix present as a regression. That is not hypothetical — an earlier partial
attempt did exactly that, turning two passing arms red.

Genuine sliding windows already carry non-negative bounds and pass through untouched.

Verification (B200, SM100)

check	result
non-CP causal through DotProductAttention	3.453125e+00, was 0.000000e+00; relerr vs float64 2.046e-03
sliding windows (256,0), (128,0), (0,0)	relerr 2.046e-03, 2.046e-03, 0.000e+00
CP with FA4 on: p2p bshd/sbhd/thd	3/3 pass — p2p failed before this change
CP with FA4 on: all_gather, a2a bshd/sbhd	4/4 pass — were passing before, and still are
CP with FA4 off (control)	3/3 pass
test_fa4_window_sentinel.py	12 passed
test_attention.py sweep (regression)	487 passed, 138 skipped, 0 failed

Numerics are graded against a float64 reference, never against another backend. That matters
here: the defect survived because the CP suite compares against a run of the same backend, so an
error present on both sides cancels. float64 rather than float32 because torch computes fp32
matmuls in TF32 on Ampere and newer, whose significand is 11 bits — the same as fp16 — so an fp32
reference cannot judge a bf16 kernel.

Tests

tests/pytorch/attention/test_fa4_window_sentinel.py, registered in qa/L0_pytorch_unittest:

• eleven parametrized cases for the normalizer, which run on any machine, including that a genuine
  sliding window such as (511, 0) is passed through unchanged
• one end-to-end case, skipped unless FA4 is installed on SM100, pinned to FlashAttention and
  anchored to a float64 reference. It asserts a non-zero output separately from the tolerance
  check, so the specific failure is named rather than reported as a large error.

Note on CI

Three independent reasons this was invisible, each worth addressing on its own and none of them
fixed here:

1. NVTE_FLASH_ATTN_V4=0 is exported in qa/L1_pytorch_distributed_unittest/test.sh and the L0
   suites, so FA4 is never exercised under context parallelism.
2. CI pins flash-attn-4==4.0.0b11, which predates cute: don't widen intentionally-empty offset windows to full attention Dao-AILab/flash-attention#2490.
3. run_attention_with_cp.py grades a CP run against a non-CP run of the same backend, so it
   cannot detect any error affecting both equally. Two of the three comm types reported PASS on a
   backend that was returning zeros.

[greptile-apps]

Confidence Score: 5/5

The PR appears safe to merge, with no outstanding correctness, security, or repository-rule issues.

Summary

This PR prevents FlashAttention 4 from receiving Transformer Engine’s negative unbounded-window sentinel, avoiding silent all-zero causal-attention results.

• Converts negative FA4 window bounds to None at all four FA4 entry points.
• Covers standard, variable-length, and context-parallel forward/backward paths.
• Adds normalization unit tests and an SM100 end-to-end numerical regression test.
• Pins and verifies FA4 backend selection so another backend cannot produce a misleading green test.
• Registers the regression suite in the L0 PyTorch test runner.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  TE[TE attention kwargs<br/>window bound = -1] --> Normalize[FA4 wrapper<br/>negative bound → None]
  Normalize --> Standard[flash_attn_func / varlen]
  Normalize --> CP[_flash_attn_fwd / _flash_attn_bwd]
  Standard --> FA4[FlashAttention 4]
  CP --> FA4

Loading

_{Reviews (3) · Last reviewed commit: "docs(attention): say why the FA4 window ..."}

[sudhakarsingh27]

IMO, it's unnecessary to add the test here, could we add a versioned gate on the calls to "normalize windows" method? (something like this but I'm sure that I've got the minimum version incorrect)

[nvegesna-netizen]

I checked whether a gate is needed, and I do not think it is.

I transcribed _resolve_causal_local_window from before #2490 (interface.py @ 59cf537812) and after it, verified both transcriptions against the real source, and ran every (causal, window) pair TE emits through each. TE derives causal from attn_mask_type and the window from check_set_window_size, so the two always agree and causal=False with (-1, 0) never arises.

mask type	pre as-is	pre normalized	post as-is	post normalized
causal	CAUSAL	CAUSAL	LOCAL left=-1 right=0	CAUSAL
causal_bottom_right	CAUSAL	CAUSAL	LOCAL left=-1 right=0	CAUSAL
no_mask	FULL	FULL	FULL	FULL
sliding window	LOCAL 511/0	LOCAL 511/0	LOCAL 511/0	LOCAL 511/0

Normalizing is inert on the older release, not merely harmless. The old code reaches the same resolved state from (-1, 0) by widening that the new code reaches from (None, 0) directly, via the same window_size_left is None and window_size_right == 0 branch, which is unchanged context in #2490's diff. All four entry points funnel through that one resolver on both releases.

So ungated looks safer rather than just simpler: a threshold has to be right, and any release between #2490 merging on 9 Sep and the gated version would stay broken. Reasoning is now in the docstring in 4a84a21e.

On the test, I would keep it. Eleven of the twelve cases are pure Python with no GPU, so the CI cost is small. What earns their place is that this bug was invisible three separate ways: NVTE_FLASH_ATTN_V4=0 in the distributed suite, a flash-attn-4 pin predating #2490, and the CP suite comparing a CP run against a non-CP run of the same backend, which agrees when both sides return zeros. The end to end case pins NVTE_FLASH_ATTN_V4=1 and asserts the selected backend reports version 4 before grading numbers, so it cannot pass on FA2 the way the original would have.