Thread tanh logit softcapping through FlashAttention (FA2 and FA3)

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()

tests/pytorch/attention/test_attention_with_cp.py

@@ -418,6 +418,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

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,18 @@
 
     fa_utils.set_flash_attention_3_params()
 
+    # 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
+            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
+
 # Try to import Flash Attention v4
 try:
     fa_utils.fa4_version = PkgVersion(get_pkg_version("flash-attn-4"))
@@ -435,6 +448,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,
@@ -626,6 +640,8 @@ def _forward(
         key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1)
 
         # Raw attention scores. [b * h, sq, sk]
+        # `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(
                 matmul_result,
@@ -669,9 +685,19 @@ 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
+
+        # 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)
+
+        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
@@ -894,6 +920,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 +1137,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 +1172,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 +1270,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 +1292,17 @@ def forward(
                         **fa_optional_forward_kwargs,
                     )
                 else:
+                    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:
+                        # 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"] = (
                             cu_seqlens_q[1:] - cu_seqlens_q[:-1]

transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py

@@ -1617,6 +1617,7 @@ def forward(
         deterministic,
         use_fused_attention,
         return_max_logit,
+        softcap,
         fp8,
         fp8_meta,
         cp_group,
@@ -1906,7 +1907,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]
@@ -2399,6 +2400,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
@@ -2705,7 +2707,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):
@@ -3212,6 +3214,7 @@ def backward(ctx, dout, *_args):
             None,
             None,
             None,
+            None,
         )
 
 
@@ -3289,6 +3292,7 @@ def forward(
         deterministic,
         use_fused_attention,
         return_max_logit,
+        softcap,
         window_size,
         cp_group,
         cp_stream,
@@ -3382,7 +3386,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
 
@@ -3974,6 +3978,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.use_flash_attn_4 = use_flash_attn_4
@@ -4183,7 +4188,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
                 if (
                     ctx.qkv_format == "thd"
                     and ctx.load_balancing_strategy is CPLoadBalancingStrategy.NO_LOAD_BALANCE
@@ -4571,6 +4576,7 @@ def backward(ctx, dout, *_args):
             None,
             None,
             None,
+            None,
         )
 
 
@@ -4602,6 +4608,7 @@ def forward(
         deterministic,
         use_fused_attention,
         return_max_logit,
+        softcap,
         window_size,
         fp8,
         fp8_meta,
@@ -4703,7 +4710,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__
@@ -5028,6 +5035,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
@@ -5178,7 +5186,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:
@@ -5406,6 +5414,7 @@ def backward(ctx, dout, *_args):
             None,
             None,
             None,
+            None,
             d_softmax_offset,
             None,
         )
@@ -5435,6 +5444,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,
@@ -5621,6 +5631,7 @@ def attn_forward_func_with_cp(
         deterministic,
         use_fused_attention,
         return_max_logit,
+        softcap,
     ]
 
     if cp_comm_type in ["p2p", "a2a+p2p"]: