Thread tanh logit softcapping through FlashAttention (FA2 and FA3)

[nvegesna-netizen]

Adds a softcap kwarg to DotProductAttention so models with attention-logit soft-capping (cap·tanh(x/cap), e.g. Gemma2) can run on TE's fused flash-attention kernels instead of falling back to an unfused/non-TE path.

Companion PRs (needed together for an end-to-end model to pick this up):

• Megatron-LM: exposes TransformerConfig.attn_logit_softcapping and maps it to this softcap kwarg
• Megatron-Bridge: adds an opt-in Gemma2 attention path that sets attn_logit_softcapping

What changed

• DotProductAttention (init + forward) and AttentionParams gain a softcap: float = 0.0 kwarg. softcap=0.0 is a no-op — existing behavior is unchanged.
• Threaded through the FA2 non-CP path and all three context-parallel autograd functions (forward + ctx-saved backward).
• UnfusedDotProductAttention applies cap * tanh(scores / cap) to the already-scaled logits, so it serves as the in-tree reference implementation (and the numerical reference the tests compare the flash backends against). With qk-layer-scaling the layer_number factor is divided out of the cap, since that path defers scaling to the softmax.
• get_attention_backend: when softcap != 0, FusedAttention (cuDNN), FA4, and FA3 (unless the checks below pass) are disqualified and selection steers to FA2 or unfused (also disqualifies FA2 builds too old to carry the softcap kernel). This is a deliberate safety net — softcap must never be silently dropped, nor hit a NotImplementedError at runtime.
• FA3 (Hopper only): gated on a build-capability probe (softcap present on all three FA3 entry points TE dispatches to), max(head_dim_qk, head_dim_v) <= 256, and non-CP; forward threads softcap into FA3's kwargs, backward is handled automatically by the existing Hopper autograd function. No new env var — FA3 eligibility is governed by the existing NVTE_FLASH_ATTN_V3 (default 1). Since TE already prefers FA3 over FA2 on sm90, FA3 is the default softcap backend on Hopper when a softcap-capable FA3 build is installed; that is intentional. NVTE_FLASH_ATTN_V3=0 steers to FA2.
• Tests: three tests in tests/pytorch/attention/test_attention.py. test_dpa_softcap sweeps the available backends through the existing test_dot_product_attention harness (forward + backward parity against the unfused reference); softcap always disqualifies FusedAttention, so it opts out of the harness's fused-unavailable fallback to keep the dQ/dK/dV comparison. test_dpa_softcap_zero_backend_selection asserts softcap=0.0 leaves FusedAttention selectable and a nonzero cap does not. test_dpa_softcap_vs_reference compares forward and dQ/dK/dV against a closed-form pure-PyTorch oracle one backend at a time, so UnfusedDotProductAttention stays covered on machines without flash-attn; it uses its own randn inputs because the shared harness's 0.1 * randn puts logits at O(1e-2), where a cap of 50 moves the output by ~1e-8 and a dropped cap would be undetectable.

Supported configurations

• FA2, non-CP — supported with flash-attn >= 2.6.0 (first FA2 release exposing a softcap kwarg). The default path wherever FA3 is not eligible (i.e. off sm90, or no softcap-capable FA3 build). Requires zero attention dropout while training — see below.
• FA2 + context parallelism — supported for all four cp_comm_type values: p2p, all_gather, a2a, a2a+p2p (p2p and a2a+p2p share one autograd function). Same flash-attn >= 2.6.0 requirement; forward and backward both carry softcap.
• FA3, non-CP — requires an FA3 build whose flash_attn_func, flash_attn_varlen_func and flash_attn_with_kvcache all expose softcap (signature probe, fail-closed) and max(head_dim_qk, head_dim_v) <= 256. FA3 is Hopper (sm90)-only upstream and governed by the existing NVTE_FLASH_ATTN_V3 (default 1). When eligible it takes precedence over FA2 on Hopper — deliberate; both paths were exercised on Hopper (see Validation). Any check failing steers to FA2 (>= 2.6.0) or unfused.
• UnfusedDotProductAttention, non-CP — supported; the reference/fallback path, used when no flash backend is eligible (no version floor). Unfused attention does not support context parallelism at all, independent of softcap.

Not supported:

