[JAX] Return max_logit from TE fused attn

[KshitijLakhani]

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 time added for non CP is <1 min and for CP <3 mins

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

Please list the changes introduced in this PR:

• Change A
• Change B

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]

Greptile Summary

The PR adds an opt-in JAX fused-attention auxiliary result containing the global per-head maximum logit.

• Extends the JAX, Flax, primitive, and native FFI interfaces with return_max_logit.
• Reduces cuDNN’s raw Max tensor locally and across context-parallel mesh axes.
• Adds forward, backward, layout, masking, and distributed coverage for the new result.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains in the eligible follow-up review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/jax/attention.py	Adds the public option and custom-VJP plumbing for returning max-logit telemetry while preserving the default array-only contract.
transformer_engine/jax/cpp_extensions/attention.py	Adds raw Max output metadata, per-head reductions, sharding rules, and context-parallel aggregation.
transformer_engine/jax/csrc/extensions/attention.cpp	Extends fused-attention workspace and FFI output handling to expose cuDNN’s raw Max tensor.
transformer_engine/jax/csrc/extensions.h	Updates native declarations to carry the new max-logit configuration and output contract.
transformer_engine/jax/flax/transformer.py	Exposes the opt-in auxiliary result through the Flax fused-attention path.
tests/jax/test_fused_attn.py	Adds reference comparisons and backward-path coverage across representative layouts, masks, and bias configurations.
tests/jax/test_distributed_fused_attn.py	Adds context-parallel max-logit coverage for all-gather and ring strategies.
tests/jax/test_fused_attn_score_mod.py	Updates the score-mod test double for the expanded fused-attention signature.

Sequence Diagram

sequenceDiagram
    participant User as JAX/Flax caller
    participant API as fused_attn
    participant Primitive as JAX fused-attention primitive
    participant FFI as C++/cuDNN FFI
    User->>API: "return_max_logit=True"
    API->>Primitive: attention inputs and configuration
    Primitive->>FFI: fused-attention forward
    FFI-->>Primitive: output, softmax aux, raw Max, RNG state
    Primitive->>Primitive: reduce raw Max to global per-head maximum
    Primitive-->>API: output and max_logit
    API-->>User: "(output, {max_logit})"

Loading

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