[JAX] Return max_logit from TE fused attn by KshitijLakhani · Pull Request #3112 · NVIDIA/TransformerEngine

KshitijLakhani · 2026-06-10T06:17:50Z

Description

Return the max_logit from TE JAX fused attn to maintain parity with TE PyT fused attn.
Returning the softmax aux stats is a separate request and a bit mroe involved, so it'll be addressed separately.
Exposing the max_logit gives sufficient telemetry insight to users, for starters.

Fixes #2945

Test timing added

Non-CP

   Level              Tests added    Approx. time added
  ━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━
   L0       41 (23 pass, 18 skip)                  ~23s
  ───────  ───────────────────────  ────────────────────
   L2       41 (23 pass, 18 skip)                  ~23s

CP


   Level                Tests added    Approx. time added
  ━━━━━━━  ━━━━━━━━━━━━━━━━━━━━━━━━━  ━━━━━━━━━━━━━━━━━━━━
   L1       180 (74 pass, 106 skip)                 ~147s
  ───────  ─────────────────────────  ────────────────────
   L2         60 (16 pass, 44 skip)                  ~26s

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Checklist:

I have read and followed the contributing guidelines
The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

greptile-apps · 2026-08-21T06:52:37Z

Greptile Summary

This PR adds an opt-in per-head max_logit result to JAX fused attention while preserving the existing output-only contract by default.

Propagates the option through the public JAX and Flax APIs.
Extends the JAX primitive and native FFI result contracts.
Reduces local statistics across batches, context-parallel steps, and relevant mesh axes.
Adds single-device and distributed coverage across layouts, masks, GQA, and context-parallel strategies.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/jax/attention.py	Adds the public option, tuple-return contract, kernel-availability propagation, and custom-VJP handling for max-logit telemetry.
transformer_engine/jax/cpp_extensions/attention.py	Extends fused-attention primitive results and implements per-head reduction, batching, sharding, and context-parallel aggregation.
transformer_engine/jax/csrc/extensions/attention.cpp	Extends the native fused-attention forward interface with an optional cuDNN Max output buffer aligned with the Python primitive contract.
transformer_engine/jax/csrc/extensions.h	Updates native extension declarations for the optional max-logit result.
transformer_engine/jax/flax/transformer.py	Exposes the fused-only telemetry option through Flax while retaining plain-array returns when disabled.
tests/jax/test_fused_attn.py	Adds non-context-parallel numerical and backward coverage for the optional auxiliary result.
tests/jax/test_distributed_fused_attn.py	Adds distributed max-logit coverage for supported all-gather and ring context-parallel configurations.

Sequence Diagram

sequenceDiagram
    participant App as JAX/Flax caller
    participant API as fused_attn
    participant Primitive as JAX fused-attention primitive
    participant Native as Native TE/cuDNN
    participant Mesh as Mesh reduction
    App->>API: "QKV, return_max_logit=True"
    API->>Primitive: Bind fused forward
    Primitive->>Native: Output and Max buffers
    Native-->>Primitive: Attention output and row-wise Max
    Primitive->>Primitive: Reduce Max to per-head values
    Primitive->>Mesh: pmax across non-head mesh axes
    Mesh-->>API: Global per-head max_logit
    API-->>App: (output, max_logit)

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

jberchtold-nvidia · 2026-08-24T17:13:31Z

                                       NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend,
-                                       void *softmax_buf, void *rng_state_buf = nullptr,
-                                       void *bias_buf = nullptr,
+                                       void *softmax_buf, void *max_buf = nullptr,


Can we rename max_buf to max_logits_buf

Addressed in SHA 704957a

jberchtold-nvidia · 2026-08-24T17:15:36Z

+      nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data);
+    }
+
+    NVTETensor &rng_state_aux = tensor_pack->tensors[size++];


Do we need to check the capacity of tensor_pack->tensors here?

I was not too concerned about it here, because I think the max size is 10 and with max_logit we will have 5 tensors only: 1 softmax + 1 max + 1 RNG + 1 bias + 1 offset = 5 maximum so I wasn't concerned about exceeding.

However, that's a good point for the future when someone is packing in more tensors in here. Let me push a change not only for max_logit but also for the other 4 tensors packed in sothat it serves as a template for any tensors being packed in the future

Addressed in SHA 2e71bc8

jberchtold-nvidia · 2026-08-24T20:00:16Z

                        (batch, head, q.shape[1], 1), -jnp.inf, dtype=jnp.float32
                    )
-                    return output_per_step, softmax_aux_per_step
+                    max_logit_per_step = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(


Why does this instance use empty max logits?

That's a good question @jberchtold-nvidia !
TL;DR: So, this is the no-compute branch for a fully causal masked Ring step and hence attention, softmax, isn't calculated in this branch and we need a neutral placeholder (-inf/0) because veery lax.cond branch must return the same shapes .

Let me try expanding on that:
With causal Ring attention, every rank cycles through KV blocks from all CP ranks. Some received KV blocks are entirely in the future relative to that rank’s query block and because causal attention must not attend to any token in such a block, TE skips the cuDNN attention call completely. So skip_compute does not run attention, softmax, or max-logit computation. It returns identity values so the JAX control flow and subsequent accumulation remain valid.

When max logits are requested, the helper returns [H] filled with -inf, so the skipped step cannot affect the running maximum. FWIW, it returns (0,) only when max-logit return is disabled.

jberchtold-nvidia · 2026-08-24T20:07:14Z

@@ -399,6 +399,12 @@ def test_softcap_score_mod_with_aux_params_backward(
    ),
 ]



