[JAX] Return max_logit from TE fused attn - #3112

Merged

KshitijLakhani merged 22 commits into

NVIDIA:mainfrom

KshitijLakhani:klakhani/feat/return-softmax-stats

Aug 25, 2026

Commits on Aug 24, 2026

Add support in lower level JAX API for returning max logit and softmax aux to the user from TE JAX fused attn output
KshitijLakhani
committed
3ab4725
Add support for returning reduced per head max logit. Plumb max logit and softmax through the JAX fused attn primitives
KshitijLakhani
committed
84f180a
Add max logit to JAX fused attn FFI and set it in the workspace
KshitijLakhani
committed
8f565cf
Add first pass tests for max logit and softmax aux tensor outputs in JAX fused attn tests
KshitijLakhani
committed
31ede58
Reject aux returns with score_mod
KshitijLakhani
committed
dde413b
Handle SM120 max-logit layout
KshitijLakhani
committed
9da2226
Drop softmax aux return
KshitijLakhani
committed
38ce6e3
Modify static args in fused attn tests for jax
KshitijLakhani
committed
3198726
[pre-commit.ci] auto fixes from pre-commit.com hooks

pre-commit-ci[bot]
authored and
KshitijLakhani
committed
e56b971
nit: Inline the choice of what is to be returned and remove redundant function for it
KshitijLakhani
committed
cbe18ac
Expose JAX max logit
KshitijLakhani
committed
7f9977e
Support CP max logit
KshitijLakhani
committed
de929ba
Expand JAX max logit tests
KshitijLakhani
committed
78c8828
Remove JAX max logit integration tests
KshitijLakhani
committed
e36a4a0
Broaden JAX CP max logit tests
KshitijLakhani
committed
9b537ba
Reduce JAX max logit across DP
KshitijLakhani
committed
bb7ce9b
Simplify JAX max logit return
KshitijLakhani
committed
c093329
Document JAX max logit reductions
KshitijLakhani
committed
0f7bb0b
Refine JAX max logit tests
KshitijLakhani
committed
90a3eaf
Rename JAX max logit buffer
KshitijLakhani
committed
704957a
Check JAX attention tensor pack capacity
KshitijLakhani
committed
2e71bc8
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed
d7bf4cf

Navigation Menu

Uh oh!

[JAX] Return max_logit from TE fused attn - #3112

[JAX] Return max_logit from TE fused attn#3112
KshitijLakhani merged 22 commits into
NVIDIA:mainfrom
KshitijLakhani:klakhani/feat/return-softmax-stats

Commits on Aug 24, 2026

Add support in lower level JAX API for returning max logit and softmax aux to the user from TE JAX fused attn output

Add support for returning reduced per head max logit. Plumb max logit and softmax through the JAX fused attn primitives

Add max logit to JAX fused attn FFI and set it in the workspace

Add first pass tests for max logit and softmax aux tensor outputs in JAX fused attn tests

Reject aux returns with score_mod

Handle SM120 max-logit layout

Drop softmax aux return

Modify static args in fused attn tests for jax

[pre-commit.ci] auto fixes from pre-commit.com hooks

nit: Inline the choice of what is to be returned and remove redundant function for it

Expose JAX max logit

Support CP max logit

Expand JAX max logit tests

Remove JAX max logit integration tests

Broaden JAX CP max logit tests

Reduce JAX max logit across DP

Simplify JAX max logit return

Document JAX max logit reductions

Refine JAX max logit tests

Rename JAX max logit buffer

Check JAX attention tensor pack capacity

[pre-commit.ci] auto fixes from pre-commit.com hooks