• FusedAttention (cuDNN) — disqualified in get_attention_backend when softcap != 0.0, so selection steers to FA2 (or unfused) rather than silently dropping the cap.
• FA3 + CP — raises NotImplementedError (use FA2).
• FA4 — disqualified in get_attention_backend when softcap != 0.0. FA4 does implement softcap, but TE does not plumb the cap into the FA4 call path (fa_4_optional_forward_kwargs never carries it), so without this filter FA4 would be selected on SM100 (NVTE_FLASH_ATTN_V4 defaults to 1) and silently drop it. See the follow-up section below.
• FA2 + nonzero attention dropout, while training — flash-attn rejects a nonzero softcap combined with nonzero dropout at dispatch ("Softcapping does not support dropout for now", csrc/flash_attn/flash_api.cpp), so FA2 is disqualified and selection steers to unfused. Dropout reaches the kernel as 0.0 in eval, so inference configs are unaffected. Since unfused does not support CP, CP + softcap + dropout while training has no eligible backend and raises rather than crashing inside flash-attn.

ONNX export force-selects the unfused backend, which now honors softcap via torch.tanh (exportable as the ONNX Tanh op), so export is expected to work — but there is no ONNX softcap test in this PR.

softcap = 0.0 (the default) disables softcapping: backend selection and numerics are identical to today.

Validation

• Unit: test_dpa_softcap — forward and backward parity against the unfused reference, across whichever backends are available on the test machine.
• Unit: test_dpa_softcap_vs_reference — forward and dQ/dK/dV against a closed-form reference, for softcap in {0.0, 0.5}, with logits at O(1) so tanh runs in its saturating region. softcap=0.0 compares against a reference that never applies tanh, and the nonzero case asserts the cap moves the reference output by more than the comparison tolerance, so the test cannot pass an implementation that drops the cap.
• Unit: test_dpa_softcap_zero_backend_selection — the softcap=0.0 no-op claim for backend selection, which is the half the filter actually changed.
• End-to-end: exercised via the companion Megatron-LM/Megatron-Bridge changes above — multi-step training runs on both a Hopper and a Blackwell target, confirmed numerically consistent with the prior unfused path, and confirmed via kernel-level profiling that the fused flash kernels (not the fallback path) were selected at runtime.

On the Hopper target both the FA2 and the FA3 softcap paths were exercised. FA3 taking precedence over FA2 there is an intentional design choice.

Follow-up, not in this PR: FA4

Wiring softcap through to FA4 is left to a separate PR rather than folded in here. FA4 exposes a softcap argument on the non-CP entry points TE calls and implements it for forward and backward, so this is a plumbing task on the TE side rather than a missing kernel. Three Blackwell shapes are excluded: the SM100 forward and the SM100 backward at head_dim=256, and the MLA path with head_dim_v=512. Until that plumbing lands, get_attention_backend disqualifies FA4 whenever the cap is nonzero, so it can never be selected and silently drop the cap.

[greptile-apps]

Greptile Summary

Adds tanh attention-logit softcapping throughout TransformerLayer, MultiheadAttention, and DotProductAttention, with backend capability filtering and reference coverage.

• Threads softcap through FA2, eligible FA3 paths, and all context-parallel FlashAttention autograd functions.
• Implements softcapping in the unfused backend, including ONNX-compatible behavior and qk-layer-scaling handling.
• Prevents unsupported fused, FA4, legacy FA2, FA3-context-parallel, and training-dropout configurations from silently ignoring the cap.
• Adds backend-selection, numerical, gradient, context-parallel, and high-level plumbing tests.

Confidence Score: 5/5

The PR appears safe to merge; no actionable regressions remain in the reviewed changes.

No new correctness, security, or repository-rule violations were identified. All previous findings were manually resolved, and the latest change consistently rejects negative softcap values before backend dispatch.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Implements unfused softcapping and passes the cap to supported FA2 and FA3 entry points.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Carries softcap through the three context-parallel FA2 autograd implementations in forward and backward.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds constructor and forward softcap plumbing with non-negative input validation.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds softcap to backend parameters and filters backends that cannot honor it.
transformer_engine/pytorch/attention/multi_head_attention.py	Propagates constructor and per-forward softcap values into DotProductAttention.
transformer_engine/pytorch/transformer.py	Propagates softcap to both self-attention and decoder cross-attention.
tests/pytorch/attention/test_attention.py	Adds numerical, gradient, backend-selection, bias-ordering, layer-scaling, and high-level plumbing coverage.
tests/pytorch/attention/test_attention_with_cp.py	Adds context-parallel softcap coverage across the distinct FA2 communication implementations.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[DotProductAttention with softcap] --> B{softcap is zero?}
    B -->|Yes| C[Existing backend selection]
    B -->|No| D{Eligible FA3 build and shape and non-CP?}
    D -->|Yes| E[FlashAttention 3]
    D -->|No| F{FA2 >= 2.6 and supported dropout mode?}
    F -->|Yes| G[FlashAttention 2]
    F -->|No| H{Context parallel?}
    H -->|No| I[Unfused attention]
    H -->|Yes| J[No eligible backend: raise]