Can you run L1 CI on this PR? I don't see a CI run

I had run the test suite locally only. Launched CI right now. Thanks for the reminder !

…x aux to the user from TE JAX fused attn output Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

… and softmax through the JAX fused attn primitives Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

…JAX fused attn tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

… function for it Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani · 2026-08-24T21:28:05Z

/te-ci jax L0 L1 L2

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

for more information, see https://pre-commit.ci

KshitijLakhani self-assigned this Jun 10, 2026

KshitijLakhani added the 2.17 label Jun 10, 2026

KshitijLakhani added 2.18 and removed 2.17 labels Jun 25, 2026

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch 2 times, most recently from 4f3fbb3 to 9a56117 Compare July 8, 2026 23:28

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch 2 times, most recently from ca8551c to fc151f9 Compare July 27, 2026 00:57

KshitijLakhani added 2.19 and removed 2.18 labels Jul 30, 2026

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch 2 times, most recently from fd2fff3 to f31d5de Compare August 21, 2026 06:49

KshitijLakhani marked this pull request as ready for review August 21, 2026 06:49

KshitijLakhani requested a review from jberchtold-nvidia as a code owner August 21, 2026 06:49

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch from cc33228 to 8dabb52 Compare August 21, 2026 07:45

KshitijLakhani requested a review from cyanguwa August 21, 2026 07:47

KshitijLakhani changed the title ~~[JAX] Return max_logit and softmax aux stats from TE JAX fused attn~~ [JAX] Return max_logit from TE JAX fused attn Aug 21, 2026

KshitijLakhani changed the title ~~[JAX] Return max_logit from TE JAX fused attn~~ [JAX] Return max_logit from TE fused attn Aug 21, 2026

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch 2 times, most recently from 2cb2b8d to 032eeea Compare August 24, 2026 16:21

jberchtold-nvidia reviewed Aug 24, 2026

View reviewed changes

KshitijLakhani added 8 commits August 24, 2026 14:15

Add support in lower level JAX API for returning max logit and softma…

3ab4725

…x aux to the user from TE JAX fused attn output Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Add support for returning reduced per head max logit. Plumb max logit…

84f180a

… and softmax through the JAX fused attn primitives Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

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

8f565cf

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

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

31ede58

…JAX fused attn tests Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Reject aux returns with score_mod

dde413b

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Handle SM120 max-logit layout

9da2226

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Drop softmax aux return

38ce6e3

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Modify static args in fused attn tests for jax

3198726

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

pre-commit-ci Bot and others added 11 commits August 24, 2026 14:15

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

e56b971

for more information, see https://pre-commit.ci

nit: Inline the choice of what is to be returned and remove redundant…

cbe18ac

… function for it Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Expose JAX max logit

7f9977e

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Support CP max logit

de929ba

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Expand JAX max logit tests

78c8828

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Remove JAX max logit integration tests

e36a4a0

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Broaden JAX CP max logit tests

9b537ba

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Reduce JAX max logit across DP

bb7ce9b

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Simplify JAX max logit return

c093329

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Document JAX max logit reductions

0f7bb0b

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Refine JAX max logit tests

90a3eaf

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch from b305c49 to 90a3eaf Compare August 24, 2026 21:16

KshitijLakhani added 2 commits August 24, 2026 14:56

Rename JAX max logit buffer

704957a

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

Check JAX attention tensor pack capacity

2e71bc8

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>

KshitijLakhani force-pushed the klakhani/feat/return-softmax-stats branch from 3e247dd to 2e71bc8 Compare August 24, 2026 21:59

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

d7bf4cf

for more information, see https://pre-commit.ci

Navigation Menu

Uh oh!

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

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

KshitijLakhani commented Jun 10, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Aug 21, 2026 •

edited

Loading

Uh oh!

jberchtold-nvidia Aug 24, 2026

Uh oh!

KshitijLakhani Aug 24, 2026

Uh oh!

jberchtold-nvidia Aug 24, 2026

Uh oh!

KshitijLakhani Aug 24, 2026

Uh oh!

KshitijLakhani Aug 24, 2026

Uh oh!

jberchtold-nvidia Aug 24, 2026

Uh oh!

KshitijLakhani Aug 24, 2026

Uh oh!

jberchtold-nvidia Aug 24, 2026

Uh oh!

KshitijLakhani Aug 24, 2026 •

edited

Loading

Uh oh!

KshitijLakhani commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		@@ -399,6 +399,12 @@ def test_softcap_score_mod_with_aux_params_backward(
		),
		]

Uh oh!

Conversation

KshitijLakhani commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test timing added

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

KshitijLakhani commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

KshitijLakhani commented Jun 10, 2026 •

edited

Loading

greptile-apps Bot commented Aug 21, 2026 •

edited

Loading

KshitijLakhani Aug 24, 2026 •

edited

Loading