diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a5d2de9a4c5..fc632e9e9ed 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/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}" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 9a339233a42..097e1bde38c 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,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()