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).
• get_attention_backend: when softcap != 0, fused/unfused attention and FA3/FA4 are disqualified and selection steers to FA2 (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, opt-in (Hopper only): gated behind an env flag plus a build-capability probe and a head_dim <= 256 check; forward threads softcap into FA3's kwargs, backward is handled automatically by the existing Hopper autograd function. Default behavior (flag unset) is unchanged — steers to FA2.
• Tests: forward/backward parity vs. a pure-PyTorch reference implementation.

Validation

Exercised end-to-end 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) are actually selected at runtime.

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 tanh attention-logit softcapping to DotProductAttention, routes it through supported FlashAttention paths, and fails closed for unsupported backend and export combinations.

• Adds softcap to the public attention configuration and backend-selection cache key.
• Threads softcap through FA2 forward/backward, including context-parallel variants.
• Adds opt-in FA3 support with capability, shape, and context-parallel guards.
• Adds forward/backward numerical parity coverage and invokes it from L0 CI.

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/dot_product_attention.py	Resolves and propagates softcap, includes it in backend parameters, and explicitly rejects unsupported ONNX export.
transformer_engine/pytorch/attention/dot_product_attention/utils.py	Adds softcap to backend-cache inputs and filters unsupported backends, including FA3 under context parallelism.
transformer_engine/pytorch/attention/dot_product_attention/backends.py	Probes FA3 softcap capability and forwards softcap through supported FA2 and opt-in FA3 execution.
transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py	Threads softcap through all three FA2 context-parallel forward and backward autograd paths.
tests/pytorch/attention/test_softcap.py	Adds FA2 forward and gradient parity coverage across dtypes, masks, GQA, and disabled/enabled softcapping.
qa/L0_pytorch_unittest/test.sh	Explicitly adds the softcap numerical test to the L0 PyTorch test job.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A[DotProductAttention softcap] --> B{ONNX export?}
  B -->|Yes, nonzero softcap| C[Raise ValueError]
  B -->|No| D{Context parallel?}
  D -->|Yes| E[Select supported FA2 path]
  D -->|No| F{FA3 opt-in and capable?}
  F -->|Yes| G[FA3 forward and backward]
  F -->|No| E

Loading

_{Reviews (4): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}