Thread tanh logit softcapping through FlashAttention (FA2, opt-in FA3) - #3391

Open

nvegesna-netizen wants to merge 10 commits into

NVIDIA:mainfrom

nvegesna-netizen:nvegesna/gemma2-softcap-core

Commits on Aug 27, 2026

feat: thread tanh logit softcapping through FlashAttention (FA2, opt-in FA3)

nvegesna-netizen
and
claude
committed
69b4ffb
fix: use raise instead of assert for the ONNX+softcap guard

nvegesna-netizen
and
claude
committed
9839ae0
[pre-commit.ci] auto fixes from pre-commit.com hooks
nvegesna-netizen
committed
eb215b5
refactor: move softcap reference into UnfusedDotProductAttention and fold test into test_attention.py
nvegesna-netizen
committed
2475521
fix(pytorch): gate FA3 softcap on existing NVTE_FLASH_ATTN_V3
nvegesna-netizen
committed
3c5eb4a
fix(pytorch): disable FlashAttention 4 for softcap
nvegesna-netizen
committed
900371a
fix(pytorch): disable FlashAttention 2 for softcap with dropout
nvegesna-netizen
committed
5ecabac
Merge branch 'main' into nvegesna/gemma2-softcap-core
nvegesna-netizen
authored
818ce26
test: restore softcap dQ/dK/dV parity in the shared DPA harness
nvegesna-netizen
committed
4cc2e9c
test: add softcap no-op and closed-form reference coverage
nvegesna-netizen
committed
13d65a1