Loading

_{Reviews (36): Last reviewed commit: "fix(pytorch): reject a negative softcap ..." | Re-trigger Greptile}

[cyanguwa]

Please follow this link to fix the DCO of this PR as well. Thanks!

https://github.com/NVIDIA/TransformerEngine/pull/3391/checks?check_run_id=97902991433

[cyanguwa]

/te-ci pytorch L0 L1 L3

[nvegesna-netizen]

Triaged the te-ci failures. One was ours; the rest reproduce on unrelated PRs. Fix pushed in 8e49940e.

Ours: L0_pytorch_unittest--A100_1GPU (and part of --L40_1GPU)

FAILED test_attention.py::test_dpa_softcap_zero_backend_selection[softcap_1_0-...]
>   assert fused_off, "softcap=0.0 must not disable FusedAttention"
E   AssertionError: softcap=0.0 must not disable FusedAttention
Error: sub-test failed: NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py

Only the deterministic pass fails; the normal test_attention.py run passes on the same machine (5865 passed). The test asserted that FusedAttention is available at softcap=0.0, but availability depends on arch and on NVTE_ALLOW_NONDETERMINISTIC_ALGO, not on softcap — cuDNN offers no fused backend for this config on A100/L40 in deterministic mode, so the assert fires. It passes on H100 and B200 because a fused backend exists there.

That half is now a pytest.skip rather than an assert. The direction the test actually exists to check — a nonzero softcap must disable FusedAttention — still runs wherever a fused backend is available, and skipping prevents the comparison from being silently vacuous where one isn't.

Not ours

job	failure	evidence
L0_pytorch_unittest--L40_1GPU	test_numerics.py::test_layernorm_linear_accuracy[...] numeric drift (e.g. -2.5 vs -2.765625)	also fails on #3466, #3465, #3457
L1_pytorch_distributed_unittest--{H100_4GPU,B200_8GPU}	test_fusible_ops.py::test_distributed_fuser_ops[2]	also fails on #3466, #3457
L1_pytorch_mcore_fsdp_integration--B200_8GPU	No route to host on the c10d TCPStore rendezvous	infrastructure
L0_pytorch_lint	<EMPTY LOG>	no output produced

None of those touch attention.

[nvegesna-netizen]

/te-ci pytorch L0 L1 L3

[cyanguwa]

softcap is only exposed through DotProductAttention right now. Could you please also thread it through MultiheadAttention and TransformerLayer? Thanks.

[nvegesna-netizen]

Added to both. Each takes a softcap constructor argument stored on the module plus a forward override that falls back to it, matching how window_size is threaded. TransformerLayer passes it to both self_attention and inter_attention, so decoder cross attention is capped too. Resolution sits outside the thd_attention_policies branch at both levels, since softcap is not mask specific. Defaults are 0.0 and None, so existing callers are unaffected.

Two things worth flagging.

I appended softcap to the end of both forward signatures rather than grouping it with the other attention arguments. Inserting it mid signature shifted encoder_output and is_first_microbatch, so a downstream caller passing those positionally would have silently bound the wrong value. Nothing in the repo does that, but both classes are public API. Constructors keep it grouped with window_size, since it sits past every positional tuple used to build these modules.

I also added test_transformer_layer_softcap_plumbing, which hooks both core_attention modules in a decoder layer and asserts each receives the cap, from the constructor and from a forward override. Writing it showed that passing softcap through common_attention_kwargs was partly masking rather than redundant: it pre-set each MultiheadAttention's self.softcap, so deleting the cross attention pass through still satisfied the constructor assertion, since the fallback picked the value back up. The forward override assertion still failed, so the test caught it either way, but only on one of its two halves. Removing that entry makes both halves load bearing, and leaves softcap matching window_size exactly.

