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 carries it consistently through backend selection, unfused attention, FlashAttention 2 context-parallel execution, and eligible FlashAttention 3 execution.

• Adds softcap to the public dot-product-attention request and backend capability model.
• Implements softcapping in the unfused fallback, including ONNX-export execution.
• Routes supported requests to FA2 or FA3 while rejecting incompatible fused, FA4, dropout, and context-parallel combinations.
• Adds backend, forward, and gradient comparisons against a closed-form reference.

Confidence Score: 5/5

The PR appears safe to merge with no blocking failure remaining from the previously reported issues.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds softcap-aware capability filtering, including the corrected FA3 context-parallel exclusion and FA2 version/dropout requirements.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Implements unfused softcapping and forwards the option into supported FA2 and FA3 entry points.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Threads softcap through FA2 context-parallel forward state and backward calls across communication strategies.
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	Adds the public constructor and forward override and passes the resolved value consistently into selection and execution.
tests/pytorch/attention/test_attention.py	Adds collected forward and gradient parity tests, zero-cap behavior checks, and a closed-form numerical reference.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention request] --> B{softcap nonzero?}
  B -- No --> C[Existing backend selection]
  B -- Yes --> D[Disable FusedAttention and FA4]
  D --> E{FA3 capable, non-CP, head dimensions <= 256?}
  E -- Yes --> F[FA3 with softcap]
  E -- No --> G{FA2 >= 2.6 and supported dropout mode?}
  G -- Yes --> H[FA2 with softcap]
  G -- No --> I{Context parallel?}
  I -- No --> J[Unfused attention with softcap]
  I -- Yes --> K[No eligible backend: raise]

Loading

_{Reviews (12): Last reviewed commit: "test: add softcap no-op and closed-form ..." | 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