Skip to content

[JAX] Return max_logit and softmax aux stats from TE JAX fused attn#3112

Draft
KshitijLakhani wants to merge 4 commits into
NVIDIA:mainfrom
KshitijLakhani:klakhani/feat/return-softmax-stats
Draft

[JAX] Return max_logit and softmax aux stats from TE JAX fused attn#3112
KshitijLakhani wants to merge 4 commits into
NVIDIA:mainfrom
KshitijLakhani:klakhani/feat/return-softmax-stats

Add first pass tests for max logit and softmax aux tensor outputs in …

f8dbc15
Select commit
Loading
Failed to load commit list.
DCO / DCO succeeded Jun 10, 2026 in 1s

DCO

All commits are signed off!

View more details on DCO