From 5917f0d39c6552699b57e17d23e35de662e64529 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 17 Aug 2026 14:20:36 -0700 Subject: [PATCH 1/3] 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 5ae46ce59545d610e7842b3427ca8e5499b617e9 Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 17 Aug 2026 14:34:54 -0700 Subject: [PATCH 2/3] 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 af4587630530709c20c61f911a104f8c3478afec Mon Sep 17 00:00:00 2001 From: Nitin Vegesna Date: Mon, 17 Aug 2026 14:35:45 -0700 Subject: [PATCH 3/3] [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) --- 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."), ]