The test is more than you asked for and there is no precedent for it here, so easy to drop. My reason for adding it is that a merge from main during this branch's life already silently dropped softcap from attention_params_kwargs, and only a manual check caught it.

Validated on an H100 since writing the above. With the softcap suite green as the baseline, 26 passed, deleting the cross attention pass through fails test_transformer_layer_softcap_plumbing, and restoring it returns the suite to green. The other two arms behave the same way: capping the bias fails the bias_outside_cap variant, and dropping the cap / layer_number division fails the qk_layer_scaling variant. So this is measured rather than inferred.

[cyanguwa]

Thanks for thoroughly designing the tests, but I wonder if we could consolidate them a little bit, without losing meaningful test signals. How about this:

• keep only 2 configs in model_configs_softcap: softcap_1_0 for the high softcap value and to test flash_attn_func API; softcap_3_1 (with "padding_causal" mask though) for a low softcap value and to test flash_attn_varlen_func API
• fold model_configs_softcap_reference from three configs to two, by merging GQA into the causal config; they are orthogonal to softcapping
• merge test_dpa_softcap_vs_reference, test_dpa_softcap_bias_ordering and test_dpa_softcap_qk_layer_scaling into one test, parametrized over variants {"plain", "bias_outside_cap", "qk_layer_scaling"}; we can extract a helper function to do the repeated work: force a backend via env vars, build randn q/k/v, construct the module with the cap, compute a right and a wrong variant, assert the two differ by more than 10 * atol, assert TE matches the right one, return the DPA kwargs, forward kwargs and the two closures
• drop the causal variant from the bias-ordering and qk-layer-scaling cases; the other tests should be sufficient for testing whether the additive bias sits outside the tanh and whether the cap is divided by layer_number

Thanks!

[nvegesna-netizen]

Done in a7b5160, following all four bullets.

model_configs_softcap is down to softcap_1_0 and softcap_3_1, the latter switched to padding_causal as you suggested so it goes through flash_attn_varlen_func. model_configs_softcap_reference is down to two, with GQA merged into the causal config. The three tests are now one test_dpa_softcap_vs_reference parametrized over {plain, bias_outside_cap, qk_layer_scaling}, with a _softcap_variant_spec helper returning the DPA kwargs, the forward kwargs and a right and a wrong reference closure. The shared work of forcing the backend, building inputs, and asserting the two references differ by more than 10 * atol now happens once. The bias variant reuses _softcap_reference_attention via new bias and cap_includes_bias arguments instead of duplicating its GQA and causal handling, which is what made the merged version shorter than the three it replaces. Causal is dropped from both special variants, as you asked.

The matrix is 48 cases on a bf16 capable GPU, since param_types picks up bf16 when available, with 19 running and 29 skipping and each special variant running once on the non causal reference config. qk_layer_scaling is fp16 only. On an fp16 only machine it is 24, 10 and 14.

Two things I checked before making the change, since consolidating tests can quietly remove the signal they exist for. Both variants still fail if their bug is reintroduced: capping the bias separates the two references by about 0.98, and dropping the cap / layer_number division by about 0.44, against the 10 * atol threshold, which is 0.20 in fp16 and 0.40 in bf16. Those margins hold because the surviving softcap_ref_1_0 is the same shape the originals ran on. I also hoisted layer_number to _SOFTCAP_QK_LAYER_NUMBER = 8 with a note that 3 is too small to separate, because that value is load bearing and would otherwise look arbitrary to the next person.

On head dimensions, the two surviving configs cover 64 and 128: dbb36c33 moves softcap_3_1 to head_dim 128, since FA2 and FA3 compile a separate softcap kernel per head_dim and the input scale is head_dim invariant.

I kept backward on the plain variant only, matching what runs today. Extending it to all three would be more coverage, but it is behaviour that has never run on hardware and I would rather not fold that in under a consolidation.

[nvegesna-netizen]

/te-ci pytorch

[nvegesna-netizen]

Tightened the softcap comments in backends.py and utils.py in cef71d85. Comment lines this PR adds to those two files go from 35 to 15, though total added lines only fall from 102 to 82, and the cut is uneven: utils.py loses about three quarters of its comment lines, backends.py about a third. I kept the two facts that are not inferable from the code, that FA2 adds ALiBi outside the cap and why the cap is divided by layer_number.

Two corrections to things this PR previously asserted, and one further fix.

