diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 037d071c1d..8d1b870d3a 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.py b/tests/pytorch/attention/test_attention.py index 0f5cdfa647..b9637b0601 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, @@ -645,6 +650,362 @@ 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) + # High cap, no padding -> flash_attn_func. + "softcap_1_0": ModelConfig(4, 128, 16, 64, softcap=50.0), + # 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. + # 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), +} + + +@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. + + 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( + 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) + + 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" + + +def _softcap_reference_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + 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]: + 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 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( + 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_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_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) +# 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"]) +@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. 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, + 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" + if variant == "qk_layer_scaling": + os.environ["NVTE_APPLY_QK_LAYER_SCALING"] = "1" + _attention_backends["backend_selection_requires_update"] = True + + 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") + + dpa_kwargs, forward_kwargs, right_fn, wrong_fn = _softcap_variant_spec( + variant, config, dtype, softcap, softmax_scale, q_ref, k_ref, v_ref + ) + + 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, + **dpa_kwargs, + ).to(dtype=dtype, device="cuda") + 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 = right_fn() + + tols = dict(atol=2e-2, rtol=2e-2) + 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) + + +@pytest.mark.parametrize("dtype", param_types) +def test_transformer_layer_softcap_plumbing(dtype): + """Test that TransformerLayer forwards softcap to both of its attention modules. + + 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. + """ + 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 + + # 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 + + 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" + ), + ) + + # 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" + 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 = { # test: ModelConfig(b, sq, hq, dqk) "mla_1_0": ModelConfig(8, 128, 16, 64, head_dim_v=128), @@ -1181,6 +1542,7 @@ def make_dot_product_attention( attention_type=config.attn_type if attention_type is None else attention_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_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index af22c42016..4d6369dbee 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/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 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 bca1f3200d..47f3c9933a 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 @@ -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] 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 e9948cb9d8..dd3d06b64e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/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"]: 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 de082c8fae..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 @@ -745,6 +745,11 @@ 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. 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 @@ -844,6 +849,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, @@ -880,6 +886,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: @@ -1924,6 +1931,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, @@ -2096,6 +2104,10 @@ 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. 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. @@ -2377,6 +2389,16 @@ def forward( }: bottom_right_diagonal = True + # softcap is not mask-specific: resolve it outside the thd_mask_policies branch so the + # 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: qkv_format = self.qkv_format @@ -2750,6 +2772,7 @@ def forward( "max_seqlen_kv": max_seqlen_kv, "head_dim_qk": head_dim_qk, "head_dim_v": head_dim_v, + "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, @@ -2910,6 +2933,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, @@ -3044,6 +3068,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, @@ -3068,6 +3093,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 33d612b4f4..13ec638d69 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -148,6 +148,8 @@ class FlashAttentionUtils: v4_is_installed = False fa4_version = PkgVersion("0") use_v4 = False + # 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]""" v4_warning_printed = False @@ -229,6 +231,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, 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 @@ -289,6 +294,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 +439,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 +771,36 @@ def _disable_all_flash_attention() -> None: use_unfused_attention = False logger.debug("Disabling all backends for max_logit with FP8 attention") + # Filter: softcap + # 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: + if FlashAttentionUtils.v4_is_installed: + # 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 + ): + logger.debug( + "Disabling FlashAttention 3 for softcap (requires softcap-capable FA3 build, " + "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 + if use_flash_attention_2 and attention_dropout != 0.0 and 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 + # 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") diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index 82221f0c83..e07e203ca7 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 @@ -754,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. @@ -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..db9a352881 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 @@ -677,6 +685,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) @@ -711,6 +720,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 +831,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 +917,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 +956,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,