Thread tanh logit softcapping through FlashAttention (FA2, opt-in FA3)

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,
@@ -645,6 +650,197 @@ 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),
+    # 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.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,
+        fwd_only_without_fused_attn=False,
+    )
+
+
+@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),
@@ -1447,6 +1643,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()

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,

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
@@ -166,6 +167,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 head_dim <= 256 + non-CP 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"))
@@ -435,6 +449,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,
@@ -673,6 +688,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(
@@ -894,6 +916,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,
@@ -1110,6 +1133,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,
@@ -1140,6 +1168,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,
@@ -1237,6 +1266,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"] = (
@@ -1257,9 +1288,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) + 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]