FA4 does implement softcap. Both non-CP entry points TE calls expose it, lowered to a score mod for forward and backward. It is excluded on three Blackwell shapes rather than one: SM100 forward and SM100 backward at head_dim=256, and the MLA head_dim_v=512 path. TE disqualifies FA4 for softcap because fa_4_optional_forward_kwargs never carries the cap, so the filter is correct, but the log line attributed the limitation to the library and I have fixed it. Also, TE does not pin an FA4 version; 4.0.0b11 is only what the install hint suggests, and backends.py has explicit handling for 4.0.0b24.

The FA3 probe establishes API presence, not build configuration. It reads the Python wrapper, so a FLASHATTENTION_DISABLE_SOFTCAP build still passes it and rejects a nonzero cap at dispatch instead.

The probe now covers all three FA3 entry points (8805255b). The cap is injected into fa_3_optional_forward_kwargs without regard to which entry point was selected, and flash_attn_with_kvcache_v3 is chosen on the padding or thd branch when inference_params is set, so it was previously validated only by proxy. No shipped build changes behaviour: all three gained softcap in 2.7.3/hopper and none had it before. For precision, utils.py notes that 2.7.3+ has different APIs but enforces no FA3 minimum, so that is a documented API break rather than a version floor TE checks.

The two numerical variants and the plumbing test are validated on an H100 by mutation, with the suite green as the baseline: capping the bias, removing the cap / layer_number division, and dropping the cross attention pass through each fail their corresponding test, and restoring returns the suite to green. Both numerical variants run on UnfusedDotProductAttention only, so they pin TE's reference arithmetic rather than the flash kernels.

---

Collecting the open items in one place, since they are spread across several threads.

1. FA3 is the default softcap path on Hopper, where FA3 is sm90 only. Dropping NVTE_FA3_SOFTCAP at your suggestion means FA3 is preferred over FA2 whenever the build exposes softcap, max(head_dim_qk, head_dim_v) <= 256, and context parallelism is off. Under CP it still falls back to FA2. The title said "opt-in FA3" and no longer matched, so I corrected it. One consequence worth your view: on a FLASHATTENTION_DISABLE_SOFTCAP FA3 build there is now no escape hatch short of NVTE_FLASH_ATTN_V3=0, because the probe cannot detect that build. Happy to reinstate a subordinate opt-in if you would rather have one.

2. The FA2 dropout gate should stay. I offered on 08-27, in the thread on 5ecabac, to drop it if you wanted this PR narrower. Withdrawing that offer. Before this PR, softcap was hardcoded to 0.0 at every FA call site, so softcap combined with dropout was unreachable. The gate closes a failure mode this PR itself creates rather than an unrelated one. Without it, that configuration raises RuntimeError("Softcapping does not support dropout for now") from a host side TORCH_CHECK in flash_api.cpp, on a combination get_attention_backend reports as supported.

   A related gap I should flag rather than leave implicit: FA2 has no build probe at all, unlike FA3. A FLASHATTENTION_DISABLE_SOFTCAP FA2 build passes the v2_6_0_plus check and then rejects at dispatch. Happy to add a probe there for symmetry if you want it.

3. test_transformer_layer_softcap_plumbing is more than you asked for, and easy to drop. Removing the common_attention_kwargs entry is what makes both of its assertions load bearing, matching how window_size is threaded. Note the test asserts on captured kwargs rather than numerics, so it is hardware independent. I have also wrapped it in try/finally in 590c7ca0 to restore the backend env vars, which it was leaking into the rest of the suite.

4. Head dimension, already committed in dbb36c33, so revert it if you disagree. I widened softcap_3_1 from 64 to 128 rather than adding a third config, so the count stays at the two you named and it costs no new cases. Being precise, this is a swap rather than a pure gain: model_configs_softcap now has no head_dim 64 varlen case. I think 128 is the better of the two, since FA2 and FA3 instantiate separate kernels per head dim and 128 is what Gemma 2 27B and Grok 1 use, whereas nothing that softcaps uses 64. Your call. Neither of your justifications for the pair changes: head_dim plays no part in the flash_attn_func versus flash_attn_varlen_func dispatch, and the 0.01 cap rationale is head dim invariant because softmax_scale cancels the sqrt(d) exactly.

   head_dim 256, which Gemma 2 2B and 9B use, stays uncovered. A third config would cost two cases on bf16 capable hardware. I am not proposing it, since the end to end Megatron runs exercise 256 on both FA2 and FA3, and it would take the dict back to three configs after you asked for two.