From 69b4ffb1a8c4ea00599d4d611e1102534ecc040b Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 17 Aug 2026 14:20:36 -0700 Subject: [PATCH 01/28] feat: thread tanh logit softcapping through FlashAttention (FA2, opt-in FA3) Add a user `softcap` value (tanh logit softcapping, `softcap*tanh(x/softcap)`) to DotProductAttention so models like Gemma2 can run on the fused flash path instead of an unfused/FlexAttention kernel. - Add `softcap` to DotProductAttention (init+forward) and AttentionParams; thread it into the FA2 non-CP kwargs and all three context-parallel autograd functions (forward + ctx-saved backward). softcap=0.0 reproduces prior behavior. - get_attention_backend: when softcap != 0, disable FusedAttention/unfused and steer to FA2 -- disable FA3/FA4, and disable FA2 < 2.6.0 -- so the cap is never silently dropped (FA2 < 2.6.0) or hit at runtime as NotImplementedError (FA3/FA4). Also disable FA3 under context parallelism (its CP path hard-rejects nonzero softcap) so CP+softcap steers to FA2, which supports it, instead of crashing. - FA3 softcap opt-in: NVTE_FA3_SOFTCAP=1, Hopper (sm90) hd<=256, non-CP only, gated on a fail-closed signature probe (fa3_supports_softcap). Forward threads softcap into fa_3_optional_forward_kwargs; the existing Hopper autograd function carries it into backward automatically. Default off; unchanged behavior steers to FA2. - ONNX export: fail loudly (assert) rather than silently drop softcap -- export unconditionally force-selects UnfusedDotProductAttention, which has no softcap support, so this previously exported models with softcapping silently omitted. - Tests: test_softcap.py (FA2 fwd/bwd parity vs pure-PyTorch reference), wired into qa/L0_pytorch_unittest/test.sh. FA4 softcap opt-in is deliberately NOT included here -- see follow-up PR. On Blackwell (SM100), FA4's dedicated head_dim=256 forward kernel has no score_mod support at all (kernel constructor asserts `score_mod is None`), so there is currently no FA4 kernel path this could opt into; adding the scaffolding now would just be inert code with nothing to exercise. Addresses review findings: CP+FA3 softcap selection crash, ONNX silent drop, and the missing CI wiring for test_softcap.py. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/attention/test_softcap.py | 149 ++++++++++++++++++ .../dot_product_attention/backends.py | 38 +++++ .../dot_product_attention/context_parallel.py | 23 ++- .../dot_product_attention.py | 26 +++ .../attention/dot_product_attention/utils.py | 47 ++++++ 6 files changed, 278 insertions(+), 6 deletions(-) create mode 100644 tests/pytorch/attention/test_softcap.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 14a5f4fe3d..482642e5e1 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -67,6 +67,7 @@ NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" 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_softcap.xml $TE_PATH/tests/pytorch/attention/test_softcap.py || test_fail "test_softcap.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_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 diff --git a/tests/pytorch/attention/test_softcap.py b/tests/pytorch/attention/test_softcap.py new file mode 100644 index 0000000000..b0b94bb4e8 --- /dev/null +++ b/tests/pytorch/attention/test_softcap.py @@ -0,0 +1,149 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Isolation numerics test for tanh logit softcapping in DotProductAttention. + +The reference implements softcapping in pure PyTorch: + + scores = (Q @ K^T) * scale + scores = softcap * tanh(scores / softcap) # only when softcap != 0.0 + scores = scores + mask + attn = softmax(scores) + out = attn @ V + +and is compared against ``DotProductAttention(..., softcap=...)`` forced onto the +FlashAttention backend, for both the forward output and the input gradients +(dQ/dK/dV obtained via autograd). +""" + +import sys +import pathlib + +import pytest +import torch +from packaging.version import Version as PkgVersion + +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends + +_current_file = pathlib.Path(__file__).resolve() +sys.path = [str(_current_file.parent.parent)] + sys.path +from utils import reset_rng_states # pylint: disable=wrong-import-position + + +def _flash_attn_2_6_available() -> bool: + """Whether flash-attn >= 2.6.0 (the first version exposing ``softcap``) is installed.""" + try: + import flash_attn # pylint: disable=import-outside-toplevel + except ImportError: + return False + return PkgVersion(flash_attn.__version__) >= PkgVersion("2.6.0") + + +# Softcapping through DotProductAttention is only wired through the FlashAttention 2 +# backend (>= 2.6.0), and requires CUDA tensors. +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required."), + pytest.mark.skipif( + not _flash_attn_2_6_available(), reason="flash-attn >= 2.6.0 is required." + ), +] + + +def _force_flash_backend() -> None: + """Force DotProductAttention to select the FlashAttention backend.""" + import os # pylint: disable=import-outside-toplevel + + os.environ["NVTE_FLASH_ATTN"] = "1" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "0" + _attention_backends["backend_selection_requires_update"] = True + + +def _reference_attention(q, k, v, scale, softcap, causal): + """Pure-PyTorch reference for softcapped scaled dot product attention. + + q, k, v are in ``bshd`` layout. GQA is supported: ``k``/``v`` may have fewer + heads than ``q``. + """ + # bshd -> bhsd + qt = q.transpose(1, 2).float() + kt = k.transpose(1, 2).float() + vt = v.transpose(1, 2).float() + + num_heads = qt.shape[1] + num_gqa_groups = kt.shape[1] + if num_heads != num_gqa_groups: + assert num_heads % num_gqa_groups == 0 + repeats = num_heads // num_gqa_groups + kt = kt.repeat_interleave(repeats, dim=1) + vt = vt.repeat_interleave(repeats, dim=1) + + scores = torch.matmul(qt, kt.transpose(-2, -1)) * scale + if softcap != 0.0: + scores = softcap * torch.tanh(scores / softcap) + if causal: + sq, skv = scores.shape[-2], scores.shape[-1] + mask = torch.triu( + torch.ones(sq, skv, dtype=torch.bool, device=scores.device), + diagonal=1 + skv - sq, + ) + scores = scores.masked_fill(mask, float("-inf")) + attn = torch.softmax(scores, dim=-1) + out = torch.matmul(attn, vt) + # bhsd -> bshd + return out.transpose(1, 2) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("softcap", [0.0, 50.0]) +@pytest.mark.parametrize("num_gqa_groups", [4, 2]) +@pytest.mark.parametrize("causal", [False, True]) +def test_softcap_numerics(dtype, softcap, num_gqa_groups, causal): + """FlashAttention softcap forward + grads match a pure-PyTorch reference. + + ``softcap == 0.0`` additionally proves that softcapping is a no-op relative to + the plain (no-softcap) reference, i.e. today's behavior is reproduced exactly. + """ + reset_rng_states() + + batch_size = 2 + max_seqlen = 32 + num_heads = 4 + head_dim = 64 + scale = 1.0 / (head_dim**0.5) + + q_shape = (batch_size, max_seqlen, num_heads, head_dim) + kv_shape = (batch_size, max_seqlen, num_gqa_groups, head_dim) + + q = (0.5 * torch.randn(q_shape, dtype=dtype, device="cuda")).requires_grad_() + k = (0.5 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() + v = (0.5 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() + q_ref, k_ref, v_ref = [x.detach().clone().requires_grad_() for x in (q, k, v)] + + grad_output = torch.randn(q_shape, dtype=dtype, device="cuda") + + _force_flash_backend() + dpa = DotProductAttention( + num_heads, + head_dim, + num_gqa_groups=num_gqa_groups, + qkv_format="bshd", + attn_mask_type="causal" if causal else "no_mask", + softmax_scale=scale, + softcap=softcap, + layer_number=1, + ).to(dtype=dtype, device="cuda") + + out = dpa(q, k, v) + out.backward(grad_output) + + out_ref = _reference_attention(q_ref, k_ref, v_ref, scale, softcap, causal) + out_ref.backward(grad_output.float()) + + atol, rtol = (2e-2, 2e-2) if dtype == torch.float16 else (3.5e-2, 3.5e-2) + + torch.testing.assert_close(out.float(), out_ref.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(q.grad.float(), q_ref.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(k.grad.float(), k_ref.grad.float(), atol=atol, rtol=rtol) + torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), atol=atol, rtol=rtol) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 8a219a6a4d..12a2ed1492 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -7,6 +7,7 @@ from contextlib import nullcontext from importlib.metadata import version as get_pkg_version from importlib.metadata import PackageNotFoundError +import inspect import os from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings @@ -165,6 +166,19 @@ fa_utils.set_flash_attention_3_params() + # Probe whether this FA3 build exposes a `softcap` parameter on BOTH entry points. FA3's Hopper + # (sm90) kernels DO implement tanh logit softcapping in fwd AND bwd (dedicated + # flash_{fwd,bwd}_hdim256_bf16_softcap_sm90 instantiations, off only behind a compile-time + # DISABLE_SOFTCAP flag), so this is a mature path. Still fail-closed and additionally + # gated on opt-in (NVTE_FA3_SOFTCAP) + head_dim <= 256 in get_attention_backend. + try: + fa_utils.fa3_supports_softcap = ( + "softcap" in inspect.signature(flash_attn_func_v3).parameters + and "softcap" in inspect.signature(flash_attn_varlen_func_v3).parameters + ) + except (ValueError, TypeError): + fa_utils.fa3_supports_softcap = False + # Try to import Flash Attention v4 try: fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-4")) @@ -885,6 +899,7 @@ def forward( max_seqlen_kv: Optional[int] = None, attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, + softcap: float = 0.0, alibi_slopes: Optional[torch.Tensor] = None, cp_group: Optional[Union[dist_group_type, List[dist_group_type]]] = None, cp_global_ranks: List[int] = None, @@ -1100,6 +1115,11 @@ def forward( assert ( alibi_slopes is None ), "Alibi slope bias addition is not supported with context parallelism." + if use_flash_attn_3 and softcap != 0.0: + raise NotImplementedError( + "softcap is not supported by the FlashAttention 3 backend in context " + "parallel. Please use FlashAttention 2 (>= 2.6.0) for softcap support." + ) with self.attention_dropout_ctx(): output = attn_forward_func_with_cp( self.training, @@ -1130,6 +1150,7 @@ def forward( attn_mask_type=attn_mask_type, deterministic=self.deterministic, window_size=window_size, + softcap=softcap, quantizers=quantizers, pad_between_seqs=pad_between_seqs, use_flash_attn_3=use_flash_attn_3, @@ -1215,6 +1236,8 @@ def forward( fa_optional_forward_kwargs["alibi_slopes"] = alibi_slopes if fa_utils.v2_4_1_plus: fa_optional_forward_kwargs["deterministic"] = self.deterministic + if fa_utils.v2_6_0_plus: + fa_optional_forward_kwargs["softcap"] = softcap if inference_params is not None: # use block_table kwarg to support thd_2bshd for non-paged fa_optional_forward_kwargs["block_table"] = ( @@ -1235,9 +1258,24 @@ def forward( **fa_optional_forward_kwargs, ) else: + # Fail-loud net: get_attention_backend only keeps FA3 for softcap on a + # softcap-capable build (signature probe) + opt-in (NVTE_FA3_SOFTCAP) + Hopper + # (FA3 is sm90-only upstream) + head_dim <= 256. If FA3 is still reached with + # softcap while the build lacks support (force-selected / regressed path), raise + # rather than silently drop the cap. The non-CP FA3 entry points + # (flash_attn_func_v3 / flash_attn_varlen_func_v3) are self-contained autograd + # functions, so threading `softcap` into the forward call also drives the + # matching FA3 softcap backward kernel. (CP + FA3 + softcap stays blocked above.) + if softcap != 0.0 and not fa_utils.fa3_supports_softcap: + raise NotImplementedError( + "softcap is not supported by the installed FlashAttention 3 build. " + "Please use FlashAttention 2 (>= 2.6.0) for softcap support." + ) fa_3_optional_forward_kwargs = {} fa_3_optional_forward_kwargs["window_size"] = window_size fa_3_optional_forward_kwargs["num_splits"] = num_splits + if softcap != 0.0 and fa_utils.fa3_supports_softcap: + fa_3_optional_forward_kwargs["softcap"] = softcap if pad_between_seqs: fa_3_optional_forward_kwargs["seqused_q"] = ( cu_seqlens_q[1:] - cu_seqlens_q[:-1] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index ea89ca97eb..3484d4e9cd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1387,6 +1387,7 @@ def forward( deterministic, use_fused_attention, return_max_logit, + softcap, fp8, fp8_meta, cp_group, @@ -1664,7 +1665,7 @@ def forward( if fa_utils.v2_5_7_plus and qkv_format == "thd": fa_forward_kwargs["block_table"] = None if fa_utils.v2_6_0_plus: - fa_forward_kwargs["softcap"] = 0.0 + fa_forward_kwargs["softcap"] = softcap # set up inputs for forward q_inputs = [None, None] @@ -2156,6 +2157,7 @@ def forward( ctx.attn_bias_type = attn_bias_type ctx.attn_bias_shape = None if attn_bias is None else attn_bias.shape ctx.deterministic = deterministic + ctx.softcap = softcap ctx.use_fused_attention = use_fused_attention ctx.pad_between_seqs = pad_between_seqs ctx.softmax_lse_in_packed_format = softmax_lse_in_packed_format @@ -2454,7 +2456,7 @@ def backward(ctx, dout, *_args): if fa_utils.v2_4_1_plus: fa_backward_kwargs["deterministic"] = ctx.deterministic if fa_utils.v2_6_0_plus: - fa_backward_kwargs["softcap"] = 0.0 + fa_backward_kwargs["softcap"] = ctx.softcap send_recv_reqs = [] for i in range(cp_size): @@ -2970,6 +2972,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -3047,6 +3050,7 @@ def forward( deterministic, use_fused_attention, return_max_logit, + softcap, window_size, cp_group, cp_stream, @@ -3128,7 +3132,7 @@ def forward( if fa_utils.v2_5_7_plus and qkv_format == "thd": fa_forward_kwargs["block_table"] = None if fa_utils.v2_6_0_plus: - fa_forward_kwargs["softcap"] = 0.0 + fa_forward_kwargs["softcap"] = softcap qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format @@ -3644,6 +3648,7 @@ def forward( ctx.attn_bias_type = attn_bias_type ctx.attn_mask_type = attn_mask_type ctx.deterministic = deterministic + ctx.softcap = softcap ctx.use_fused_attention = use_fused_attention ctx.use_flash_attn_3 = use_flash_attn_3 ctx.pad_between_seqs = pad_between_seqs @@ -3840,7 +3845,7 @@ def backward(ctx, dout, *_args): if fa_utils.v2_4_1_plus: fa_backward_kwargs["deterministic"] = ctx.deterministic if fa_utils.v2_6_0_plus: - fa_backward_kwargs["softcap"] = 0.0 + fa_backward_kwargs["softcap"] = ctx.softcap local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1] for i in range(len(local_seq_chunk_ids) + 1): @@ -4164,6 +4169,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -4195,6 +4201,7 @@ def forward( deterministic, use_fused_attention, return_max_logit, + softcap, window_size, fp8, fp8_meta, @@ -4284,7 +4291,7 @@ def forward( if fa_utils.v2_5_7_plus and qkv_format == "thd": fa_forward_kwargs["block_table"] = None if fa_utils.v2_6_0_plus: - fa_forward_kwargs["softcap"] = 0.0 + fa_forward_kwargs["softcap"] = softcap assert isinstance(k, q.__class__) and isinstance( v, q.__class__ @@ -4585,6 +4592,7 @@ def forward( ctx.attn_mask_type = attn_mask_type ctx.attn_bias_type = attn_bias_type ctx.deterministic = deterministic + ctx.softcap = softcap ctx.window_size = window_size ctx.use_fused_attention = use_fused_attention ctx.fp8_meta = fp8_meta @@ -4725,7 +4733,7 @@ def backward(ctx, dout, *_args): if fa_utils.v2_4_1_plus: fa_backward_kwargs["deterministic"] = ctx.deterministic if fa_utils.v2_6_0_plus: - fa_backward_kwargs["softcap"] = 0.0 + fa_backward_kwargs["softcap"] = ctx.softcap dq_fp8, dk_fp8, dv_fp8 = None, None, None if ctx.use_fused_attention: @@ -4916,6 +4924,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, d_softmax_offset, None, ) @@ -4945,6 +4954,7 @@ def attn_forward_func_with_cp( deterministic=False, use_fused_attention=False, window_size=None, + softcap=0.0, fp8=False, fp8_meta=None, quantizers=None, @@ -5091,6 +5101,7 @@ def attn_forward_func_with_cp( deterministic, use_fused_attention, return_max_logit, + softcap, ] if cp_comm_type in ["p2p", "a2a+p2p"]: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index d5adbbcadf..fad2c51951 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -578,6 +578,12 @@ def nvfp4_linear_mxfp8_dpa_factory(role): or bottom right (`True`) corner of the softmax matrix in the encoder. If `None`, it will be set to `False` for `attn_mask_type` = {'causal', 'padding_causal'} and `True` for other mask types. + softcap : float, default = 0.0 + tanh logit softcapping value applied to the attention scores as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables + softcapping. Softcapping is only supported by the FlashAttention + backend. Similar to :attr:`window_size`, ``softcap`` can be + overridden by :attr:`softcap` in ``forward`` as well. attention_type : str, default = "self" type of attention, either ``"self"`` and ``"cross"``. layer_number : int, default = None @@ -677,6 +683,7 @@ def __init__( attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: float = 0.0, sequence_parallel: bool = False, tp_size: int = 1, get_rng_state_tracker: Optional[Callable] = None, @@ -713,6 +720,7 @@ def __init__( self.attn_mask_type = attn_mask_type self.window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size) self.bottom_right_diagonal = bottom_right_diagonal + self.softcap = softcap if tp_group is None: self.tp_size = tp_size if tp_size == 1: @@ -1393,6 +1401,7 @@ def forward( attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: Optional[float] = None, checkpoint_core_attention: bool = False, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, @@ -1563,6 +1572,11 @@ def forward( causal masks are aligned to the bottom right corner. window_size: Optional[Tuple[int, int]], default = None Sliding window size for local attention. + softcap: Optional[float], default = None + tanh logit softcapping value applied to the attention scores as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables + softcapping. When `None`, the value passed to the constructor is used. + Softcapping is only supported by the FlashAttention backend. bottom_right_diagonal: Optional[bool], default = None Align sliding window and ALiBi diagonal to the top left (`False`) or bottom right (`True`) corner of the softmax matrix in the encoder. @@ -1743,6 +1757,8 @@ def forward( if window_size is None: window_size = self.window_size window_size = dpa_utils.check_set_window_size(attn_mask_type, window_size) + if softcap is None: + softcap = self.softcap if bottom_right_diagonal is None: bottom_right_diagonal = self.bottom_right_diagonal if attn_mask_type in {"causal", "padding_causal"}: @@ -2025,6 +2041,14 @@ def forward( else: pad_between_seqs = False + # softcap is not supported by UnfusedDotProductAttention, which is the backend ONNX + # export unconditionally force-selects further down (bypassing get_attention_backend's + # softcap-aware filter). Fail loudly rather than silently export a model that omits + # softcapping. + assert ( + softcap == 0.0 or not is_in_onnx_export_mode() + ), "Attention logit softcapping (softcap != 0.0) is not supported with ONNX export!" + # Validate experimental Flex Attention API inputs that backend selection # cannot represent. if score_mod is None: @@ -2074,6 +2098,7 @@ def forward( attn_mask_type=attn_mask_type, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, + softcap=softcap, alibi_slopes_shape=alibi_slopes.shape if alibi_slopes is not None else None, core_attention_bias_type=core_attention_bias_type, core_attention_bias_shape=core_attention_bias_shape, @@ -2205,6 +2230,7 @@ def forward( cu_seqlens_kv=cu_seqlens_kv, attn_mask_type=attn_mask_type, window_size=window_size, + softcap=softcap, alibi_slopes=alibi_slopes, cp_group=self.cp_group, cp_global_ranks=self.cp_global_ranks, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ba049c9aef..aeab47007e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -148,6 +148,11 @@ class FlashAttentionUtils: v4_is_installed = False fa4_version = PkgVersion("0") use_v4 = False + # True only if the installed FA3 build exposes a `softcap` parameter (signature probe in + # backends.py, fail-closed default False). Necessary-but-not-sufficient: FA3 softcap is also + # gated on opt-in (NVTE_FA3_SOFTCAP=1) and head_dim <= 256 in get_attention_backend. FA3 is + # already restricted to Hopper (sm90) upstream, where its softcap fwd+bwd kernels are mature. + fa3_supports_softcap = False v4_installation_steps = """\ pip install flash-attn-4==4.0.0b11 nvidia-cutlass-dsl[cu13]""" v4_warning_printed = False @@ -229,6 +234,9 @@ class AttentionParams: bottom_right_diagonal: bool, default = `None` Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. + softcap : float, default = 0.0 + Tanh logit softcapping value applied to the attention scores. A value of + ``0.0`` disables softcapping. Only supported by the FlashAttention backend. alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias @@ -289,6 +297,7 @@ class AttentionParams: attn_mask_type: str = "no_mask" window_size: Union[Tuple[int, int], None] = None bottom_right_diagonal: bool = True + softcap: float = 0.0 alibi_slopes_shape: Union[torch.Size, List, None] = None core_attention_bias_type: str = "no_bias" core_attention_bias_shape: str = "1hss" @@ -433,6 +442,7 @@ def get_attention_backend( attn_mask_type = attention_params.attn_mask_type window_size = attention_params.window_size bottom_right_diagonal = attention_params.bottom_right_diagonal + softcap = attention_params.softcap alibi_slopes_shape = attention_params.alibi_slopes_shape core_attention_bias_type = attention_params.core_attention_bias_type core_attention_bias_shape = attention_params.core_attention_bias_shape @@ -764,6 +774,43 @@ def _disable_all_flash_attention() -> None: use_unfused_attention = False logger.debug("Disabling all backends for max_logit with FP8 attention") + # Filter: softcap + # The scalar `softcap` kwarg (tanh logit softcapping) is plumbed to the FlashAttention 2 + # backend (>= 2.6.0) by default, and to FA3 only behind an explicit opt-in gate below. + # FusedAttention/unfused don't take the scalar kwarg (cuDNN can softcap via score_mod, but that + # path is not used here). Steer selection to FA2 rather than (a) hitting a runtime + # NotImplementedError when an unwired backend is selected, or (b) silently dropping the cap. + if softcap != 0.0: + if use_fused_attention: + logger.debug("Disabling FusedAttention as it does not support softcap") + use_fused_attention = False + if use_unfused_attention: + logger.debug("Disabling UnfusedDotProductAttention as it does not support softcap") + use_unfused_attention = False + if use_flash_attention_3 and not ( + FlashAttentionUtils.fa3_supports_softcap + and os.getenv("NVTE_FA3_SOFTCAP", "0") == "1" + and max(head_dim_qk, head_dim_v) <= 256 + and not context_parallel + ): + # FA3 softcap is opt-in (NVTE_FA3_SOFTCAP=1) and requires a softcap-capable FA3 build, + # head_dim <= 256 (the range FA3's sm90 softcap kernels are instantiated for), and no + # context parallelism -- FA3's CP path hard-rejects nonzero softcap (backends.py), so + # selecting it here would just crash at dispatch instead of steering to FA2, which does + # support CP+softcap via context_parallel.py's autograd threading. FA3 is already + # Hopper-only upstream. FA3's non-CP softcap fwd+bwd is mature, so no arch/beta caveat is + # needed beyond the build probe; keep it opt-in to preserve FA2 as the default (unchanged + # behavior) and allow a clean FA2-vs-FA3 comparison. When all conditions hold, FA3 + # survives and the softcap kwarg is threaded in backends.py. + logger.debug( + "Disabling FlashAttention 3 for softcap (requires softcap-capable FA3 build, " + "NVTE_FA3_SOFTCAP=1, head_dim <= 256, and no context parallelism)" + ) + use_flash_attention_3 = False + if use_flash_attention_2 and not FlashAttentionUtils.v2_6_0_plus: + logger.debug("Disabling FlashAttention 2 for softcap (requires flash-attn >= 2.6.0)") + use_flash_attention_2 = False + # Filter: score_mod if has_score_mod_bprop and not has_score_mod: logger.debug("Disabling all backends because score_mod_bprop requires score_mod") From 9839ae08ac71d6bfa7bbda5752eb62334685fd21 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 17 Aug 2026 14:34:54 -0700 Subject: [PATCH 02/28] fix: use raise instead of assert for the ONNX+softcap guard python -O / PYTHONOPTIMIZE strips assert statements, which would silently reopen the ONNX export softcap-drop bug the previous commit fixed (ONNX mode would again force-select UnfusedDotProductAttention with softcap silently omitted, with no error). Switch to an explicit if/raise ValueError, which survives optimized execution. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- .../dot_product_attention/dot_product_attention.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index fad2c51951..8ca8fbef3a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -2044,10 +2044,13 @@ def forward( # softcap is not supported by UnfusedDotProductAttention, which is the backend ONNX # export unconditionally force-selects further down (bypassing get_attention_backend's # softcap-aware filter). Fail loudly rather than silently export a model that omits - # softcapping. - assert ( - softcap == 0.0 or not is_in_onnx_export_mode() - ), "Attention logit softcapping (softcap != 0.0) is not supported with ONNX export!" + # softcapping. Uses an explicit raise (not assert) so the check survives python -O / + # PYTHONOPTIMIZE, which strips asserts and would otherwise silently re-open this gap. + if softcap != 0.0 and is_in_onnx_export_mode(): + raise ValueError( + "Attention logit softcapping (softcap != 0.0) is not supported with " + "ONNX export!" + ) # Validate experimental Flex Attention API inputs that backend selection # cannot represent. From eb215b51decb54891895d92c516360e60ce377fc Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 17 Aug 2026 14:35:45 -0700 Subject: [PATCH 03/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci (reapplied after a force-push rebase clobbered pre-commit.ci's original 19a21eb7 commit; same content, restored by hand) Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_softcap.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/pytorch/attention/test_softcap.py b/tests/pytorch/attention/test_softcap.py index b0b94bb4e8..78a8cf2829 100644 --- a/tests/pytorch/attention/test_softcap.py +++ b/tests/pytorch/attention/test_softcap.py @@ -44,9 +44,7 @@ def _flash_attn_2_6_available() -> bool: # backend (>= 2.6.0), and requires CUDA tensors. pytestmark = [ pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required."), - pytest.mark.skipif( - not _flash_attn_2_6_available(), reason="flash-attn >= 2.6.0 is required." - ), + pytest.mark.skipif(not _flash_attn_2_6_available(), reason="flash-attn >= 2.6.0 is required."), ] From 2475521e4a2b284bd7d44ff98f18f40688c3811b Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 27 Aug 2026 09:22:15 -0700 Subject: [PATCH 04/28] refactor: move softcap reference into UnfusedDotProductAttention and fold test into test_attention.py UnfusedDotProductAttention now applies softcap * tanh(scores / softcap) to the already-scaled logits, matching how FlashAttention folds softmax_scale into its tanh argument, so it can serve as the softcap reference backend. Backend selection therefore no longer disqualifies unfused attention for softcap, and the ONNX-export guard is dropped since the export path force-selects unfused and torch.tanh is exportable. test_softcap.py is replaced by a model_configs_softcap dict and test_dpa_softcap in test_attention.py, which reuses test_dot_product_attention for backend sweeping. Signed-off-by: Nitin Vegesna --- qa/L0_pytorch_unittest/test.sh | 1 - tests/pytorch/attention/test_attention.py | 26 ++++ tests/pytorch/attention/test_softcap.py | 147 ------------------ tests/pytorch/utils.py | 3 + .../dot_product_attention/backends.py | 8 + .../dot_product_attention.py | 21 +-- .../attention/dot_product_attention/utils.py | 14 +- 7 files changed, 50 insertions(+), 170 deletions(-) delete mode 100644 tests/pytorch/attention/test_softcap.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 482642e5e1..14a5f4fe3d 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -67,7 +67,6 @@ NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py" 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_softcap.xml $TE_PATH/tests/pytorch/attention/test_softcap.py || test_fail "test_softcap.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_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 diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index bfd2cdf9fd..8e8e557405 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -645,6 +645,31 @@ def test_dpa_softmax_thd(dtype, model_configs, model): test_dot_product_attention(dtype, model_configs, model, True, "thd_thd_thd", False, False) +model_configs_softcap = { + # test: ModelConfig(b, sq, hq, dqk) + "softcap_1_0": ModelConfig(4, 128, 16, 64, softcap=50.0), + "softcap_1_1": ModelConfig(4, 128, 16, 64, num_gqa_groups=4, softcap=50.0), + "softcap_2_0": ModelConfig(2, 512, 16, 64, attn_mask_type="causal", softcap=50.0), + "softcap_2_1": ModelConfig(2, 512, 24, 128, attn_mask_type="padding_causal", softcap=50.0), + # 0.01 is on the order of the logits these inputs produce, so tanh runs in its nonlinear + # region instead of acting as a no-op, and a misapplied softmax_scale or a missing outer + # softcap factor changes the output. + "softcap_3_0": ModelConfig(4, 128, 16, 64, softcap=0.01), + "softcap_3_1": ModelConfig(2, 512, 16, 64, attn_mask_type="causal", softcap=0.01), +} + + +@pytest.mark.skipif( + not FlashAttentionUtils.v2_6_0_plus, reason="flash-attn 2.6.0+ is required for softcap." +) +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_softcap]) +@pytest.mark.parametrize("model", model_configs_softcap.keys()) +def test_dpa_softcap(dtype, model_configs, model): + """Test DotProductAttention module with tanh logit softcapping""" + test_dot_product_attention(dtype, model_configs, model, False, "bshd_bshd_bshd", False, False) + + model_configs_mla = { # test: ModelConfig(b, sq, hq, dqk) "mla_1_0": ModelConfig(8, 128, 16, 64, head_dim_v=128), @@ -1447,6 +1472,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attention_type=config.attn_type, softmax_type=config.softmax_type, return_max_logit=config.return_max_logit, + softcap=config.softcap, ).to(dtype=dtype, device="cuda") if not is_training: block = block.eval() diff --git a/tests/pytorch/attention/test_softcap.py b/tests/pytorch/attention/test_softcap.py deleted file mode 100644 index 78a8cf2829..0000000000 --- a/tests/pytorch/attention/test_softcap.py +++ /dev/null @@ -1,147 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. -"""Isolation numerics test for tanh logit softcapping in DotProductAttention. - -The reference implements softcapping in pure PyTorch: - - scores = (Q @ K^T) * scale - scores = softcap * tanh(scores / softcap) # only when softcap != 0.0 - scores = scores + mask - attn = softmax(scores) - out = attn @ V - -and is compared against ``DotProductAttention(..., softcap=...)`` forced onto the -FlashAttention backend, for both the forward output and the input gradients -(dQ/dK/dV obtained via autograd). -""" - -import sys -import pathlib - -import pytest -import torch -from packaging.version import Version as PkgVersion - -from transformer_engine.pytorch import DotProductAttention -from transformer_engine.pytorch.attention.dot_product_attention import _attention_backends - -_current_file = pathlib.Path(__file__).resolve() -sys.path = [str(_current_file.parent.parent)] + sys.path -from utils import reset_rng_states # pylint: disable=wrong-import-position - - -def _flash_attn_2_6_available() -> bool: - """Whether flash-attn >= 2.6.0 (the first version exposing ``softcap``) is installed.""" - try: - import flash_attn # pylint: disable=import-outside-toplevel - except ImportError: - return False - return PkgVersion(flash_attn.__version__) >= PkgVersion("2.6.0") - - -# Softcapping through DotProductAttention is only wired through the FlashAttention 2 -# backend (>= 2.6.0), and requires CUDA tensors. -pytestmark = [ - pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required."), - pytest.mark.skipif(not _flash_attn_2_6_available(), reason="flash-attn >= 2.6.0 is required."), -] - - -def _force_flash_backend() -> None: - """Force DotProductAttention to select the FlashAttention backend.""" - import os # pylint: disable=import-outside-toplevel - - os.environ["NVTE_FLASH_ATTN"] = "1" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "0" - _attention_backends["backend_selection_requires_update"] = True - - -def _reference_attention(q, k, v, scale, softcap, causal): - """Pure-PyTorch reference for softcapped scaled dot product attention. - - q, k, v are in ``bshd`` layout. GQA is supported: ``k``/``v`` may have fewer - heads than ``q``. - """ - # bshd -> bhsd - qt = q.transpose(1, 2).float() - kt = k.transpose(1, 2).float() - vt = v.transpose(1, 2).float() - - num_heads = qt.shape[1] - num_gqa_groups = kt.shape[1] - if num_heads != num_gqa_groups: - assert num_heads % num_gqa_groups == 0 - repeats = num_heads // num_gqa_groups - kt = kt.repeat_interleave(repeats, dim=1) - vt = vt.repeat_interleave(repeats, dim=1) - - scores = torch.matmul(qt, kt.transpose(-2, -1)) * scale - if softcap != 0.0: - scores = softcap * torch.tanh(scores / softcap) - if causal: - sq, skv = scores.shape[-2], scores.shape[-1] - mask = torch.triu( - torch.ones(sq, skv, dtype=torch.bool, device=scores.device), - diagonal=1 + skv - sq, - ) - scores = scores.masked_fill(mask, float("-inf")) - attn = torch.softmax(scores, dim=-1) - out = torch.matmul(attn, vt) - # bhsd -> bshd - return out.transpose(1, 2) - - -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -@pytest.mark.parametrize("softcap", [0.0, 50.0]) -@pytest.mark.parametrize("num_gqa_groups", [4, 2]) -@pytest.mark.parametrize("causal", [False, True]) -def test_softcap_numerics(dtype, softcap, num_gqa_groups, causal): - """FlashAttention softcap forward + grads match a pure-PyTorch reference. - - ``softcap == 0.0`` additionally proves that softcapping is a no-op relative to - the plain (no-softcap) reference, i.e. today's behavior is reproduced exactly. - """ - reset_rng_states() - - batch_size = 2 - max_seqlen = 32 - num_heads = 4 - head_dim = 64 - scale = 1.0 / (head_dim**0.5) - - q_shape = (batch_size, max_seqlen, num_heads, head_dim) - kv_shape = (batch_size, max_seqlen, num_gqa_groups, head_dim) - - q = (0.5 * torch.randn(q_shape, dtype=dtype, device="cuda")).requires_grad_() - k = (0.5 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() - v = (0.5 * torch.randn(kv_shape, dtype=dtype, device="cuda")).requires_grad_() - q_ref, k_ref, v_ref = [x.detach().clone().requires_grad_() for x in (q, k, v)] - - grad_output = torch.randn(q_shape, dtype=dtype, device="cuda") - - _force_flash_backend() - dpa = DotProductAttention( - num_heads, - head_dim, - num_gqa_groups=num_gqa_groups, - qkv_format="bshd", - attn_mask_type="causal" if causal else "no_mask", - softmax_scale=scale, - softcap=softcap, - layer_number=1, - ).to(dtype=dtype, device="cuda") - - out = dpa(q, k, v) - out.backward(grad_output) - - out_ref = _reference_attention(q_ref, k_ref, v_ref, scale, softcap, causal) - out_ref.backward(grad_output.float()) - - atol, rtol = (2e-2, 2e-2) if dtype == torch.float16 else (3.5e-2, 3.5e-2) - - torch.testing.assert_close(out.float(), out_ref.float(), atol=atol, rtol=rtol) - torch.testing.assert_close(q.grad.float(), q_ref.grad.float(), atol=atol, rtol=rtol) - torch.testing.assert_close(k.grad.float(), k_ref.grad.float(), atol=atol, rtol=rtol) - torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), atol=atol, rtol=rtol) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 21601d8cdd..0002bcef2c 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -282,6 +282,7 @@ def __init__( alibi_type: str = "none", bias_shape: str = "1hss", window_size: Tuple[int, int] = (-1, -1), + softcap: float = 0.0, context_parallel: bool = False, cp_comm_type: str = "p2p", return_max_logit=False, @@ -312,6 +313,7 @@ def __init__( self.attn_type = "self" if (self.max_seqlen_q == self.max_seqlen_kv) else "cross" self.bias_shape = bias_shape self.window_size = check_set_window_size(self.attn_mask_type, window_size) + self.softcap = softcap self.context_parallel = context_parallel self.cp_comm_type = cp_comm_type self.return_max_logit = return_max_logit @@ -390,6 +392,7 @@ def test(): head_dim_v=config.head_dim_v, attn_mask_type=config.attn_mask_type, window_size=config.window_size, + softcap=config.softcap, alibi_slopes_shape=alibi_slopes_shape, core_attention_bias_type=config.attn_bias_type, core_attention_bias_shape=core_attention_bias_shape, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 12a2ed1492..c74e0f1f04 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -440,6 +440,7 @@ def _forward( attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: float = 0.0, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, alibi_slopes: Optional[torch.Tensor] = None, @@ -678,6 +679,13 @@ def _forward( dtype=query_layer.dtype ) + # Cap the scaled logits -- softcap * tanh(scores * scale / softcap) -- matching how + # FlashAttention folds softmax_scale into its tanh argument. qk layer scaling defers the + # layer_number factor to the softmax below, so it is divided out of the cap here. + if softcap != 0.0: + cap = softcap / self.layer_number if apply_qk_layer_scaling else softcap + matmul_result = cap * torch.tanh(matmul_result / cap) + if fp8: # quantize and dequantize dP to emulate FP8 matmul_result, *_ = FP8EmulationFunc.apply( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 8ca8fbef3a..2b51ed8f2c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -582,8 +582,9 @@ def nvfp4_linear_mxfp8_dpa_factory(role): tanh logit softcapping value applied to the attention scores as ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. Softcapping is only supported by the FlashAttention - backend. Similar to :attr:`window_size`, ``softcap`` can be - overridden by :attr:`softcap` in ``forward`` as well. + and UnfusedDotProductAttention backends. Similar to + :attr:`window_size`, ``softcap`` can be overridden by + :attr:`softcap` in ``forward`` as well. attention_type : str, default = "self" type of attention, either ``"self"`` and ``"cross"``. layer_number : int, default = None @@ -1576,7 +1577,8 @@ def forward( tanh logit softcapping value applied to the attention scores as ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. When `None`, the value passed to the constructor is used. - Softcapping is only supported by the FlashAttention backend. + Softcapping is only supported by the FlashAttention and + UnfusedDotProductAttention backends. bottom_right_diagonal: Optional[bool], default = None Align sliding window and ALiBi diagonal to the top left (`False`) or bottom right (`True`) corner of the softmax matrix in the encoder. @@ -2041,17 +2043,6 @@ def forward( else: pad_between_seqs = False - # softcap is not supported by UnfusedDotProductAttention, which is the backend ONNX - # export unconditionally force-selects further down (bypassing get_attention_backend's - # softcap-aware filter). Fail loudly rather than silently export a model that omits - # softcapping. Uses an explicit raise (not assert) so the check survives python -O / - # PYTHONOPTIMIZE, which strips asserts and would otherwise silently re-open this gap. - if softcap != 0.0 and is_in_onnx_export_mode(): - raise ValueError( - "Attention logit softcapping (softcap != 0.0) is not supported with " - "ONNX export!" - ) - # Validate experimental Flex Attention API inputs that backend selection # cannot represent. if score_mod is None: @@ -2365,6 +2356,7 @@ def forward( attention_mask=attention_mask, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, + softcap=softcap, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, alibi_slopes=alibi_slopes, @@ -2389,6 +2381,7 @@ def forward( attention_mask=attention_mask, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, + softcap=softcap, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, alibi_slopes=alibi_slopes, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index aeab47007e..94c2fc4d99 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -236,7 +236,8 @@ class AttentionParams: of the softmax matrix. softcap : float, default = 0.0 Tanh logit softcapping value applied to the attention scores. A value of - ``0.0`` disables softcapping. Only supported by the FlashAttention backend. + ``0.0`` disables softcapping. Only supported by the FlashAttention and + UnfusedDotProductAttention backends. alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias @@ -776,17 +777,14 @@ def _disable_all_flash_attention() -> None: # Filter: softcap # The scalar `softcap` kwarg (tanh logit softcapping) is plumbed to the FlashAttention 2 - # backend (>= 2.6.0) by default, and to FA3 only behind an explicit opt-in gate below. - # FusedAttention/unfused don't take the scalar kwarg (cuDNN can softcap via score_mod, but that - # path is not used here). Steer selection to FA2 rather than (a) hitting a runtime - # NotImplementedError when an unwired backend is selected, or (b) silently dropping the cap. + # backend (>= 2.6.0) and to UnfusedDotProductAttention by default, and to FA3 only behind an + # explicit opt-in gate below. FusedAttention does not take the scalar kwarg (cuDNN can softcap + # via score_mod, but that path is not used here), so disable it rather than silently dropping + # the cap. if softcap != 0.0: if use_fused_attention: logger.debug("Disabling FusedAttention as it does not support softcap") use_fused_attention = False - if use_unfused_attention: - logger.debug("Disabling UnfusedDotProductAttention as it does not support softcap") - use_unfused_attention = False if use_flash_attention_3 and not ( FlashAttentionUtils.fa3_supports_softcap and os.getenv("NVTE_FA3_SOFTCAP", "0") == "1" From 3c5eb4a254b616f00a0421b7ca672d5144f70a99 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 27 Aug 2026 09:28:45 -0700 Subject: [PATCH 05/28] fix(pytorch): gate FA3 softcap on existing NVTE_FLASH_ATTN_V3 Drop the redundant NVTE_FA3_SOFTCAP opt-in. `use_flash_attention_3` already derives from NVTE_FLASH_ATTN_V3, so the existing flag governs the FA3 softcap path and NVTE_FLASH_ATTN_V3=0 disables it. Correctness stays established by the build-capability probe, head_dim <= 256, and the non-CP requirement. Signed-off-by: Nitin Vegesna --- .../dot_product_attention/backends.py | 10 +++---- .../attention/dot_product_attention/utils.py | 26 ++++++++----------- 2 files changed, 16 insertions(+), 20 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index c74e0f1f04..d2efd9048d 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -170,7 +170,7 @@ # (sm90) kernels DO implement tanh logit softcapping in fwd AND bwd (dedicated # flash_{fwd,bwd}_hdim256_bf16_softcap_sm90 instantiations, off only behind a compile-time # DISABLE_SOFTCAP flag), so this is a mature path. Still fail-closed and additionally - # gated on opt-in (NVTE_FA3_SOFTCAP) + head_dim <= 256 in get_attention_backend. + # gated on head_dim <= 256 + non-CP in get_attention_backend. try: fa_utils.fa3_supports_softcap = ( "softcap" in inspect.signature(flash_attn_func_v3).parameters @@ -1267,10 +1267,10 @@ def forward( ) else: # Fail-loud net: get_attention_backend only keeps FA3 for softcap on a - # softcap-capable build (signature probe) + opt-in (NVTE_FA3_SOFTCAP) + Hopper - # (FA3 is sm90-only upstream) + head_dim <= 256. If FA3 is still reached with - # softcap while the build lacks support (force-selected / regressed path), raise - # rather than silently drop the cap. The non-CP FA3 entry points + # softcap-capable build (signature probe) + Hopper (FA3 is sm90-only upstream) + # + head_dim <= 256. If FA3 is still reached with softcap while the build lacks + # support (force-selected / regressed path), raise rather than silently drop the + # cap. The non-CP FA3 entry points # (flash_attn_func_v3 / flash_attn_varlen_func_v3) are self-contained autograd # functions, so threading `softcap` into the forward call also drives the # matching FA3 softcap backward kernel. (CP + FA3 + softcap stays blocked above.) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 94c2fc4d99..ea0329f557 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -150,8 +150,8 @@ class FlashAttentionUtils: use_v4 = False # True only if the installed FA3 build exposes a `softcap` parameter (signature probe in # backends.py, fail-closed default False). Necessary-but-not-sufficient: FA3 softcap is also - # gated on opt-in (NVTE_FA3_SOFTCAP=1) and head_dim <= 256 in get_attention_backend. FA3 is - # already restricted to Hopper (sm90) upstream, where its softcap fwd+bwd kernels are mature. + # gated on head_dim <= 256 and non-CP in get_attention_backend. FA3 is already restricted to + # Hopper (sm90) upstream, where its softcap fwd+bwd kernels are mature. fa3_supports_softcap = False v4_installation_steps = """\ pip install flash-attn-4==4.0.0b11 nvidia-cutlass-dsl[cu13]""" @@ -777,8 +777,8 @@ def _disable_all_flash_attention() -> None: # Filter: softcap # The scalar `softcap` kwarg (tanh logit softcapping) is plumbed to the FlashAttention 2 - # backend (>= 2.6.0) and to UnfusedDotProductAttention by default, and to FA3 only behind an - # explicit opt-in gate below. FusedAttention does not take the scalar kwarg (cuDNN can softcap + # backend (>= 2.6.0) and to UnfusedDotProductAttention by default, and to FA3 subject to the + # build/shape checks below. FusedAttention does not take the scalar kwarg (cuDNN can softcap # via score_mod, but that path is not used here), so disable it rather than silently dropping # the cap. if softcap != 0.0: @@ -787,22 +787,18 @@ def _disable_all_flash_attention() -> None: use_fused_attention = False if use_flash_attention_3 and not ( FlashAttentionUtils.fa3_supports_softcap - and os.getenv("NVTE_FA3_SOFTCAP", "0") == "1" and max(head_dim_qk, head_dim_v) <= 256 and not context_parallel ): - # FA3 softcap is opt-in (NVTE_FA3_SOFTCAP=1) and requires a softcap-capable FA3 build, - # head_dim <= 256 (the range FA3's sm90 softcap kernels are instantiated for), and no - # context parallelism -- FA3's CP path hard-rejects nonzero softcap (backends.py), so - # selecting it here would just crash at dispatch instead of steering to FA2, which does - # support CP+softcap via context_parallel.py's autograd threading. FA3 is already - # Hopper-only upstream. FA3's non-CP softcap fwd+bwd is mature, so no arch/beta caveat is - # needed beyond the build probe; keep it opt-in to preserve FA2 as the default (unchanged - # behavior) and allow a clean FA2-vs-FA3 comparison. When all conditions hold, FA3 - # survives and the softcap kwarg is threaded in backends.py. + # FA3 softcap requires a softcap-capable FA3 build, head_dim <= 256 (the range FA3's + # sm90 softcap kernels are instantiated for), and no context parallelism -- FA3's CP + # path hard-rejects nonzero softcap (backends.py), so selecting it here would just + # crash at dispatch instead of steering to FA2, which does support CP+softcap via + # context_parallel.py's autograd threading. Whether FA3 is eligible at all is governed + # by NVTE_FLASH_ATTN_V3 through use_flash_attention_3. logger.debug( "Disabling FlashAttention 3 for softcap (requires softcap-capable FA3 build, " - "NVTE_FA3_SOFTCAP=1, head_dim <= 256, and no context parallelism)" + "head_dim <= 256, and no context parallelism)" ) use_flash_attention_3 = False if use_flash_attention_2 and not FlashAttentionUtils.v2_6_0_plus: From 900371a61dbc0a9573e79532738f293f5706b8bf Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 27 Aug 2026 09:45:02 -0700 Subject: [PATCH 06/28] fix(pytorch): disable FlashAttention 4 for softcap FA4 exposes no softcap kwarg and its head_dim=256 kernel asserts score_mod is None, so there is no kernel to route the cap through. The FA4 call path in backends.py passes no softcap, so an FA4 selection with a nonzero softcap silently dropped the cap instead of failing closed. NVTE_FLASH_ATTN_V4 defaults to enabled, so this was reachable on SM100+ with flash-attn v4 installed and no context parallelism. Signed-off-by: Nitin Vegesna --- .../pytorch/attention/dot_product_attention/utils.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ea0329f557..75007e60b3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -779,12 +779,19 @@ def _disable_all_flash_attention() -> None: # The scalar `softcap` kwarg (tanh logit softcapping) is plumbed to the FlashAttention 2 # backend (>= 2.6.0) and to UnfusedDotProductAttention by default, and to FA3 subject to the # build/shape checks below. FusedAttention does not take the scalar kwarg (cuDNN can softcap - # via score_mod, but that path is not used here), so disable it rather than silently dropping - # the cap. + # via score_mod, but that path is not used here), and FA4 has no softcap kernel to call, so + # disable both rather than silently dropping the cap. if softcap != 0.0: if use_fused_attention: logger.debug("Disabling FusedAttention as it does not support softcap") use_fused_attention = False + if use_flash_attention_4: + # FA4 exposes no softcap kwarg and its head_dim=256 kernel asserts score_mod is None, + # so there is no kernel to route the cap through, and the FA4 call path in backends.py + # passes no softcap -- selecting it here would silently drop the cap. + if FlashAttentionUtils.v4_is_installed: + logger.debug("Disabling FlashAttention 4 as it does not support softcap") + use_flash_attention_4 = False if use_flash_attention_3 and not ( FlashAttentionUtils.fa3_supports_softcap and max(head_dim_qk, head_dim_v) <= 256 From 5ecabacd816ec754bb1cc9f389e8dceb22a8a93f Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 27 Aug 2026 09:45:46 -0700 Subject: [PATCH 07/28] fix(pytorch): disable FlashAttention 2 for softcap with dropout flash-attn rejects a nonzero softcap combined with nonzero dropout at dispatch: "Softcapping does not support dropout for now" in csrc/flash_attn/flash_api.cpp, present in mha_fwd and mha_varlen_fwd from v2.6.0 (the earliest version TE allows softcap on) onwards. Backend selection did not model this, so a softcap + attention-dropout config passed selection, routed to FA2, and crashed inside flash-attn. Dropout only reaches the kernel while training, since backends.py passes `self.attention_dropout if self.training else 0.0`, so the gate is on `attention_dropout != 0.0 and is_training` to avoid blocking valid inference configs. UnfusedDotProductAttention supports both softcap and dropout and stays available, so this steers rather than hard-fails. Signed-off-by: Nitin Vegesna --- .../pytorch/attention/dot_product_attention/utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 75007e60b3..6c0ac25006 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -811,6 +811,12 @@ def _disable_all_flash_attention() -> None: if use_flash_attention_2 and not FlashAttentionUtils.v2_6_0_plus: logger.debug("Disabling FlashAttention 2 for softcap (requires flash-attn >= 2.6.0)") use_flash_attention_2 = False + if use_flash_attention_2 and attention_dropout != 0.0 and is_training: + # FA2 hard-rejects a nonzero softcap combined with nonzero dropout at dispatch + # ("Softcapping does not support dropout for now", flash_api.cpp). Dropout only reaches + # the kernel while training -- backends.py passes 0.0 in eval -- hence the is_training. + logger.debug("Disabling FlashAttention 2 for softcap with dropout") + use_flash_attention_2 = False # Filter: score_mod if has_score_mod_bprop and not has_score_mod: From 4cc2e9c773dc10a444768b8b05a781f8f49425fb Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 27 Aug 2026 10:16:16 -0700 Subject: [PATCH 08/28] test: restore softcap dQ/dK/dV parity in the shared DPA harness test_dot_product_attention forced is_training=False whenever FusedAttention could not train a config, so that backends only available for inference could still be compared. softcap always disables FusedAttention, so test_dpa_softcap silently degraded to a forward-only comparison and the PR's backward-parity claim -- the FA2 softcap backward kernel included -- went untested. Add fwd_only_without_fused_attn (default True, so every other caller is byte-for-byte unchanged) and opt test_dpa_softcap out, which pairs FlashAttention against UnfusedDotProductAttention with is_training=True and restores the dgrad comparison. Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 9e3a078817..b9e446cd37 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -170,6 +170,7 @@ def test_dot_product_attention( pad_between_seqs, declarative_packed=False, is_training=True, + fwd_only_without_fused_attn=True, ): """Test DotProductAttention module""" @@ -222,7 +223,11 @@ def test_dot_product_attention( ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - if not fused_attn_supported: + # Some backends are only available in inference mode, so when FusedAttention cannot train this + # config the query is repeated forward-only to recover enough backends to compare. Callers + # whose backward-capable pair does not include FusedAttention -- softcap, where + # get_attention_backend always disables FusedAttention -- opt out to keep dgrad coverage. + if not fused_attn_supported and fwd_only_without_fused_attn: is_training = False available_backends, _, fused_attn_backends = get_available_attention_backends( config, @@ -667,7 +672,16 @@ def test_dpa_softmax_thd(dtype, model_configs, model): @pytest.mark.parametrize("model", model_configs_softcap.keys()) def test_dpa_softcap(dtype, model_configs, model): """Test DotProductAttention module with tanh logit softcapping""" - test_dot_product_attention(dtype, model_configs, model, False, "bshd_bshd_bshd", False, False) + test_dot_product_attention( + dtype, + model_configs, + model, + False, + "bshd_bshd_bshd", + False, + False, + fwd_only_without_fused_attn=False, + ) model_configs_mla = { From 13d65a10929168792ab849e21ba17fe73138dbd6 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Thu, 27 Aug 2026 10:16:30 -0700 Subject: [PATCH 09/28] test: add softcap no-op and closed-form reference coverage Two gaps remained after folding test_softcap.py into test_attention.py. softcap=0.0 no-op: every model_configs_softcap entry uses a nonzero cap, so nothing asserted the backward-compatibility claim. The half that the PR actually changed is backend selection, and a filter that fired at 0.0 would silently remove FusedAttention and FA4 from other tests rather than fail one. test_dpa_softcap_zero_backend_selection asserts FusedAttention survives softcap=0.0 and is disabled by a nonzero cap. Unfused coverage and tanh's nonlinear region: test_dpa_softcap needs two TE backends, so it skips entirely without flash-attn even though UnfusedDotProductAttention now implements softcap and is the reference for everything else. It also cannot detect a dropped cap at all: 0.1 * randn inputs put the logits at O(1e-2), where the reference output moves by 9e-9 at cap=50 and 2e-4 at cap=0.01. test_dpa_softcap_vs_reference compares forward and dQ/dK/dV against a pure-PyTorch oracle one backend at a time, so it runs with unfused alone, and uses randn inputs so the cap moves the output by O(1). An assertion on that displacement keeps the test from going vacuous if the config drifts. Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 169 +++++++++++++++++++++- 1 file changed, 163 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index b9e446cd37..0ad3c90e1b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -656,17 +656,17 @@ def test_dpa_softmax_thd(dtype, model_configs, model): "softcap_1_1": ModelConfig(4, 128, 16, 64, num_gqa_groups=4, softcap=50.0), "softcap_2_0": ModelConfig(2, 512, 16, 64, attn_mask_type="causal", softcap=50.0), "softcap_2_1": ModelConfig(2, 512, 24, 128, attn_mask_type="padding_causal", softcap=50.0), - # 0.01 is on the order of the logits these inputs produce, so tanh runs in its nonlinear - # region instead of acting as a no-op, and a misapplied softmax_scale or a missing outer - # softcap factor changes the output. + # The shared harness feeds 0.1 * randn, which puts the logits at O(1e-2) whatever the head + # dim, so tanh is numerically linear at a Gemma-sized cap. A cap of 0.01 is the one regime + # these inputs can distinguish: dropping the outer softcap factor would leave logits of + # O(1) instead of O(1e-2) and move the output well past the tolerance. Softcapping in + # tanh's saturating region is covered by test_dpa_softcap_vs_reference, which uses its own + # inputs. "softcap_3_0": ModelConfig(4, 128, 16, 64, softcap=0.01), "softcap_3_1": ModelConfig(2, 512, 16, 64, attn_mask_type="causal", softcap=0.01), } -@pytest.mark.skipif( - not FlashAttentionUtils.v2_6_0_plus, reason="flash-attn 2.6.0+ is required for softcap." -) @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model_configs", [model_configs_softcap]) @pytest.mark.parametrize("model", model_configs_softcap.keys()) @@ -684,6 +684,163 @@ def test_dpa_softcap(dtype, model_configs, model): ) +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_softcap]) +@pytest.mark.parametrize("model", ["softcap_1_0"]) +def test_dpa_softcap_zero_backend_selection(dtype, model_configs, model): + """Test that softcap=0.0 leaves backend selection untouched. + + The softcap filter in get_attention_backend disables FusedAttention (and FA4) whenever the + cap is nonzero. If it also fired at 0.0, those backends would silently drop out of every + other test in this file rather than failing one, so assert both halves here. + """ + config = copy.deepcopy(model_configs[model]) + query = dict( + qkv_dtype=dtype, + qkv_layout="bshd_bshd_bshd", + is_training=True, + deterministic=_deterministic, + ) + + config.softcap = 0.0 + (_, fused_off, unfused_off), _, _ = get_available_attention_backends(config, **query) + config.softcap = 50.0 + (_, fused_on, unfused_on), _, _ = get_available_attention_backends(config, **query) + + assert fused_off, "softcap=0.0 must not disable FusedAttention" + assert not fused_on, "a nonzero softcap must disable FusedAttention" + assert unfused_off and unfused_on, "UnfusedDotProductAttention must support softcap" + + +def _softcap_reference_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + softmax_scale: float, + softcap: float, + causal: bool, +) -> torch.Tensor: + """Closed-form softcapped attention in bshd layout, computed in fp32. + + scores = softcap * tanh(Q @ K^T * softmax_scale / softcap), with the tanh skipped entirely + when softcap == 0.0, so this doubles as the reference for the no-op claim. GQA is supported. + """ + q, k, v = (x.transpose(1, 2).float() for x in (q, k, v)) + if q.shape[1] != k.shape[1]: + repeats = q.shape[1] // k.shape[1] + k = k.repeat_interleave(repeats, dim=1) + v = v.repeat_interleave(repeats, dim=1) + scores = torch.matmul(q, k.transpose(-2, -1)) * softmax_scale + if softcap != 0.0: + scores = softcap * torch.tanh(scores / softcap) + if causal: + max_seqlen_q, max_seqlen_kv = scores.shape[-2], scores.shape[-1] + mask = torch.triu( + torch.ones(max_seqlen_q, max_seqlen_kv, dtype=torch.bool, device=scores.device), + diagonal=1 + max_seqlen_kv - max_seqlen_q, + ) + scores = scores.masked_fill(mask, float("-inf")) + return torch.matmul(torch.softmax(scores, dim=-1), v).transpose(1, 2) + + +model_configs_softcap_reference = { + # test: ModelConfig(b, sq, hq, dqk) + "softcap_ref_1_0": ModelConfig(2, 128, 8, 64), + "softcap_ref_1_1": ModelConfig(2, 128, 8, 64, num_gqa_groups=2), + "softcap_ref_2_0": ModelConfig(2, 128, 8, 64, attn_mask_type="causal"), +} + + +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_softcap_reference]) +@pytest.mark.parametrize("model", model_configs_softcap_reference.keys()) +@pytest.mark.parametrize("softcap", [0.0, 0.5]) +@pytest.mark.parametrize("backend", ["UnfusedDotProductAttention", "FlashAttention"]) +def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend): + """Test softcap forward and dQ/dK/dV against a closed-form reference, one backend at a time. + + This needs only one TE backend, so UnfusedDotProductAttention -- the reference + implementation for every other softcap test -- stays covered on machines without + flash-attn. softcap=0.0 checks against a reference that never applies tanh, which is the + numerical half of the no-op claim. + """ + config = copy.deepcopy(model_configs[model]) + config.softcap = softcap + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=dtype, + qkv_layout="bshd_bshd_bshd", + is_training=True, + deterministic=_deterministic, + ) + supported = dict( + zip(["FlashAttention", "FusedAttention", "UnfusedDotProductAttention"], available_backends) + ) + if not supported[backend]: + pytest.skip(f"{backend} is unavailable for this config.") + + reset_rng_states() + os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FlashAttention" else "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention" else "0" + _attention_backends["backend_selection_requires_update"] = True + + causal = "causal" in config.attn_mask_type + softmax_scale = 1.0 / config.head_dim_qk**0.5 + q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) + k_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_qk) + v_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_v) + out_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_v) + # randn puts the logits at O(1), so a cap of 0.5 lands in tanh's saturating region and moves + # the output by O(1). The shared harness uses 0.1 * randn, where the logits are O(1e-2) and + # no cap value is distinguishable from no cap at all. + q, k, v = ( + torch.randn(shape, dtype=dtype, device="cuda").requires_grad_() + for shape in (q_shape, k_shape, v_shape) + ) + q_ref, k_ref, v_ref = (x.detach().clone().requires_grad_() for x in (q, k, v)) + # DotProductAttention merges the head and head-dim axes of its output. + d_out = torch.randn(out_shape, dtype=dtype, device="cuda") + + block = DotProductAttention( + config.num_heads, + (config.head_dim_qk, config.head_dim_v), + num_gqa_groups=config.num_gqa_groups, + qkv_format="bshd", + attn_mask_type=config.attn_mask_type, + softmax_scale=softmax_scale, + softcap=softcap, + layer_number=1, + ).to(dtype=dtype, device="cuda") + out = block(q, k, v).view(out_shape) + out.backward(d_out) + + out_ref = _softcap_reference_attention(q_ref, k_ref, v_ref, softmax_scale, softcap, causal) + out_ref.backward(d_out.float()) + + tols = dict(atol=2e-2, rtol=2e-2) + if dtype == torch.bfloat16: + tols = dict(atol=4e-2, rtol=4e-2) + + if softcap != 0.0: + # Without this the test could be vacuous: a backend that dropped softcap on the floor + # would still match a reference whose tanh is numerically the identity. + out_ref_uncapped = _softcap_reference_attention( + q_ref.detach(), k_ref.detach(), v_ref.detach(), softmax_scale, 0.0, causal + ) + cap_effect = (out_ref.detach() - out_ref_uncapped).abs().max().item() + assert cap_effect > 10 * tols["atol"], ( + f"softcap={softcap} moves the reference output by only {cap_effect:.2e}; this config" + " would pass even if the backend ignored softcap" + ) + + torch.testing.assert_close(out.float(), out_ref, **tols) + torch.testing.assert_close(q.grad.float(), q_ref.grad.float(), **tols) + torch.testing.assert_close(k.grad.float(), k_ref.grad.float(), **tols) + torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), **tols) + + model_configs_mla = { # test: ModelConfig(b, sq, hq, dqk) "mla_1_0": ModelConfig(8, 128, 16, 64, head_dim_v=128), From 9ba7803d2871ecadd6ed2bf469c343b6faa1a3ee Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 31 Aug 2026 11:19:06 -0700 Subject: [PATCH 10/28] fix(pytorch): align CP autograd backward arity with the softcap forward input Threading `softcap` through the three context-parallel autograd functions added a forward input to each without adding the matching gradient slot to the corresponding backward return tuple, leaving every CP backward one gradient short of its forward inputs. Because `softcap` sits mid-signature, the omission also shifted every later slot in AttnFuncWithCPAndQKVOA2A: `d_softmax_offset` was being returned in `softmax_type`'s position. Insert the missing slot at the `softcap` position in all three tuples. Signed-off-by: Nitin Vegesna --- .../attention/dot_product_attention/context_parallel.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index ad3b45d2d2..dd3d06b64e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3214,6 +3214,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -4575,6 +4576,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, ) @@ -5412,6 +5414,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, d_softmax_offset, None, ) From db6a119809bc1d229d6313adf8496c2980a57efb Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 31 Aug 2026 11:34:23 -0700 Subject: [PATCH 11/28] test(pytorch): cover context parallelism with softcap No test exercised CP together with softcap, which is why a backward that returned one gradient fewer than its forward inputs went unnoticed in all three CP autograd functions. Thread softcap through the CP runner so it reaches DotProductAttention, and add one case per CP autograd function -- p2p, all_gather and a2a -- checking the softcapped forward and dgrad against the non-CP reference. The cap sits in tanh's nonlinear region so a path that dropped it diverges rather than matching a numerically linear reference. Signed-off-by: Nitin Vegesna --- .../attention/run_attention_with_cp.py | 3 ++ .../attention/test_attention_with_cp.py | 39 +++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index b620d7e4fe..ff2ff4cc15 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -234,6 +234,7 @@ def run_dpa_with_cp( fa_pad_between_seqs="False", deterministic="False", load_balancing_strategy="DUAL_CHUNK_SWAP", + softcap="0.0", log_level=logging.WARNING, ): """Test DotProductAttention module with context parallelism""" @@ -281,6 +282,7 @@ def run_dpa_with_cp( config.attn_mask_type = "padding_causal" else: config.attn_mask_type = "padding" + config.softcap = float(softcap) # set up distributed group rank = int(os.getenv("RANK", "0")) @@ -342,6 +344,7 @@ def run_dpa_with_cp( qkv_format=qkv_format, attn_mask_type=config.attn_mask_type, window_size=config.window_size, + softcap=config.softcap, softmax_type=config.softmax_type, return_max_logit=config.return_max_logit, ).cuda() diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 3b44c31a2f..25caca1054 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -405,6 +405,45 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type ) +@pytest.mark.skipif( + not FlashAttentionUtils.v2_6_0_plus, reason="CP softcap requires flash-attn 2.6.0+." +) +@pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.") +@pytest.mark.parametrize("cp_comm_type", ["p2p", "all_gather", "a2a"]) +def test_cp_with_flash_attention_softcap(cp_pool, cp_comm_type): + """Check softcap forward and dgrad against the non-CP reference. + + One case per CP autograd function, since P2P, all-gather and A2A each thread softcap + through their own forward inputs and gradient slots. + """ + config = copy.deepcopy(model_configs_flash_attn["cp_2_0"]) + config.context_parallel = True + config.cp_comm_type = cp_comm_type + # The runner's clamped-randn inputs put the scaled logits at O(1), so this cap sits in + # tanh's nonlinear region and a path that dropped it would diverge from the reference. + config.softcap = 0.5 + available_backends, _, _ = get_available_attention_backends( + config, + qkv_dtype=torch.bfloat16, + qkv_layout="bshd_bshd_bshd", + is_training=True, + deterministic=_deterministic, + ) + if not available_backends[0]: + pytest.skip("FlashAttention is unavailable.") + _submit( + cp_pool(2), + dtype="bf16", + model="cp_2_0", + qkv_format="bshd", + kernel_backend="FlashAttention", + cp_comm_type=cp_comm_type, + softcap=config.softcap, + deterministic=_deterministic, + log_level=pytest_logging_level, + ) + + model_configs_fused_attn = { # test: ModelConfig(b, sq, hq, dqk) "cp_1_0": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", return_max_logit=True), # MHA From 49ede92c9c4aebb51e6d3505389704ca8d991af7 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Tue, 1 Sep 2026 14:03:43 -0700 Subject: [PATCH 12/28] fix(pytorch): apply softcap before the additive bias, matching FlashAttention UnfusedDotProductAttention applied the tanh cap after adding post_scale_bias / ALiBi, computing cap(scale*QK + bias). FlashAttention-2 computes cap(scale*QK) + bias: its kernel softcaps immediately after the QK^T gemm and only then adds ALiBi, and it pre-divides alibi_slope by scale_softmax -- which softcapping sets to `softcap` -- so the bias deliberately lands outside the tanh (csrc/flash_attn/src/flash_fwd_kernel.h, mask.h, flash_api.cpp). ALiBi is the one bias type flash supports, so with softcap + ALiBi the unfused and flash paths returned different numerics depending only on whether a suitable flash-attn was installed. Defer the additive bias until after the cap so the two agree. pre_scale_bias is folded in before the scaling by construction and stays inside the cap; flash does not support it. softcap = 0.0 remains a bit-exact no-op for every bias type. Add two tests, both forcing UnfusedDotProductAttention: - test_dpa_softcap_bias_ordering pins cap(scores) + bias against cap(scores + bias), using post_scale_bias to drive the same branch ALiBi uses without needing slope machinery in the reference. - test_dpa_softcap_qk_layer_scaling covers softcap under NVTE_APPLY_QK_LAYER_SCALING, where the cap must be divided by layer_number; omitting that leaves an effective cap of softcap * layer_number. Both carry anti-vacuity asserts, and both were verified by mutation on an H100: reintroducing either bug makes the corresponding test fail, and the existing softcap suite still passes. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/pytorch/attention/test_attention.py | 161 ++++++++++++++++++ .../dot_product_attention/backends.py | 22 ++- 2 files changed, 178 insertions(+), 5 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 0ad3c90e1b..181b3c8bea 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -841,6 +841,167 @@ def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend) torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), **tols) +model_configs_softcap_bias = { + # test: ModelConfig(b, sq, hq, dqk) + "softcap_bias_1_0": ModelConfig(2, 128, 8, 64, attn_bias_type="post_scale_bias"), + "softcap_bias_2_0": ModelConfig( + 2, 128, 8, 64, attn_mask_type="causal", attn_bias_type="post_scale_bias" + ), +} + + +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("model_configs", [model_configs_softcap_bias]) +@pytest.mark.parametrize("model", model_configs_softcap_bias.keys()) +def test_dpa_softcap_bias_ordering(dtype, model_configs, model): + """An additive bias is added *after* the cap, not capped together with the logits. + + FlashAttention softcaps immediately after the QK^T gemm and only then adds ALiBi: its + alibi_slope is pre-divided by scale_softmax, which softcapping sets to `softcap`, so the bias + term lands outside the tanh. UnfusedDotProductAttention serves ALiBi and post_scale_bias from + the same branch, so capping the sum there would silently disagree with the flash backends on + softcap + ALiBi. post_scale_bias drives that same branch without needing ALiBi slope + machinery in the reference. + """ + config = copy.deepcopy(model_configs[model]) + softcap = 0.5 + + reset_rng_states() + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + + causal = "causal" in config.attn_mask_type + softmax_scale = 1.0 / config.head_dim_qk**0.5 + q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) + kv_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_qk) + out_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_v) + q, k, v = ( + torch.randn(shape, dtype=dtype, device="cuda") for shape in (q_shape, kv_shape, kv_shape) + ) + # O(1) against a cap of 0.5, so capping the bias too is clearly visible in the output while + # the softmax stays well conditioned (a much larger bias drives it to one-hot, which only + # sharpens fp16 rounding against the tolerance without adding signal). + bias = torch.randn( + 1, config.num_heads, config.max_seqlen_q, config.max_seqlen_kv, dtype=dtype, device="cuda" + ) + + block = DotProductAttention( + config.num_heads, + (config.head_dim_qk, config.head_dim_v), + num_gqa_groups=config.num_gqa_groups, + qkv_format="bshd", + attn_mask_type=config.attn_mask_type, + softmax_scale=softmax_scale, + softcap=softcap, + layer_number=1, + ).to(dtype=dtype, device="cuda") + out = block( + q, k, v, core_attention_bias_type="post_scale_bias", core_attention_bias=bias + ).view(out_shape) + + def _reference(cap_includes_bias): + q_f, k_f, v_f = (x.transpose(1, 2).float() for x in (q, k, v)) + scores = torch.matmul(q_f, k_f.transpose(-2, -1)) * softmax_scale + bias_f = bias.float() + if cap_includes_bias: + scores = softcap * torch.tanh((scores + bias_f) / softcap) + else: + scores = softcap * torch.tanh(scores / softcap) + bias_f + if causal: + max_seqlen_q, max_seqlen_kv = scores.shape[-2], scores.shape[-1] + mask = torch.triu( + torch.ones(max_seqlen_q, max_seqlen_kv, dtype=torch.bool, device=scores.device), + diagonal=1 + max_seqlen_kv - max_seqlen_q, + ) + scores = scores.masked_fill(mask, float("-inf")) + return torch.matmul(torch.softmax(scores, dim=-1), v_f).transpose(1, 2) + + out_ref = _reference(cap_includes_bias=False) + out_capped_bias = _reference(cap_includes_bias=True) + + tols = dict(atol=2e-2, rtol=2e-2) + if dtype == torch.bfloat16: + tols = dict(atol=4e-2, rtol=4e-2) + + # Without this the test could be vacuous: the two orderings have to be distinguishable at + # this cap and bias magnitude for the comparison below to mean anything. + ordering_effect = (out_ref - out_capped_bias).abs().max().item() + assert ordering_effect > 10 * tols["atol"], ( + f"the two bias orderings differ by only {ordering_effect:.2e}; this config would pass" + " whichever one the backend implements" + ) + torch.testing.assert_close(out.float(), out_ref, **tols) + + +@pytest.mark.parametrize("model_configs", [model_configs_softcap_reference]) +@pytest.mark.parametrize("model", ["softcap_ref_1_0", "softcap_ref_2_0"]) +def test_dpa_softcap_qk_layer_scaling(model_configs, model): + """softcap survives NVTE_APPLY_QK_LAYER_SCALING, which defers a layer_number factor. + + With qk layer scaling UnfusedDotProductAttention scales the logits by + softmax_scale / layer_number and lets the softmax multiply them back by layer_number, so the + cap must be divided by layer_number to land on the true logits. Dropping that division leaves + an effective cap of softcap * layer_number, which this test pins down. fp16 only: the backend + gates qk layer scaling on the key dtype. + """ + dtype = torch.float16 + config = copy.deepcopy(model_configs[model]) + softcap = 0.5 + # The undivided cap would be softcap * layer_number; layer_number is chosen large enough that + # the two caps are far enough apart to separate well beyond the comparison tolerance (the + # assert below pins this down), since the logits here are O(1). + layer_number = 8 + + reset_rng_states() + os.environ["NVTE_FLASH_ATTN"] = "0" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "1" + os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "1" + _attention_backends["backend_selection_requires_update"] = True + + causal = "causal" in config.attn_mask_type + softmax_scale = 1.0 / config.head_dim_qk**0.5 + q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) + kv_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_qk) + out_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_v) + q, k, v = ( + torch.randn(shape, dtype=dtype, device="cuda") for shape in (q_shape, kv_shape, kv_shape) + ) + + try: + block = DotProductAttention( + config.num_heads, + (config.head_dim_qk, config.head_dim_v), + num_gqa_groups=config.num_gqa_groups, + qkv_format="bshd", + attn_mask_type=config.attn_mask_type, + softmax_scale=softmax_scale, + softcap=softcap, + layer_number=layer_number, + ).to(dtype=dtype, device="cuda") + out = block(q, k, v).view(out_shape) + finally: + os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "0" + _attention_backends["backend_selection_requires_update"] = True + + out_ref = _softcap_reference_attention(q, k, v, softmax_scale, softcap, causal) + # Omitting the cap / layer_number division caps the reduced logits instead, which after the + # softmax's layer_number factor is exactly a softcap * layer_number cap on the true logits. + out_undivided = _softcap_reference_attention( + q, k, v, softmax_scale, softcap * layer_number, causal + ) + + tols = dict(atol=2e-2, rtol=2e-2) + scaling_effect = (out_ref - out_undivided).abs().max().item() + assert scaling_effect > 10 * tols["atol"], ( + f"dividing the cap by layer_number moves the output by only {scaling_effect:.2e}; this" + " config would pass even without the division" + ) + torch.testing.assert_close(out.float(), out_ref, **tols) + + model_configs_mla = { # test: ModelConfig(b, sq, hq, dqk) "mla_1_0": ModelConfig(8, 128, 16, 64, head_dim_v=128), diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 0904c1a605..4304f93ebd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -641,6 +641,9 @@ def _forward( key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1) # Raw attention scores. [b * h, sq, sk] + # An additive `post_scale_bias`/ALiBi term is deferred until after the softcap below, so + # that the cap applies to the bare scaled logits (see the softcap comment for why). + deferred_bias = None if core_attention_bias_type == "no_bias": matmul_result = torch.baddbmm( matmul_result, @@ -684,17 +687,26 @@ def _forward( beta=0.0, alpha=scale, ) - matmul_result = (matmul_result.view(*output_size) + core_attention_bias).to( - dtype=query_layer.dtype - ) + matmul_result = matmul_result.view(*output_size) + deferred_bias = core_attention_bias # Cap the scaled logits -- softcap * tanh(scores * scale / softcap) -- matching how - # FlashAttention folds softmax_scale into its tanh argument. qk layer scaling defers the - # layer_number factor to the softmax below, so it is divided out of the cap here. + # FlashAttention folds softmax_scale into its tanh argument. The cap is applied to the + # bare scaled logits, before any additive bias: FA2 softcaps immediately after the QK^T + # gemm and only then adds ALiBi (its alibi_slope is pre-divided by scale_softmax, which + # softcapping sets to `softcap`, so the bias lands outside the tanh). Capping the bias + # too would silently diverge from FA2, which is selectable alongside this backend for + # ALiBi -- the one bias type flash supports (pre/post_scale_bias disable it outright). + # `pre_scale_bias` is folded in before the scaling by construction, so it is necessarily + # inside the cap. qk layer scaling defers the layer_number factor to the softmax below, + # so it is divided out of the cap here. if softcap != 0.0: cap = softcap / self.layer_number if apply_qk_layer_scaling else softcap matmul_result = cap * torch.tanh(matmul_result / cap) + if deferred_bias is not None: + matmul_result = (matmul_result + deferred_bias).to(dtype=query_layer.dtype) + if fp8: # quantize and dequantize dP to emulate FP8 matmul_result, *_ = FP8EmulationFunc.apply( From 18f67b73f56b67a4c8b1411376dee5916d0135dd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Sep 2026 21:06:55 +0000 Subject: [PATCH 13/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_attention.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 181b3c8bea..81d5b7898b 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -897,9 +897,9 @@ def test_dpa_softcap_bias_ordering(dtype, model_configs, model): softcap=softcap, layer_number=1, ).to(dtype=dtype, device="cuda") - out = block( - q, k, v, core_attention_bias_type="post_scale_bias", core_attention_bias=bias - ).view(out_shape) + out = block(q, k, v, core_attention_bias_type="post_scale_bias", core_attention_bias=bias).view( + out_shape + ) def _reference(cap_includes_bias): q_f, k_f, v_f = (x.transpose(1, 2).float() for x in (q, k, v)) From 8e49940ed4b74ad5b0554d253fc76aa1a54ce8e2 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 2 Sep 2026 23:36:15 -0700 Subject: [PATCH 14/28] fix(test): skip the softcap no-op check when FusedAttention is unavailable test_dpa_softcap_zero_backend_selection asserted that FusedAttention is selectable at softcap=0.0. Whether it is available at all depends on the arch and on NVTE_ALLOW_NONDETERMINISTIC_ALGO, not on softcap: the A100 and L40 CI jobs fail the assert in their deterministic pass, where cuDNN offers no fused backend for this config, while the same test passes on H100 and B200. The test is about the softcap filter, so make the availability half a skip. The meaningful direction -- a nonzero softcap must disable FusedAttention -- still runs wherever a fused backend exists, and skipping keeps the comparison from being silently vacuous where one does not. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/pytorch/attention/test_attention.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 81d5b7898b..caff965d2d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -694,6 +694,10 @@ def test_dpa_softcap_zero_backend_selection(dtype, model_configs, model): The softcap filter in get_attention_backend disables FusedAttention (and FA4) whenever the cap is nonzero. If it also fired at 0.0, those backends would silently drop out of every other test in this file rather than failing one, so assert both halves here. + + Whether FusedAttention is available at all is arch- and mode-dependent (cuDNN support, + NVTE_ALLOW_NONDETERMINISTIC_ALGO=0), and is not what this test is about, so that half is a + skip rather than an assert. """ config = copy.deepcopy(model_configs[model]) query = dict( @@ -708,7 +712,12 @@ def test_dpa_softcap_zero_backend_selection(dtype, model_configs, model): config.softcap = 50.0 (_, fused_on, unfused_on), _, _ = get_available_attention_backends(config, **query) - assert fused_off, "softcap=0.0 must not disable FusedAttention" + if not fused_off: + pytest.skip( + "FusedAttention is unavailable for this config irrespective of softcap (no cuDNN" + " support for this arch/shape, or deterministic mode), so the softcap filter has" + " nothing to disable and the comparison below would be vacuous." + ) assert not fused_on, "a nonzero softcap must disable FusedAttention" assert unfused_off and unfused_on, "UnfusedDotProductAttention must support softcap" From 42975ac44d5aa19576372835ae7181c630754395 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 2 Sep 2026 23:51:56 -0700 Subject: [PATCH 15/28] fix(lint): silence missing-function-docstring on _IdentityWithMaskedGradient The Lint job fails on this branch with dot_product_attention.py:243:4: C0116: Missing function or method docstring dot_product_attention.py:248:4: C0116: Missing function or method docstring _IdentityWithMaskedGradient came from #3274 and is byte-identical to main; the same two errors fail Lint on that PR's own branch (desh/mixed-thd-pr-minimal), which merged before the check went green, so this branch inherited a red Lint by merging main. Not introduced here, but it blocks this PR. Use the same `# pylint: disable=missing-function-docstring` the other autograd Functions in this package use (e.g. context_parallel.py) rather than inventing docstrings. Drop this commit if #3274 is fixed upstream first. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- .../attention/dot_product_attention/dot_product_attention.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 201ee71cfd..0c132a32a2 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -241,11 +241,13 @@ class _IdentityWithMaskedGradient(torch.autograd.Function): @staticmethod def forward(ctx, tensor: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + # pylint: disable=missing-function-docstring ctx.save_for_backward(mask) return tensor @staticmethod def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None]: + # pylint: disable=missing-function-docstring (mask,) = ctx.saved_tensors return torch.where(mask, grad_output, 0.0), None From 67277a03b485ae6eef421aded636285649a079ae Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Fri, 4 Sep 2026 12:36:26 -0700 Subject: [PATCH 16/28] feat: expose softcap on MultiheadAttention and TransformerLayer; trim comments Review feedback from #3391. Thread softcap through MultiheadAttention and TransformerLayer, following the window_size pattern: a constructor argument stored on the module and a forward argument that falls back to it. TransformerLayer passes softcap via common_attention_kwargs, so it reaches cross-attention as well as self-attention, matching how softmax_type (also not mask-specific) is handled; scoping it to self-attention would silently leave decoder cross-attention uncapped. Resolution sits outside the thd_attention_policies branch at both levels, since softcap is not mask-specific. Drop "only supported by the FlashAttention and UnfusedDotProductAttention backends" from the softcap docstrings. It is inaccurate, since support depends on the FlashAttention version, and get_attention_backend() is the single place that should describe backend eligibility. The same wording appeared in the DotProductAttention class and forward docstrings, so all three now describe only what softcap does. Shorten the softcap comment in UnfusedDotProductAttention, keeping why the cap precedes the additive bias and why it is divided by layer_number, as each of those was a real bug. Remove the fail-loud block comment in the FA3 branch, whose selection criteria get_attention_backend() already documents, and keep a single line noting that the FA3 entry points are autograd functions and so the forward kwarg drives the backward too. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- .../dot_product_attention/backends.py | 24 +++++-------------- .../dot_product_attention.py | 8 ++----- .../attention/dot_product_attention/utils.py | 5 ++-- .../pytorch/attention/multi_head_attention.py | 17 +++++++++++++ transformer_engine/pytorch/transformer.py | 20 ++++++++++++++++ 5 files changed, 47 insertions(+), 27 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 4304f93ebd..582a8a2d44 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -690,16 +690,11 @@ def _forward( matmul_result = matmul_result.view(*output_size) deferred_bias = core_attention_bias - # Cap the scaled logits -- softcap * tanh(scores * scale / softcap) -- matching how - # FlashAttention folds softmax_scale into its tanh argument. The cap is applied to the - # bare scaled logits, before any additive bias: FA2 softcaps immediately after the QK^T - # gemm and only then adds ALiBi (its alibi_slope is pre-divided by scale_softmax, which - # softcapping sets to `softcap`, so the bias lands outside the tanh). Capping the bias - # too would silently diverge from FA2, which is selectable alongside this backend for - # ALiBi -- the one bias type flash supports (pre/post_scale_bias disable it outright). - # `pre_scale_bias` is folded in before the scaling by construction, so it is necessarily - # inside the cap. qk layer scaling defers the layer_number factor to the softmax below, - # so it is divided out of the cap here. + # Cap the scaled logits: softcap * tanh(scores * scale / softcap), matching how + # FlashAttention folds softmax_scale into its tanh argument. The cap must land on the + # bare scaled logits, before any additive bias: FA2 caps right after the QK^T gemm and + # adds ALiBi afterwards, so capping the bias too would diverge from it. qk layer scaling + # defers the layer_number factor to the softmax below, so it is divided out of the cap. if softcap != 0.0: cap = softcap / self.layer_number if apply_qk_layer_scaling else softcap matmul_result = cap * torch.tanh(matmul_result / cap) @@ -1300,14 +1295,6 @@ def forward( **fa_optional_forward_kwargs, ) else: - # Fail-loud net: get_attention_backend only keeps FA3 for softcap on a - # softcap-capable build (signature probe) + Hopper (FA3 is sm90-only upstream) - # + head_dim <= 256. If FA3 is still reached with softcap while the build lacks - # support (force-selected / regressed path), raise rather than silently drop the - # cap. The non-CP FA3 entry points - # (flash_attn_func_v3 / flash_attn_varlen_func_v3) are self-contained autograd - # functions, so threading `softcap` into the forward call also drives the - # matching FA3 softcap backward kernel. (CP + FA3 + softcap stays blocked above.) if softcap != 0.0 and not fa_utils.fa3_supports_softcap: raise NotImplementedError( "softcap is not supported by the installed FlashAttention 3 build. " @@ -1317,6 +1304,7 @@ def forward( fa_3_optional_forward_kwargs["window_size"] = window_size fa_3_optional_forward_kwargs["num_splits"] = num_splits if softcap != 0.0 and fa_utils.fa3_supports_softcap: + # FA3 entry points are autograd functions, so this drives the backward too. fa_3_optional_forward_kwargs["softcap"] = softcap if pad_between_seqs: fa_3_optional_forward_kwargs["seqused_q"] = ( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 3746410a7b..c3eebfcc6b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -748,10 +748,8 @@ def nvfp4_linear_mxfp8_dpa_factory(role): softcap : float, default = 0.0 tanh logit softcapping value applied to the attention scores as ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables - softcapping. Softcapping is only supported by the FlashAttention - and UnfusedDotProductAttention backends. Similar to - :attr:`window_size`, ``softcap`` can be overridden by - :attr:`softcap` in ``forward`` as well. + softcapping. Similar to :attr:`window_size`, ``softcap`` can be + overridden by :attr:`softcap` in ``forward`` as well. attention_type : str, default = "self" type of attention, either ``"self"`` and ``"cross"``. layer_number : int, default = None @@ -2110,8 +2108,6 @@ def forward( tanh logit softcapping value applied to the attention scores as ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. When `None`, the value passed to the constructor is used. - Softcapping is only supported by the FlashAttention and - UnfusedDotProductAttention backends. bottom_right_diagonal: Optional[bool], default = None Align sliding window and ALiBi diagonal to the top left (`False`) or bottom right (`True`) corner of the softmax matrix in the encoder. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 5e7f323b0f..fa70884f93 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -235,9 +235,8 @@ class AttentionParams: Whether to align sliding window and ALiBi diagonal to the bottom right corner of the softmax matrix. softcap : float, default = 0.0 - Tanh logit softcapping value applied to the attention scores. A value of - ``0.0`` disables softcapping. Only supported by the FlashAttention and - UnfusedDotProductAttention backends. + Tanh logit softcapping value applied to the attention scores, as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. alibi_slopes_shape : Optional[Union[torch.Size, List]], default = None Tensor shape of :attr:`alibi_slopes` in `DotProductAttention`. core_attention_bias_type : str, default = no_bias diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 82221f0c83..e689da7313 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -100,6 +100,11 @@ class MultiheadAttention(torch.nn.Module): or bottom right (`True`) corner of the softmax matrix in the encoder. If `None`, it will be set to `False` for `attn_mask_type` = {`causal`, `padding_causal`} and `True` for other mask types. + softcap : float, default = 0.0 + tanh logit softcapping value applied to the attention scores as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. + Similar to :attr:`window_size`, ``softcap`` can be overridden by + :attr:`softcap` in :meth:`forward` as well. num_gqa_groups : int, default = None number of GQA groups in the transformer layer. Grouped Query Attention is described in @@ -256,6 +261,7 @@ def __init__( attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: float = 0.0, tp_group: Optional[dist_group_type] = None, tp_size: int = 1, num_gqa_groups: Optional[int] = None, @@ -295,6 +301,7 @@ def __init__( self.attn_mask_type = attn_mask_type self.window_size = window_size self.bottom_right_diagonal = bottom_right_diagonal + self.softcap = softcap self.layer_number = 1 if layer_number is None else layer_number self.input_layernorm = input_layernorm self.attention_type = attention_type @@ -737,6 +744,7 @@ def forward( attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: Optional[float] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[InferenceParams] = None, @@ -790,6 +798,10 @@ def forward( or bottom right (`True`) corner of the softmax matrix in the encoder. If `None`, it will be set to `False` for `attn_mask_type` = {`causal`, `padding_causal`} and `True` for other mask types. + softcap: Optional[float], default = None + tanh logit softcapping value applied to the attention scores as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables + softcapping. When `None`, the value passed to the constructor is used. thd_attention_policies: Optional[List[Dict[str, Any]]], default = None Per-sequence policies for packed THD attention. Passed through to :class:`DotProductAttention`; do not also pass :attr:`attn_mask_type` @@ -873,6 +885,10 @@ def forward( }: bottom_right_diagonal = True + # softcap is not mask-specific, so resolve it outside the policy branch above. + if softcap is None: + softcap = self.softcap + if ( thd_attention_policies is None and "padding" in attn_mask_type @@ -1240,6 +1256,7 @@ def forward( attention_mask=attention_mask, attn_mask_type=attn_mask_type, window_size=window_size, + softcap=softcap, bottom_right_diagonal=bottom_right_diagonal, thd_attention_policies=thd_attention_policies, thd_attention_policy_dispatch=thd_attention_policy_dispatch, diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 664683b555..0d24f082a8 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -154,6 +154,12 @@ class TransformerLayer(torch.nn.Module): or bottom right (`True`) corner of the softmax matrix in the encoder. If `None`, it will be set to `False` for `self_attn_mask_type` = {`causal`, `padding_causal`} and `True` for other mask types. + softcap : float, default = 0.0 + tanh logit softcapping value applied to the attention scores as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. + Applied to both self-attention and, in decoder layers, cross-attention. + Similar to :attr:`window_size`, ``softcap`` can be overridden by + :attr:`softcap` in :meth:`forward` as well. enc_dec_attn_mask_type : {'no_mask', 'causal', 'padding', 'padding_causal', 'arbitrary'}, default = "no_mask" type of attention mask passed into softmax operation for decoder. @@ -314,6 +320,7 @@ def __init__( self_attn_mask_type: str = "causal", window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: float = 0.0, enc_dec_attn_mask_type: str = "no_mask", enc_dec_bottom_right_diagonal: Optional[bool] = None, enc_dec_window_size: Optional[Tuple[int, int]] = None, @@ -358,6 +365,7 @@ def __init__( self.self_attn_mask_type = self_attn_mask_type self.window_size = window_size self.bottom_right_diagonal = bottom_right_diagonal + self.softcap = softcap self.enc_dec_attn_mask_type = enc_dec_attn_mask_type self.enc_dec_window_size = enc_dec_window_size self.enc_dec_bottom_right_diagonal = enc_dec_bottom_right_diagonal @@ -460,6 +468,7 @@ def __init__( "seq_length": seq_length, "micro_batch_size": micro_batch_size, "softmax_type": self.softmax_type, + "softcap": self.softcap, } self.self_attention = MultiheadAttention( @@ -655,6 +664,7 @@ def forward( self_attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, + softcap: Optional[float] = None, encoder_output: Optional[torch.Tensor] = None, enc_dec_attn_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, enc_dec_attn_mask_type: Optional[str] = None, @@ -711,6 +721,10 @@ def forward( or bottom right (`True`) corner of the softmax matrix in the encoder. If `None`, it will be set to `False` for `self_attn_mask_type` = {`causal`, `padding_causal`} and `True` for other mask types. + softcap: Optional[float], default = None + tanh logit softcapping value applied to the attention scores as + ``softcap * tanh(scores / softcap)``. A value of ``0.0`` disables softcapping. + When `None`, the value passed to the constructor is used. thd_attention_policies: Optional[List[Dict[str, Any]]], default = None Per-sequence policies for packed THD self-attention. Passed through to :class:`MultiheadAttention`; do not also pass :attr:`self_attn_mask_type` @@ -818,6 +832,10 @@ def forward( }: bottom_right_diagonal = True + # softcap is not mask-specific, so resolve it outside the policy branch above. + if softcap is None: + softcap = self.softcap + if enc_dec_attn_mask_type is None: enc_dec_attn_mask_type = self.enc_dec_attn_mask_type if enc_dec_window_size is None: @@ -900,6 +918,7 @@ def forward( attn_mask_type=self_attn_mask_type, window_size=window_size, bottom_right_diagonal=bottom_right_diagonal, + softcap=softcap, thd_attention_policies=thd_attention_policies, thd_attention_policy_dispatch=thd_attention_policy_dispatch, inference_params=inference_params, @@ -938,6 +957,7 @@ def forward( attn_mask_type=enc_dec_attn_mask_type, window_size=enc_dec_window_size, bottom_right_diagonal=enc_dec_bottom_right_diagonal, + softcap=softcap, encoder_output=encoder_output, inference_params=inference_params, is_first_microbatch=is_first_microbatch, From 33fe7329a1eeb497015c65c7ba026c326edec594 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Fri, 4 Sep 2026 14:34:48 -0700 Subject: [PATCH 17/28] fix: append softcap to the MHA and TransformerLayer forward signatures softcap was inserted mid-signature in MultiheadAttention.forward and TransformerLayer.forward, next to the other attention arguments. That shifts every parameter after it, so an external caller passing encoder_output or is_first_microbatch positionally would silently bind it to softcap instead. Both accept arbitrary objects, so it would be a wrong answer rather than a TypeError. Nothing in this repo calls either forward with enough positional arguments to hit it, but both classes are public API. Move softcap to the end of both forward signatures. Diffing against main confirms no existing parameter changes index in either. The constructors are unaffected and keep softcap grouped with window_size and bottom_right_diagonal, since it sits past every positional tuple used to construct these modules. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- transformer_engine/pytorch/attention/multi_head_attention.py | 2 +- transformer_engine/pytorch/transformer.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index e689da7313..e07e203ca7 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -744,7 +744,6 @@ def forward( attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, - softcap: Optional[float] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[InferenceParams] = None, @@ -762,6 +761,7 @@ def forward( pad_between_seqs: Optional[bool] = None, thd_attention_policies: Optional[List[Dict[str, Any]]] = None, thd_attention_policy_dispatch: str = "auto", + softcap: Optional[float] = None, ) -> Tuple[Union[torch.Tensor, None], ...]: r""" Forward propagation for MultiheadAttention layer. diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 0d24f082a8..5493a18111 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -664,7 +664,6 @@ def forward( self_attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, bottom_right_diagonal: Optional[bool] = None, - softcap: Optional[float] = None, encoder_output: Optional[torch.Tensor] = None, enc_dec_attn_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, enc_dec_attn_mask_type: Optional[str] = None, @@ -687,6 +686,7 @@ def forward( pad_between_seqs: Optional[bool] = None, thd_attention_policies: Optional[List[Dict[str, Any]]] = None, thd_attention_policy_dispatch: str = "auto", + softcap: Optional[float] = None, ) -> torch.Tensor: r""" Transformer Layer: attention block and a feedforward network (MLP) From a7b51600c94110c426beeee5ebdb720ccbecfd9c Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Fri, 4 Sep 2026 14:40:12 -0700 Subject: [PATCH 18/28] test: consolidate the softcap tests Review feedback from #3391. Cut model_configs_softcap from six configs to two: softcap_1_0 for a high cap through flash_attn_func, and softcap_3_1 with padding_causal for a low cap through flash_attn_varlen_func. Fold model_configs_softcap_reference from three to two by merging GQA into the causal config, since GQA is orthogonal to softcapping. Merge test_dpa_softcap_vs_reference, test_dpa_softcap_bias_ordering and test_dpa_softcap_qk_layer_scaling into one test parametrized over {plain, bias_outside_cap, qk_layer_scaling}. A helper returns the DPA kwargs, the forward kwargs and a right and a wrong reference closure per variant, so the shared work of forcing a backend, building inputs and asserting the two references differ by more than the tolerance happens once. The bias variant reuses _softcap_reference_attention through new bias and cap_includes_bias arguments rather than duplicating its GQA and causal handling. Both mutation signals are preserved. The special variants run on softcap_ref_1_0, the same shape the original margins were measured on, so capping the bias still separates the two references by about 0.98 and dropping the cap / layer_number division by about 0.44, against a threshold of 0.20. layer_number is hoisted to _SOFTCAP_QK_LAYER_NUMBER with a note that 3 is too small, since that value is load-bearing and would otherwise look arbitrary. Backward stays on the plain variant only, matching the previous coverage. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/pytorch/attention/test_attention.py | 308 ++++++++-------------- 1 file changed, 113 insertions(+), 195 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 48615ee8ba..55e1ed8d27 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -652,18 +652,13 @@ def test_dpa_softmax_thd(dtype, model_configs, model): model_configs_softcap = { # test: ModelConfig(b, sq, hq, dqk) + # High cap, no padding -> flash_attn_func. "softcap_1_0": ModelConfig(4, 128, 16, 64, softcap=50.0), - "softcap_1_1": ModelConfig(4, 128, 16, 64, num_gqa_groups=4, softcap=50.0), - "softcap_2_0": ModelConfig(2, 512, 16, 64, attn_mask_type="causal", softcap=50.0), - "softcap_2_1": ModelConfig(2, 512, 24, 128, attn_mask_type="padding_causal", softcap=50.0), - # The shared harness feeds 0.1 * randn, which puts the logits at O(1e-2) whatever the head - # dim, so tanh is numerically linear at a Gemma-sized cap. A cap of 0.01 is the one regime - # these inputs can distinguish: dropping the outer softcap factor would leave logits of - # O(1) instead of O(1e-2) and move the output well past the tolerance. Softcapping in - # tanh's saturating region is covered by test_dpa_softcap_vs_reference, which uses its own - # inputs. - "softcap_3_0": ModelConfig(4, 128, 16, 64, softcap=0.01), - "softcap_3_1": ModelConfig(2, 512, 16, 64, attn_mask_type="causal", softcap=0.01), + # Low cap, padding -> flash_attn_varlen_func. The shared harness feeds 0.1 * randn, putting + # logits at O(1e-2) whatever the head dim, so tanh is numerically linear at a Gemma-sized + # cap. A cap of 0.01 is the one regime these inputs can distinguish. Softcapping in tanh's + # saturating region is covered by test_dpa_softcap_vs_reference, which uses its own inputs. + "softcap_3_1": ModelConfig(2, 512, 16, 64, attn_mask_type="padding_causal", softcap=0.01), } @@ -729,11 +724,15 @@ def _softcap_reference_attention( softmax_scale: float, softcap: float, causal: bool, + bias: torch.Tensor = None, + cap_includes_bias: bool = False, ) -> torch.Tensor: """Closed-form softcapped attention in bshd layout, computed in fp32. scores = softcap * tanh(Q @ K^T * softmax_scale / softcap), with the tanh skipped entirely when softcap == 0.0, so this doubles as the reference for the no-op claim. GQA is supported. + An additive `bias` lands outside the tanh unless `cap_includes_bias`, which builds the + reference an implementation that capped the bias along with the logits would produce. """ q, k, v = (x.transpose(1, 2).float() for x in (q, k, v)) if q.shape[1] != k.shape[1]: @@ -741,8 +740,12 @@ def _softcap_reference_attention( k = k.repeat_interleave(repeats, dim=1) v = v.repeat_interleave(repeats, dim=1) scores = torch.matmul(q, k.transpose(-2, -1)) * softmax_scale + if bias is not None and cap_includes_bias: + scores = scores + bias.float() if softcap != 0.0: scores = softcap * torch.tanh(scores / softcap) + if bias is not None and not cap_includes_bias: + scores = scores + bias.float() if causal: max_seqlen_q, max_seqlen_kv = scores.shape[-2], scores.shape[-1] mask = torch.triu( @@ -756,25 +759,93 @@ def _softcap_reference_attention( model_configs_softcap_reference = { # test: ModelConfig(b, sq, hq, dqk) "softcap_ref_1_0": ModelConfig(2, 128, 8, 64), - "softcap_ref_1_1": ModelConfig(2, 128, 8, 64, num_gqa_groups=2), - "softcap_ref_2_0": ModelConfig(2, 128, 8, 64, attn_mask_type="causal"), + "softcap_ref_2_0": ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal"), } +# "plain" checks the cap itself. The other two pin down *where* the cap is applied inside +# UnfusedDotProductAttention, so they run on that backend only, at a cap that saturates. +softcap_variants = ["plain", "bias_outside_cap", "qk_layer_scaling"] + +# Large enough that a cap of softcap * layer_number is far from a cap of softcap for O(1) +# logits; at layer_number=3 the two references differ by less than the tolerance below. +_SOFTCAP_QK_LAYER_NUMBER = 8 + + +def _softcap_variant_spec(variant, config, dtype, softcap, softmax_scale, q, k, v): + """Return (dpa_kwargs, forward_kwargs, right_fn, wrong_fn) for one softcap variant. + + `wrong_fn` is the reference that an implementation carrying the bug this variant guards + against would produce, or None when there is no distinguishable wrong answer. + """ + causal = "causal" in config.attn_mask_type + + def _ref(cap, **kwargs): + return _softcap_reference_attention(q, k, v, softmax_scale, cap, causal, **kwargs) + + if variant == "plain": + # A backend that ignored the cap entirely would land on the uncapped reference. + wrong_fn = None if softcap == 0.0 else (lambda: _ref(0.0)) + return dict(layer_number=1), {}, (lambda: _ref(softcap)), wrong_fn + + if variant == "bias_outside_cap": + # O(1) against the cap, so capping the bias too is visible in the output while the + # softmax stays well conditioned; a larger bias only sharpens fp16 rounding. + bias = torch.randn( + 1, + config.num_heads, + config.max_seqlen_q, + config.max_seqlen_kv, + dtype=dtype, + device="cuda", + ) + forward_kwargs = dict(core_attention_bias_type="post_scale_bias", core_attention_bias=bias) + return ( + dict(layer_number=1), + forward_kwargs, + (lambda: _ref(softcap, bias=bias)), + (lambda: _ref(softcap, bias=bias, cap_includes_bias=True)), + ) + + if variant == "qk_layer_scaling": + # Omitting the cap / layer_number division caps the reduced logits instead, which after + # the softmax's layer_number factor is exactly a softcap * layer_number cap. + layer_number = _SOFTCAP_QK_LAYER_NUMBER + return ( + dict(layer_number=layer_number), + {}, + (lambda: _ref(softcap)), + (lambda: _ref(softcap * layer_number)), + ) + + raise ValueError(f"Unknown softcap variant {variant}!") + @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model_configs", [model_configs_softcap_reference]) @pytest.mark.parametrize("model", model_configs_softcap_reference.keys()) @pytest.mark.parametrize("softcap", [0.0, 0.5]) @pytest.mark.parametrize("backend", ["UnfusedDotProductAttention", "FlashAttention"]) -def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend): - """Test softcap forward and dQ/dK/dV against a closed-form reference, one backend at a time. +@pytest.mark.parametrize("variant", softcap_variants) +def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend, variant): + """Test softcap against a closed-form reference, one backend and one variant at a time. This needs only one TE backend, so UnfusedDotProductAttention -- the reference implementation for every other softcap test -- stays covered on machines without flash-attn. softcap=0.0 checks against a reference that never applies tanh, which is the - numerical half of the no-op claim. + numerical half of the no-op claim. Every variant with a distinguishable wrong answer + asserts the two references are further apart than the tolerance, so a backend that + implemented the wrong one could not pass. """ config = copy.deepcopy(model_configs[model]) + causal = "causal" in config.attn_mask_type + if variant != "plain": + # These pin down UnfusedDotProductAttention's own arithmetic and need one saturating + # cap on one mask type; "plain" carries the backend and mask coverage. + if backend != "UnfusedDotProductAttention" or softcap == 0.0 or causal: + pytest.skip(f"{variant} is covered once, on the non-causal config with a nonzero cap") + if variant == "qk_layer_scaling" and dtype != torch.float16: + pytest.skip("qk layer scaling is gated on fp16 keys") + config.softcap = softcap available_backends, _, _ = get_available_attention_backends( config, @@ -793,9 +864,10 @@ def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend) os.environ["NVTE_FLASH_ATTN"] = "1" if backend == "FlashAttention" else "0" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "1" if backend == "UnfusedDotProductAttention" else "0" + if variant == "qk_layer_scaling": + os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "1" _attention_backends["backend_selection_requires_update"] = True - causal = "causal" in config.attn_mask_type softmax_scale = 1.0 / config.head_dim_qk**0.5 q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) k_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_qk) @@ -812,171 +884,8 @@ def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend) # DotProductAttention merges the head and head-dim axes of its output. d_out = torch.randn(out_shape, dtype=dtype, device="cuda") - block = DotProductAttention( - config.num_heads, - (config.head_dim_qk, config.head_dim_v), - num_gqa_groups=config.num_gqa_groups, - qkv_format="bshd", - attn_mask_type=config.attn_mask_type, - softmax_scale=softmax_scale, - softcap=softcap, - layer_number=1, - ).to(dtype=dtype, device="cuda") - out = block(q, k, v).view(out_shape) - out.backward(d_out) - - out_ref = _softcap_reference_attention(q_ref, k_ref, v_ref, softmax_scale, softcap, causal) - out_ref.backward(d_out.float()) - - tols = dict(atol=2e-2, rtol=2e-2) - if dtype == torch.bfloat16: - tols = dict(atol=4e-2, rtol=4e-2) - - if softcap != 0.0: - # Without this the test could be vacuous: a backend that dropped softcap on the floor - # would still match a reference whose tanh is numerically the identity. - out_ref_uncapped = _softcap_reference_attention( - q_ref.detach(), k_ref.detach(), v_ref.detach(), softmax_scale, 0.0, causal - ) - cap_effect = (out_ref.detach() - out_ref_uncapped).abs().max().item() - assert cap_effect > 10 * tols["atol"], ( - f"softcap={softcap} moves the reference output by only {cap_effect:.2e}; this config" - " would pass even if the backend ignored softcap" - ) - - torch.testing.assert_close(out.float(), out_ref, **tols) - torch.testing.assert_close(q.grad.float(), q_ref.grad.float(), **tols) - torch.testing.assert_close(k.grad.float(), k_ref.grad.float(), **tols) - torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), **tols) - - -model_configs_softcap_bias = { - # test: ModelConfig(b, sq, hq, dqk) - "softcap_bias_1_0": ModelConfig(2, 128, 8, 64, attn_bias_type="post_scale_bias"), - "softcap_bias_2_0": ModelConfig( - 2, 128, 8, 64, attn_mask_type="causal", attn_bias_type="post_scale_bias" - ), -} - - -@pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("model_configs", [model_configs_softcap_bias]) -@pytest.mark.parametrize("model", model_configs_softcap_bias.keys()) -def test_dpa_softcap_bias_ordering(dtype, model_configs, model): - """An additive bias is added *after* the cap, not capped together with the logits. - - FlashAttention softcaps immediately after the QK^T gemm and only then adds ALiBi: its - alibi_slope is pre-divided by scale_softmax, which softcapping sets to `softcap`, so the bias - term lands outside the tanh. UnfusedDotProductAttention serves ALiBi and post_scale_bias from - the same branch, so capping the sum there would silently disagree with the flash backends on - softcap + ALiBi. post_scale_bias drives that same branch without needing ALiBi slope - machinery in the reference. - """ - config = copy.deepcopy(model_configs[model]) - softcap = 0.5 - - reset_rng_states() - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "1" - _attention_backends["backend_selection_requires_update"] = True - - causal = "causal" in config.attn_mask_type - softmax_scale = 1.0 / config.head_dim_qk**0.5 - q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) - kv_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_qk) - out_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_v) - q, k, v = ( - torch.randn(shape, dtype=dtype, device="cuda") for shape in (q_shape, kv_shape, kv_shape) - ) - # O(1) against a cap of 0.5, so capping the bias too is clearly visible in the output while - # the softmax stays well conditioned (a much larger bias drives it to one-hot, which only - # sharpens fp16 rounding against the tolerance without adding signal). - bias = torch.randn( - 1, config.num_heads, config.max_seqlen_q, config.max_seqlen_kv, dtype=dtype, device="cuda" - ) - - block = DotProductAttention( - config.num_heads, - (config.head_dim_qk, config.head_dim_v), - num_gqa_groups=config.num_gqa_groups, - qkv_format="bshd", - attn_mask_type=config.attn_mask_type, - softmax_scale=softmax_scale, - softcap=softcap, - layer_number=1, - ).to(dtype=dtype, device="cuda") - out = block(q, k, v, core_attention_bias_type="post_scale_bias", core_attention_bias=bias).view( - out_shape - ) - - def _reference(cap_includes_bias): - q_f, k_f, v_f = (x.transpose(1, 2).float() for x in (q, k, v)) - scores = torch.matmul(q_f, k_f.transpose(-2, -1)) * softmax_scale - bias_f = bias.float() - if cap_includes_bias: - scores = softcap * torch.tanh((scores + bias_f) / softcap) - else: - scores = softcap * torch.tanh(scores / softcap) + bias_f - if causal: - max_seqlen_q, max_seqlen_kv = scores.shape[-2], scores.shape[-1] - mask = torch.triu( - torch.ones(max_seqlen_q, max_seqlen_kv, dtype=torch.bool, device=scores.device), - diagonal=1 + max_seqlen_kv - max_seqlen_q, - ) - scores = scores.masked_fill(mask, float("-inf")) - return torch.matmul(torch.softmax(scores, dim=-1), v_f).transpose(1, 2) - - out_ref = _reference(cap_includes_bias=False) - out_capped_bias = _reference(cap_includes_bias=True) - - tols = dict(atol=2e-2, rtol=2e-2) - if dtype == torch.bfloat16: - tols = dict(atol=4e-2, rtol=4e-2) - - # Without this the test could be vacuous: the two orderings have to be distinguishable at - # this cap and bias magnitude for the comparison below to mean anything. - ordering_effect = (out_ref - out_capped_bias).abs().max().item() - assert ordering_effect > 10 * tols["atol"], ( - f"the two bias orderings differ by only {ordering_effect:.2e}; this config would pass" - " whichever one the backend implements" - ) - torch.testing.assert_close(out.float(), out_ref, **tols) - - -@pytest.mark.parametrize("model_configs", [model_configs_softcap_reference]) -@pytest.mark.parametrize("model", ["softcap_ref_1_0", "softcap_ref_2_0"]) -def test_dpa_softcap_qk_layer_scaling(model_configs, model): - """softcap survives NVTE_APPLY_QK_LAYER_SCALING, which defers a layer_number factor. - - With qk layer scaling UnfusedDotProductAttention scales the logits by - softmax_scale / layer_number and lets the softmax multiply them back by layer_number, so the - cap must be divided by layer_number to land on the true logits. Dropping that division leaves - an effective cap of softcap * layer_number, which this test pins down. fp16 only: the backend - gates qk layer scaling on the key dtype. - """ - dtype = torch.float16 - config = copy.deepcopy(model_configs[model]) - softcap = 0.5 - # The undivided cap would be softcap * layer_number; layer_number is chosen large enough that - # the two caps are far enough apart to separate well beyond the comparison tolerance (the - # assert below pins this down), since the logits here are O(1). - layer_number = 8 - - reset_rng_states() - os.environ["NVTE_FLASH_ATTN"] = "0" - os.environ["NVTE_FUSED_ATTN"] = "0" - os.environ["NVTE_UNFUSED_ATTN"] = "1" - os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "1" - _attention_backends["backend_selection_requires_update"] = True - - causal = "causal" in config.attn_mask_type - softmax_scale = 1.0 / config.head_dim_qk**0.5 - q_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_qk) - kv_shape = (config.batch_size, config.max_seqlen_kv, config.num_gqa_groups, config.head_dim_qk) - out_shape = (config.batch_size, config.max_seqlen_q, config.num_heads, config.head_dim_v) - q, k, v = ( - torch.randn(shape, dtype=dtype, device="cuda") for shape in (q_shape, kv_shape, kv_shape) + dpa_kwargs, forward_kwargs, right_fn, wrong_fn = _softcap_variant_spec( + variant, config, dtype, softcap, softmax_scale, q_ref, k_ref, v_ref ) try: @@ -988,28 +897,37 @@ def test_dpa_softcap_qk_layer_scaling(model_configs, model): attn_mask_type=config.attn_mask_type, softmax_scale=softmax_scale, softcap=softcap, - layer_number=layer_number, + **dpa_kwargs, ).to(dtype=dtype, device="cuda") - out = block(q, k, v).view(out_shape) + out = block(q, k, v, **forward_kwargs).view(out_shape) finally: os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "0" _attention_backends["backend_selection_requires_update"] = True - out_ref = _softcap_reference_attention(q, k, v, softmax_scale, softcap, causal) - # Omitting the cap / layer_number division caps the reduced logits instead, which after the - # softmax's layer_number factor is exactly a softcap * layer_number cap on the true logits. - out_undivided = _softcap_reference_attention( - q, k, v, softmax_scale, softcap * layer_number, causal - ) + out_ref = right_fn() tols = dict(atol=2e-2, rtol=2e-2) - scaling_effect = (out_ref - out_undivided).abs().max().item() - assert scaling_effect > 10 * tols["atol"], ( - f"dividing the cap by layer_number moves the output by only {scaling_effect:.2e}; this" - " config would pass even without the division" - ) + if dtype == torch.bfloat16: + tols = dict(atol=4e-2, rtol=4e-2) + + if wrong_fn is not None: + # Without this the test could be vacuous: the right and wrong references have to be + # distinguishable at this cap for the comparison below to mean anything. + variant_effect = (out_ref.detach() - wrong_fn().detach()).abs().max().item() + assert variant_effect > 10 * tols["atol"], ( + f"{variant} moves the reference output by only {variant_effect:.2e}; this config" + " would pass even if the backend implemented the wrong variant" + ) + torch.testing.assert_close(out.float(), out_ref, **tols) + if variant == "plain": + out.backward(d_out) + out_ref.backward(d_out.float()) + torch.testing.assert_close(q.grad.float(), q_ref.grad.float(), **tols) + torch.testing.assert_close(k.grad.float(), k_ref.grad.float(), **tols) + torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), **tols) + model_configs_mla = { # test: ModelConfig(b, sq, hq, dqk) From 022a37b86d135dd0d1b7357b652d123dca67671d Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Fri, 4 Sep 2026 15:14:32 -0700 Subject: [PATCH 19/28] test: assert TransformerLayer forwards softcap to both attention modules Adds test_transformer_layer_softcap_plumbing. It builds a decoder TransformerLayer with a softcap, hooks both self_attention.core_attention and inter_attention.core_attention with a forward pre-hook, and asserts each sees the value, once from the constructor and once from a forward override. Only arrival is checked; the numerics are covered by the DotProductAttention tests. Cross-attention is the half worth asserting. It is reached through a separate call site from self-attention, so a refactor can drop the cap there while self-attention keeps working, and nothing else in the suite would notice. Main restructured both forward methods during this branch's life, so that is a live risk rather than a hypothetical one. Writing the test showed that passing softcap through common_attention_kwargs was not merely redundant but harmful: it pre-set each MultiheadAttention's self.softcap, so deleting the cross-attention pass-through still produced the right answer on the constructor path and the assertion could not fail. Removing that entry makes both assertions catch the regression, and leaves softcap matching how window_size is threaded, stored on self and passed per call rather than at construction. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/pytorch/attention/test_attention.py | 57 +++++++++++++++++++++++ transformer_engine/pytorch/transformer.py | 1 - 2 files changed, 57 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 55e1ed8d27..08a1987e79 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -929,6 +929,63 @@ def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend, torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), **tols) +@pytest.mark.parametrize("dtype", param_types) +def test_transformer_layer_softcap_plumbing(dtype): + """Test that TransformerLayer forwards softcap to both of its attention modules. + + The value only has to arrive; the numerics are covered by the DotProductAttention tests + above. Cross-attention is the half worth asserting: it is reached through a separate call + site from self-attention, so a refactor can drop the cap there while self-attention keeps + working, and nothing else in the suite would notice. + """ + # head_dim 64, matching the other softcap configs, so every backend can serve the shape. + hidden_size, num_heads, seqlen, batch_size = 256, 4, 32, 2 + seen = {} + + def _record(name): + def hook(_module, _args, kwargs): + seen[name] = kwargs.get("softcap") + + return hook + + block = TransformerLayer( + hidden_size, + 4 * hidden_size, + num_heads, + layer_type="decoder", + softcap=50.0, + params_dtype=dtype, + device="cuda", + ) + block.self_attention.core_attention.register_forward_pre_hook( + _record("self"), with_kwargs=True + ) + block.inter_attention.core_attention.register_forward_pre_hook( + _record("cross"), with_kwargs=True + ) + + hidden_states = torch.randn( + seqlen, batch_size, hidden_size, dtype=dtype, device="cuda", requires_grad=True + ) + forward_kwargs = dict( + encoder_output=hidden_states, + enc_dec_attn_mask=torch.zeros( + batch_size, 1, 1, seqlen, dtype=torch.bool, device="cuda" + ), + ) + + # The constructor value reaches both attention modules. + block(hidden_states, **forward_kwargs) + assert seen["self"] == 50.0, f"self-attention saw softcap={seen['self']}, expected 50.0" + assert seen["cross"] == 50.0, f"cross-attention saw softcap={seen['cross']}, expected 50.0" + + # A forward override wins over the constructor, for both. + seen.clear() + block(hidden_states, softcap=10.0, **forward_kwargs) + assert seen["self"] == 10.0, f"self-attention saw softcap={seen['self']}, expected 10.0" + assert seen["cross"] == 10.0, f"cross-attention saw softcap={seen['cross']}, expected 10.0" + + model_configs_mla = { # test: ModelConfig(b, sq, hq, dqk) "mla_1_0": ModelConfig(8, 128, 16, 64, head_dim_v=128), diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 5493a18111..db9a352881 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -468,7 +468,6 @@ def __init__( "seq_length": seq_length, "micro_batch_size": micro_batch_size, "softmax_type": self.softmax_type, - "softcap": self.softcap, } self.self_attention = MultiheadAttention( From 2df366826a3d2abf8547184c50d6b2bab491f847 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 4 Sep 2026 22:15:43 +0000 Subject: [PATCH 20/28] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/attention/test_attention.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 08a1987e79..a45bca56c0 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -957,9 +957,7 @@ def hook(_module, _args, kwargs): params_dtype=dtype, device="cuda", ) - block.self_attention.core_attention.register_forward_pre_hook( - _record("self"), with_kwargs=True - ) + block.self_attention.core_attention.register_forward_pre_hook(_record("self"), with_kwargs=True) block.inter_attention.core_attention.register_forward_pre_hook( _record("cross"), with_kwargs=True ) @@ -969,9 +967,7 @@ def hook(_module, _args, kwargs): ) forward_kwargs = dict( encoder_output=hidden_states, - enc_dec_attn_mask=torch.zeros( - batch_size, 1, 1, seqlen, dtype=torch.bool, device="cuda" - ), + enc_dec_attn_mask=torch.zeros(batch_size, 1, 1, seqlen, dtype=torch.bool, device="cuda"), ) # The constructor value reaches both attention modules. From 443f692144f2c4db56a9fefab1e9fc9cca1811a0 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Fri, 4 Sep 2026 15:20:47 -0700 Subject: [PATCH 21/28] test: set explicit backend state in the softcap plumbing test The test inherited NVTE_FLASH_ATTN / NVTE_FUSED_ATTN / NVTE_UNFUSED_ATTN from whichever softcap test ran before it, since none of them restore those vars. A nonzero cap also drops FusedAttention and FA4 in the backend filter, so a leftover NVTE_UNFUSED_ATTN=0 could leave no eligible backend at all and the forward would raise "No dot product attention backend is available" rather than fail an assertion. The outcome also depended on which parametrization of the preceding test happened to run last. Set the three variables explicitly, matching the convention of the other tests in this file, leaving flash and unfused both enabled so at least one backend is available with or without flash-attn. The assertions stay backend agnostic because the hook fires on DotProductAttention itself. Signed-off-by: Nitin Vegesna Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/pytorch/attention/test_attention.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index a45bca56c0..0493d12ace 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -948,6 +948,17 @@ def hook(_module, _args, kwargs): return hook + # Every softcap test above sets these and none restore them, so inherit nothing: with a + # nonzero cap the filter also drops FusedAttention and FA4, and a leftover + # NVTE_UNFUSED_ATTN=0 would leave no backend at all on a machine without flash-attn. + # Leaving flash and unfused both enabled keeps this backend-agnostic; the hook fires on + # DotProductAttention regardless of which one is selected. + reset_rng_states() + os.environ["NVTE_FLASH_ATTN"] = "1" + os.environ["NVTE_FUSED_ATTN"] = "0" + os.environ["NVTE_UNFUSED_ATTN"] = "1" + _attention_backends["backend_selection_requires_update"] = True + block = TransformerLayer( hidden_size, 4 * hidden_size, From cef71d859186caa0657219c09edeeb65fd11f558 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Sat, 5 Sep 2026 23:58:50 -0700 Subject: [PATCH 22/28] chore(pytorch): tighten softcap comments and correct the FA4 attribution Halve the comment volume added to backends.py and utils.py, keeping only the facts that are not inferable from the code: that FA2 adds ALiBi outside the cap, and why the cap is divided by layer_number. Which backend supports what is already expressed by get_attention_backend and its debug messages. Also correct two inaccuracies found while auditing the surviving text. FA4 does implement softcap on both entry points TE calls, forward and backward, so it is disqualified because TE does not plumb the cap to it, not because the library lacks it; the log message said otherwise. And the FA3 signature probe detects whether the installed release exposes the kwarg at all, not whether the build set FLASHATTENTION_DISABLE_SOFTCAP, which still exposes it and rejects a nonzero cap at dispatch. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 15 ++++------ .../dot_product_attention/backends.py | 20 ++++++------- .../attention/dot_product_attention/utils.py | 28 ++++--------------- 3 files changed, 19 insertions(+), 44 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 0493d12ace..f67d840cbc 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -933,12 +933,10 @@ def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend, def test_transformer_layer_softcap_plumbing(dtype): """Test that TransformerLayer forwards softcap to both of its attention modules. - The value only has to arrive; the numerics are covered by the DotProductAttention tests - above. Cross-attention is the half worth asserting: it is reached through a separate call - site from self-attention, so a refactor can drop the cap there while self-attention keeps - working, and nothing else in the suite would notice. + Numerics are covered above; this only checks the value arrives. Cross-attention is + reached through a separate call site, so a refactor can drop the cap there while + self-attention keeps working and nothing else in the suite would notice. """ - # head_dim 64, matching the other softcap configs, so every backend can serve the shape. hidden_size, num_heads, seqlen, batch_size = 256, 4, 32, 2 seen = {} @@ -948,11 +946,8 @@ def hook(_module, _args, kwargs): return hook - # Every softcap test above sets these and none restore them, so inherit nothing: with a - # nonzero cap the filter also drops FusedAttention and FA4, and a leftover - # NVTE_UNFUSED_ATTN=0 would leave no backend at all on a machine without flash-attn. - # Leaving flash and unfused both enabled keeps this backend-agnostic; the hook fires on - # DotProductAttention regardless of which one is selected. + # Set explicitly rather than inheriting: the tests above leave these set, and a stale + # NVTE_UNFUSED_ATTN=0 would leave no eligible backend once softcap drops the fused ones. reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 582a8a2d44..346167883a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -167,11 +167,9 @@ fa_utils.set_flash_attention_3_params() - # Probe whether this FA3 build exposes a `softcap` parameter on BOTH entry points. FA3's Hopper - # (sm90) kernels DO implement tanh logit softcapping in fwd AND bwd (dedicated - # flash_{fwd,bwd}_hdim256_bf16_softcap_sm90 instantiations, off only behind a compile-time - # DISABLE_SOFTCAP flag), so this is a mature path. Still fail-closed and additionally - # gated on head_dim <= 256 + non-CP in get_attention_backend. + # Older FA3 releases expose no `softcap` kwarg, so probe the API rather than the version. + # This cannot see a FLASHATTENTION_DISABLE_SOFTCAP build: that still exposes the kwarg and + # rejects a nonzero cap at dispatch. try: fa_utils.fa3_supports_softcap = ( "softcap" in inspect.signature(flash_attn_func_v3).parameters @@ -641,8 +639,7 @@ def _forward( key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1) # Raw attention scores. [b * h, sq, sk] - # An additive `post_scale_bias`/ALiBi term is deferred until after the softcap below, so - # that the cap applies to the bare scaled logits (see the softcap comment for why). + # `post_scale_bias`/ALiBi are deferred until after the softcap below; see the cap. deferred_bias = None if core_attention_bias_type == "no_bias": matmul_result = torch.baddbmm( @@ -690,11 +687,10 @@ def _forward( matmul_result = matmul_result.view(*output_size) deferred_bias = core_attention_bias - # Cap the scaled logits: softcap * tanh(scores * scale / softcap), matching how - # FlashAttention folds softmax_scale into its tanh argument. The cap must land on the - # bare scaled logits, before any additive bias: FA2 caps right after the QK^T gemm and - # adds ALiBi afterwards, so capping the bias too would diverge from it. qk layer scaling - # defers the layer_number factor to the softmax below, so it is divided out of the cap. + # The cap lands on the scaled logits before `post_scale_bias`/ALiBi: FA2 caps right + # after the QK^T gemm and adds ALiBi afterwards, so capping those would diverge from it. + # `pre_scale_bias` is folded in before the scaling, so it stays inside the cap. qk layer + # scaling defers the layer_number factor to the softmax below, so divide it out here. if softcap != 0.0: cap = softcap / self.layer_number if apply_qk_layer_scaling else softcap matmul_result = cap * torch.tanh(matmul_result / cap) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index fa70884f93..13ec638d69 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -148,10 +148,7 @@ class FlashAttentionUtils: v4_is_installed = False fa4_version = PkgVersion("0") use_v4 = False - # True only if the installed FA3 build exposes a `softcap` parameter (signature probe in - # backends.py, fail-closed default False). Necessary-but-not-sufficient: FA3 softcap is also - # gated on head_dim <= 256 and non-CP in get_attention_backend. FA3 is already restricted to - # Hopper (sm90) upstream, where its softcap fwd+bwd kernels are mature. + # Set by a signature probe in backends.py; fail-closed default. fa3_supports_softcap = False v4_installation_steps = """\ pip install flash-attn-4==4.0.0b11 nvidia-cutlass-dsl[cu13]""" @@ -775,33 +772,21 @@ def _disable_all_flash_attention() -> None: logger.debug("Disabling all backends for max_logit with FP8 attention") # Filter: softcap - # The scalar `softcap` kwarg (tanh logit softcapping) is plumbed to the FlashAttention 2 - # backend (>= 2.6.0) and to UnfusedDotProductAttention by default, and to FA3 subject to the - # build/shape checks below. FusedAttention does not take the scalar kwarg (cuDNN can softcap - # via score_mod, but that path is not used here), and FA4 has no softcap kernel to call, so - # disable both rather than silently dropping the cap. + # Disable any backend that would not honour a nonzero cap, rather than silently dropping it. if softcap != 0.0: if use_fused_attention: logger.debug("Disabling FusedAttention as it does not support softcap") use_fused_attention = False if use_flash_attention_4: - # FA4 exposes no softcap kwarg and its head_dim=256 kernel asserts score_mod is None, - # so there is no kernel to route the cap through, and the FA4 call path in backends.py - # passes no softcap -- selecting it here would silently drop the cap. if FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 as it does not support softcap") + # FA4 implements softcap; TE does not plumb it to the FA4 call path yet. + logger.debug("Disabling FlashAttention 4 as TE does not pass it softcap") use_flash_attention_4 = False if use_flash_attention_3 and not ( FlashAttentionUtils.fa3_supports_softcap and max(head_dim_qk, head_dim_v) <= 256 and not context_parallel ): - # FA3 softcap requires a softcap-capable FA3 build, head_dim <= 256 (the range FA3's - # sm90 softcap kernels are instantiated for), and no context parallelism -- FA3's CP - # path hard-rejects nonzero softcap (backends.py), so selecting it here would just - # crash at dispatch instead of steering to FA2, which does support CP+softcap via - # context_parallel.py's autograd threading. Whether FA3 is eligible at all is governed - # by NVTE_FLASH_ATTN_V3 through use_flash_attention_3. logger.debug( "Disabling FlashAttention 3 for softcap (requires softcap-capable FA3 build, " "head_dim <= 256, and no context parallelism)" @@ -811,9 +796,8 @@ def _disable_all_flash_attention() -> None: logger.debug("Disabling FlashAttention 2 for softcap (requires flash-attn >= 2.6.0)") use_flash_attention_2 = False if use_flash_attention_2 and attention_dropout != 0.0 and is_training: - # FA2 hard-rejects a nonzero softcap combined with nonzero dropout at dispatch - # ("Softcapping does not support dropout for now", flash_api.cpp). Dropout only reaches - # the kernel while training -- backends.py passes 0.0 in eval -- hence the is_training. + # FA2 rejects softcap with dropout at dispatch (flash_api.cpp). Dropout only + # reaches the kernel while training, hence the is_training guard. logger.debug("Disabling FlashAttention 2 for softcap with dropout") use_flash_attention_2 = False From 8805255bc07113d62d95cbe78fccd8acab94122f Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Sun, 6 Sep 2026 00:05:44 -0700 Subject: [PATCH 23/28] fix(pytorch): probe softcap on all three FA3 entry points softcap is injected into fa_3_optional_forward_kwargs without regard to which entry point was selected, so the probe should cover every function the cap can reach. flash_attn_with_kvcache is the one chosen when inference_params is set, and it was validated only by proxy through the other two. No shipped FA3 build changes behaviour: all three entry points gained softcap in 2.7.3/hopper, which is TE's documented FA3 floor. Before that none of them had it. The symbol is already imported and None-guarded, so the probe cannot fire on an uninstalled FA3. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- .../pytorch/attention/dot_product_attention/backends.py | 1 + 1 file changed, 1 insertion(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 346167883a..47f3c9933a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -174,6 +174,7 @@ fa_utils.fa3_supports_softcap = ( "softcap" in inspect.signature(flash_attn_func_v3).parameters and "softcap" in inspect.signature(flash_attn_varlen_func_v3).parameters + and "softcap" in inspect.signature(flash_attn_with_kvcache_v3).parameters ) except (ValueError, TypeError): fa_utils.fa3_supports_softcap = False From dbb36c335636ba44b2ede6beb38736800abc773b Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 7 Sep 2026 10:14:39 -0700 Subject: [PATCH 24/28] test: exercise softcap at head_dim 128, not only 64 Cutting model_configs_softcap to two configs dropped head_dim 128 as a side effect, leaving every softcap config at 64. FA2 and FA3 compile a separate softcap kernel per head_dim bucket, and above 128 FA2 additionally forces the non-even-MN predicated path, so a bug in those kernels is invisible at 64. Widen the existing padding config instead of adding a third, so the count stays at two. Neither justification for the pair changes: head_dim plays no part in the flash_attn_func vs flash_attn_varlen_func dispatch, and the 0.01 cap rationale is head_dim invariant because softmax_scale cancels the sqrt(d) growth of the dot product. 128 rather than 256 because get_attention_backend disables FA3 for deterministic backward at head_dim >= 256, which would leave only UnfusedDotProductAttention and skip the comparison in the NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 CI pass. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index f67d840cbc..e5fb8e22cf 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -658,7 +658,9 @@ def test_dpa_softmax_thd(dtype, model_configs, model): # logits at O(1e-2) whatever the head dim, so tanh is numerically linear at a Gemma-sized # cap. A cap of 0.01 is the one regime these inputs can distinguish. Softcapping in tanh's # saturating region is covered by test_dpa_softcap_vs_reference, which uses its own inputs. - "softcap_3_1": ModelConfig(2, 512, 16, 64, attn_mask_type="padding_causal", softcap=0.01), + # head_dim 128 rather than 64: FA2 and FA3 compile a separate softcap kernel per head_dim, + # and the logit scale above is head_dim invariant since softmax_scale cancels the sqrt(d). + "softcap_3_1": ModelConfig(2, 512, 16, 128, attn_mask_type="padding_causal", softcap=0.01), } From 590c7ca05d3006b212c69825262e76806a54a7c0 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 7 Sep 2026 10:32:10 -0700 Subject: [PATCH 25/28] test: restore backend env vars in the softcap plumbing test The test set NVTE_FLASH_ATTN, NVTE_FUSED_ATTN and NVTE_UNFUSED_ATTN and never restored them, leaking NVTE_FUSED_ATTN=0 into every test that ran after it. Its own comment names that hazard as the reason it sets them explicitly, then reproduces it. Wrap the body in try/finally, restore the previous values, and invalidate the backend selection cache on the way out. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 74 ++++++++++++++--------- 1 file changed, 45 insertions(+), 29 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index e5fb8e22cf..b9637b0601 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -766,6 +766,8 @@ def _softcap_reference_attention( # "plain" checks the cap itself. The other two pin down *where* the cap is applied inside # UnfusedDotProductAttention, so they run on that backend only, at a cap that saturates. +_SOFTCAP_BACKEND_ENV = ("NVTE_FLASH_ATTN", "NVTE_FUSED_ATTN", "NVTE_UNFUSED_ATTN") + softcap_variants = ["plain", "bias_outside_cap", "qk_layer_scaling"] # Large enough that a cap of softcap * layer_number is far from a cap of softcap for O(1) @@ -950,44 +952,58 @@ def hook(_module, _args, kwargs): # Set explicitly rather than inheriting: the tests above leave these set, and a stale # NVTE_UNFUSED_ATTN=0 would leave no eligible backend once softcap drops the fused ones. + # Restored in the finally below so this test does not do to others what they did to it. + backend_env = {k: os.environ.get(k) for k in _SOFTCAP_BACKEND_ENV} reset_rng_states() os.environ["NVTE_FLASH_ATTN"] = "1" os.environ["NVTE_FUSED_ATTN"] = "0" os.environ["NVTE_UNFUSED_ATTN"] = "1" _attention_backends["backend_selection_requires_update"] = True - block = TransformerLayer( - hidden_size, - 4 * hidden_size, - num_heads, - layer_type="decoder", - softcap=50.0, - params_dtype=dtype, - device="cuda", - ) - block.self_attention.core_attention.register_forward_pre_hook(_record("self"), with_kwargs=True) - block.inter_attention.core_attention.register_forward_pre_hook( - _record("cross"), with_kwargs=True - ) + try: + block = TransformerLayer( + hidden_size, + 4 * hidden_size, + num_heads, + layer_type="decoder", + softcap=50.0, + params_dtype=dtype, + device="cuda", + ) + block.self_attention.core_attention.register_forward_pre_hook( + _record("self"), with_kwargs=True + ) + block.inter_attention.core_attention.register_forward_pre_hook( + _record("cross"), with_kwargs=True + ) - hidden_states = torch.randn( - seqlen, batch_size, hidden_size, dtype=dtype, device="cuda", requires_grad=True - ) - forward_kwargs = dict( - encoder_output=hidden_states, - enc_dec_attn_mask=torch.zeros(batch_size, 1, 1, seqlen, dtype=torch.bool, device="cuda"), - ) + hidden_states = torch.randn( + seqlen, batch_size, hidden_size, dtype=dtype, device="cuda", requires_grad=True + ) + forward_kwargs = dict( + encoder_output=hidden_states, + enc_dec_attn_mask=torch.zeros( + batch_size, 1, 1, seqlen, dtype=torch.bool, device="cuda" + ), + ) - # The constructor value reaches both attention modules. - block(hidden_states, **forward_kwargs) - assert seen["self"] == 50.0, f"self-attention saw softcap={seen['self']}, expected 50.0" - assert seen["cross"] == 50.0, f"cross-attention saw softcap={seen['cross']}, expected 50.0" + # The constructor value reaches both attention modules. + block(hidden_states, **forward_kwargs) + assert seen["self"] == 50.0, f"self-attention saw softcap={seen['self']}, expected 50.0" + assert seen["cross"] == 50.0, f"cross-attention saw softcap={seen['cross']}, expected 50.0" - # A forward override wins over the constructor, for both. - seen.clear() - block(hidden_states, softcap=10.0, **forward_kwargs) - assert seen["self"] == 10.0, f"self-attention saw softcap={seen['self']}, expected 10.0" - assert seen["cross"] == 10.0, f"cross-attention saw softcap={seen['cross']}, expected 10.0" + # A forward override wins over the constructor, for both. + seen.clear() + block(hidden_states, softcap=10.0, **forward_kwargs) + assert seen["self"] == 10.0, f"self-attention saw softcap={seen['self']}, expected 10.0" + assert seen["cross"] == 10.0, f"cross-attention saw softcap={seen['cross']}, expected 10.0" + finally: + for key, value in backend_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + _attention_backends["backend_selection_requires_update"] = True model_configs_mla = { From f0cd6720ac10b34345b8545b4fa15cb2998069de Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 7 Sep 2026 11:07:42 -0700 Subject: [PATCH 26/28] fix(pytorch): reject a negative softcap instead of diverging silently tanh is odd, so UnfusedDotProductAttention's `cap * tanh(x / cap)` gives a bit-identical result for -c and +c, i.e. it caps at the absolute value. FlashAttention only enables capping under `softcap > 0.0`, so it applies no cap at all. get_attention_backend gates on `softcap != 0.0`, so a negative value passed every filter and the two backends silently disagreed, which is the exact failure mode the rest of this PR exists to prevent. Raise on a negative cap where softcap is resolved, which is the choke point every caller reaches, including MultiheadAttention and TransformerLayer. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- .../attention/dot_product_attention/dot_product_attention.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index c3eebfcc6b..b4b86f4c02 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -2393,6 +2393,11 @@ def forward( # packed-THD policy path gets the constructor value too, rather than a bare None. if softcap is None: softcap = self.softcap + # A negative cap is silently inconsistent rather than harmless: tanh is odd, so + # UnfusedDotProductAttention's `cap * tanh(x / cap)` treats it as its absolute value, + # while FlashAttention only caps when `softcap > 0` and so applies none at all. + if softcap < 0.0: + raise ValueError(f"softcap must be non-negative, got {softcap}.") # checks for qkv_format if qkv_format is None: From 138ec5756055ec540841f9b8c6d8c6de5253bf70 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 9 Sep 2026 15:04:27 -0700 Subject: [PATCH 27/28] fix(pytorch): reject a non-finite softcap as well as a negative one The existing guard exists to stop the backends disagreeing silently about a cap they cannot all honour, but NaN slipped through it and caused exactly that disagreement: UnfusedDotProductAttention computes cap * tanh(x / cap) and returns NaN for every score, while FA2 and FA3 gate their kernels on softcap > 0.0, which is false for NaN, so they drop the cap and return ordinary uncapped attention. Infinity produces NaN on every backend instead, which at least fails loudly. Neither is a legitimate way to spell a disabled cap. The limit of c * tanh(x / c) as c grows is the identity, but evaluated at infinity it is inf * 0, which is NaN, so the intuition does not survive floating point. Fold both into the existing check with math.isfinite, matching how ScaledTanhSReLU and ScaledSiTUGLU validate the same shape of parameter. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- .../dot_product_attention/dot_product_attention.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index b4b86f4c02..18a9ad85ed 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -2393,11 +2393,13 @@ def forward( # packed-THD policy path gets the constructor value too, rather than a bare None. if softcap is None: softcap = self.softcap - # A negative cap is silently inconsistent rather than harmless: tanh is odd, so - # UnfusedDotProductAttention's `cap * tanh(x / cap)` treats it as its absolute value, - # while FlashAttention only caps when `softcap > 0` and so applies none at all. - if softcap < 0.0: - raise ValueError(f"softcap must be non-negative, got {softcap}.") + # A cap that is negative or non-finite is silently inconsistent rather than + # harmless, because the backends disagree about it. tanh is odd, so + # UnfusedDotProductAttention's `cap * tanh(x / cap)` applies a negative cap as its + # absolute value and yields NaN for a non-finite one, while FlashAttention caps only + # when `softcap > 0` and so drops both without a word. + if not math.isfinite(softcap) or softcap < 0.0: + raise ValueError(f"softcap must be finite and non-negative, got {softcap}.") # checks for qkv_format if qkv_format is None: From 458c03fd1b2afc2e9bf7d2e7883ba299218e965d Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Wed, 9 Sep 2026 15:23:02 -0700 Subject: [PATCH 28/28] test: pin the softcap validity guard The guard had no coverage, for either the negative case or the non-finite one. Parametrised over -1.0, inf and nan, exercising both the constructor value and the forward override, since both resolve on the same path. Mirrors test_scaled_tanh_srelu_rejects_bad_clamp_scale, which pins the same shape of check on ScaledTanhSReLU. Co-Authored-By: Claude Opus 5 Signed-off-by: Nitin Vegesna --- tests/pytorch/attention/test_attention.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index b9637b0601..6bdbcd023d 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -933,6 +933,26 @@ def test_dpa_softcap_vs_reference(dtype, model_configs, model, softcap, backend, torch.testing.assert_close(v.grad.float(), v_ref.grad.float(), **tols) +@pytest.mark.parametrize("softcap", (-1.0, float("inf"), float("nan"))) +def test_dpa_softcap_rejects_invalid(softcap): + """A cap that is negative or non-finite must raise rather than diverge by backend. + + The backends disagree about these values: unfused applies a negative cap as its + absolute value and returns NaN for a non-finite one, while flash caps only when + softcap > 0 and so drops both silently. + """ + block = DotProductAttention(4, 64, qkv_format="bshd", softcap=softcap).to(device="cuda") + q, k, v = (torch.randn(2, 32, 4, 64, dtype=torch.float16, device="cuda") for _ in range(3)) + + with pytest.raises(ValueError, match="softcap"): + block(q, k, v) + + # A forward override is validated on the same path as the constructor value. + ok = DotProductAttention(4, 64, qkv_format="bshd").to(device="cuda") + with pytest.raises(ValueError, match="softcap"): + ok(q, k, v, softcap=softcap) + + @pytest.mark.parametrize("dtype", param_types) def test_transformer_layer_softcap_plumbing(dtype): """Test that TransformerLayer forwards softcap to both of its attention modules.