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

tests/pytorch/attention/test_attention.py

@@ -3894,3 +3894,139 @@ def forward(
                 self.quantizers,
             )
         return out
+
+
+# FlashAttention 4 must never be handed TE's -1 window sentinel: FA4 widens a window only when
+# both bounds are negative, so causal (-1, 0) reaches the kernel as an empty band and returns an
+# all-zero output with nothing raised.
+
+
+from transformer_engine.pytorch.attention.dot_product_attention.backends import (
+    _fa4_with_none_window_sentinel,
+)
+
+
+def _fa4_normalized_kwargs(kwargs):
+    """Return the kwargs FA4 would actually receive, through the real wrapper."""
+    received = {}
+    _fa4_with_none_window_sentinel(received.update)(**kwargs)
+    return received
+
+
+@pytest.mark.parametrize(
+    "sent,expected",
+    [
+        # TE's causal encoding: the pair that produced zeros.
+        ({"window_size": (-1, 0)}, {"window_size": (None, 0)}),
+        # TE's no-mask encoding. FA4 widens this one itself, but one rule for every window is
+        # simpler to reason about.
+        ({"window_size": (-1, -1)}, {"window_size": (None, None)}),
+        # A genuine sliding window must survive untouched.
+        ({"window_size": (511, 0)}, {"window_size": (511, 0)}),
+        ({"window_size": (511, -1)}, {"window_size": (511, None)}),
+        # The private entry points take the bounds separately.
+        (
+            {"window_size_left": -1, "window_size_right": 0},
+            {"window_size_left": None, "window_size_right": 0},
+        ),
+        (
+            {"window_size_left": -1, "window_size_right": -1},
+            {"window_size_left": None, "window_size_right": None},
+        ),
+        (
+            {"window_size_left": 256, "window_size_right": 0},
+            {"window_size_left": 256, "window_size_right": 0},
+        ),
+        # Already idiomatic, and absent: both must be left alone.
+        ({"window_size": (None, None)}, {"window_size": (None, None)}),
+        ({"window_size": None}, {"window_size": None}),
+        ({}, {}),
+    ],
+)
+def test_fa4_window_sentinel_normalization(sent, expected):
+    """The -1 sentinel becomes None; every other value is passed through unchanged."""
+    assert _fa4_normalized_kwargs(sent) == expected
+
+
+def _fa4_causal_unavailable():
+    """Why this machine cannot exercise FA4 causal attention, or None if it can."""
+    if not torch.cuda.is_available():
+        return "no CUDA device"
+    from transformer_engine.pytorch.attention.dot_product_attention.backends import (
+        flash_attn_func_v4,
+    )
+
+    if flash_attn_func_v4 is None:
+        return "flash-attn-4 is not installed"
+    if torch.cuda.get_device_capability() != (10, 0):
+        return "the FA4 CuTe kernels under test are SM100"
+    return None
+
+
+_SKIP = _fa4_causal_unavailable()
+
+
+@pytest.mark.skipif(_SKIP is not None, reason=str(_SKIP))
+def test_fa4_causal_attention_is_not_all_zeros():
+    """End to end: FA4 causal through DotProductAttention must match a float64 reference.
+
+    All four backend switches are pinned, since each can turn this green while measuring something
+    else: cuDNN FusedAttention wins selection for this shape, unfused would serve it too, and
+    qa/L0_pytorch_unittest exports NVTE_FLASH_ATTN_V4=0, which would leave FA2 answering a test
+    named for FA4. The selected backend is asserted afterwards rather than assumed.
+    """
+    import os
+
+    from transformer_engine.pytorch import DotProductAttention
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+
+    b, h, s, d = 2, 8, 1024, 128
+    dtype = torch.bfloat16
+    torch.manual_seed(0)
+    q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=dtype) for _ in range(3))
+
+    pinned = {
+        "NVTE_FLASH_ATTN": "1",
+        "NVTE_FUSED_ATTN": "0",
+        "NVTE_UNFUSED_ATTN": "0",
+        "NVTE_FLASH_ATTN_V4": "1",
+    }
+    saved = {key: os.environ.get(key) for key in pinned}
+    os.environ.update(pinned)
+    _attention_backends["backend_selection_requires_update"] = True
+    try:
+        dpa = DotProductAttention(
+            h, d, qkv_format="bshd", attn_mask_type="causal", attention_dropout=0.0
+        ).to(dtype=dtype, device="cuda")
+        out = dpa(q, k, v).view(b, s, h, d)
+    finally:
+        for key, value in saved.items():
+            if value is None:
+                os.environ.pop(key, None)
+            else:
+                os.environ[key] = value
+        _attention_backends["backend_selection_requires_update"] = True
+
+    # Prove FA4 actually answered before grading the numbers. Without this the test passes on FA2
+    # wherever FA4 is disabled, which is the default in at least one existing CI lane.
+    selected = _attention_backends["flash_attention_backend"]
+    assert selected is not None and str(selected).startswith("4"), (
+        f"expected FlashAttention 4 to serve this config, got {selected}; the test would otherwise"
+        " grade a backend it is not written for"
+    )
+
+    # Stated separately from the tolerance check: an all-zero output is the specific failure this
+    # test exists for, and it should be reported as such rather than as a large error.
+    assert out.abs().max() > 0, "FA4 returned an all-zero output for causal attention"
+
+    qs, ks, vs = (t[0, :, 0].double() for t in (q, k, v))
+    scores = (qs @ ks.transpose(-1, -2)) * (d**-0.5)
+    scores = scores.masked_fill(
+        torch.ones(s, s, dtype=torch.bool, device=qs.device).triu(1), float("-inf")
+    )
+    reference = torch.softmax(scores, dim=-1) @ vs
+
+    error = (out[0, :, 0].double() - reference).abs().max() / reference.abs().max()
+    assert error < 2e-2, f"FA4 causal attention differs from the float64 reference: {error:.3e}"

