From 1e703bb152d4ace8558593709629b84e03dde3b4 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 17 Sep 2026 00:15:37 -0700 Subject: [PATCH 1/9] fix(attention): stop handing FlashAttention 4 the -1 window sentinel FA4 returns an all-zero output for causal attention on SM100, silently. No exception, no warning, finite values, every element zero. Measured against a float64 reference on B200, with no context parallelism involved: flash_attn_func(q, k, v, causal=True, window_size=(-1, 0)) out.abs().max() = 0.0 relerr = 1.000e+00 flash_attn_func(q, k, v, causal=True, window_size=(None, None)) out.abs().max() = 3.453125 relerr = 2.046e-03 TE normalizes an unbounded window side to -1; FA4 spells it None. Until flash-attention #2490 a shim widened any pair summing below zero to full attention, so TE's (-1, 0) was always erased and causal= 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. The kernel returns zeros and an all -inf LSE. Normalizing negative bounds to None at the four FA4 entry points, rather than at the ~20 sites that build these kwargs: those sites are shared with FA2 >= 2.7 and FA3, where -1 is the correct spelling, so fixing them in place would mean forking every one on use_flash_attn_4. All four entry points, not just the two used by context parallelism. The non-CP path calls flash_attn_func/flash_attn_varlen_func; CP calls _flash_attn_fwd/_flash_attn_bwd. Covering 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 zero the comparison agrees, and correcting one side alone makes the fix present as a regression. That is not hypothetical -- it is what an earlier partial attempt did. Three reasons CI could not see this, all still worth addressing separately: NVTE_FLASH_ATTN_V4=0 in the distributed suite, a flash-attn-4 pin predating #2490, and the self-referential comparison above, which cannot detect any error affecting CP and non-CP equally. The new test anchors to a float64 reference rather than to another backend, and pins FlashAttention -- cuDNN FusedAttention wins selection for the shape under test and would mask the defect entirely. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- qa/L0_pytorch_unittest/test.sh | 1 + .../attention/test_fa4_window_sentinel.py | 139 ++++++++++++++++++ .../dot_product_attention/backends.py | 59 +++++++- 3 files changed, 195 insertions(+), 4 deletions(-) create mode 100644 tests/pytorch/attention/test_fa4_window_sentinel.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a78a99d7f95..c7e912c4ee2 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -70,6 +70,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_upro python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mixed_thd_attention.xml $TE_PATH/tests/pytorch/attention/test_mixed_thd_attention.py || test_fail "test_mixed_thd_attention.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fa4_window_sentinel.xml $TE_PATH/tests/pytorch/attention/test_fa4_window_sentinel.py || test_fail "test_fa4_window_sentinel.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then diff --git a/tests/pytorch/attention/test_fa4_window_sentinel.py b/tests/pytorch/attention/test_fa4_window_sentinel.py new file mode 100644 index 00000000000..7206584be1f --- /dev/null +++ b/tests/pytorch/attention/test_fa4_window_sentinel.py @@ -0,0 +1,139 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""FlashAttention 4 must never be handed TE's -1 window sentinel. + +TE spells an unbounded window side as -1; FA4 spells it None. Since flash-attention #2490 a negative +bound is honoured arithmetically rather than widened to full attention, so the causal encoding +(-1, 0) describes the band [row + 1, row] -- empty. FA4 then returns an all-zero output and an +all -inf LSE and raises nothing, which is a silent wrong answer rather than a crash. + +The numerical test below is anchored to a float64 reference rather than to another backend. That +matters here: the bug survived because tests/pytorch/attention/run_attention_with_cp.py grades a CP +run against a non-CP run *of the same backend*, so an error present on both sides cancels and the +comparison passes while measuring nothing. + +float64 and not float32 -- 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. +""" + +import pytest +import torch + +from transformer_engine.pytorch.attention.dot_product_attention.backends import ( + _fa4_normalized_window_kwargs, +) + + +@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 widened this one correctly on its own, since #2490 widens when + # both bounds are negative -- it is normalized anyway so one rule covers every case. + ({"window_size": (-1, -1)}, {"window_size": (None, None)}), + # A genuine sliding window must survive untouched, bounds and all. + ({"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): + """Every negative bound becomes None; every other value is passed through unchanged.""" + assert _fa4_normalized_window_kwargs(dict(sent)) == expected + + +def test_fa4_window_normalization_preserves_other_kwargs(): + """The normalizer must not disturb anything else it is handed.""" + kwargs = {"causal": True, "window_size": (-1, 0), "softmax_scale": 0.125, "num_splits": 1} + out = _fa4_normalized_window_kwargs(dict(kwargs)) + assert out["window_size"] == (None, 0) + assert out["causal"] is True and out["softmax_scale"] == 0.125 and out["num_splits"] == 1 + + +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. + + Pinned to FlashAttention because cuDNN FusedAttention wins backend selection for this shape and + would mask the defect entirely -- the first attempt at reproducing this measured a correct + result for exactly that reason. + """ + import os + + from transformer_engine.pytorch import DotProductAttention + from transformer_engine.pytorch.attention.dot_product_attention.utils 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)) + + saved = {key: os.environ.get(key) for key in ("NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN")} + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + _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 + + # 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}" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 9a339233a42..1da1e8dd73c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/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,54 @@ except (ValueError, TypeError): fa_utils.fa3_supports_softcap = False + +def _fa4_normalized_window_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: + """Replace the -1 "unbounded" window sentinel with the ``None`` FlashAttention 4 expects. + + TE normalizes an unbounded side to ``-1`` (``check_set_window_size``); FA4 spells it ``None``. + Since flash-attention #2490 a negative bound is honoured arithmetically instead of being widened + to full attention, and ``flash_attn/cute/mask.py`` guards on ``is not None``, so the causal + encoding ``(-1, 0)`` describes the band ``[row + 1, row]`` -- empty. The kernel then returns an + all-zero output and an all ``-inf`` LSE without raising, which is a silent wrong answer. + + Genuine sliding windows already carry non-negative bounds and pass through untouched. + """ + window = kwargs.get("window_size") + if window is not None: + kwargs["window_size"] = tuple( + None if bound is not None and bound < 0 else bound for bound in window + ) + for bound_name in ("window_size_left", "window_size_right"): + bound = kwargs.get(bound_name) + if bound is not None and bound < 0: + kwargs[bound_name] = None + return kwargs + + +def _fa4_with_none_window_sentinel(func: Callable) -> Callable: + """Wrap an FA4 entry point so it can never receive a negative window bound. + + Done here rather than at the ~20 call sites that build these kwargs, because those sites are + shared with FA2 >= 2.7 and FA3, for which ``-1`` is the correct spelling; rewriting them there + would have to fork every one on ``use_flash_attn_4``. + + All four entry points are wrapped. The non-CP path calls ``flash_attn_func`` / + ``flash_attn_varlen_func`` while context parallelism calls ``_flash_attn_fwd`` / + ``_flash_attn_bwd``, so covering only one pair leaves the other returning zeros -- and that is + worse than leaving both broken, because ``run_attention_with_cp.py`` grades a CP run against a + non-CP run of the same backend. With both sides zero the comparison agrees and passes; correcting + one side alone makes the fix present as a regression. + + Every TE call site passes these by keyword. A positional caller would bypass this. + """ + + @functools.wraps(func) + def wrapper(*args, **kwargs): + return func(*args, **_fa4_normalized_window_kwargs(kwargs)) + + return wrapper + + # Try to import Flash Attention v4 try: fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-4")) @@ -222,10 +271,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() From dcc69aa912d52808dcd4589f7bbd968798a27a3c Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 17 Sep 2026 00:54:58 -0700 Subject: [PATCH 2/9] fix(tests): import _attention_backends from where it is defined It lives in dot_product_attention.py, not utils.py, so the GPU arm of the new test failed at import on B200 while the eleven parametrized cases passed. The behaviour it checks was separately confirmed in the same run -- a non-CP forward through DotProductAttention returned 3.453125e+00 against 0.000000e+00 before the fix, relerr 2.046e-03 -- so this is the test reaching the assertion, not a change in what is being asserted. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_fa4_window_sentinel.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_fa4_window_sentinel.py b/tests/pytorch/attention/test_fa4_window_sentinel.py index 7206584be1f..04be5230e7a 100644 --- a/tests/pytorch/attention/test_fa4_window_sentinel.py +++ b/tests/pytorch/attention/test_fa4_window_sentinel.py @@ -98,7 +98,7 @@ def test_fa4_causal_attention_is_not_all_zeros(): import os from transformer_engine.pytorch import DotProductAttention - from transformer_engine.pytorch.attention.dot_product_attention.utils import ( + from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import ( _attention_backends, ) From 1b6c23d4c45b8962bf31835a16b326ee73ae602f Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 17 Sep 2026 01:30:54 -0700 Subject: [PATCH 3/9] test(attention): make the FA4 test prove FA4 actually answered qa/L0_pytorch_unittest exports NVTE_FLASH_ATTN_V4=0, and the skip condition checked flash_attn_func_v4 is None -- an import binding that env var does not affect. So on an SM100 L0 job the test would not skip: it would run, TE would decline FA4, fall through to FA2 (the same lane sets NVTE_FLASH_ATTN_V2=1), and pass while measuring a backend the test is not written for. A green test that exercises the wrong backend is the failure mode this fix exists to close, so relying on the lane to have FA4 enabled is the wrong dependency. The test now pins NVTE_FLASH_ATTN_V4=1 alongside the fused and unfused disables, and asserts the selected backend reports version 4 before grading any numbers. That makes it correct in any lane and stops a future lane quietly opting it out. Reported by greptile-apps on #3532, which flagged that the L0 runner disables FA4; the fall-through to FA2 is the part that makes it a false green rather than a skip. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- .../attention/test_fa4_window_sentinel.py | 23 ++++++++++++++----- 1 file changed, 17 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/attention/test_fa4_window_sentinel.py b/tests/pytorch/attention/test_fa4_window_sentinel.py index 04be5230e7a..7b426ff57eb 100644 --- a/tests/pytorch/attention/test_fa4_window_sentinel.py +++ b/tests/pytorch/attention/test_fa4_window_sentinel.py @@ -91,9 +91,12 @@ def _fa4_causal_unavailable(): def test_fa4_causal_attention_is_not_all_zeros(): """End to end: FA4 causal through DotProductAttention must match a float64 reference. - Pinned to FlashAttention because cuDNN FusedAttention wins backend selection for this shape and - would mask the defect entirely -- the first attempt at reproducing this measured a correct - result for exactly that reason. + Pins the backend three ways, because every one of them can silently turn this green while + measuring something else: cuDNN FusedAttention wins selection for this shape (the first attempt + at reproducing the bug measured a correct result for exactly that reason), the unfused path + would serve it too, and NVTE_FLASH_ATTN_V4=0 is exported by qa/L0_pytorch_unittest -- which + would leave FA2 answering a test named for FA4. The selected backend is asserted afterwards + rather than assumed, so a future lane cannot quietly opt this out again. """ import os @@ -107,9 +110,9 @@ def test_fa4_causal_attention_is_not_all_zeros(): torch.manual_seed(0) q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=dtype) for _ in range(3)) - saved = {key: os.environ.get(key) for key in ("NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN")} - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" + pinned = {"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( @@ -124,6 +127,14 @@ def test_fa4_causal_attention_is_not_all_zeros(): 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" From 962b9022e899a640534a6f18ab22057f360f2a75 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 22 Sep 2026 06:04:45 -0700 Subject: [PATCH 4/9] docs(attention): say why the FA4 window normalization is unconditional Review asked whether the -1 to None translation should sit behind a version gate keyed on the FA4 release that changed the sentinel. It does not need one, and the reason is worth recording where the next reader will look. Transcribing flash-attention's _resolve_causal_local_window from before #2490 (interface.py @ 59cf537812) and after it, and running every (causal, window) pair TE emits through both: mask type b11 as-is b11 normalized post-2490 as-is causal CAUSAL CAUSAL LOCAL left=-1 right=0 causal_bottom_right CAUSAL CAUSAL LOCAL left=-1 right=0 no_mask FULL FULL FULL sliding window LOCAL 511/0 LOCAL 511/0 LOCAL 511/0 Normalizing changes nothing on the older release, because the old code reaches the same resolved state from (-1, 0) by widening that the new code reaches from (None, 0) directly. So the translation is inert before the change and corrective after it, with no version boundary to get wrong. All four entry points funnel through that one resolver, on both releases, so the trace covers them. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- .../pytorch/attention/dot_product_attention/backends.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 1da1e8dd73c..fdb0572d269 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -193,6 +193,13 @@ def _fa4_normalized_window_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: all-zero output and an all ``-inf`` LSE without raising, which is a silent wrong answer. Genuine sliding windows already carry non-negative bounds and pass through untouched. + + Applied unconditionally rather than behind a version gate, because it is inert on releases + predating #2490 rather than merely harmless there. Both the old and new + ``_resolve_causal_local_window`` reduce ``(None, 0)`` with ``causal=True`` to plain causal via + the same ``window_size_left is None and window_size_right == 0`` branch; the old one reaches + the identical state from ``(-1, 0)`` by widening. So every window TE emits resolves the same + way before the change and correctly after it, and there is no version boundary to track. """ window = kwargs.get("window_size") if window is not None: From e718fc8203dde00953a2e9f70e6d475cc7f65562 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 1 Oct 2026 00:16:01 -0700 Subject: [PATCH 5/9] fix(attention): match only the -1 window sentinel at the FA4 boundary FA4 widens a window to full attention only when both bounds are negative, so a lone negative bound is a valid request for an empty window. Matching on `bound < 0` would silently turn one of those into full attention; `== -1` translates exactly TE's sentinel and leaves every other value with its FA4 meaning. It also matches how TE spells this sentinel elsewhere, including utils.py:1435, which already treats a non-(-1) negative as a genuine sliding window. Folds the normalizer and the decorator into one function, and drops the kwargs-passthrough test, which checked nothing about windows. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- .../attention/test_fa4_window_sentinel.py | 32 +++++----- .../dot_product_attention/backends.py | 61 ++++++------------- 2 files changed, 35 insertions(+), 58 deletions(-) diff --git a/tests/pytorch/attention/test_fa4_window_sentinel.py b/tests/pytorch/attention/test_fa4_window_sentinel.py index 7b426ff57eb..9247f144e6c 100644 --- a/tests/pytorch/attention/test_fa4_window_sentinel.py +++ b/tests/pytorch/attention/test_fa4_window_sentinel.py @@ -4,10 +4,10 @@ """FlashAttention 4 must never be handed TE's -1 window sentinel. -TE spells an unbounded window side as -1; FA4 spells it None. Since flash-attention #2490 a negative -bound is honoured arithmetically rather than widened to full attention, so the causal encoding -(-1, 0) describes the band [row + 1, row] -- empty. FA4 then returns an all-zero output and an -all -inf LSE and raises nothing, which is a silent wrong answer rather than a crash. +TE spells an unbounded window side as -1; FA4 spells it None. FA4 widens a window to full +attention only when both bounds are negative, so the causal encoding (-1, 0) reaches the kernel as +the band [row + 1, row] -- empty. FA4 then returns an all-zero output and an all -inf LSE and +raises nothing, which is a wrong answer rather than a crash. The numerical test below is anchored to a float64 reference rather than to another backend. That matters here: the bug survived because tests/pytorch/attention/run_attention_with_cp.py grades a CP @@ -22,17 +22,24 @@ import torch from transformer_engine.pytorch.attention.dot_product_attention.backends import ( - _fa4_normalized_window_kwargs, + _fa4_with_none_window_sentinel, ) +def _normalized(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 widened this one correctly on its own, since #2490 widens when - # both bounds are negative -- it is normalized anyway so one rule covers every case. + # TE's no-mask encoding. FA4 already widens this one itself, both bounds being negative, + # but rewriting it anyway keeps a single rule for every window TE emits. ({"window_size": (-1, -1)}, {"window_size": (None, None)}), # A genuine sliding window must survive untouched, bounds and all. ({"window_size": (511, 0)}, {"window_size": (511, 0)}), @@ -57,16 +64,9 @@ ], ) def test_fa4_window_sentinel_normalization(sent, expected): - """Every negative bound becomes None; every other value is passed through unchanged.""" - assert _fa4_normalized_window_kwargs(dict(sent)) == expected - + """The -1 sentinel becomes None; every other value is passed through unchanged.""" + assert _normalized(sent) == expected -def test_fa4_window_normalization_preserves_other_kwargs(): - """The normalizer must not disturb anything else it is handed.""" - kwargs = {"causal": True, "window_size": (-1, 0), "softmax_scale": 0.125, "num_splits": 1} - out = _fa4_normalized_window_kwargs(dict(kwargs)) - assert out["window_size"] == (None, 0) - assert out["causal"] is True and out["softmax_scale"] == 0.125 and out["num_splits"] == 1 def _fa4_causal_unavailable(): diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index fdb0572d269..e383615f6bd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -183,56 +183,33 @@ fa_utils.fa3_supports_softcap = False -def _fa4_normalized_window_kwargs(kwargs: Dict[str, Any]) -> Dict[str, Any]: - """Replace the -1 "unbounded" window sentinel with the ``None`` FlashAttention 4 expects. - - TE normalizes an unbounded side to ``-1`` (``check_set_window_size``); FA4 spells it ``None``. - Since flash-attention #2490 a negative bound is honoured arithmetically instead of being widened - to full attention, and ``flash_attn/cute/mask.py`` guards on ``is not None``, so the causal - encoding ``(-1, 0)`` describes the band ``[row + 1, row]`` -- empty. The kernel then returns an - all-zero output and an all ``-inf`` LSE without raising, which is a silent wrong answer. - - Genuine sliding windows already carry non-negative bounds and pass through untouched. - - Applied unconditionally rather than behind a version gate, because it is inert on releases - predating #2490 rather than merely harmless there. Both the old and new - ``_resolve_causal_local_window`` reduce ``(None, 0)`` with ``causal=True`` to plain causal via - the same ``window_size_left is None and window_size_right == 0`` branch; the old one reaches - the identical state from ``(-1, 0)`` by widening. So every window TE emits resolves the same - way before the change and correctly after it, and there is no version boundary to track. - """ - window = kwargs.get("window_size") - if window is not None: - kwargs["window_size"] = tuple( - None if bound is not None and bound < 0 else bound for bound in window - ) - for bound_name in ("window_size_left", "window_size_right"): - bound = kwargs.get(bound_name) - if bound is not None and bound < 0: - kwargs[bound_name] = None - return kwargs - - def _fa4_with_none_window_sentinel(func: Callable) -> Callable: - """Wrap an FA4 entry point so it can never receive a negative window bound. + """Rewrite TE's ``-1`` unbounded-window sentinel to the ``None`` FlashAttention 4 expects. - Done here rather than at the ~20 call sites that build these kwargs, because those sites are - shared with FA2 >= 2.7 and FA3, for which ``-1`` is the correct spelling; rewriting them there - would have to fork every one on ``use_flash_attn_4``. + TE spells an unbounded side ``-1``; FA4 spells it ``None``. FA4 widens a window to full + attention only when *both* bounds are negative, so TE's causal ``(-1, 0)`` reaches the kernel + as the band ``[row + 1, row]`` -- empty. The output is then all zeros and the LSE all ``-inf``, + with nothing raised: a wrong answer rather than a failure. - All four entry points are wrapped. The non-CP path calls ``flash_attn_func`` / - ``flash_attn_varlen_func`` while context parallelism calls ``_flash_attn_fwd`` / - ``_flash_attn_bwd``, so covering only one pair leaves the other returning zeros -- and that is - worse than leaving both broken, because ``run_attention_with_cp.py`` grades a CP run against a - non-CP run of the same backend. With both sides zero the comparison agrees and passes; correcting - one side alone makes the fix present as a regression. + Only ``-1`` is rewritten. A lone negative bound is a valid FA4 request for an empty window, so + widening every negative here would silently turn one of those into full attention. - Every TE call site passes these by keyword. A positional caller would bypass this. + Wrapped at the entry points rather than at the ~20 call sites that build these kwargs, since + those are shared with FA2 >= 2.7 and FA3, where ``-1`` is the correct spelling. All four are + wrapped together because ``run_attention_with_cp.py`` grades a CP run against a non-CP run of + the same backend: with both sides empty the comparison agrees, so fixing one pair alone would + present as a regression. Keyword arguments only; a positional caller bypasses this. """ @functools.wraps(func) def wrapper(*args, **kwargs): - return func(*args, **_fa4_normalized_window_kwargs(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 From 7b412e4c11f7e1ca72a186caf783f8c721660ff6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 1 Oct 2026 07:18:47 +0000 Subject: [PATCH 6/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_fa4_window_sentinel.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/pytorch/attention/test_fa4_window_sentinel.py b/tests/pytorch/attention/test_fa4_window_sentinel.py index 9247f144e6c..4551ea3e640 100644 --- a/tests/pytorch/attention/test_fa4_window_sentinel.py +++ b/tests/pytorch/attention/test_fa4_window_sentinel.py @@ -68,7 +68,6 @@ def test_fa4_window_sentinel_normalization(sent, expected): assert _normalized(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(): From c88e28d83724a6f472bb17409abc7e61752f3d5f Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 1 Oct 2026 14:24:58 -0700 Subject: [PATCH 7/9] test(attention): fold the FA4 window sentinel tests into test_attention.py A file and an L0 entry for two tests on a one-line fix is not worth it. The directory already carries 14 test files, and the rest each cover a backend or subsystem rather than a bug fix. Nothing is dropped: the module docstring's reasoning moves into a section comment above the tests, since it explains why the numerical test is anchored to a float64 reference rather than to another backend, which is the thing that made the bug detectable at all. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- qa/L0_pytorch_unittest/test.sh | 1 - tests/pytorch/attention/test_attention.py | 145 +++++++++++++++++ .../attention/test_fa4_window_sentinel.py | 149 ------------------ 3 files changed, 145 insertions(+), 150 deletions(-) delete mode 100644 tests/pytorch/attention/test_fa4_window_sentinel.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 773c6129c0e..ea2d58825f0 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -72,7 +72,6 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_upro python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mixed_thd_attention.xml $TE_PATH/tests/pytorch/attention/test_mixed_thd_attention.py || test_fail "test_mixed_thd_attention.py" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fa4_window_sentinel.xml $TE_PATH/tests/pytorch/attention/test_fa4_window_sentinel.py || test_fail "test_fa4_window_sentinel.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py" export NVTE_TEST_CHECKPOINT_ARTIFACT_PATH=$TE_PATH/artifacts/tests/pytorch/test_checkpoint if [ ! -d "$NVTE_TEST_CHECKPOINT_ARTIFACT_PATH" ]; then diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a5d2de9a4c5..6fbb3c0cb74 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3894,3 +3894,148 @@ def forward( self.quantizers, ) return out + + +# ---------------------------------------------------------------------------------------------- +# FlashAttention 4 must never be handed TE's -1 window sentinel. +# +# TE spells an unbounded window side as -1; FA4 spells it None. FA4 widens a window to full +# attention only when both bounds are negative, so the causal encoding (-1, 0) reaches the kernel +# as the band [row + 1, row] -- empty. FA4 then returns an all-zero output and an all -inf LSE and +# raises nothing, which is a wrong answer rather than a crash. +# +# The numerical test is anchored to a float64 reference rather than to another backend. That +# matters here: the bug survived because run_attention_with_cp.py grades a CP run against a non-CP +# run *of the same backend*, so an error present on both sides cancels and the comparison passes +# while measuring nothing. float64 and not float32 -- 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. +# ---------------------------------------------------------------------------------------------- + + +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 already widens this one itself, both bounds being negative, + # but rewriting it anyway keeps a single rule for every window TE emits. + ({"window_size": (-1, -1)}, {"window_size": (None, None)}), + # A genuine sliding window must survive untouched, bounds and all. + ({"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. + + Pins the backend three ways, because every one of them can silently turn this green while + measuring something else: cuDNN FusedAttention wins selection for this shape (the first attempt + at reproducing the bug measured a correct result for exactly that reason), the unfused path + would serve it too, and NVTE_FLASH_ATTN_V4=0 is exported by qa/L0_pytorch_unittest -- which + would leave FA2 answering a test named for FA4. The selected backend is asserted afterwards + rather than assumed, so a future lane cannot quietly opt this out again. + """ + 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_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}" diff --git a/tests/pytorch/attention/test_fa4_window_sentinel.py b/tests/pytorch/attention/test_fa4_window_sentinel.py deleted file mode 100644 index 4551ea3e640..00000000000 --- a/tests/pytorch/attention/test_fa4_window_sentinel.py +++ /dev/null @@ -1,149 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""FlashAttention 4 must never be handed TE's -1 window sentinel. - -TE spells an unbounded window side as -1; FA4 spells it None. FA4 widens a window to full -attention only when both bounds are negative, so the causal encoding (-1, 0) reaches the kernel as -the band [row + 1, row] -- empty. FA4 then returns an all-zero output and an all -inf LSE and -raises nothing, which is a wrong answer rather than a crash. - -The numerical test below is anchored to a float64 reference rather than to another backend. That -matters here: the bug survived because tests/pytorch/attention/run_attention_with_cp.py grades a CP -run against a non-CP run *of the same backend*, so an error present on both sides cancels and the -comparison passes while measuring nothing. - -float64 and not float32 -- 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. -""" - -import pytest -import torch - -from transformer_engine.pytorch.attention.dot_product_attention.backends import ( - _fa4_with_none_window_sentinel, -) - - -def _normalized(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 already widens this one itself, both bounds being negative, - # but rewriting it anyway keeps a single rule for every window TE emits. - ({"window_size": (-1, -1)}, {"window_size": (None, None)}), - # A genuine sliding window must survive untouched, bounds and all. - ({"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 _normalized(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. - - Pins the backend three ways, because every one of them can silently turn this green while - measuring something else: cuDNN FusedAttention wins selection for this shape (the first attempt - at reproducing the bug measured a correct result for exactly that reason), the unfused path - would serve it too, and NVTE_FLASH_ATTN_V4=0 is exported by qa/L0_pytorch_unittest -- which - would leave FA2 answering a test named for FA4. The selected backend is asserted afterwards - rather than assumed, so a future lane cannot quietly opt this out again. - """ - 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_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}" From 771f85209a088b53bd55b1f445af5acad3eed63b Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 1 Oct 2026 14:31:40 -0700 Subject: [PATCH 8/9] test(attention): pin NVTE_FLASH_ATTN alongside the v4 switch Consequence of moving this test into test_attention.py: it now shares a process with tests that set NVTE_FLASH_ATTN=0, and several of those do not restore it. utils.py derives use_flash_attention_4 as (NVTE_FLASH_ATTN and NVTE_FLASH_ATTN_V4), so an inherited 0 disables FA4 whatever the v4 switch says, and the backend assertion then fails rather than grading the fix. The save and restore iterate the pinned dict, so the new key is covered by both. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 6fbb3c0cb74..d6115acb5c7 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3982,7 +3982,7 @@ def _fa4_causal_unavailable(): def test_fa4_causal_attention_is_not_all_zeros(): """End to end: FA4 causal through DotProductAttention must match a float64 reference. - Pins the backend three ways, because every one of them can silently turn this green while + Pins the backend four ways, because every one of them can silently turn this green while measuring something else: cuDNN FusedAttention wins selection for this shape (the first attempt at reproducing the bug measured a correct result for exactly that reason), the unfused path would serve it too, and NVTE_FLASH_ATTN_V4=0 is exported by qa/L0_pytorch_unittest -- which @@ -4001,7 +4001,12 @@ def test_fa4_causal_attention_is_not_all_zeros(): torch.manual_seed(0) q, k, v = (torch.randn(b, s, h, d, device="cuda", dtype=dtype) for _ in range(3)) - pinned = {"NVTE_FUSED_ATTN": "0", "NVTE_UNFUSED_ATTN": "0", "NVTE_FLASH_ATTN_V4": "1"} + 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 From c0d1c37d787a9eb832fa4f6cbdd4fcffff1e5932 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 6 Oct 2026 17:39:22 -0700 Subject: [PATCH 9/9] refactor(attention): condense the FA4 sentinel comments The new prose ran long for its surroundings. backends.py has a median docstring of 1 line and only one of 28 reaches 10; this helper's was 16, the second longest in the file. test_attention.py has a median of 1 and a max of 9, and carried a 15-line banner. Kept the facts a reader cannot derive: that FA4 widens a window only when both bounds are negative, which turns TE's causal (-1, 0) into an empty band and an all-zero output with nothing raised; that only -1 is rewritten because a lone negative is a valid empty-window request; and why all four backend switches are pinned in the end-to-end test. Cut the rationale the PR description already carries. Helper docstring 16 lines to 6, test banner 15 to 3, end-to-end docstring 8 to 5. No code changed: verified by comparing ASTs with docstrings stripped. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 34 ++++++------------- .../dot_product_attention/backends.py | 17 +++------- 2 files changed, 14 insertions(+), 37 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index d6115acb5c7..fc632e9e9ed 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -3896,21 +3896,9 @@ def forward( return out -# ---------------------------------------------------------------------------------------------- -# FlashAttention 4 must never be handed TE's -1 window sentinel. -# -# TE spells an unbounded window side as -1; FA4 spells it None. FA4 widens a window to full -# attention only when both bounds are negative, so the causal encoding (-1, 0) reaches the kernel -# as the band [row + 1, row] -- empty. FA4 then returns an all-zero output and an all -inf LSE and -# raises nothing, which is a wrong answer rather than a crash. -# -# The numerical test is anchored to a float64 reference rather than to another backend. That -# matters here: the bug survived because run_attention_with_cp.py grades a CP run against a non-CP -# run *of the same backend*, so an error present on both sides cancels and the comparison passes -# while measuring nothing. float64 and not float32 -- 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. -# ---------------------------------------------------------------------------------------------- +# 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 ( @@ -3930,10 +3918,10 @@ def _fa4_normalized_kwargs(kwargs): [ # TE's causal encoding: the pair that produced zeros. ({"window_size": (-1, 0)}, {"window_size": (None, 0)}), - # TE's no-mask encoding. FA4 already widens this one itself, both bounds being negative, - # but rewriting it anyway keeps a single rule for every window TE emits. + # 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, bounds and all. + # 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. @@ -3982,12 +3970,10 @@ def _fa4_causal_unavailable(): def test_fa4_causal_attention_is_not_all_zeros(): """End to end: FA4 causal through DotProductAttention must match a float64 reference. - Pins the backend four ways, because every one of them can silently turn this green while - measuring something else: cuDNN FusedAttention wins selection for this shape (the first attempt - at reproducing the bug measured a correct result for exactly that reason), the unfused path - would serve it too, and NVTE_FLASH_ATTN_V4=0 is exported by qa/L0_pytorch_unittest -- which - would leave FA2 answering a test named for FA4. The selected backend is asserted afterwards - rather than assumed, so a future lane cannot quietly opt this out again. + 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 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index e383615f6bd..097e1bde38c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -186,19 +186,10 @@ def _fa4_with_none_window_sentinel(func: Callable) -> Callable: """Rewrite TE's ``-1`` unbounded-window sentinel to the ``None`` FlashAttention 4 expects. - TE spells an unbounded side ``-1``; FA4 spells it ``None``. FA4 widens a window to full - attention only when *both* bounds are negative, so TE's causal ``(-1, 0)`` reaches the kernel - as the band ``[row + 1, row]`` -- empty. The output is then all zeros and the LSE all ``-inf``, - with nothing raised: a wrong answer rather than a failure. - - Only ``-1`` is rewritten. A lone negative bound is a valid FA4 request for an empty window, so - widening every negative here would silently turn one of those into full attention. - - Wrapped at the entry points rather than at the ~20 call sites that build these kwargs, since - those are shared with FA2 >= 2.7 and FA3, where ``-1`` is the correct spelling. All four are - wrapped together because ``run_attention_with_cp.py`` grades a CP run against a non-CP run of - the same backend: with both sides empty the comparison agrees, so fixing one pair alone would - present as a regression. Keyword arguments only; a positional caller bypasses this. + 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)