Thread tanh logit softcapping through FlashAttention (FA2, opt-in 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 in both FA3 entry points), 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 both 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 — no FA4 kernel can carry the cap, so FA4 is disqualified in get_attention_backend when softcap != 0.0; without that it would be selected on SM100 (NVTE_FLASH_ATTN_V4 defaults to 1) and silently drop the cap. 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

FA4 softcap support is deliberately excluded here rather than included as dead scaffolding. On Blackwell (SM100), FA4's dedicated head_dim=256 forward kernel has no score_mod/softcap fusion logic in it at all — the kernel constructor asserts score_mod is None. So there's currently no FA4 kernel path capable of serving this shape; adding an opt-in flag now would just be inert code with nothing to opt into. This follows as its own PR (stacked on this branch) once — or if — an FA4 kernel with softcap fusion for head_dim=256 lands.

[greptile-apps]

Greptile Summary

The PR adds configurable tanh attention-logit softcapping and threads it through FlashAttention 2, eligible FlashAttention 3 configurations, context-parallel execution, and the unfused fallback.

• Adds capability-aware backend filtering so unsupported fused implementations cannot silently discard softcapping.
• Preserves softcap through context-parallel forward and backward execution.
• Adds numerical, backend-selection, bias-ordering, layer-scaling, and distributed coverage.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds softcap to attention parameters and aligns backend eligibility with FA2, FA3, context-parallel, dropout, and head-dimension capabilities.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Probes FA3 softcap support, forwards the option into flash kernels, and implements equivalent score capping in the unfused backend.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Carries softcap through the three context-parallel autograd implementations and their FA2 forward and backward calls.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Exposes constructor and forward overrides and consistently forwards the resolved value into selection and backend dispatch.
tests/pytorch/attention/test_attention.py	Adds reference-based numerical, backend-selection, bias-ordering, and QK-layer-scaling coverage.
tests/pytorch/attention/test_attention_with_cp.py	Adds distributed forward and gradient comparisons for each distinct context-parallel autograd implementation.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention request] --> B{softcap nonzero?}
  B -- No --> C[Existing backend selection]
  B -- Yes --> D{Supported configuration}
  D -->|FA3 capable, non-CP, head dim <= 256| E[FlashAttention 3]
  D -->|FA2 >= 2.6, supported dropout mode| F[FlashAttention 2]
  D -->|Context parallel| F
  D -->|No eligible flash backend| G[Unfused attention]
  E --> H[Softcapped attention output]
  F --> H
  G --> H

Loading

_{Reviews (18): Last reviewed commit: "Merge branch 'main' into nvegesna/gemma2..." | 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

[nvegesna-netizen]

Pushed one more fix found while re-reviewing the UnfusedDotProductAttention refactor from the last round.

The bug. Softcap was applied after the additive bias, i.e. cap(scale*QK + bias). FA2 computes cap(scale*QK) + bias — its kernel softcaps immediately after the QK^T gemm and only then adds ALiBi inside apply_mask, and it pre-divides alibi_slope by params.scale_softmax, which set_params_fprop sets to softcap when capping is enabled. So ALiBi deliberately lands outside the tanh.

ALiBi is the one bias type flash supports (pre/post_scale_bias disable it), so softcap + ALiBi produced different numerics from the unfused vs. flash path depending only on whether a suitable flash-attn was installed. Since this PR positions unfused as the in-tree reference the flash backends are compared against, that seemed worth correcting rather than documenting.

The fix defers the additive bias until after the cap. pre_scale_bias is folded in before the scaling by construction, so it necessarily stays inside the cap; flash doesn't support that bias type. softcap = 0.0 remains a bit-exact no-op for all four bias types.

Two tests added, both forcing UnfusedDotProductAttention:

• test_dpa_softcap_bias_ordering — pins cap(scores) + bias against cap(scores + bias). Uses post_scale_bias, which drives the same code path as ALiBi without needing slope machinery in the reference.
• test_dpa_softcap_qk_layer_scaling — covers softcap under NVTE_APPLY_QK_LAYER_SCALING, where the cap has to be divided by layer_number. Dropping that division leaves an effective cap of softcap * layer_number. This path had no coverage.

Verified by mutation on an H100, not just by passing: reintroducing either bug makes the corresponding test fail, and restoring makes it pass. The existing softcap suite (including the CP tests) still passes alongside. Both tests carry anti-vacuity asserts — the first draft of the layer-scaling test had too small a margin and would have failed, which the mutation run is what caught.