Follow-up: opt-in FA4 softcap scaffolding (Blackwell, currently non-functional) - #1
Open
nvegesna-netizen wants to merge 4 commits into
Open
Conversation
…in FA3) Add a user `softcap` value (tanh logit softcapping, `softcap*tanh(x/softcap)`) to DotProductAttention so models like Gemma2 can run on the fused flash path instead of an unfused/FlexAttention kernel. - Add `softcap` to DotProductAttention (init+forward) and AttentionParams; thread it into the FA2 non-CP kwargs and all three context-parallel autograd functions (forward + ctx-saved backward). softcap=0.0 reproduces prior behavior. - get_attention_backend: when softcap != 0, disable FusedAttention/unfused and steer to FA2 -- disable FA3/FA4, and disable FA2 < 2.6.0 -- so the cap is never silently dropped (FA2 < 2.6.0) or hit at runtime as NotImplementedError (FA3/FA4). Also disable FA3 under context parallelism (its CP path hard-rejects nonzero softcap) so CP+softcap steers to FA2, which supports it, instead of crashing. - FA3 softcap opt-in: NVTE_FA3_SOFTCAP=1, Hopper (sm90) hd<=256, non-CP only, gated on a fail-closed signature probe (fa3_supports_softcap). Forward threads softcap into fa_3_optional_forward_kwargs; the existing Hopper autograd function carries it into backward automatically. Default off; unchanged behavior steers to FA2. - ONNX export: fail loudly (assert) rather than silently drop softcap -- export unconditionally force-selects UnfusedDotProductAttention, which has no softcap support, so this previously exported models with softcapping silently omitted. - Tests: test_softcap.py (FA2 fwd/bwd parity vs pure-PyTorch reference), wired into qa/L0_pytorch_unittest/test.sh. FA4 softcap opt-in is deliberately NOT included here -- see follow-up PR. On Blackwell (SM100), FA4's dedicated head_dim=256 forward kernel has no score_mod support at all (kernel constructor asserts `score_mod is None`), so there is currently no FA4 kernel path this could opt into; adding the scaffolding now would just be inert code with nothing to exercise. Addresses review findings: CP+FA3 softcap selection crash, ONNX silent drop, and the missing CI wiring for test_softcap.py. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
nvegesna-netizen
force-pushed
the
nvegesna/gemma2-softcap-core
branch
from
August 17, 2026 21:21
a6a793b to
5917f0d
Compare
nvegesna-netizen
force-pushed
the
nvegesna/gemma2-fa4-softcap-followup
branch
from
August 17, 2026 21:23
b809133 to
3d761b1
Compare
python -O / PYTHONOPTIMIZE strips assert statements, which would silently reopen the ONNX export softcap-drop bug the previous commit fixed (ONNX mode would again force-select UnfusedDotProductAttention with softcap silently omitted, with no error). Switch to an explicit if/raise ValueError, which survives optimized execution. Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
nvegesna-netizen
force-pushed
the
nvegesna/gemma2-softcap-core
branch
from
August 17, 2026 21:35
19a21eb to
5ae46ce
Compare
for more information, see https://pre-commit.ci (reapplied after a force-push rebase clobbered pre-commit.ci's original 19a21eb commit; same content, restored by hand)
…onal) Follow-up to the FA2/FA3 softcap PR. Adds the same fail-closed opt-in pattern for FA4 (NVTE_FA4_SOFTCAP=1, Blackwell/sm100 gate, signature capability probe) plus an on-device verification harness (verify_fa4_softcap.py) asserting FA4 was actually selected as the backend. Known limitation, not fixed by this PR: FA4's dedicated head_dim=256 forward kernel on SM100 has no score_mod/softcap fusion logic in it at all -- the kernel constructor asserts `score_mod is None`. So this opt-in currently has no working kernel path to exercise; it exists so the plumbing is ready and gated the moment a softcap-capable FA4 build exists, without a runtime NotImplementedError surprising anyone who sets the env var today (the opt-in's capability probe fails closed and the code stays on FA2 unless the build genuinely supports it). Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
nvegesna-netizen
force-pushed
the
nvegesna/gemma2-fa4-softcap-followup
branch
from
August 17, 2026 21:36
3d761b1 to
ea55d69
Compare
nvegesna-netizen
force-pushed
the
nvegesna/gemma2-softcap-core
branch
from
August 27, 2026 16:57
82a2bf4 to
5ecabac
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-up to NVIDIA#3391 (base branch of this PR — stacked, review that one first). Opened against my own fork's core branch rather than upstream
mainsince it depends on NVIDIA#3391 merging first; retarget toNVIDIA/TransformerEngine:mainonce that lands.Adds the same fail-closed opt-in pattern as the FA3 change in NVIDIA#3391, but for FA4: an env-gated opt-in, a Blackwell (sm100) gate, and a signature capability probe, plus an on-device verification harness asserting FA4 was actually selected as the backend.
Why this is a follow-up, not part of the base PR
FA4's dedicated
head_dim=256forward kernel on SM100 has noscore_mod/softcap fusion logic in it at all — the kernel constructor assertsscore_mod is None. So this opt-in currently has no working kernel path to exercise. It exists so the plumbing is ready and gated the moment a softcap-capable FA4 build exists, without a runtimeNotImplementedErrorsurprising anyone who sets the env var today — the opt-in's capability probe fails closed and the code stays on FA2 unless the installed build genuinely supports it.Given it's inert until that kernel support lands, keeping it out of the base PR avoids review noise on code with nothing to actually test end-to-end yet.