transformer_engine/pytorch/attention/dot_product_attention/backends.py

@@ -8,6 +8,7 @@
 from dataclasses import dataclass
 from importlib.metadata import version as get_pkg_version
 from importlib.metadata import PackageNotFoundError
+import functools
 import inspect
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -181,6 +182,29 @@
     except (ValueError, TypeError):
         fa_utils.fa3_supports_softcap = False
 
+
+def _fa4_with_none_window_sentinel(func: Callable) -> Callable:
+    """Rewrite TE's ``-1`` unbounded-window sentinel to the ``None`` FlashAttention 4 expects.
+
+    FA4 widens a window to full attention only when *both* bounds are negative, so TE's causal
+    ``(-1, 0)`` reaches the kernel as the empty band ``[row + 1, row]``: all-zero output and an
+    all ``-inf`` LSE, with nothing raised. Only ``-1`` is rewritten, since a lone negative bound
+    is a valid FA4 request for an empty window. Keyword arguments only.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        window = kwargs.get("window_size")
+        if window is not None:
+            kwargs["window_size"] = tuple(None if bound == -1 else bound for bound in window)
+        for bound_name in ("window_size_left", "window_size_right"):
+            if kwargs.get(bound_name) == -1:
+                kwargs[bound_name] = None
+        return func(*args, **kwargs)
+
+    return wrapper
+
+
 # Try to import Flash Attention v4
 try:
     fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-4"))
@@ -222,10 +246,12 @@
     else:
         # Unlike versions 2 and 3, FlashAttention 4 registers no custom ops: it builds
         # its kernels through the CUTLASS DSL as it runs. Keep it an eager island.
-        flash_attn_func_v4 = no_torch_dynamo()(_flash_attn_func_v4)
-        flash_attn_varlen_func_v4 = no_torch_dynamo()(_flash_attn_varlen_func_v4)
-        _flash_attn_fwd_v4 = no_torch_dynamo()(_flash_attn_fwd_v4)
-        _flash_attn_bwd_v4 = no_torch_dynamo()(_flash_attn_bwd_v4)
+        flash_attn_func_v4 = no_torch_dynamo()(_fa4_with_none_window_sentinel(_flash_attn_func_v4))
+        flash_attn_varlen_func_v4 = no_torch_dynamo()(
+            _fa4_with_none_window_sentinel(_flash_attn_varlen_func_v4)
+        )
+        _flash_attn_fwd_v4 = no_torch_dynamo()(_fa4_with_none_window_sentinel(_flash_attn_fwd_v4))
+        _flash_attn_bwd_v4 = no_torch_dynamo()(_fa4_with_none_window_sentinel(_flash_attn_bwd_v4))
 
         fa_utils.v4_validate_head_dims = _fa4_validate_head_dims
         fa_utils.set_flash_attention_4_params()