From 3ab472581add8832f0f0024cbce095f8d04fb45f Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Tue, 9 Jun 2026 23:08:02 -0700
Subject: [PATCH 01/22] Add support in lower level JAX API for returning max
 logit and softmax aux to the user from TE JAX fused attn output

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 transformer_engine/jax/attention.py | 74 +++++++++++++++++++++++++++--
 1 file changed, 69 insertions(+), 5 deletions(-)

diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index ecca4a3871..9d3aaae90d 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -339,6 +339,7 @@ def is_fused_attn_kernel_available(
     head_dim_qk,
     head_dim_v,
     window_size: Optional[Tuple[int, int]] = None,
+    return_max_logit: bool = False,
 ):
     """
     To check whether the fused attention kernel is supported
@@ -362,6 +363,7 @@ def make_helper(attn_mask_type):
             head_dim_qk,
             head_dim_v,
             window_size_tuple,
+            return_max_logit,
         )
 
     return make_helper(attn_mask_type).is_fused_attn_kernel_available()
@@ -1053,6 +1055,8 @@ def _legacy_fused_attn(
     context_parallel_causal_load_balanced: bool = False,
     context_parallel_axis: str = "",
     softmax_offset: Optional[jnp.ndarray] = None,
+    return_max_logit: bool = False,
+    return_softmax_aux: bool = False,
 ):
     """
     Perform non-THD (non-packed) cuDNN fused attention.
@@ -1084,8 +1088,18 @@ def _legacy_fused_attn(
         context_parallel_causal_load_balanced (bool):
             Indicates the sequences are ordered for causal mask load balancing when running context parallelism.
         context_parallel_axis (str): The name of the context parallel axis.
+        softmax_offset (Optional[jnp.ndarray]): An optional learnable softmax offset tensor with shape
+            [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX.
+        return_max_logit (bool): If True, also return per-head maximum attention logits
+            in an auxiliary dictionary under ``"max_logit"``.
+        return_softmax_aux (bool): If True, also return backend-specific softmax statistics
+            in an auxiliary dictionary under ``"softmax_aux"``.
     Returns:
-        (jnp.ndarray): The output tensor from the fused attention.
+        jnp.ndarray:
+            Attention output when neither ``return_max_logit`` nor ``return_softmax_aux`` is True.
+        tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
+            ``(output, aux)`` when either flag is True. ``aux`` may contain:
+            ``"max_logit"`` (shape ``[h]``) and/or ``"softmax_aux"`` (float32).
     """
     assert (
         not qkv_layout.is_thd()
@@ -1139,6 +1153,8 @@ def _legacy_fused_attn(
         context_parallel_strategy=context_parallel_strategy,
         context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
         context_parallel_axis=context_parallel_axis,
+        return_max_logit=return_max_logit,
+        return_softmax_aux=return_softmax_aux,
     )
 
     return output
@@ -1164,6 +1180,8 @@ def fused_attn_thd(
     context_parallel_causal_load_balanced: bool = False,
     context_parallel_axis: str = "",
     softmax_offset: Optional[jnp.ndarray] = None,
+    return_max_logit: bool = False,
+    return_softmax_aux: bool = False,
 ):
     """
     Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor
@@ -1218,12 +1236,17 @@ def fused_attn_thd(
         context_parallel_strategy=context_parallel_strategy,
         context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
         context_parallel_axis=context_parallel_axis,
+        return_max_logit=return_max_logit,
+        return_softmax_aux=return_softmax_aux,
     )
 
     return output
 
 
-@partial(jax.custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18))
+@partial(
+    jax.custom_vjp,
+    nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20),
+)
 def _fused_attn(
     qkv: Tuple[jnp.ndarray, ...],
     bias: Optional[jnp.ndarray],
@@ -1244,6 +1267,8 @@ def _fused_attn(
     context_parallel_axis: str,
     context_checkpoint_name: str = "context",
     stripe_size: int | None = None,
+    return_max_logit: bool = False,
+    return_softmax_aux: bool = False,
 ):
     output, _ = _fused_attn_fwd_rule(
         qkv,
@@ -1265,6 +1290,8 @@ def _fused_attn(
         context_parallel_axis,
         context_checkpoint_name=context_checkpoint_name,
         stripe_size=stripe_size,
+        return_max_logit=return_max_logit,
+        return_softmax_aux=return_softmax_aux,
     )
     return output
 
@@ -1289,8 +1316,10 @@ def _fused_attn_fwd_rule(
     context_parallel_axis,
     context_checkpoint_name,
     stripe_size,
+    return_max_logit,
+    return_softmax_aux,
 ):
-    output, softmax_aux, rng_state = tex.fused_attn_fwd(
+    output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd(
         qkv,
         bias,
         softmax_offset,
@@ -1309,11 +1338,16 @@ def _fused_attn_fwd_rule(
         context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
         context_parallel_axis=context_parallel_axis,
         stripe_size=stripe_size,
+        return_max_logit=return_max_logit,
     )
     output = checkpoint_name(output, context_checkpoint_name)
     softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name)
     rng_state = checkpoint_name(rng_state, context_checkpoint_name)
-    return output, (
+    max_logit = checkpoint_name(max_logit, context_checkpoint_name)
+    attn_output = _resolve_fused_attn_output(
+        output, max_logit, softmax_aux, return_max_logit, return_softmax_aux
+    )
+    return attn_output, (
         qkv,
         bias,
         sequence_descriptor,
@@ -1339,10 +1373,14 @@ def _fused_attn_bwd_rule(
     context_parallel_axis,
     context_checkpoint_name,
     stripe_size,
+    return_max_logit,
+    return_softmax_aux,
     ctx,
     dz,
 ):
     del context_checkpoint_name
+    if return_max_logit or return_softmax_aux:
+        dz, _ = dz
     (
         qkv,
         bias,
@@ -1388,6 +1426,18 @@ def _fused_attn_bwd_rule(
     )
 
 
+def _resolve_fused_attn_output(output, max_logit, softmax_aux, return_max_logit, return_softmax_aux):
+    if not return_max_logit and not return_softmax_aux:
+        return output
+
+    aux = {}
+    if return_max_logit:
+        aux["max_logit"] = max_logit
+    if return_softmax_aux:
+        aux["softmax_aux"] = softmax_aux
+    return output, aux
+
+
 _fused_attn.defvjp(_fused_attn_fwd_rule, _fused_attn_bwd_rule)
 
 
@@ -1468,6 +1518,8 @@ def fused_attn(
     score_mod_bprop: Optional[Callable] = None,
     score_mod_tensors: Optional[Mapping[str, Any]] = None,
     score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None,
+    return_max_logit: bool = False,
+    return_softmax_aux: bool = False,
 ):
     """
     Perform cuDNN fused attention.
@@ -1524,8 +1576,16 @@ def fused_attn(
             non-differentiable auxiliary inputs.
         score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or
             Python/NumPy scalars made available to `score_mod_bprop`.
+        return_max_logit (bool): If True, also return per-head maximum attention logits
+            in an auxiliary dictionary under ``"max_logit"``.
+        return_softmax_aux (bool): If True, also return backend-specific softmax statistics
+            in an auxiliary dictionary under ``"softmax_aux"``.
     Returns:
-        (jnp.ndarray): The output tensor from the fused attention.
+        jnp.ndarray:
+            Attention output when neither ``return_max_logit`` nor ``return_softmax_aux`` is True.
+        tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
+            ``(output, aux)`` when either flag is True. ``aux`` may contain:
+            ``"max_logit"`` (shape ``[h]``) and/or ``"softmax_aux"`` (float32).
 
     Examples (non-THD, also known as non-packed):
         >>> #  q_segment_ids = [[1, 1, 1, 0], [1, 1, 0, 0]], 0 means padded tokens
@@ -1628,6 +1688,8 @@ def fused_attn(
             context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
             context_parallel_axis=context_parallel_axis,
             softmax_offset=softmax_offset,
+            return_max_logit=return_max_logit,
+            return_softmax_aux=return_softmax_aux,
         )
     if max_segments_per_seq > 1 and not qkv_layout.is_thd():
         warnings.warn(
@@ -1658,5 +1720,7 @@ def fused_attn(
         context_parallel_axis=context_parallel_axis,
         context_checkpoint_name=context_checkpoint_name,
         stripe_size=stripe_size,
+        return_max_logit=return_max_logit,
+        return_softmax_aux=return_softmax_aux,
     )
     return output

From 84f180a45f01d8acc90c5af1a2525aec81b47781 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Tue, 9 Jun 2026 23:10:57 -0700
Subject: [PATCH 02/22] Add support for returning reduced per head max logit.
 Plumb max logit and softmax through the JAX fused attn primitives

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 .../jax/cpp_extensions/attention.py           | 129 ++++++++++++++----
 1 file changed, 104 insertions(+), 25 deletions(-)

diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 489bfde997..e129997b85 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -75,6 +75,7 @@
         "cp_axis",
         "cp_striped_window_size",
         "stripe_size",
+        "return_max_logit",
     ],
 )
 @dataclass(frozen=True)
@@ -99,6 +100,7 @@ class _FusedAttnConfig:
     stripe_size: (
         int | None
     )  # Only for CP + Striped. For Ring P2P, stripe_size=1 only.For AG, stripe_size>=1.
+    return_max_logit: bool = False
 
 
 @dataclass(frozen=True)
@@ -122,6 +124,7 @@ class FusedAttnHelper:
     head_dim_qk: int
     head_dim_v: int
     window_size: Tuple[int, int]
+    return_max_logit: bool = False
 
     def is_fused_attn_kernel_available(self):
         """Check if there is available fused attention kernel"""
@@ -146,6 +149,7 @@ def get_fused_attn_backend(self):
             self.head_dim_v,
             self.window_size[0],
             self.window_size[1],
+            self.return_max_logit,
             not self.is_non_deterministic_allowed(),
         )
 
@@ -351,6 +355,7 @@ def abstract(
             q_head_dim,
             v_head_dim,
             config.window_size,
+            config.return_max_logit,
         ).get_fused_attn_backend()
 
         if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen:
@@ -371,6 +376,14 @@ def abstract(
         else:
             raise ValueError(f"Unsupported {backend=}")
         softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype)
+        if config.return_max_logit:
+            if config.qkv_layout.is_thd() and get_cudnn_version() >= (9, 6, 0):
+                max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1)
+            else:
+                max_tensor_shape = (*batch_shape, attn_heads, q_max_seqlen, 1)
+        else:
+            max_tensor_shape = (0,)
+        max_tensor_aval = q_aval.update(shape=max_tensor_shape, dtype=softmax_dtype)
 
         # JAX does not enable 64-bit int by default so we get XLA to allocate x8 memory with
         # 32-bit unsigned int to get the buffer size we need in the C++ kernel
@@ -417,6 +430,7 @@ def abstract(
             config.max_segments_per_seq,
             config.window_size[0],
             config.window_size[1],
+            config.return_max_logit,
             bottom_right_diagonal,
         )
         wkspace_aval = q_aval.update(
@@ -437,17 +451,19 @@ def abstract(
                 f" {softmax_offset_aval.shape}"
             )
 
-        return out_aval, softmax_aux_aval, rng_state_aval, wkspace_aval
+        return out_aval, softmax_aux_aval, max_tensor_aval, rng_state_aval, wkspace_aval
 
     @staticmethod
     def outer_abstract(*args, **kwargs):
         """
         Fused attention fwd outer primitive abstract
         """
-        out_aval, softmax_aux_aval, rng_state_aval, _ = FusedAttnFwdPrimitive.abstract(
+        out_aval, softmax_aux_aval, _, rng_state_aval, _ = FusedAttnFwdPrimitive.abstract(
             *args, **kwargs
         )
-        return out_aval, softmax_aux_aval, rng_state_aval
+        max_logit_shape = (out_aval.shape[-2],) if kwargs["config"].return_max_logit else (0,)
+        max_logit_aval = out_aval.update(shape=max_logit_shape, dtype=out_aval.dtype)
+        return out_aval, softmax_aux_aval, rng_state_aval, max_logit_aval
 
     @staticmethod
     def lowering(
@@ -531,6 +547,7 @@ def lowering(
             mask_type=int(config.attn_mask_type.value),
             qkv_layout=int(config.qkv_layout.value),
             is_training=config.is_training,
+            return_max_logit=config.return_max_logit,
             deterministic=not FusedAttnHelper.is_non_deterministic_allowed(),
             window_size_left=window_size_left,
             window_size_right=window_size_right,
@@ -574,6 +591,9 @@ def impl(
                 config.max_segments_per_seq,
             )
         )
+        raw_q_seqlen = q_seqlen
+        raw_q_seq_offsets = q_seq_offsets
+
         if config.qkv_layout.is_thd():
 
             def _fix_len_take(x, condition, fill_value=-1):
@@ -630,7 +650,7 @@ def convert_to_2d(offsets, batch, max_seqlen):
         q_cu_seqlen = generate_cu_seqlen(q_seqlen.flatten())
         kv_cu_seqlen = generate_cu_seqlen(kv_seqlen.flatten())
 
-        output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind(
+        output, softmax_aux, max_tensor, rng_state, _ = FusedAttnFwdPrimitive.inner_primitive.bind(
             q,
             k,
             v,
@@ -647,7 +667,39 @@ def convert_to_2d(offsets, batch, max_seqlen):
             _kv_segment_pos,
             config=config,
         )
-        return output, softmax_aux, rng_state
+        max_logit = FusedAttnFwdPrimitive._reduce_max_logit(
+            max_tensor, output, raw_q_seqlen, raw_q_seq_offsets, config
+        )
+        return output, softmax_aux, rng_state, max_logit
+
+    @staticmethod
+    def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config):
+        """Reduce cuDNN's raw Max tensor to PyTorch-compatible per-head max_logit."""
+        if not config.return_max_logit:
+            return jnp.zeros((0,), dtype=output.dtype)
+
+        if config.qkv_layout.is_thd() and max_tensor.ndim == 4:
+            q_seqlen = jnp.where(q_seqlen > 0, q_seqlen, 0)
+            q_seq_offsets = jnp.where(q_seq_offsets >= 0, q_seq_offsets, -1)
+            token_idx = jnp.arange(output.shape[-3], dtype=q_seq_offsets.dtype)
+            valid = jnp.any(
+                (q_seq_offsets[..., :-1, None] >= 0)
+                & (token_idx >= q_seq_offsets[..., :-1, None])
+                & (token_idx < (q_seq_offsets[..., :-1, None] + q_seqlen[..., None])),
+                axis=-2,
+            )
+            if max_tensor.shape[1] == output.shape[-3]:
+                max_tensor = jnp.where(valid[:, :, None, None], max_tensor, -jnp.inf)
+            else:
+                max_tensor = jnp.where(valid[:, None, :, None], max_tensor, -jnp.inf)
+
+        if max_tensor.ndim == 3:
+            amax_dims = (0, 2)
+        elif config.qkv_layout.is_thd() and max_tensor.shape[1] == output.shape[-3]:
+            amax_dims = (0, 1, 3)
+        else:
+            amax_dims = (0, 2, 3)
+        return jnp.max(max_tensor, axis=amax_dims).astype(output.dtype)
 
     @staticmethod
     def batcher(batched_args, batch_dims, *, config):
@@ -659,7 +711,8 @@ def batcher(batched_args, batch_dims, *, config):
         q_bdim, _, _, _, _, seed_bdim, *_ = batch_dims
         # Pass through; segment_ids/segment_pos may have different batch dims (e.g. vmapped ids,
         # replicated pos). get_seqlens_and_offsets() in attention.py handles conversion without expanding.
-        out_bdims = q_bdim, q_bdim, seed_bdim
+        max_logit_bdim = q_bdim if config.return_max_logit else None
+        out_bdims = q_bdim, q_bdim, seed_bdim, max_logit_bdim
         return (
             FusedAttnFwdPrimitive.outer_primitive.bind(*batched_args, config=config),
             out_bdims,
@@ -713,12 +766,16 @@ def infer_sharding_from_operands(config, mesh, arg_infos, result_infos):
             raise ValueError(f"Unsupported {config.qkv_layout=}")
 
         rng_state_sharding = NamedSharding(mesh, PartitionSpec(get_all_mesh_axes(), None))
-        return (out_sharding, softmax_aux_sharding, rng_state_sharding)
+        max_logit_sharding = NamedSharding(
+            mesh, PartitionSpec(q_spec[-2] if config.return_max_logit else None)
+        )
+        return (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
 
     @staticmethod
     def partition(config, mesh, arg_infos, result_infos):
         out_sharding = result_infos[0].sharding
         softmax_aux_sharding = result_infos[1].sharding
+        max_logit_sharding = result_infos[3].sharding
         rng_state_sharding = seed_sharding = NamedSharding(
             mesh, PartitionSpec(get_all_mesh_axes(), None)
         )
@@ -727,7 +784,7 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-1] = arg_shardings[-3]
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
-        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding)
+        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
         impl = partial(FusedAttnFwdPrimitive.impl, config=config)
         return mesh, impl, out_shardings, arg_shardings
 
@@ -756,8 +813,10 @@ def shardy_sharding_rule(config, mesh, value_types, result_types):
         else:
             softmax_aux_sharding = ("…0", "head", "seqlen", "i")
 
+        max_logit_sharding = ("head",) if config.return_max_logit else ("max_logit",)
         return SdyShardingRule(
-            tuple(input_spec), (out_sharding, softmax_aux_sharding, rng_sharding)
+            tuple(input_spec),
+            (out_sharding, softmax_aux_sharding, rng_sharding, max_logit_sharding),
         )
 
 
@@ -1810,19 +1869,22 @@ def partition(config, mesh, arg_infos, result_infos):
         ), "Sliding window attention is not supported when context parallelism is enabled"
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
+        if config.return_max_logit:
+            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
 
         helper = _FusedAttnCPWithAllGatherHelper(mesh, config)
         helper.check_supported()
 
         out_sharding = result_infos[0].sharding
         softmax_aux_sharding = result_infos[1].sharding
+        max_logit_sharding = result_infos[3].sharding
         rng_state_sharding = seed_sharding = NamedSharding(
             mesh, PartitionSpec(get_all_mesh_axes(), None)
         )
         arg_shardings = [arg_i.sharding for arg_i in arg_infos]
         arg_shardings[5] = seed_sharding
         arg_shardings = tuple(arg_shardings)
-        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding)
+        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
 
         def impl(
             q,
@@ -1870,7 +1932,7 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
                     q_seqlen_for_step = q_seqlen / (cp_size * 2)
                     num_kv_chunks = kv_max_seqlen // kv_seqlens_for_rank[sub_idx]
                     kv_seqlen_for_step = (kv_seqlen / (cp_size * 2)) * num_kv_chunks
-                    output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl(
+                    output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.impl(
                         q_split[sub_idx],
                         k_unmasked,
                         v_unmasked,
@@ -1892,8 +1954,9 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
                 output = jnp.concatenate((results[0][0], results[1][0]), axis=1)
                 softmax_aux = jnp.concatenate((results[0][1], results[1][1]), axis=2)
                 rng_state = results[1][2]  # Use the final RNG state
+                max_logit = jnp.zeros((0,), dtype=output.dtype)
 
-                return output, softmax_aux, rng_state
+                return output, softmax_aux, rng_state, max_logit
 
             k_ag, v_ag = helper.all_gather_kv(k, v)
 
@@ -2103,19 +2166,22 @@ def partition(config, mesh, arg_infos, result_infos):
         is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
+        if config.return_max_logit:
+            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
 
         helper = _FusedAttnCPWithAllGatherHelper(mesh, config)
         helper.check_supported()
 
         out_sharding = result_infos[0].sharding
         softmax_aux_sharding = result_infos[1].sharding
+        max_logit_sharding = result_infos[3].sharding
         rng_state_sharding = seed_sharding = NamedSharding(
             mesh, PartitionSpec(get_all_mesh_axes(), None)
         )
         arg_shardings = [arg_i.sharding for arg_i in arg_infos]
         arg_shardings[5] = seed_sharding
         arg_shardings = tuple(arg_shardings)
-        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding)
+        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
 
         def impl(
             q,
@@ -2179,7 +2245,7 @@ def _cross_attn(
                     max_segments_per_seq=adjusted_max_segments_per_seq,
                 )
 
-                output, softmax_aux, rng_state = FusedAttnFwdPrimitive.impl(
+                output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.impl(
                     q,  # sharded for rank
                     k,  # ag
                     v,  # ag
@@ -2198,7 +2264,8 @@ def _cross_attn(
                         max_seqlen=kv_max_seqlen, cp_size=cp_size
                     ),
                 )
-                return output, softmax_aux, rng_state
+                max_logit = jnp.zeros((0,), dtype=output.dtype)
+                return output, softmax_aux, rng_state, max_logit
 
             # AG the k, v, kv_segment_ids and kv_segment_pos
             k_ag, v_ag = helper.all_gather_kv(k, v)
@@ -2553,12 +2620,15 @@ def partition(config, mesh, arg_infos, result_infos):
         ), "Sliding window attention is not supported when context parallelism is enabled"
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
+        if config.return_max_logit:
+            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
 
         helper = _FusedAttnCPWithP2PHelper(mesh, config)
         helper.check_supported()
 
         out_sharding = result_infos[0].sharding
         softmax_aux_sharding = result_infos[1].sharding
+        max_logit_sharding = result_infos[3].sharding
         rng_state_sharding = seed_sharding = NamedSharding(
             mesh, PartitionSpec(get_all_mesh_axes(), None)
         )
@@ -2568,7 +2638,7 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-1] = arg_shardings[-3]
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
-        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding)
+        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
 
         def ring_attn_fwd_impl(
             q,
@@ -2616,7 +2686,7 @@ def scan_kv_block(idx, carry):
                 def mask_compute(attn_mask_type):
                     q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx)
                     kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx)
-                    output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl(
+                    output_per_step, softmax_aux_per_step, _, _ = FusedAttnFwdPrimitive.impl(
                         q,
                         kv,
                         _not_used,
@@ -2642,7 +2712,7 @@ def half_kv_no_mask_compute():
                     q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx)
                     kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) // 2
                     kv_part = lax.slice_in_dim(kv, 0, kv.shape[1] // 2, axis=1)
-                    output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl(
+                    output_per_step, softmax_aux_per_step, _, _ = FusedAttnFwdPrimitive.impl(
                         q,
                         kv_part,
                         _not_used,
@@ -2665,7 +2735,7 @@ def half_q_no_mask_compute():
                     q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) // 2
                     kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx)
                     q_part = lax.slice_in_dim(q, q_max_seqlen // 2, q_max_seqlen, axis=1)
-                    output_per_step, softmax_aux_per_step, _ = FusedAttnFwdPrimitive.impl(
+                    output_per_step, softmax_aux_per_step, _, _ = FusedAttnFwdPrimitive.impl(
                         q_part,
                         kv,
                         _not_used,
@@ -2747,7 +2817,8 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
             (kv, output, softmax_aux) = carry
 
             output = output.astype(q.dtype)
-            return output, softmax_aux, rng_state
+            max_logit = jnp.zeros((0,), dtype=output.dtype)
+            return output, softmax_aux, rng_state, max_logit
 
         return mesh, ring_attn_fwd_impl, out_shardings, arg_shardings
 
@@ -3059,12 +3130,15 @@ def partition(config, mesh, arg_infos, result_infos):
         is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
+        if config.return_max_logit:
+            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
 
         helper = _FusedAttnCPWithP2PHelper(mesh, config)
         helper.check_supported()
 
         out_sharding = result_infos[0].sharding
         softmax_aux_sharding = result_infos[1].sharding
+        max_logit_sharding = result_infos[3].sharding
         rng_state_sharding = seed_sharding = NamedSharding(
             mesh, PartitionSpec(get_all_mesh_axes(), None)
         )
@@ -3074,7 +3148,7 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-1] = arg_shardings[-3]
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
-        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding)
+        out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
 
         def fwd_impl(
             q,
@@ -3157,7 +3231,7 @@ def compute(config):
                     )
                 else:
                     current_config = subblock_config
-                output_per_step, softmax_aux_per_step, _ = compute(current_config)
+                output_per_step, softmax_aux_per_step, _, _ = compute(current_config)
 
                 softmax_aux_per_step = softmax_aux_per_step.reshape((batch, q_max_seqlen, head, 1))
 
@@ -3194,7 +3268,9 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
                     carry = scan_kv_block(i, carry)
             (_, _, _, output, softmax_aux) = carry
 
-            return output.astype(q.dtype), softmax_aux, rng_state
+            output = output.astype(q.dtype)
+            max_logit = jnp.zeros((0,), dtype=output.dtype)
+            return output, softmax_aux, rng_state, max_logit
 
         return mesh, fwd_impl, out_shardings, arg_shardings
 
@@ -3375,6 +3451,7 @@ def fused_attn_fwd(
     context_parallel_causal_load_balanced: bool = False,
     context_parallel_axis: str = "",
     stripe_size: int | None = None,
+    return_max_logit: bool = False,
 ) -> jnp.ndarray:
     """
     Perform the forward pass of with cuDNN fused attention implementations.
@@ -3414,6 +3491,7 @@ def fused_attn_fwd(
             Indicates the sequences are ordered for causal mask load balancing when running context parallelism.
         context_parallel_axis (str): The name of the context parallel axis.
         stripe_size (int | None): Indicates the striping height to be used for ReorderStrategy.Striped Load Balancing
+        return_max_logit (bool): Whether to return the per-head maximum attention logit.
     Returns:
         (jnp.ndarray): The output tensor from the fused attention.
     """
@@ -3489,6 +3567,7 @@ def fused_attn_fwd(
         cp_axis=_maybe_context_parallel_axis(context_parallel_axis),
         cp_striped_window_size=None,
         stripe_size=stripe_size,
+        return_max_logit=return_max_logit,
     )
 
     primitive = None
@@ -3506,7 +3585,7 @@ def fused_attn_fwd(
                 primitive = FusedRingAttnFwdPrimitive.outer_primitive
 
     seq_desc_flatten, _ = jax.tree.flatten(sequence_descriptor)
-    output, softmax_aux, rng_state = primitive.bind(
+    output, softmax_aux, rng_state, max_logit = primitive.bind(
         *qkv_for_primitive,
         bias,
         softmax_offset,
@@ -3515,7 +3594,7 @@ def fused_attn_fwd(
         config=fused_config,
     )
     rng_state = with_sharding_constraint(rng_state, PartitionSpec(get_all_mesh_axes(), None))
-    return (output, softmax_aux, rng_state)
+    return (output, softmax_aux, rng_state, max_logit)
 
 
 def fused_attn_bwd(

From 8f565cf0227496ebd8ac7a90f5c4224b20ced6b5 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Tue, 9 Jun 2026 23:13:11 -0700
Subject: [PATCH 03/22] Add max logit to JAX fused attn FFI and set it in the
 workspace

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 transformer_engine/jax/csrc/extensions.h      |  4 +-
 .../jax/csrc/extensions/attention.cpp         | 80 ++++++++++++-------
 2 files changed, 53 insertions(+), 31 deletions(-)

diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h
index 9bd0940c4b..580219baf2 100644
--- a/transformer_engine/jax/csrc/extensions.h
+++ b/transformer_engine/jax/csrc/extensions.h
@@ -156,7 +156,7 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend(
     NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type,
     float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen,
     size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left,
-    int64_t window_size_right, bool deterministic);
+    int64_t window_size_right, bool return_max_logit, bool deterministic);
 
 pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
     size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
@@ -164,7 +164,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
     size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
     NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout,
     DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left,
-    int64_t window_size_right, bool bottom_right_diagonal);
+    int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal);
 
 pybind11::tuple GetFusedAttnBackwardWorkspaceSizes(
     size_t input_batch, size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen,
diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp
index 3fd6780d6d..9728f2cbcf 100644
--- a/transformer_engine/jax/csrc/extensions/attention.cpp
+++ b/transformer_engine/jax/csrc/extensions/attention.cpp
@@ -29,12 +29,12 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend(
     NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type,
     float dropout_probability, size_t q_attn_heads, size_t kv_attn_heads, size_t q_max_seqlen,
     size_t kv_max_seqlen, size_t qk_head_dim, size_t v_head_dim, int64_t window_size_left,
-    int64_t window_size_right, bool deterministic) {
+    int64_t window_size_right, bool return_max_logit, bool deterministic) {
   auto backend = nvte_get_fused_attn_backend(
       is_training, static_cast<NVTEDType>(q_dtype), static_cast<NVTEDType>(kv_dtype), qkv_layout,
       bias_type, mask_type, softmax_type, dropout_probability, q_attn_heads, kv_attn_heads,
       q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right,
-      false, false, deterministic);
+      return_max_logit, false, deterministic);
   return backend;
 }
 
@@ -48,8 +48,8 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
                                        const size_t bias_heads, const size_t q_max_seqlen,
                                        const size_t kv_max_seqlen, DType dtype,
                                        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,
+                                       void *rng_state_buf = nullptr, void *bias_buf = nullptr,
                                        void *softmax_offset_buf = nullptr) {
   // all backends need softmax but expect different shapes/dtypes
   tensor_pack->size = 1;
@@ -65,8 +65,23 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
 
   // arbitrary sequence length backend needs the RNG state and a different shape/dtype softmax
   if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) {
-    tensor_pack->size = 2;
-    NVTETensor &rng_state_aux = tensor_pack->tensors[1];
+    int size = 1;  // Start after softmax.
+
+    if (max_buf != nullptr) {
+      NVTETensor &max_aux = tensor_pack->tensors[size++];
+      NVTEBasicTensor max_aux_data;
+      max_aux_data.data_ptr = max_buf;
+      max_aux_data.shape = {};
+      max_aux_data.shape.ndim = 4;
+      max_aux_data.shape.data[0] = input_batch;
+      max_aux_data.shape.data[1] = attn_heads;
+      max_aux_data.shape.data[2] = q_max_seqlen;
+      max_aux_data.shape.data[3] = 1;
+      max_aux_data.dtype = static_cast<NVTEDType>(DType::kFloat32);
+      nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data);
+    }
+
+    NVTETensor &rng_state_aux = tensor_pack->tensors[size++];
     NVTEBasicTensor rng_state_aux_data;
     rng_state_aux_data.data_ptr = rng_state_buf;
     rng_state_aux_data.shape = {};
@@ -77,8 +92,6 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
     softmax_aux_data.shape.data[3] = 1;  // {B,H,Qs,Ks} -> {B,H,Qs,1}
     softmax_aux_data.dtype = static_cast<NVTEDType>(DType::kFloat32);
 
-    int size = 2;  // Start at 2 (we have softmax and rng_state at indices 0, 1)
-
     // include bias if enabled
     if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) {
       NVTETensor &bias_aux = tensor_pack->tensors[size];
@@ -136,7 +149,7 @@ void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_
   auto dummy_backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen;
   PrepareFusedAttnForwardAuxTensors(tensor_pack, input_batch, bias_batch, attn_heads, bias_heads,
                                     q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type,
-                                    dummy_backend, softmax_buf, rng_state_buf, bias_buf,
+                                    dummy_backend, softmax_buf, nullptr, rng_state_buf, bias_buf,
                                     softmax_offset_buf);
 }
 
@@ -146,7 +159,7 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
     size_t v_head_dim, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
     NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout,
     DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left,
-    int64_t window_size_right, bool bottom_right_diagonal) {
+    int64_t window_size_right, bool return_max_logit, bool bottom_right_diagonal) {
   auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD;
   auto q_shape = is_ragged
                      ? std::vector<size_t>{input_batch * q_max_seqlen, attn_heads, qk_head_dim}
@@ -200,8 +213,8 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
         dummy_softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors,
         q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), ragged_offset_tensor.data(),
         ragged_offset_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(),
-        dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false,
-        scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout),
+        dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit,
+        false, scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout),
         NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left,
         window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), nullptr);
   }
@@ -242,13 +255,14 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes(
 static void FusedAttnForwardImpl(
     cudaStream_t stream, void *q, void *k, void *v, void *bias, void *softmax_offset, void *seed,
     void *q_cu_seqlens, void *kv_cu_seqlens, void *q_seq_offsets, void *k_seq_offsets, void *output,
-    void *softmax_aux, void *rng_state, void *workspace, size_t input_batch, size_t bias_batch,
-    size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads, size_t num_gqa_groups,
-    size_t bias_heads, size_t qk_head_dim, size_t v_head_dim, size_t max_segments_per_seq,
-    size_t wkspace_size, float scaling_factor, float dropout_probability, NVTE_Bias_Type bias_type,
-    NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout,
-    DType dtype, DType wkspace_dtype, bool is_training, bool deterministic,
-    int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) {
+    void *softmax_aux, void *max_tensor, void *rng_state, void *workspace, size_t input_batch,
+    size_t bias_batch, size_t q_max_seqlen, size_t kv_max_seqlen, size_t attn_heads,
+    size_t num_gqa_groups, size_t bias_heads, size_t qk_head_dim, size_t v_head_dim,
+    size_t max_segments_per_seq, size_t wkspace_size, float scaling_factor,
+    float dropout_probability, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type,
+    NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, DType wkspace_dtype,
+    bool is_training, bool return_max_logit, bool deterministic, int64_t window_size_left,
+    int64_t window_size_right, bool bottom_right_diagonal) {
   FUSED_ATTN_IMPL_COMMON_BLOCK;
 
   /* Input tensors */
@@ -263,6 +277,9 @@ static void FusedAttnForwardImpl(
     // Memset to 0xF0 for filling large negative numbers
     auto softmax_aux_size = input_batch * q_max_seqlen * attn_heads;
     cudaMemsetAsync(softmax_aux, 0xF0, softmax_aux_size * sizeof(float), stream);
+    if (return_max_logit) {
+      cudaMemsetAsync(max_tensor, 0xF0, softmax_aux_size * sizeof(float), stream);
+    }
   }
 
   /* Output tensors */
@@ -278,7 +295,7 @@ static void FusedAttnForwardImpl(
       is_training, static_cast<NVTEDType>(dtype), static_cast<NVTEDType>(dtype), qkv_layout,
       bias_type, mask_type, softmax_type, dropout_probability, attn_heads, num_gqa_groups,
       q_max_seqlen, kv_max_seqlen, qk_head_dim, v_head_dim, window_size_left, window_size_right,
-      false, false, deterministic);
+      return_max_logit, false, deterministic);
   nvte_populate_rng_state_async(rng_state, seed, q_max_seqlen, kv_max_seqlen, backend, stream);
 
   /* Auxiliary tensors (to be propagated to the backward pass later) */
@@ -286,7 +303,8 @@ static void FusedAttnForwardImpl(
   nvte_tensor_pack_create(&aux_output_tensors);
   PrepareFusedAttnForwardAuxTensors(&aux_output_tensors, input_batch, bias_batch, attn_heads,
                                     bias_heads, q_max_seqlen, kv_max_seqlen, dtype, bias_type,
-                                    backend, softmax_aux, softmax_offset);
+                                    backend, softmax_aux, return_max_logit ? max_tensor : nullptr,
+                                    rng_state, bias, softmax_offset);
 
   /* Call the underlying NVTE API */
   auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector<size_t>{1}, DType::kInt32);
@@ -344,7 +362,7 @@ static void FusedAttnForwardImpl(
       softmax_offset_tensor.data(), s_tensor.data(), o_tensor.data(), &aux_output_tensors,
       q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(),
       k_seq_offsets_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(),
-      rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false,
+      rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, return_max_logit, false,
       scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout),
       NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left,
       window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream);
@@ -378,6 +396,7 @@ static void FusedAttnForwardImpl(
   NVTE_QKV_Layout qkv_layout =                                                                    \
       static_cast<NVTE_QKV_Layout>(get_attr_value<int64_t>(attrs, "qkv_layout"));                 \
   bool is_training = get_attr_value<bool>(attrs, "is_training");                                  \
+  bool return_max_logit = get_attr_value_or_default<bool>(attrs, "return_max_logit", false);      \
   bool deterministic = get_attr_value<bool>(attrs, "deterministic");                              \
   auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD;                  \
   size_t wkspace_size = product(workspace_buf->dimensions());                                     \
@@ -390,8 +409,9 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty
                                Buffer_Type q_cu_seqlens_buf, Buffer_Type kv_cu_seqlens_buf,
                                Buffer_Type q_seq_offsets_buf, Buffer_Type k_seq_offsets_buf,
                                Variadic_Buffer_Type _unused_args, Result_Type output_buf,
-                               Result_Type softmax_aux_buf, Result_Type rng_state_buf,
-                               Result_Type workspace_buf, Dictionary attrs) {
+                               Result_Type softmax_aux_buf, Result_Type max_tensor_buf,
+                               Result_Type rng_state_buf, Result_Type workspace_buf,
+                               Dictionary attrs) {
   FUSED_ATTN_FFI_GET_ATTRS;
 
   FusedAttnForwardImpl(
@@ -400,11 +420,12 @@ Error_Type FusedAttnForwardFFI(cudaStream_t stream, Buffer_Type q_buf, Buffer_Ty
       q_cu_seqlens_buf.untyped_data(), kv_cu_seqlens_buf.untyped_data(),
       is_ragged ? q_seq_offsets_buf.untyped_data() : nullptr,
       is_ragged ? k_seq_offsets_buf.untyped_data() : nullptr, output_buf->untyped_data(),
-      softmax_aux_buf->untyped_data(), rng_state_buf->untyped_data(), workspace_buf->untyped_data(),
-      input_batch, bias_batch, q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads,
-      qk_head_dim, v_head_dim, max_segments_per_seq, wkspace_size, scaling_factor,
-      dropout_probability, bias_type, mask_type, softmax_type, qkv_layout, dtype, wkspace_dtype,
-      is_training, deterministic, window_size_left, window_size_right, bottom_right_diagonal);
+      softmax_aux_buf->untyped_data(), max_tensor_buf->untyped_data(),
+      rng_state_buf->untyped_data(), workspace_buf->untyped_data(), input_batch, bias_batch,
+      q_max_seqlen, kv_max_seqlen, attn_heads, num_gqa_groups, bias_heads, qk_head_dim, v_head_dim,
+      max_segments_per_seq, wkspace_size, scaling_factor, dropout_probability, bias_type, mask_type,
+      softmax_type, qkv_layout, dtype, wkspace_dtype, is_training, return_max_logit, deterministic,
+      window_size_left, window_size_right, bottom_right_diagonal);
   return ffi_with_cuda_error_check();
 }
 
@@ -424,6 +445,7 @@ XLA_FFI_DEFINE_HANDLER_SYMBOL(FusedAttnForwardHandler, FusedAttnForwardFFI,
                                   .RemainingArgs()         // _cp_aux_args unused
                                   .Ret<Buffer_Type>()      // output
                                   .Ret<Buffer_Type>()      // softmax_aux
+                                  .Ret<Buffer_Type>()      // max_tensor
                                   .Ret<Buffer_Type>()      // rng_state
                                   .Ret<Buffer_Type>()      // workspace
                                   .Attrs(),

From 31ede581e99f0e6a5530ac9cf39a706feb67ec26 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Tue, 9 Jun 2026 23:15:21 -0700
Subject: [PATCH 04/22] Add first pass tests for max logit and softmax aux
 tensor outputs in JAX fused attn tests

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn.py | 224 +++++++++++++++++++++++++++++++++--
 1 file changed, 216 insertions(+), 8 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 352ab64a0d..0e65e42a30 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -79,27 +79,23 @@ def general_dot_product_attention(
     dropout_rng: ArrayLike,
     dtype: DTypeLike,
     score_mod_reference: Optional[Callable[[Array], Array]] = None,
+    is_max_logit_enabled: bool = False,
 ) -> Array:
     """
     Similar to flax.linen.dot_product_attention but with GQA support
     """
     query, key, value, bias = promote_dtype(query, key, value, bias, dtype=dtype)
     dtype = query.dtype
-
     b, s_q, h_q, d = query.shape
     _, s_kv, h_kv, _ = key.shape
     assert (h_q % h_kv == 0) and (h_q >= h_kv)
     num_groups = h_q // h_kv
     grouped_query = jnp.reshape(query, (b, s_q, h_kv, num_groups, d))
-    # logits with shape (b, h_kv, num_groups, s_q, s_kv)
     logits = scale_factor * jnp.einsum("...qhgd,...khd->...hgqk", grouped_query, key)
 
     if bias is not None:
-        # reshape logits without groups
         logits = logits.reshape((b, h_kv * num_groups, s_q, s_kv))
-        # apply post-scale bias
         logits = logits + bias
-        # reshape logits back to original
         logits = logits.reshape((b, h_kv, num_groups, s_q, s_kv))
 
     if mask is not None:
@@ -110,6 +106,8 @@ def general_dot_product_attention(
     if score_mod_reference is not None:
         # Kernel tests use NO_MASK; fused_attn rejects mask+score_mod before this reference path.
         logits = score_mod_reference(logits.astype(jnp.float32))
+    if is_max_logit_enabled:
+        return jnp.max(logits.reshape((b, h_q, s_q, s_kv)), axis=(0, 2, 3)).astype(dtype)
 
     match softmax_type:
         case AttnSoftmaxType.VANILLA_SOFTMAX:
@@ -268,7 +266,17 @@ def _split_valid_and_invalid(primitive, reference, pad):
     return primitive_valid, primitive_invalid, reference_valid, reference_invalid
 
 
-def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs):
+def jax_dpa(
+    query,
+    key,
+    value,
+    bias,
+    softmax_offset,
+    mask,
+    dropout_rng,
+    is_max_logit_enabled=False,
+    **kwargs,
+):
     """
     JAX native dot product attention implementation
     """
@@ -308,6 +316,7 @@ def jax_dpa(query, key, value, bias, softmax_offset, mask, dropout_rng, **kwargs
         dropout_rng=dropout_rng,
         dtype=jnp.float32,
         score_mod_reference=score_mod_reference,
+        is_max_logit_enabled=is_max_logit_enabled,
     )
     return output.astype(query.dtype)
 
@@ -339,9 +348,13 @@ def customcall_fused_dpa(
             qkv_args = (query, key, value)
         case _:
             raise ValueError(f"Unsupported {qkv_layout=}")
-    return fused_attn(
+    result = fused_attn(
         qkv_args, bias, sequence_descriptor, dropout_rng, softmax_offset=softmax_offset, **kwargs
-    ).astype(query.dtype)
+    )
+    if isinstance(result, tuple):
+        output, aux = result
+        return output.astype(query.dtype), aux
+    return result.astype(query.dtype)
 
 
 def test_fused_attn_score_mod_rejects_masks_before_cudnn_frontend():
@@ -1278,6 +1291,201 @@ def check_dqkv(primitive, reference, pad, idx):
                 target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text()
             assert_equal_collectives(target_hlo, self.coll_count_ref)
 
+    def _reference_args(self):
+        return [self.q, self.k, self.v, self.bias, self.softmax_offset, self.mask, self.dropout_rng]
+
+    def _customcall_args(self):
+        return [
+            jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding),
+            jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
+            jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding),
+            jax.device_put(self.bias, self.bias_sharding),
+            jax.device_put(self.softmax_offset, self.softmax_offset_sharding),
+            jax.device_put(self.sequence_desciptor, self.seq_desc_sharding),
+            jax.device_put(self.dropout_rng, self.dropout_rng_sharding),
+        ]
+
+    def _fused_attn_kwargs(self, **overrides):
+        kwargs = {
+            "attn_bias_type": self.attn_bias_type,
+            "attn_mask_type": self.attn_mask_type,
+            "softmax_type": self.softmax_type,
+            "scaling_factor": self.scaling_factor,
+            "dropout_probability": self.dropout_prob,
+            "is_training": self.is_training,
+            "qkv_layout": self.qkv_layout,
+            "max_segments_per_seq": self._get_max_segments_per_sequence(),
+            "window_size": self.window_size,
+            "context_parallel_strategy": self.cp_strategy,
+            "context_parallel_causal_load_balanced": self.cp_load_balanced,
+            "stripe_size": self.stripe_size,
+        }
+        kwargs.update(overrides)
+        return kwargs
+
+    def test_forward_with_max_logit(self):
+        """Test forward output and returned max_logit."""
+        self._setup_inputs()
+        kwargs = self._fused_attn_kwargs()
+
+        customcall_fused_dpa_jit = jit(
+            partial(customcall_fused_dpa, return_max_logit=True, **kwargs),
+            static_argnames=kwargs.keys(),
+            in_shardings=[
+                self.qkvo_sharding,
+                self.qkvo_sharding,
+                self.qkvo_sharding,
+                self.bias_sharding,
+                self.softmax_offset_sharding,
+                self.seq_desc_sharding,
+                self.dropout_rng_sharding,
+            ],
+        )
+
+        with self.mesh, autocast(mesh_resource=self.mesh_resource):
+            primitive_out, primitive_aux = customcall_fused_dpa_jit(*self._customcall_args())
+            primitive_max_logit = primitive_aux["max_logit"]
+            primitive_out = self.cp_inverse_reorder_fn(primitive_out)
+
+        reference_out = jax_dpa(*self._reference_args(), **kwargs)
+        reference_max_logit = jax_dpa(
+            *self._reference_args(), is_max_logit_enabled=True, **kwargs
+        )
+
+        primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid(
+            primitive_out, reference_out, self.pad_q
+        )
+        assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype)
+        assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
+        assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype)
+
+    def test_forward_with_softmax_aux(self):
+        """Test optional softmax_aux return wiring."""
+        self._setup_inputs()
+        kwargs = self._fused_attn_kwargs(return_softmax_aux=True)
+
+        customcall_fused_dpa_jit = jit(
+            partial(customcall_fused_dpa, **kwargs),
+            static_argnames=kwargs.keys(),
+            in_shardings=[
+                self.qkvo_sharding,
+                self.qkvo_sharding,
+                self.qkvo_sharding,
+                self.bias_sharding,
+                self.softmax_offset_sharding,
+                self.seq_desc_sharding,
+                self.dropout_rng_sharding,
+            ],
+        )
+
+        with self.mesh, autocast(mesh_resource=self.mesh_resource):
+            output, aux = customcall_fused_dpa_jit(*self._customcall_args())
+
+        assert output.shape == self.q.shape
+        assert "softmax_aux" in aux
+        assert aux["softmax_aux"].dtype == jnp.float32
+
+    def test_backward_with_max_logit(self):
+        """Ensure aux-return cotangents do not break the fused attention backward path."""
+        self._setup_inputs()
+        kwargs = self._fused_attn_kwargs()
+
+        def loss_fn(query):
+            output, _ = customcall_fused_dpa(
+                query,
+                self.k,
+                self.v,
+                self.bias,
+                self.softmax_offset,
+                self.sequence_desciptor,
+                self.dropout_rng,
+                return_max_logit=True,
+                **kwargs,
+            )
+            return jnp.mean(output.astype(jnp.float32))
+
+        grad = jax.grad(loss_fn)(self.q)
+        assert grad.shape == self.q.shape
+
+
+@pytest.mark.parametrize(
+    "qkv_layout, seq_desc_format",
+    [
+        pytest.param(QKVLayout.BSHD_BSHD_BSHD, SeqDescFormat.Seqlens, id="BSHD_SEPARATE"),
+        pytest.param(QKVLayout.T3HD, SeqDescFormat.Seqlens, id="THD_QKV_PACKED"),
+    ],
+)
+def test_fused_attn_return_max_logit(qkv_layout, seq_desc_format):
+    """Check non-CP JAX fused attention can expose PyTorch-compatible max_logit."""
+    runner = FusedAttnRunner(
+        batch_size=2,
+        max_seqlen_q=128,
+        max_seqlen_kv=128,
+        num_heads_q=8,
+        num_heads_kv=8,
+        head_dim_qk=64,
+        head_dim_v=64,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_prob=0.0,
+        dtype=jnp.bfloat16,
+        is_training=True,
+        qkv_layout=qkv_layout,
+        bias_shape=None,
+        window_size=None,
+        seq_desc_format=seq_desc_format,
+    )
+    runner.test_forward_with_max_logit()
+
+
+def test_fused_attn_return_softmax_aux():
+    """Check the optional public softmax_aux return is wired for non-CP attention."""
+    runner = FusedAttnRunner(
+        batch_size=2,
+        max_seqlen_q=128,
+        max_seqlen_kv=128,
+        num_heads_q=8,
+        num_heads_kv=8,
+        head_dim_qk=64,
+        head_dim_v=64,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_prob=0.0,
+        dtype=jnp.bfloat16,
+        is_training=True,
+        qkv_layout=QKVLayout.BS3HD,
+        bias_shape=None,
+        window_size=None,
+        seq_desc_format=SeqDescFormat.Seqlens,
+    )
+    runner.test_forward_with_softmax_aux()
+
+
+def test_fused_attn_return_max_logit_backward_smoke():
+    """Ensure aux-return cotangents do not break the fused attention backward path."""
+    runner = FusedAttnRunner(
+        batch_size=2,
+        max_seqlen_q=128,
+        max_seqlen_kv=128,
+        num_heads_q=8,
+        num_heads_kv=8,
+        head_dim_qk=64,
+        head_dim_v=64,
+        attn_bias_type=AttnBiasType.NO_BIAS,
+        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
+        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+        dropout_prob=0.0,
+        dtype=jnp.bfloat16,
+        is_training=True,
+        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+        bias_shape=None,
+        window_size=None,
+        seq_desc_format=SeqDescFormat.Seqlens,
+    )
+    runner.test_backward_with_max_logit()
+
 
 def _get_swa_window_size_for_test(s_kv: int, attn_mask_type: AttnMaskType) -> Tuple[int, int]:
     """Pick a sliding-window size for SWA tests, gated on cuDNN version.

From dde413b29c2e98854e944a9eb7c7f5f43995d399 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Wed, 8 Jul 2026 16:00:37 -0700
Subject: [PATCH 05/22] Reject aux returns with score_mod

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 transformer_engine/jax/attention.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index 9d3aaae90d..a300fb353c 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -1629,6 +1629,18 @@ def fused_attn(
         if score_mod_only_args:
             raise ValueError(f"{', '.join(score_mod_only_args)} require score_mod to be provided.")
     else:
+        aux_return_args = [
+            name
+            for name, value in (
+                ("return_max_logit", return_max_logit),
+                ("return_softmax_aux", return_softmax_aux),
+            )
+            if value
+        ]
+        if aux_return_args:
+            raise ValueError(
+                f"{', '.join(aux_return_args)} are not supported with score_mod fused_attn."
+            )
         tex.validate_fused_attn_score_mod(
             qkv,
             bias,

From 9da2226c3c9ae898e2349f77afac30d40e7ca8d7 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Wed, 8 Jul 2026 16:01:43 -0700
Subject: [PATCH 06/22] Handle SM120 max-logit layout

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 .../jax/cpp_extensions/attention.py              | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index e129997b85..53db542798 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -377,7 +377,7 @@ def abstract(
             raise ValueError(f"Unsupported {backend=}")
         softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype)
         if config.return_max_logit:
-            if config.qkv_layout.is_thd() and get_cudnn_version() >= (9, 6, 0):
+            if FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config):
                 max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1)
             else:
                 max_tensor_shape = (*batch_shape, attn_heads, q_max_seqlen, 1)
@@ -678,6 +678,7 @@ def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config):
         if not config.return_max_logit:
             return jnp.zeros((0,), dtype=output.dtype)
 
+        uses_thd_ragged_max_tensor = FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config)
         if config.qkv_layout.is_thd() and max_tensor.ndim == 4:
             q_seqlen = jnp.where(q_seqlen > 0, q_seqlen, 0)
             q_seq_offsets = jnp.where(q_seq_offsets >= 0, q_seq_offsets, -1)
@@ -688,19 +689,28 @@ def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config):
                 & (token_idx < (q_seq_offsets[..., :-1, None] + q_seqlen[..., None])),
                 axis=-2,
             )
-            if max_tensor.shape[1] == output.shape[-3]:
+            if uses_thd_ragged_max_tensor:
                 max_tensor = jnp.where(valid[:, :, None, None], max_tensor, -jnp.inf)
             else:
                 max_tensor = jnp.where(valid[:, None, :, None], max_tensor, -jnp.inf)
 
         if max_tensor.ndim == 3:
             amax_dims = (0, 2)
-        elif config.qkv_layout.is_thd() and max_tensor.shape[1] == output.shape[-3]:
+        elif uses_thd_ragged_max_tensor:
             amax_dims = (0, 1, 3)
         else:
             amax_dims = (0, 2, 3)
         return jnp.max(max_tensor, axis=amax_dims).astype(output.dtype)
 
+    @staticmethod
+    def _uses_thd_ragged_max_tensor(config):
+        """Return whether cuDNN writes THD Max with BSH-like ragged-stats layout."""
+        return (
+            config.qkv_layout.is_thd()
+            and get_cudnn_version() >= (9, 6, 0)
+            and 120 not in get_all_device_compute_capability()
+        )
+
     @staticmethod
     def batcher(batched_args, batch_dims, *, config):
         # batch_dims: each element is the batch axis (0, ...) or None. Only 0 or None allowed.

From 38ce6e3e4356c3dcc83a866759ecca896f2d0560 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Wed, 8 Jul 2026 16:02:12 -0700
Subject: [PATCH 07/22] Drop softmax aux return

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn.py        | 50 ------------------------
 transformer_engine/jax/attention.py | 60 +++++++----------------------
 2 files changed, 14 insertions(+), 96 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 0e65e42a30..6d96604964 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1359,32 +1359,6 @@ def test_forward_with_max_logit(self):
         assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
         assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype)
 
-    def test_forward_with_softmax_aux(self):
-        """Test optional softmax_aux return wiring."""
-        self._setup_inputs()
-        kwargs = self._fused_attn_kwargs(return_softmax_aux=True)
-
-        customcall_fused_dpa_jit = jit(
-            partial(customcall_fused_dpa, **kwargs),
-            static_argnames=kwargs.keys(),
-            in_shardings=[
-                self.qkvo_sharding,
-                self.qkvo_sharding,
-                self.qkvo_sharding,
-                self.bias_sharding,
-                self.softmax_offset_sharding,
-                self.seq_desc_sharding,
-                self.dropout_rng_sharding,
-            ],
-        )
-
-        with self.mesh, autocast(mesh_resource=self.mesh_resource):
-            output, aux = customcall_fused_dpa_jit(*self._customcall_args())
-
-        assert output.shape == self.q.shape
-        assert "softmax_aux" in aux
-        assert aux["softmax_aux"].dtype == jnp.float32
-
     def test_backward_with_max_logit(self):
         """Ensure aux-return cotangents do not break the fused attention backward path."""
         self._setup_inputs()
@@ -1439,30 +1413,6 @@ def test_fused_attn_return_max_logit(qkv_layout, seq_desc_format):
     runner.test_forward_with_max_logit()
 
 
-def test_fused_attn_return_softmax_aux():
-    """Check the optional public softmax_aux return is wired for non-CP attention."""
-    runner = FusedAttnRunner(
-        batch_size=2,
-        max_seqlen_q=128,
-        max_seqlen_kv=128,
-        num_heads_q=8,
-        num_heads_kv=8,
-        head_dim_qk=64,
-        head_dim_v=64,
-        attn_bias_type=AttnBiasType.NO_BIAS,
-        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
-        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
-        dropout_prob=0.0,
-        dtype=jnp.bfloat16,
-        is_training=True,
-        qkv_layout=QKVLayout.BS3HD,
-        bias_shape=None,
-        window_size=None,
-        seq_desc_format=SeqDescFormat.Seqlens,
-    )
-    runner.test_forward_with_softmax_aux()
-
-
 def test_fused_attn_return_max_logit_backward_smoke():
     """Ensure aux-return cotangents do not break the fused attention backward path."""
     runner = FusedAttnRunner(
diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index a300fb353c..3eb41ae13f 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -1056,7 +1056,6 @@ def _legacy_fused_attn(
     context_parallel_axis: str = "",
     softmax_offset: Optional[jnp.ndarray] = None,
     return_max_logit: bool = False,
-    return_softmax_aux: bool = False,
 ):
     """
     Perform non-THD (non-packed) cuDNN fused attention.
@@ -1092,14 +1091,12 @@ def _legacy_fused_attn(
             [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX.
         return_max_logit (bool): If True, also return per-head maximum attention logits
             in an auxiliary dictionary under ``"max_logit"``.
-        return_softmax_aux (bool): If True, also return backend-specific softmax statistics
-            in an auxiliary dictionary under ``"softmax_aux"``.
     Returns:
         jnp.ndarray:
-            Attention output when neither ``return_max_logit`` nor ``return_softmax_aux`` is True.
+            Attention output when ``return_max_logit`` is False.
         tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
-            ``(output, aux)`` when either flag is True. ``aux`` may contain:
-            ``"max_logit"`` (shape ``[h]``) and/or ``"softmax_aux"`` (float32).
+            ``(output, aux)`` when ``return_max_logit`` is True. ``aux`` contains
+            ``"max_logit"`` with shape ``[h]``.
     """
     assert (
         not qkv_layout.is_thd()
@@ -1154,7 +1151,6 @@ def _legacy_fused_attn(
         context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
         context_parallel_axis=context_parallel_axis,
         return_max_logit=return_max_logit,
-        return_softmax_aux=return_softmax_aux,
     )
 
     return output
@@ -1181,7 +1177,6 @@ def fused_attn_thd(
     context_parallel_axis: str = "",
     softmax_offset: Optional[jnp.ndarray] = None,
     return_max_logit: bool = False,
-    return_softmax_aux: bool = False,
 ):
     """
     Deprecated THD fused attn, please use fusd_attn with SequenceDescriptor
@@ -1237,7 +1232,6 @@ def fused_attn_thd(
         context_parallel_causal_load_balanced=context_parallel_causal_load_balanced,
         context_parallel_axis=context_parallel_axis,
         return_max_logit=return_max_logit,
-        return_softmax_aux=return_softmax_aux,
     )
 
     return output
@@ -1245,7 +1239,7 @@ def fused_attn_thd(
 
 @partial(
     jax.custom_vjp,
-    nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20),
+    nondiff_argnums=(5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19),
 )
 def _fused_attn(
     qkv: Tuple[jnp.ndarray, ...],
@@ -1268,7 +1262,6 @@ def _fused_attn(
     context_checkpoint_name: str = "context",
     stripe_size: int | None = None,
     return_max_logit: bool = False,
-    return_softmax_aux: bool = False,
 ):
     output, _ = _fused_attn_fwd_rule(
         qkv,
@@ -1291,7 +1284,6 @@ def _fused_attn(
         context_checkpoint_name=context_checkpoint_name,
         stripe_size=stripe_size,
         return_max_logit=return_max_logit,
-        return_softmax_aux=return_softmax_aux,
     )
     return output
 
@@ -1317,7 +1309,6 @@ def _fused_attn_fwd_rule(
     context_checkpoint_name,
     stripe_size,
     return_max_logit,
-    return_softmax_aux,
 ):
     output, softmax_aux, rng_state, max_logit = tex.fused_attn_fwd(
         qkv,
@@ -1344,9 +1335,7 @@ def _fused_attn_fwd_rule(
     softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name)
     rng_state = checkpoint_name(rng_state, context_checkpoint_name)
     max_logit = checkpoint_name(max_logit, context_checkpoint_name)
-    attn_output = _resolve_fused_attn_output(
-        output, max_logit, softmax_aux, return_max_logit, return_softmax_aux
-    )
+    attn_output = _resolve_fused_attn_output(output, max_logit, return_max_logit)
     return attn_output, (
         qkv,
         bias,
@@ -1374,12 +1363,11 @@ def _fused_attn_bwd_rule(
     context_checkpoint_name,
     stripe_size,
     return_max_logit,
-    return_softmax_aux,
     ctx,
     dz,
 ):
     del context_checkpoint_name
-    if return_max_logit or return_softmax_aux:
+    if return_max_logit:
         dz, _ = dz
     (
         qkv,
@@ -1426,16 +1414,11 @@ def _fused_attn_bwd_rule(
     )
 
 
-def _resolve_fused_attn_output(output, max_logit, softmax_aux, return_max_logit, return_softmax_aux):
-    if not return_max_logit and not return_softmax_aux:
+def _resolve_fused_attn_output(output, max_logit, return_max_logit):
+    if not return_max_logit:
         return output
 
-    aux = {}
-    if return_max_logit:
-        aux["max_logit"] = max_logit
-    if return_softmax_aux:
-        aux["softmax_aux"] = softmax_aux
-    return output, aux
+    return output, {"max_logit": max_logit}
 
 
 _fused_attn.defvjp(_fused_attn_fwd_rule, _fused_attn_bwd_rule)
@@ -1519,7 +1502,6 @@ def fused_attn(
     score_mod_tensors: Optional[Mapping[str, Any]] = None,
     score_mod_bprop_tensors: Optional[Mapping[str, Any]] = None,
     return_max_logit: bool = False,
-    return_softmax_aux: bool = False,
 ):
     """
     Perform cuDNN fused attention.
@@ -1578,14 +1560,12 @@ def fused_attn(
             Python/NumPy scalars made available to `score_mod_bprop`.
         return_max_logit (bool): If True, also return per-head maximum attention logits
             in an auxiliary dictionary under ``"max_logit"``.
-        return_softmax_aux (bool): If True, also return backend-specific softmax statistics
-            in an auxiliary dictionary under ``"softmax_aux"``.
     Returns:
         jnp.ndarray:
-            Attention output when neither ``return_max_logit`` nor ``return_softmax_aux`` is True.
+            Attention output when ``return_max_logit`` is False.
         tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
-            ``(output, aux)`` when either flag is True. ``aux`` may contain:
-            ``"max_logit"`` (shape ``[h]``) and/or ``"softmax_aux"`` (float32).
+            ``(output, aux)`` when ``return_max_logit`` is True. ``aux`` contains
+            ``"max_logit"`` with shape ``[h]``.
 
     Examples (non-THD, also known as non-packed):
         >>> #  q_segment_ids = [[1, 1, 1, 0], [1, 1, 0, 0]], 0 means padded tokens
@@ -1629,18 +1609,8 @@ def fused_attn(
         if score_mod_only_args:
             raise ValueError(f"{', '.join(score_mod_only_args)} require score_mod to be provided.")
     else:
-        aux_return_args = [
-            name
-            for name, value in (
-                ("return_max_logit", return_max_logit),
-                ("return_softmax_aux", return_softmax_aux),
-            )
-            if value
-        ]
-        if aux_return_args:
-            raise ValueError(
-                f"{', '.join(aux_return_args)} are not supported with score_mod fused_attn."
-            )
+        if return_max_logit:
+            raise ValueError("return_max_logit is not supported with score_mod fused_attn.")
         tex.validate_fused_attn_score_mod(
             qkv,
             bias,
@@ -1701,7 +1671,6 @@ def fused_attn(
             context_parallel_axis=context_parallel_axis,
             softmax_offset=softmax_offset,
             return_max_logit=return_max_logit,
-            return_softmax_aux=return_softmax_aux,
         )
     if max_segments_per_seq > 1 and not qkv_layout.is_thd():
         warnings.warn(
@@ -1733,6 +1702,5 @@ def fused_attn(
         context_checkpoint_name=context_checkpoint_name,
         stripe_size=stripe_size,
         return_max_logit=return_max_logit,
-        return_softmax_aux=return_softmax_aux,
     )
     return output

From 3198726e301580fdbb6d719cbddadc8fb9ecc9f4 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Mon, 20 Jul 2026 10:29:53 -0700
Subject: [PATCH 08/22] Modify static args in fused attn tests for jax

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 6d96604964..e1a46bd018 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -64,7 +64,7 @@ def init():
     yield
 
 
-@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12))
+@partial(jax.jit, static_argnums=(6, 7, 8, 9, 11, 12, 13))
 def general_dot_product_attention(
     query: ArrayLike,
     key: ArrayLike,

From e56b9713993f69106b887641d865937204fc9589 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 22 Jul 2026 01:16:37 +0000
Subject: [PATCH 09/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/jax/test_fused_attn.py                     |  4 +---
 .../jax/cpp_extensions/attention.py              | 16 ++++++++++++----
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index e1a46bd018..d3c800962e 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1348,9 +1348,7 @@ def test_forward_with_max_logit(self):
             primitive_out = self.cp_inverse_reorder_fn(primitive_out)
 
         reference_out = jax_dpa(*self._reference_args(), **kwargs)
-        reference_max_logit = jax_dpa(
-            *self._reference_args(), is_max_logit_enabled=True, **kwargs
-        )
+        reference_max_logit = jax_dpa(*self._reference_args(), is_max_logit_enabled=True, **kwargs)
 
         primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid(
             primitive_out, reference_out, self.pad_q
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 53db542798..487199620c 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -1880,7 +1880,9 @@ def partition(config, mesh, arg_infos, result_infos):
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
         if config.return_max_logit:
-            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
+            raise NotImplementedError(
+                "return_max_logit is not yet supported with context parallelism"
+            )
 
         helper = _FusedAttnCPWithAllGatherHelper(mesh, config)
         helper.check_supported()
@@ -2177,7 +2179,9 @@ def partition(config, mesh, arg_infos, result_infos):
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
         if config.return_max_logit:
-            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
+            raise NotImplementedError(
+                "return_max_logit is not yet supported with context parallelism"
+            )
 
         helper = _FusedAttnCPWithAllGatherHelper(mesh, config)
         helper.check_supported()
@@ -2631,7 +2635,9 @@ def partition(config, mesh, arg_infos, result_infos):
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
         if config.return_max_logit:
-            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
+            raise NotImplementedError(
+                "return_max_logit is not yet supported with context parallelism"
+            )
 
         helper = _FusedAttnCPWithP2PHelper(mesh, config)
         helper.check_supported()
@@ -3141,7 +3147,9 @@ def partition(config, mesh, arg_infos, result_infos):
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
         if config.return_max_logit:
-            raise NotImplementedError("return_max_logit is not yet supported with context parallelism")
+            raise NotImplementedError(
+                "return_max_logit is not yet supported with context parallelism"
+            )
 
         helper = _FusedAttnCPWithP2PHelper(mesh, config)
         helper.check_supported()

From cbe18ac0513c0fe237d024990c4efff4ef8146a8 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Fri, 14 Aug 2026 15:02:54 -0700
Subject: [PATCH 10/22] nit: Inline the choice of what is to be returned and
 remove redundant function for it

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 transformer_engine/jax/attention.py | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index 3eb41ae13f..5cd90140df 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -1335,7 +1335,7 @@ def _fused_attn_fwd_rule(
     softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name)
     rng_state = checkpoint_name(rng_state, context_checkpoint_name)
     max_logit = checkpoint_name(max_logit, context_checkpoint_name)
-    attn_output = _resolve_fused_attn_output(output, max_logit, return_max_logit)
+    attn_output = (output, {"max_logit": max_logit}) if return_max_logit else output
     return attn_output, (
         qkv,
         bias,
@@ -1413,14 +1413,6 @@ def _fused_attn_bwd_rule(
         None,
     )
 
-
-def _resolve_fused_attn_output(output, max_logit, return_max_logit):
-    if not return_max_logit:
-        return output
-
-    return output, {"max_logit": max_logit}
-
-
 _fused_attn.defvjp(_fused_attn_fwd_rule, _fused_attn_bwd_rule)
 
 

From 7f9977ea171f624bd9a3268c22cca382d06c0657 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Wed, 19 Aug 2026 15:28:49 -0700
Subject: [PATCH 11/22] Expose JAX max logit

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn.py                  | 15 ++--
 tests/jax/test_fused_attn_score_mod.py        | 71 +++++++++++++++++++
 .../jax/cpp_extensions/attention.py           | 14 ++--
 transformer_engine/jax/flax/transformer.py    | 38 +++++++++-
 4 files changed, 124 insertions(+), 14 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index d3c800962e..254edc5e75 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1323,7 +1323,7 @@ def _fused_attn_kwargs(self, **overrides):
         kwargs.update(overrides)
         return kwargs
 
-    def test_forward_with_max_logit(self):
+    def test_forward_with_max_logit(self, check_output=True):
         """Test forward output and returned max_logit."""
         self._setup_inputs()
         kwargs = self._fused_attn_kwargs()
@@ -1347,14 +1347,15 @@ def test_forward_with_max_logit(self):
             primitive_max_logit = primitive_aux["max_logit"]
             primitive_out = self.cp_inverse_reorder_fn(primitive_out)
 
-        reference_out = jax_dpa(*self._reference_args(), **kwargs)
         reference_max_logit = jax_dpa(*self._reference_args(), is_max_logit_enabled=True, **kwargs)
 
-        primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid(
-            primitive_out, reference_out, self.pad_q
-        )
-        assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype)
-        assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
+        if check_output:
+            reference_out = jax_dpa(*self._reference_args(), **kwargs)
+            primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid(
+                primitive_out, reference_out, self.pad_q
+            )
+            assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype)
+            assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
         assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype)
 
     def test_backward_with_max_logit(self):
diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py
index b1f165f491..b08393bf24 100644
--- a/tests/jax/test_fused_attn_score_mod.py
+++ b/tests/jax/test_fused_attn_score_mod.py
@@ -426,6 +426,7 @@ def fake_fused_attn(
         score_mod_bprop=None,
         score_mod_tensors=None,
         score_mod_bprop_tensors=None,
+        return_max_logit=False,
     ):
         captured.update(
             qkv=qkv,
@@ -451,7 +452,11 @@ def fake_fused_attn(
             score_mod_bprop=score_mod_bprop,
             score_mod_tensors=score_mod_tensors,
             score_mod_bprop_tensors=score_mod_bprop_tensors,
+            return_max_logit=return_max_logit,
         )
+        if return_max_logit:
+            max_logit = jnp.arange(qkv[0].shape[-2], dtype=qkv[0].dtype)
+            return qkv[0], {"max_logit": max_logit}
         return qkv[0]
 
     monkeypatch.setattr(
@@ -536,6 +541,72 @@ def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch):
     assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD
 
 
+def test_dot_product_attention_plumbs_return_max_logit_to_fused_attn(monkeypatch):
+    """DotProductAttention forwards return_max_logit to fused_attn and returns aux data."""
+    captured = _install_fake_flax_fused_attn(monkeypatch)
+    query = jnp.ones((1, 8, 2, 16), dtype=jnp.float16)
+    key = jnp.ones((1, 8, 2, 16), dtype=jnp.float16)
+    value = jnp.ones((1, 8, 2, 16), dtype=jnp.float16)
+
+    dpa = flax_transformer.DotProductAttention(
+        head_dim=16,
+        num_attention_heads=2,
+        num_gqa_groups=2,
+        attn_mask_type="no_mask",
+        qkv_layout="bshd_bshd_bshd",
+        transpose_batch_sequence=False,
+        return_max_logit=True,
+    )
+    out, aux = dpa.apply({}, query, key, value, deterministic=True)
+
+    np.testing.assert_array_equal(out, query)
+    np.testing.assert_array_equal(aux["max_logit"], jnp.arange(2, dtype=query.dtype))
+    assert captured["return_max_logit"] is True
+    assert captured["kernel_checks"][0][1]["return_max_logit"] is True
+
+
+def test_dot_product_attention_return_max_logit_requires_fused_attn_enabled(monkeypatch):
+    """DotProductAttention rejects return_max_logit when the fused backend is disabled."""
+    monkeypatch.setenv("NVTE_FUSED_ATTN", "0")
+    query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
+    key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
+    value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
+
+    dpa = flax_transformer.DotProductAttention(
+        head_dim=16,
+        num_attention_heads=1,
+        num_gqa_groups=1,
+        attn_mask_type="no_mask",
+        qkv_layout="bshd_bshd_bshd",
+        transpose_batch_sequence=False,
+        return_max_logit=True,
+    )
+
+    with pytest.raises(ValueError, match="NVTE_FUSED_ATTN=0"):
+        dpa.apply({}, query, key, value, deterministic=True)
+
+
+def test_dot_product_attention_return_max_logit_requires_available_fused_kernel(monkeypatch):
+    """DotProductAttention rejects return_max_logit instead of falling back to unfused attention."""
+    _install_fake_flax_fused_attn(monkeypatch, kernel_available=False)
+    query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
+    key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
+    value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
+
+    dpa = flax_transformer.DotProductAttention(
+        head_dim=16,
+        num_attention_heads=1,
+        num_gqa_groups=1,
+        attn_mask_type="no_mask",
+        qkv_layout="bshd_bshd_bshd",
+        transpose_batch_sequence=False,
+        return_max_logit=True,
+    )
+
+    with pytest.raises(ValueError, match="requires a cuDNN fused attention kernel"):
+        dpa.apply({}, query, key, value, deterministic=True)
+
+
 def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch):
     """Packed QKV inputs are unpacked because score_mod requires separate Q/K/V."""
     captured = _install_fake_flax_fused_attn(monkeypatch)
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 487199620c..41412ab669 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -667,6 +667,7 @@ def convert_to_2d(offsets, batch, max_seqlen):
             _kv_segment_pos,
             config=config,
         )
+        # Reduce cuDNN's raw Max tensor to TE's public per-head [H] max_logit.
         max_logit = FusedAttnFwdPrimitive._reduce_max_logit(
             max_tensor, output, raw_q_seqlen, raw_q_seq_offsets, config
         )
@@ -674,19 +675,24 @@ def convert_to_2d(offsets, batch, max_seqlen):
 
     @staticmethod
     def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config):
-        """Reduce cuDNN's raw Max tensor to PyTorch-compatible per-head max_logit."""
+        """Reduce cuDNN's raw Max tensor to framework-compatible per-head max_logit."""
         if not config.return_max_logit:
             return jnp.zeros((0,), dtype=output.dtype)
 
         uses_thd_ragged_max_tensor = FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config)
         if config.qkv_layout.is_thd() and max_tensor.ndim == 4:
+            # Dense BSHD Max rows are expected to be masked by cuDNN before TE reduces them.
+            # THD Max can include static holes/unwritten rows, so mask valid query rows here.
             q_seqlen = jnp.where(q_seqlen > 0, q_seqlen, 0)
             q_seq_offsets = jnp.where(q_seq_offsets >= 0, q_seq_offsets, -1)
+            num_segments = min(q_seqlen.shape[-1], q_seq_offsets.shape[-1])
+            q_seqlen = q_seqlen[..., :num_segments]
+            q_seq_offsets = q_seq_offsets[..., :num_segments]
             token_idx = jnp.arange(output.shape[-3], dtype=q_seq_offsets.dtype)
             valid = jnp.any(
-                (q_seq_offsets[..., :-1, None] >= 0)
-                & (token_idx >= q_seq_offsets[..., :-1, None])
-                & (token_idx < (q_seq_offsets[..., :-1, None] + q_seqlen[..., None])),
+                (q_seq_offsets[..., None] >= 0)
+                & (token_idx >= q_seq_offsets[..., None])
+                & (token_idx < (q_seq_offsets[..., None] + q_seqlen[..., None])),
                 axis=-2,
             )
             if uses_thd_ragged_max_tensor:
diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py
index 76922d2b55..ce3dab4818 100644
--- a/transformer_engine/jax/flax/transformer.py
+++ b/transformer_engine/jax/flax/transformer.py
@@ -308,6 +308,7 @@ class _FusedDotProductAttention(nn.Module):  # pylint: disable=too-few-public-me
     score_mod: Optional[Callable] = None
     score_mod_bprop: Optional[Callable] = None
     score_mod_requested: bool = False
+    return_max_logit: bool = False
 
     @nn.compact
     def __call__(
@@ -363,6 +364,7 @@ def __call__(
             "score_mod_bprop": self.score_mod_bprop,
             "score_mod_tensors": score_mod_tensors,
             "score_mod_bprop_tensors": score_mod_bprop_tensors,
+            "return_max_logit": self.return_max_logit,
         }
 
         if self.qkv_layout.is_qkvpacked():
@@ -434,12 +436,17 @@ def __call__(
         else:
             raise ValueError(f"Unsupported {self.qkv_layout=}.")
 
+        if self.return_max_logit:
+            x, aux = x
+
         if self.transpose_batch_sequence:
             x = x.transpose([1, 0, 2, 3])
 
         assert (
             x.dtype == query.dtype
         ), f"output dtype {x.dtype} does not match query dtype {query.dtype}"
+        if self.return_max_logit:
+            return x, aux
         return x
 
 
@@ -619,6 +626,9 @@ class DotProductAttention(nn.Module):  # pylint: disable=too-few-public-methods
         argument to keep tensor operands as normal JAX inputs.
     score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None
         Additional tensors or pass-by-value scalars for ``score_mod_bprop``.
+    return_max_logit: bool, default = False
+        If True, return ``(output, aux)`` where ``aux["max_logit"]`` contains the per-head
+        maximum attention logits with shape ``[h]``. This path requires fused attention.
 
     Optimization parameters
     -----------------------
@@ -647,6 +657,7 @@ class DotProductAttention(nn.Module):  # pylint: disable=too-few-public-methods
     softmax_type: str = "vanilla"
     score_mod: Optional[Callable] = None
     score_mod_bprop: Optional[Callable] = None
+    return_max_logit: bool = False
 
     def __post_init__(self):
         # TODO(KshitijLakhani): Remove warning in TransformerEngine v2.12
@@ -717,8 +728,8 @@ def __call__(
 
         Returns
         -------
-        outputs: jax.numpy.ndarray
-            Output tensors.
+        outputs: jax.numpy.ndarray or tuple[jax.numpy.ndarray, dict[str, jax.numpy.ndarray]]
+            Output tensor, or ``(output, aux)`` when ``return_max_logit`` is enabled.
         """
         input_dtype = query.dtype
 
@@ -777,6 +788,12 @@ def __call__(
 
         # Use fused attn (if kernel check below passes) by default
         enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1"))
+        if self.return_max_logit and not enable_fused_attn:
+            raise ValueError(
+                "DotProductAttention(return_max_logit=True) requires fused attention, but "
+                "NVTE_FUSED_ATTN=0 disables it. Set NVTE_FUSED_ATTN=1 or unset the variable, "
+                "then ensure a cuDNN fused attention kernel is available for this configuration."
+            )
 
         sequence_dim = 0 if self.transpose_batch_sequence else 1
         seqlen_q = query.shape[sequence_dim]
@@ -815,11 +832,24 @@ def __call__(
             head_dim_qk,
             head_dim_v,
             self.window_size,
+            return_max_logit=self.return_max_logit,
         )
         if score_mod_requested and not has_fused_attn_kernel:
             raise ValueError(
                 "score_mod requires fused attention, but no fused attention kernel is available."
             )
+        if self.return_max_logit and not has_fused_attn_kernel:
+            raise ValueError(
+                "DotProductAttention(return_max_logit=True) requires a cuDNN fused attention "
+                "kernel, but no compatible kernel is available for this configuration. "
+                "Set NVTE_FUSED_ATTN=1 or unset it, update cuDNN/Transformer Engine if needed, "
+                "and check the configuration: "
+                f"{qkv_layout=}, {attn_bias_type=}, {attn_mask_type=}, "
+                f"{softmax_type=}, attention_dropout={self.attention_dropout}, "
+                f"num_attention_heads={self.num_attention_heads}, "
+                f"num_gqa_groups={self.num_gqa_groups}, {seqlen_q=}, {seqlen_kv=}, "
+                f"{head_dim_qk=}, {head_dim_v=}, window_size={self.window_size}."
+            )
 
         use_fused_attn = enable_fused_attn and has_fused_attn_kernel
 
@@ -916,6 +946,7 @@ def __call__(
                 score_mod=self.score_mod,
                 score_mod_bprop=self.score_mod_bprop,
                 score_mod_requested=score_mod_requested,
+                return_max_logit=self.return_max_logit,
             )(
                 query,
                 key,
@@ -927,7 +958,8 @@ def __call__(
                 score_mod_tensors=score_mod_tensors,
                 score_mod_bprop_tensors=score_mod_bprop_tensors,
             )
-        assert x.dtype == input_dtype, f"output_dtype={x.dtype}, input_dtype={input_dtype}"
+        output = x[0] if self.return_max_logit else x
+        assert output.dtype == input_dtype, f"output_dtype={output.dtype}, input_dtype={input_dtype}"
         return x
 
 

From de929bad9e02ea823591eadbf2c6ecd99ad491b1 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Wed, 19 Aug 2026 15:29:16 -0700
Subject: [PATCH 12/22] Support CP max logit

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_distributed_fused_attn.py      | 127 ++++++++-
 .../jax/cpp_extensions/attention.py           | 243 +++++++++++-------
 2 files changed, 281 insertions(+), 89 deletions(-)

diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
index b6e11b8bea..b508120362 100644
--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -399,6 +399,75 @@ def test_softcap_score_mod_with_aux_params_backward(
     ),
 ]
 
+DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CASES = [
+    pytest.param(
+        QKVLayout.BSHD_BSHD_BSHD,
+        AttnMaskType.CAUSAL_MASK,
+        CPStrategy.ALL_GATHER,
+        (-1, -1),
+        None,
+        None,
+        False,
+        True,
+        id="AG-BSHD",
+    ),
+    pytest.param(
+        QKVLayout.THD_THD_THD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        CPStrategy.ALL_GATHER,
+        (-1, -1),
+        64,
+        5,
+        False,
+        True,
+        id="AG-THD",
+    ),
+    pytest.param(
+        QKVLayout.BSHD_BSHD_BSHD,
+        AttnMaskType.CAUSAL_MASK,
+        CPStrategy.RING,
+        (-1, -1),
+        None,
+        None,
+        False,
+        True,
+        id="RING-BSHD-NO_SCAN",
+    ),
+    pytest.param(
+        QKVLayout.BSHD_BSHD_BSHD,
+        AttnMaskType.CAUSAL_MASK,
+        CPStrategy.RING,
+        (-1, -1),
+        None,
+        None,
+        True,
+        True,
+        id="RING-BSHD-SCAN",
+    ),
+    pytest.param(
+        QKVLayout.THD_THD_THD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        CPStrategy.RING,
+        (-1, -1),
+        1,
+        5,
+        False,
+        False,
+        id="RING-THD-NO_SCAN",
+    ),
+    pytest.param(
+        QKVLayout.THD_THD_THD,
+        AttnMaskType.PADDING_CAUSAL_MASK,
+        CPStrategy.RING,
+        (-1, -1),
+        1,
+        5,
+        True,
+        False,
+        id="RING-THD-SCAN",
+    ),
+]
+
 
 class TestDistributedContextParallelSelfAttn:
     # TODO(KshitijLakhani): parametrize num_segments_per_seq for all CP tests
@@ -419,6 +488,8 @@ def impl_test_context_parallel_attn(
         window_size=None,
         stripe_size=None,
         num_segments_per_seq=None,
+        return_max_logit=False,
+        check_forward_output=True,
     ):
         if qkv_layout.is_thd():
             if not load_balanced and (
@@ -513,9 +584,63 @@ def check_has_backend_for_mask(mask_type):
         if num_head % kv_groups != 0 or (num_head // kv_groups) % tp_size != 0:
             pytest.skip(f"Skipping {kv_groups=} not multiple of {data_shape=} or {tp_size=}")
 
-        runner.test_backward()
+        if return_max_logit:
+            runner.test_forward_with_max_logit(check_output=check_forward_output)
+        else:
+            runner.test_backward()
         del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"]
 
+    @pytest_parametrize_wrapper(
+        "device_count,mesh_shape,mesh_axes,mesh_resource",
+        generate_context_parallel_configs_for_attn(),
+    )
+    @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1])
+    @pytest.mark.parametrize("kv_groups", [1])
+    @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")])
+    @pytest.mark.parametrize(
+        "qkv_layout, attn_mask_type, cp_strategy, window_size, stripe_size,"
+        " num_segments_per_seq, use_scan_ring, check_forward_output",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CASES,
+    )
+    def test_context_parallel_return_max_logit(
+        self,
+        device_count,
+        mesh_shape,
+        mesh_axes,
+        mesh_resource,
+        data_shape,
+        kv_groups,
+        dtype,
+        qkv_layout,
+        attn_mask_type,
+        cp_strategy,
+        window_size,
+        stripe_size,
+        num_segments_per_seq,
+        use_scan_ring,
+        check_forward_output,
+    ):
+        """Check CP fused attention returns global per-head max_logit."""
+        self.impl_test_context_parallel_attn(
+            device_count,
+            mesh_shape,
+            mesh_axes,
+            mesh_resource,
+            data_shape,
+            kv_groups,
+            attn_mask_type,
+            dtype,
+            qkv_layout,
+            True,
+            cp_strategy,
+            use_scan_ring=use_scan_ring,
+            window_size=window_size,
+            stripe_size=stripe_size,
+            num_segments_per_seq=num_segments_per_seq,
+            return_max_logit=True,
+            check_forward_output=check_forward_output,
+        )
+
     @pytest_parametrize_wrapper(
         "device_count,mesh_shape,mesh_axes,mesh_resource",
         generate_context_parallel_configs_for_attn(),
diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 41412ab669..45ce08cb73 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -717,6 +717,36 @@ def _uses_thd_ragged_max_tensor(config):
             and 120 not in get_all_device_compute_capability()
         )
 
+    @staticmethod
+    def _empty_or_neg_inf_max_logit(head, dtype, config):
+        """Return the neutral value for per-head max_logit accumulation."""
+        if config.return_max_logit:
+            return jnp.full((head,), -jnp.inf, dtype=dtype)
+        return jnp.zeros((0,), dtype=dtype)
+
+    @staticmethod
+    def _max_logit_reduce_axes(mesh, max_logit_sharding):
+        """Return mesh axes to reduce while preserving max_logit's head sharding."""
+        # max_logit is [H], so axes that shard H (typically TP) are preserved.
+        # Axes for collapsed dimensions such as batch/sequence (DP/CP) must pmax.
+        head_axes = set()
+        for axis in max_logit_sharding.spec:
+            if axis is None:
+                continue
+            if isinstance(axis, tuple):
+                head_axes.update(axis)
+            else:
+                head_axes.add(axis)
+        return tuple(axis for axis in mesh.axis_names if axis not in head_axes)
+
+    @staticmethod
+    def _reduce_max_logit_across_mesh(max_logit, mesh, reduce_axes, config):
+        """Reduce max_logit across mesh axes absent from the [H] result."""
+        if config.return_max_logit:
+            for axis in reduce_axes:
+                max_logit = lax_paral_op(max_logit, lax.pmax, axis, mesh=mesh)
+        return max_logit
+
     @staticmethod
     def batcher(batched_args, batch_dims, *, config):
         # batch_dims: each element is the batch axis (0, ...) or None. Only 0 or None allowed.
@@ -1501,6 +1531,7 @@ def get_step_config(self) -> _FusedAttnConfig:
             cp_axis=self.config.cp_axis,
             cp_striped_window_size=None,
             stripe_size=self.config.stripe_size,
+            return_max_logit=self.config.return_max_logit,
         )
 
     def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig:
@@ -1521,6 +1552,7 @@ def get_step_config_for_striped(self, max_seqlen, cp_size) -> _FusedAttnConfig:
             cp_axis=self.config.cp_axis,
             cp_striped_window_size=None,
             stripe_size=self.config.stripe_size,
+            return_max_logit=self.config.return_max_logit,
         )
 
     def all_gather_kv(self, k, v):
@@ -1885,10 +1917,6 @@ def partition(config, mesh, arg_infos, result_infos):
         ), "Sliding window attention is not supported when context parallelism is enabled"
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
-        if config.return_max_logit:
-            raise NotImplementedError(
-                "return_max_logit is not yet supported with context parallelism"
-            )
 
         helper = _FusedAttnCPWithAllGatherHelper(mesh, config)
         helper.check_supported()
@@ -1903,6 +1931,9 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[5] = seed_sharding
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
+        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
+            mesh, max_logit_sharding
+        )
 
         def impl(
             q,
@@ -1950,7 +1981,8 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
                     q_seqlen_for_step = q_seqlen / (cp_size * 2)
                     num_kv_chunks = kv_max_seqlen // kv_seqlens_for_rank[sub_idx]
                     kv_seqlen_for_step = (kv_seqlen / (cp_size * 2)) * num_kv_chunks
-                    output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.impl(
+                    # max_logit returned here is already reduced to shape [H]
+                    output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl(
                         q_split[sub_idx],
                         k_unmasked,
                         v_unmasked,
@@ -1967,12 +1999,13 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
                         _kv_segment_pos,
                         config=helper.get_step_config(),
                     )
-                    results.append((output, softmax_aux, rng_state))
+                    results.append((output, softmax_aux, rng_state, max_logit))
 
                 output = jnp.concatenate((results[0][0], results[1][0]), axis=1)
                 softmax_aux = jnp.concatenate((results[0][1], results[1][1]), axis=2)
                 rng_state = results[1][2]  # Use the final RNG state
-                max_logit = jnp.zeros((0,), dtype=output.dtype)
+                # Rank-local [H] max across both local dual-chunk query pieces.
+                max_logit = jnp.maximum(results[0][3], results[1][3])
 
                 return output, softmax_aux, rng_state, max_logit
 
@@ -1985,7 +2018,12 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
                 for idx in range(cp_size)
             ]
 
-            return lax.switch(cp_rank, functions)
+            output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions)
+            # Reduce over non-head mesh axes to make [H] global over batch/sequence.
+            max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
+                max_logit, mesh, max_logit_reduce_axes, config
+            )
+            return output, softmax_aux, rng_state, max_logit
 
         return mesh, impl, out_shardings, arg_shardings
 
@@ -2184,10 +2222,6 @@ def partition(config, mesh, arg_infos, result_infos):
         is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
-        if config.return_max_logit:
-            raise NotImplementedError(
-                "return_max_logit is not yet supported with context parallelism"
-            )
 
         helper = _FusedAttnCPWithAllGatherHelper(mesh, config)
         helper.check_supported()
@@ -2202,6 +2236,9 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[5] = seed_sharding
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
+        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
+            mesh, max_logit_sharding
+        )
 
         def impl(
             q,
@@ -2265,7 +2302,7 @@ def _cross_attn(
                     max_segments_per_seq=adjusted_max_segments_per_seq,
                 )
 
-                output, softmax_aux, rng_state, _ = FusedAttnFwdPrimitive.impl(
+                output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl(
                     q,  # sharded for rank
                     k,  # ag
                     v,  # ag
@@ -2284,7 +2321,6 @@ def _cross_attn(
                         max_seqlen=kv_max_seqlen, cp_size=cp_size
                     ),
                 )
-                max_logit = jnp.zeros((0,), dtype=output.dtype)
                 return output, softmax_aux, rng_state, max_logit
 
             # AG the k, v, kv_segment_ids and kv_segment_pos
@@ -2306,7 +2342,11 @@ def _cross_attn(
                 )
                 for _ in range(cp_size)
             ]
-            return lax.switch(cp_rank, functions)
+            output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions)
+            max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
+                max_logit, mesh, max_logit_reduce_axes, config
+            )
+            return output, softmax_aux, rng_state, max_logit
 
         return mesh, impl, out_shardings, arg_shardings
 
@@ -2579,6 +2619,7 @@ def get_step_config(self, attn_mask_type) -> _FusedAttnConfig:
             cp_axis=self.config.cp_axis,
             cp_striped_window_size=None,
             stripe_size=self.config.stripe_size,
+            return_max_logit=self.config.return_max_logit,
         )
 
     def stack_kv(self, k, v):
@@ -2640,10 +2681,6 @@ def partition(config, mesh, arg_infos, result_infos):
         ), "Sliding window attention is not supported when context parallelism is enabled"
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
-        if config.return_max_logit:
-            raise NotImplementedError(
-                "return_max_logit is not yet supported with context parallelism"
-            )
 
         helper = _FusedAttnCPWithP2PHelper(mesh, config)
         helper.check_supported()
@@ -2661,6 +2698,9 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
+        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
+            mesh, max_logit_sharding
+        )
 
         def ring_attn_fwd_impl(
             q,
@@ -2698,9 +2738,10 @@ def ring_attn_fwd_impl(
             # support dropout currently.
             rng_state_shape = (seed.shape[0], *result_infos[2].shape[1:])
             rng_state = jnp.zeros(rng_state_shape).astype(result_infos[2].dtype)
+            max_logit = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(head, q.dtype, config)
 
             def scan_kv_block(idx, carry):
-                kv, output, softmax_aux = carry
+                kv, output, softmax_aux, max_logit = carry
 
                 # Send KV block to next step so we can overlap compute.
                 kv_next = helper.permute_kv(kv, cp_perm)
@@ -2708,24 +2749,26 @@ def scan_kv_block(idx, carry):
                 def mask_compute(attn_mask_type):
                     q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx)
                     kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx)
-                    output_per_step, softmax_aux_per_step, _, _ = FusedAttnFwdPrimitive.impl(
-                        q,
-                        kv,
-                        _not_used,
-                        bias,
-                        _softmax_offset,
-                        seed,
-                        q_seqlen_per_step,
-                        kv_seqlen_per_step,
-                        q_seq_offsets,
-                        k_seq_offsets,
-                        _q_segment_ids,
-                        _kv_segment_ids,
-                        _q_segment_pos,
-                        _kv_segment_pos,
-                        config=helper.get_step_config(attn_mask_type),
+                    output_per_step, softmax_aux_per_step, _, max_logit_per_step = (
+                        FusedAttnFwdPrimitive.impl(
+                            q,
+                            kv,
+                            _not_used,
+                            bias,
+                            _softmax_offset,
+                            seed,
+                            q_seqlen_per_step,
+                            kv_seqlen_per_step,
+                            q_seq_offsets,
+                            k_seq_offsets,
+                            _q_segment_ids,
+                            _kv_segment_ids,
+                            _q_segment_pos,
+                            _kv_segment_pos,
+                            config=helper.get_step_config(attn_mask_type),
+                        )
                     )
-                    return output_per_step, softmax_aux_per_step
+                    return output_per_step, softmax_aux_per_step, max_logit_per_step
 
                 causal_mask_compute = partial(mask_compute, AttnMaskType.CAUSAL_MASK)
                 no_mask_compute = partial(mask_compute, AttnMaskType.NO_MASK)
@@ -2734,45 +2777,49 @@ def half_kv_no_mask_compute():
                     q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx)
                     kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx) // 2
                     kv_part = lax.slice_in_dim(kv, 0, kv.shape[1] // 2, axis=1)
-                    output_per_step, softmax_aux_per_step, _, _ = FusedAttnFwdPrimitive.impl(
-                        q,
-                        kv_part,
-                        _not_used,
-                        bias,
-                        _softmax_offset,
-                        seed,
-                        q_seqlen_per_step,
-                        kv_seqlen_per_step,
-                        q_seq_offsets,
-                        k_seq_offsets,
-                        _q_segment_ids,
-                        _kv_segment_ids,
-                        _q_segment_pos,
-                        _kv_segment_pos,
-                        config=helper.get_step_config(AttnMaskType.NO_MASK),
+                    output_per_step, softmax_aux_per_step, _, max_logit_per_step = (
+                        FusedAttnFwdPrimitive.impl(
+                            q,
+                            kv_part,
+                            _not_used,
+                            bias,
+                            _softmax_offset,
+                            seed,
+                            q_seqlen_per_step,
+                            kv_seqlen_per_step,
+                            q_seq_offsets,
+                            k_seq_offsets,
+                            _q_segment_ids,
+                            _kv_segment_ids,
+                            _q_segment_pos,
+                            _kv_segment_pos,
+                            config=helper.get_step_config(AttnMaskType.NO_MASK),
+                        )
                     )
-                    return output_per_step, softmax_aux_per_step
+                    return output_per_step, softmax_aux_per_step, max_logit_per_step
 
                 def half_q_no_mask_compute():
                     q_seqlen_per_step = helper.adjust_seqlen(q_seqlen, q_max_seqlen, idx) // 2
                     kv_seqlen_per_step = helper.adjust_seqlen(kv_seqlen, kv_max_seqlen, idx)
                     q_part = lax.slice_in_dim(q, q_max_seqlen // 2, q_max_seqlen, axis=1)
-                    output_per_step, softmax_aux_per_step, _, _ = FusedAttnFwdPrimitive.impl(
-                        q_part,
-                        kv,
-                        _not_used,
-                        bias,
-                        _softmax_offset,
-                        seed,
-                        q_seqlen_per_step,
-                        kv_seqlen_per_step,
-                        q_seq_offsets,
-                        k_seq_offsets,
-                        _q_segment_ids,
-                        _kv_segment_ids,
-                        _q_segment_pos,
-                        _kv_segment_pos,
-                        config=helper.get_step_config(AttnMaskType.NO_MASK),
+                    output_per_step, softmax_aux_per_step, _, max_logit_per_step = (
+                        FusedAttnFwdPrimitive.impl(
+                            q_part,
+                            kv,
+                            _not_used,
+                            bias,
+                            _softmax_offset,
+                            seed,
+                            q_seqlen_per_step,
+                            kv_seqlen_per_step,
+                            q_seq_offsets,
+                            k_seq_offsets,
+                            _q_segment_ids,
+                            _kv_segment_ids,
+                            _q_segment_pos,
+                            _kv_segment_pos,
+                            config=helper.get_step_config(AttnMaskType.NO_MASK),
+                        )
                     )
                     output_per_step = jnp.concat([jnp.zeros_like(q_part), output_per_step], axis=1)
                     softmax_aux_per_step = jnp.concat(
@@ -2782,14 +2829,17 @@ def half_q_no_mask_compute():
                         ],
                         axis=2,
                     )
-                    return output_per_step, softmax_aux_per_step
+                    return output_per_step, softmax_aux_per_step, max_logit_per_step
 
                 def skip_compute():
                     output_per_step = jnp.zeros_like(q)
                     softmax_aux_per_step = jnp.full(
                         (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(
+                        head, q.dtype, config
+                    )
+                    return output_per_step, softmax_aux_per_step, max_logit_per_step
 
                 if config.attn_mask_type == AttnMaskType.CAUSAL_MASK:
                     # This is for nested jax.lax.cond
@@ -2800,11 +2850,11 @@ def jax_cond_wrap():
                             )
                         return lax.cond((idx <= cp_rank), no_mask_compute, skip_compute)
 
-                    output_per_step, softmax_aux_per_step = lax.cond(
+                    output_per_step, softmax_aux_per_step, max_logit_per_step = lax.cond(
                         idx == 0, causal_mask_compute, jax_cond_wrap
                     )
                 else:
-                    output_per_step, softmax_aux_per_step = no_mask_compute()
+                    output_per_step, softmax_aux_per_step, max_logit_per_step = no_mask_compute()
 
                 def skip_correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
                     # No correction done here but we cast outputs to float32 and perform reduction
@@ -2827,19 +2877,23 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
                     output_per_step,
                     softmax_aux_per_step,
                 )
+                # Running per-head max over all ring steps for this rank.
+                max_logit = jnp.maximum(max_logit, max_logit_per_step)
 
-                return (kv_next, output, softmax_aux)
+                return (kv_next, output, softmax_aux, max_logit)
 
-            carry = (kv, output, softmax_aux)
+            carry = (kv, output, softmax_aux, max_logit)
             if helper.use_scanloop():
                 carry = lax.fori_loop(0, cp_size, scan_kv_block, carry)
             else:
                 for i in range(0, cp_size):
                     carry = scan_kv_block(i, carry)
-            (kv, output, softmax_aux) = carry
+            (kv, output, softmax_aux, max_logit) = carry
 
             output = output.astype(q.dtype)
-            max_logit = jnp.zeros((0,), dtype=output.dtype)
+            max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
+                max_logit, mesh, max_logit_reduce_axes, config
+            )
             return output, softmax_aux, rng_state, max_logit
 
         return mesh, ring_attn_fwd_impl, out_shardings, arg_shardings
@@ -3152,10 +3206,6 @@ def partition(config, mesh, arg_infos, result_infos):
         is_context_parallel = get_mesh_axis_size(config.cp_axis, mesh) > 1
         if not is_context_parallel:
             return FusedAttnFwdPrimitive.partition(config, mesh, arg_infos, result_infos)
-        if config.return_max_logit:
-            raise NotImplementedError(
-                "return_max_logit is not yet supported with context parallelism"
-            )
 
         helper = _FusedAttnCPWithP2PHelper(mesh, config)
         helper.check_supported()
@@ -3173,6 +3223,9 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
+        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
+            mesh, max_logit_sharding
+        )
 
         def fwd_impl(
             q,
@@ -3215,9 +3268,10 @@ def fwd_impl(
             # support dropout currently.
             rng_state_shape = (seed.shape[0], *result_infos[2].shape[1:])
             rng_state = jnp.zeros(rng_state_shape).astype(result_infos[2].dtype)
+            max_logit = FusedAttnFwdPrimitive._empty_or_neg_inf_max_logit(head, q.dtype, config)
 
             def scan_kv_block(idx, carry):
-                kv, kv_segment_ids, kv_segment_pos, output, softmax_aux = carry
+                kv, kv_segment_ids, kv_segment_pos, output, softmax_aux, max_logit = carry
 
                 # TODO(rewang): To check whether we need special handle for the last idx
                 # Send KV block to next step so we can overlap compute.
@@ -3255,7 +3309,9 @@ def compute(config):
                     )
                 else:
                     current_config = subblock_config
-                output_per_step, softmax_aux_per_step, _, _ = compute(current_config)
+                output_per_step, softmax_aux_per_step, _, max_logit_per_step = compute(
+                    current_config
+                )
 
                 softmax_aux_per_step = softmax_aux_per_step.reshape((batch, q_max_seqlen, head, 1))
 
@@ -3281,19 +3337,30 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
                     output_per_step,
                     softmax_aux_per_step,
                 )
+                # Running per-head max over all ring steps for this rank.
+                max_logit = jnp.maximum(max_logit, max_logit_per_step)
 
-                return (kv_next, kv_segment_ids_next, kv_segment_pos_next, output, softmax_aux)
+                return (
+                    kv_next,
+                    kv_segment_ids_next,
+                    kv_segment_pos_next,
+                    output,
+                    softmax_aux,
+                    max_logit,
+                )
 
-            carry = (kv, kv_segment_ids, kv_segment_pos, output, softmax_aux)
+            carry = (kv, kv_segment_ids, kv_segment_pos, output, softmax_aux, max_logit)
             if helper.use_scanloop():
                 carry = lax.fori_loop(0, cp_size, scan_kv_block, carry)
             else:
                 for i in range(0, cp_size):
                     carry = scan_kv_block(i, carry)
-            (_, _, _, output, softmax_aux) = carry
+            (_, _, _, output, softmax_aux, max_logit) = carry
 
             output = output.astype(q.dtype)
-            max_logit = jnp.zeros((0,), dtype=output.dtype)
+            max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
+                max_logit, mesh, max_logit_reduce_axes, config
+            )
             return output, softmax_aux, rng_state, max_logit
 
         return mesh, fwd_impl, out_shardings, arg_shardings

From 78c8828ea5f673ec5bc21a81b6f9bd9389690eb4 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Thu, 20 Aug 2026 18:19:14 -0700
Subject: [PATCH 13/22] Expand JAX max logit tests

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn.py | 308 +++++++++++++++++------------------
 1 file changed, 151 insertions(+), 157 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 254edc5e75..19aea7f58b 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -949,7 +949,7 @@ def to_dp_shardings(x):
         self.seq_length_offset_pspec = PartitionSpec(self.mesh_resource.dp_resource, None)
         self.seq_length_offset_sharding = NamedSharding(self.mesh, self.seq_length_offset_pspec)
 
-    def test_forward(self):
+    def test_forward(self, return_max_logit=False, check_output=True):
         """
         Test forward with JITted primitive and unJITted reference
         """
@@ -997,6 +997,7 @@ def test_forward(self):
             "score_mod_bprop": self.score_mod_bprop,
             "score_mod_tensors": self.score_mod_tensors,
             "score_mod_bprop_tensors": self.score_mod_bprop_tensors,
+            "return_max_logit": return_max_logit,
         }
         reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference}
 
@@ -1016,31 +1017,40 @@ def test_forward(self):
 
         with self.mesh, autocast(mesh_resource=self.mesh_resource):
             primitive_out = customcall_fused_dpa_jit(*customcall_args)
+            if return_max_logit:
+                primitive_out, primitive_aux = primitive_out
+                primitive_max_logit = primitive_aux["max_logit"]
             primitive_out = self.cp_inverse_reorder_fn(primitive_out)
 
-        reference_out = jax_dpa(*args, **reference_kwargs)
+        if return_max_logit:
+            reference_max_logit = jax_dpa(
+                *args, is_max_logit_enabled=True, **reference_kwargs
+            )
 
-        if self.is_training and self.dropout_prob > 0.0:
-            return
+        if check_output and not (self.is_training and self.dropout_prob > 0.0):
+            reference_out = jax_dpa(*args, **reference_kwargs)
 
-        primitive_valid, primitive_invalid, reference_valid, reference_invalid = (
-            _split_valid_and_invalid(primitive_out, reference_out, self.pad_q)
-        )
+            primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid(
+                primitive_out, reference_out, self.pad_q
+            )
 
-        assert_allclose(
-            primitive_invalid,
-            jnp.zeros_like(primitive_invalid),
-            rtol=self.rtol,
-            atol=self.atol,
-            dtype=self.dtype,
-        )
-        assert_allclose(
-            primitive_valid,
-            reference_valid,
-            rtol=self.rtol,
-            atol=self.atol,
-            dtype=self.dtype,
-        )
+            assert_allclose(
+                primitive_invalid,
+                jnp.zeros_like(primitive_invalid),
+                rtol=self.rtol,
+                atol=self.atol,
+                dtype=self.dtype,
+            )
+            assert_allclose(
+                primitive_valid,
+                reference_valid,
+                rtol=self.rtol,
+                atol=self.atol,
+                dtype=self.dtype,
+            )
+
+        if return_max_logit:
+            assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype)
 
         if self.coll_count_ref is not None:
             with self.mesh, autocast(mesh_resource=self.mesh_resource):
@@ -1049,7 +1059,7 @@ def test_forward(self):
                 )
             assert_equal_collectives(target_hlo, self.coll_count_ref)
 
-    def test_backward(self):
+    def test_backward(self, return_max_logit=False):
         """
         Test value_and_grad with JIT, which includes both forward and backward.
 
@@ -1077,6 +1087,8 @@ def grad_func(
             if self.attn_mask_type.is_causal():
                 gradient_multiplier /= 10
             output = func(q, k, v, bias, softmax_offset, sequence_descriptor, dropout_rng, **kwargs)
+            if isinstance(output, tuple):
+                output, _ = output
             if cp_reverse_out:
                 output = self.cp_inverse_reorder_fn(output)
             # Keep only valid result for the gradient
@@ -1142,6 +1154,7 @@ def grad_func(
             "score_mod_bprop": self.score_mod_bprop,
             "score_mod_tensors": self.score_mod_tensors,
             "score_mod_bprop_tensors": self.score_mod_bprop_tensors,
+            "return_max_logit": return_max_logit,
         }
         reference_kwargs = {**kwargs, "score_mod_reference": self.score_mod_reference}
 
@@ -1291,149 +1304,130 @@ def check_dqkv(primitive, reference, pad, idx):
                 target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text()
             assert_equal_collectives(target_hlo, self.coll_count_ref)
 
-    def _reference_args(self):
-        return [self.q, self.k, self.v, self.bias, self.softmax_offset, self.mask, self.dropout_rng]
-
-    def _customcall_args(self):
-        return [
-            jax.device_put(self.cp_reorder_fn(self.q), self.qkvo_sharding),
-            jax.device_put(self.cp_reorder_fn(self.k), self.qkvo_sharding),
-            jax.device_put(self.cp_reorder_fn(self.v), self.qkvo_sharding),
-            jax.device_put(self.bias, self.bias_sharding),
-            jax.device_put(self.softmax_offset, self.softmax_offset_sharding),
-            jax.device_put(self.sequence_desciptor, self.seq_desc_sharding),
-            jax.device_put(self.dropout_rng, self.dropout_rng_sharding),
-        ]
-
-    def _fused_attn_kwargs(self, **overrides):
-        kwargs = {
-            "attn_bias_type": self.attn_bias_type,
-            "attn_mask_type": self.attn_mask_type,
-            "softmax_type": self.softmax_type,
-            "scaling_factor": self.scaling_factor,
-            "dropout_probability": self.dropout_prob,
-            "is_training": self.is_training,
-            "qkv_layout": self.qkv_layout,
-            "max_segments_per_seq": self._get_max_segments_per_sequence(),
-            "window_size": self.window_size,
-            "context_parallel_strategy": self.cp_strategy,
-            "context_parallel_causal_load_balanced": self.cp_load_balanced,
-            "stripe_size": self.stripe_size,
-        }
-        kwargs.update(overrides)
-        return kwargs
-
     def test_forward_with_max_logit(self, check_output=True):
         """Test forward output and returned max_logit."""
-        self._setup_inputs()
-        kwargs = self._fused_attn_kwargs()
-
-        customcall_fused_dpa_jit = jit(
-            partial(customcall_fused_dpa, return_max_logit=True, **kwargs),
-            static_argnames=kwargs.keys(),
-            in_shardings=[
-                self.qkvo_sharding,
-                self.qkvo_sharding,
-                self.qkvo_sharding,
-                self.bias_sharding,
-                self.softmax_offset_sharding,
-                self.seq_desc_sharding,
-                self.dropout_rng_sharding,
-            ],
-        )
-
-        with self.mesh, autocast(mesh_resource=self.mesh_resource):
-            primitive_out, primitive_aux = customcall_fused_dpa_jit(*self._customcall_args())
-            primitive_max_logit = primitive_aux["max_logit"]
-            primitive_out = self.cp_inverse_reorder_fn(primitive_out)
-
-        reference_max_logit = jax_dpa(*self._reference_args(), is_max_logit_enabled=True, **kwargs)
-
-        if check_output:
-            reference_out = jax_dpa(*self._reference_args(), **kwargs)
-            primitive_valid, primitive_invalid, reference_valid, _ = _split_valid_and_invalid(
-                primitive_out, reference_out, self.pad_q
-            )
-            assert_allclose(primitive_invalid, jnp.zeros_like(primitive_invalid), dtype=self.dtype)
-            assert_allclose(primitive_valid, reference_valid, dtype=self.dtype)
-        assert_allclose(primitive_max_logit, reference_max_logit, dtype=self.dtype)
+        self.test_forward(return_max_logit=True, check_output=check_output)
 
     def test_backward_with_max_logit(self):
         """Ensure aux-return cotangents do not break the fused attention backward path."""
-        self._setup_inputs()
-        kwargs = self._fused_attn_kwargs()
-
-        def loss_fn(query):
-            output, _ = customcall_fused_dpa(
-                query,
-                self.k,
-                self.v,
-                self.bias,
-                self.softmax_offset,
-                self.sequence_desciptor,
-                self.dropout_rng,
-                return_max_logit=True,
-                **kwargs,
-            )
-            return jnp.mean(output.astype(jnp.float32))
+        self.test_backward(return_max_logit=True)
+
+
+FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS = [
+    pytest.param(
+        QKVLayout.BSHD_BSHD_BSHD,
+        8,
+        8,
+        id="BSHD_SEPARATE",
+    ),
+    pytest.param(
+        QKVLayout.BS3HD,
+        8,
+        8,
+        id="BS3HD",
+    ),
+    pytest.param(
+        QKVLayout.BSHD_BS2HD,
+        8,
+        4,
+        id="BSHD_KV_PACKED-GQA",
+    ),
+    pytest.param(
+        QKVLayout.T3HD,
+        8,
+        8,
+        id="THD_QKV_PACKED",
+    ),
+    pytest.param(
+        QKVLayout.THD_THD_THD,
+        8,
+        8,
+        id="THD_SEPARATE",
+    ),
+]
+
+
+class TestFusedAttnMaxLogit:
+    """Targeted non-CP max_logit coverage."""
 
-        grad = jax.grad(loss_fn)(self.q)
-        assert grad.shape == self.q.shape
-
-
-@pytest.mark.parametrize(
-    "qkv_layout, seq_desc_format",
-    [
-        pytest.param(QKVLayout.BSHD_BSHD_BSHD, SeqDescFormat.Seqlens, id="BSHD_SEPARATE"),
-        pytest.param(QKVLayout.T3HD, SeqDescFormat.Seqlens, id="THD_QKV_PACKED"),
-    ],
-)
-def test_fused_attn_return_max_logit(qkv_layout, seq_desc_format):
-    """Check non-CP JAX fused attention can expose PyTorch-compatible max_logit."""
-    runner = FusedAttnRunner(
-        batch_size=2,
-        max_seqlen_q=128,
-        max_seqlen_kv=128,
-        num_heads_q=8,
-        num_heads_kv=8,
-        head_dim_qk=64,
-        head_dim_v=64,
-        attn_bias_type=AttnBiasType.NO_BIAS,
-        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
-        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
-        dropout_prob=0.0,
-        dtype=jnp.bfloat16,
-        is_training=True,
-        qkv_layout=qkv_layout,
-        bias_shape=None,
-        window_size=None,
-        seq_desc_format=seq_desc_format,
+    @staticmethod
+    @pytest.mark.parametrize(
+        "qkv_layout, num_heads_q, num_heads_kv",
+        FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS,
     )
-    runner.test_forward_with_max_logit()
-
-
-def test_fused_attn_return_max_logit_backward_smoke():
-    """Ensure aux-return cotangents do not break the fused attention backward path."""
-    runner = FusedAttnRunner(
-        batch_size=2,
-        max_seqlen_q=128,
-        max_seqlen_kv=128,
-        num_heads_q=8,
-        num_heads_kv=8,
-        head_dim_qk=64,
-        head_dim_v=64,
-        attn_bias_type=AttnBiasType.NO_BIAS,
-        attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
-        softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
-        dropout_prob=0.0,
-        dtype=jnp.bfloat16,
-        is_training=True,
-        qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
-        bias_shape=None,
-        window_size=None,
-        seq_desc_format=SeqDescFormat.Seqlens,
+    @pytest.mark.parametrize(
+        "attn_bias_type, bias_shape",
+        [
+            pytest.param(AttnBiasType.NO_BIAS, None, id="NO_BIAS"),
+            pytest.param(
+                AttnBiasType.POST_SCALE_BIAS,
+                BiasShape._1HSS,
+                id="POST_SCALE_BIAS-1HSS",
+            ),
+        ],
     )
-    runner.test_backward_with_max_logit()
+    @pytest.mark.parametrize(
+        "attn_mask_type",
+        [
+            pytest.param(AttnMaskType.NO_MASK, id="NO_MASK"),
+            pytest.param(AttnMaskType.PADDING_MASK, id="PADDING_MASK"),
+            pytest.param(AttnMaskType.CAUSAL_MASK, id="CAUSAL_MASK"),
+            pytest.param(AttnMaskType.PADDING_CAUSAL_MASK, id="PADDING_CAUSAL_MASK"),
+        ],
+    )
+    def test_forward(
+        qkv_layout,
+        num_heads_q,
+        num_heads_kv,
+        attn_bias_type,
+        bias_shape,
+        attn_mask_type,
+    ):
+        """Check non-CP JAX fused attention can expose framework-compatible max_logit."""
+        runner = FusedAttnRunner(
+            batch_size=2,
+            max_seqlen_q=128,
+            max_seqlen_kv=128,
+            num_heads_q=num_heads_q,
+            num_heads_kv=num_heads_kv,
+            head_dim_qk=64,
+            head_dim_v=64,
+            attn_bias_type=attn_bias_type,
+            attn_mask_type=attn_mask_type,
+            softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+            dropout_prob=0.0,
+            dtype=jnp.bfloat16,
+            is_training=True,
+            qkv_layout=qkv_layout,
+            bias_shape=bias_shape,
+            window_size=None,
+            seq_desc_format=SeqDescFormat.Seqlens,
+        )
+        runner.test_forward(return_max_logit=True)
+
+    @staticmethod
+    def test_backward():
+        """Ensure aux-return cotangents do not break the fused attention backward path."""
+        runner = FusedAttnRunner(
+            batch_size=2,
+            max_seqlen_q=128,
+            max_seqlen_kv=128,
+            num_heads_q=8,
+            num_heads_kv=8,
+            head_dim_qk=64,
+            head_dim_v=64,
+            attn_bias_type=AttnBiasType.NO_BIAS,
+            attn_mask_type=AttnMaskType.PADDING_CAUSAL_MASK,
+            softmax_type=AttnSoftmaxType.VANILLA_SOFTMAX,
+            dropout_prob=0.0,
+            dtype=jnp.bfloat16,
+            is_training=True,
+            qkv_layout=QKVLayout.BSHD_BSHD_BSHD,
+            bias_shape=None,
+            window_size=None,
+            seq_desc_format=SeqDescFormat.Seqlens,
+        )
+        runner.test_backward(return_max_logit=True)
 
 
 def _get_swa_window_size_for_test(s_kv: int, attn_mask_type: AttnMaskType) -> Tuple[int, int]:

From e36a4a087668c92ae11ed22c085aa4a3aa4a5208 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Thu, 20 Aug 2026 23:41:48 -0700
Subject: [PATCH 14/22] Remove JAX max logit integration tests

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn_score_mod.py | 70 --------------------------
 1 file changed, 70 deletions(-)

diff --git a/tests/jax/test_fused_attn_score_mod.py b/tests/jax/test_fused_attn_score_mod.py
index b08393bf24..7d1137b21f 100644
--- a/tests/jax/test_fused_attn_score_mod.py
+++ b/tests/jax/test_fused_attn_score_mod.py
@@ -452,11 +452,7 @@ def fake_fused_attn(
             score_mod_bprop=score_mod_bprop,
             score_mod_tensors=score_mod_tensors,
             score_mod_bprop_tensors=score_mod_bprop_tensors,
-            return_max_logit=return_max_logit,
         )
-        if return_max_logit:
-            max_logit = jnp.arange(qkv[0].shape[-2], dtype=qkv[0].dtype)
-            return qkv[0], {"max_logit": max_logit}
         return qkv[0]
 
     monkeypatch.setattr(
@@ -541,72 +537,6 @@ def test_dot_product_attention_plumbs_score_mod_to_fused_attn(monkeypatch):
     assert captured["kernel_checks"][0][0][3] is QKVLayout.BSHD_BSHD_BSHD
 
 
-def test_dot_product_attention_plumbs_return_max_logit_to_fused_attn(monkeypatch):
-    """DotProductAttention forwards return_max_logit to fused_attn and returns aux data."""
-    captured = _install_fake_flax_fused_attn(monkeypatch)
-    query = jnp.ones((1, 8, 2, 16), dtype=jnp.float16)
-    key = jnp.ones((1, 8, 2, 16), dtype=jnp.float16)
-    value = jnp.ones((1, 8, 2, 16), dtype=jnp.float16)
-
-    dpa = flax_transformer.DotProductAttention(
-        head_dim=16,
-        num_attention_heads=2,
-        num_gqa_groups=2,
-        attn_mask_type="no_mask",
-        qkv_layout="bshd_bshd_bshd",
-        transpose_batch_sequence=False,
-        return_max_logit=True,
-    )
-    out, aux = dpa.apply({}, query, key, value, deterministic=True)
-
-    np.testing.assert_array_equal(out, query)
-    np.testing.assert_array_equal(aux["max_logit"], jnp.arange(2, dtype=query.dtype))
-    assert captured["return_max_logit"] is True
-    assert captured["kernel_checks"][0][1]["return_max_logit"] is True
-
-
-def test_dot_product_attention_return_max_logit_requires_fused_attn_enabled(monkeypatch):
-    """DotProductAttention rejects return_max_logit when the fused backend is disabled."""
-    monkeypatch.setenv("NVTE_FUSED_ATTN", "0")
-    query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
-    key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
-    value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
-
-    dpa = flax_transformer.DotProductAttention(
-        head_dim=16,
-        num_attention_heads=1,
-        num_gqa_groups=1,
-        attn_mask_type="no_mask",
-        qkv_layout="bshd_bshd_bshd",
-        transpose_batch_sequence=False,
-        return_max_logit=True,
-    )
-
-    with pytest.raises(ValueError, match="NVTE_FUSED_ATTN=0"):
-        dpa.apply({}, query, key, value, deterministic=True)
-
-
-def test_dot_product_attention_return_max_logit_requires_available_fused_kernel(monkeypatch):
-    """DotProductAttention rejects return_max_logit instead of falling back to unfused attention."""
-    _install_fake_flax_fused_attn(monkeypatch, kernel_available=False)
-    query = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
-    key = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
-    value = jnp.ones((1, 8, 1, 16), dtype=jnp.float16)
-
-    dpa = flax_transformer.DotProductAttention(
-        head_dim=16,
-        num_attention_heads=1,
-        num_gqa_groups=1,
-        attn_mask_type="no_mask",
-        qkv_layout="bshd_bshd_bshd",
-        transpose_batch_sequence=False,
-        return_max_logit=True,
-    )
-
-    with pytest.raises(ValueError, match="requires a cuDNN fused attention kernel"):
-        dpa.apply({}, query, key, value, deterministic=True)
-
-
 def test_dot_product_attention_unpacks_packed_score_mod_to_separate_layout(monkeypatch):
     """Packed QKV inputs are unpacked because score_mod requires separate Q/K/V."""
     captured = _install_fake_flax_fused_attn(monkeypatch)

From 9b537babc6ff8cb2e4299fd70c454316cde3ffa3 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Fri, 21 Aug 2026 00:44:55 -0700
Subject: [PATCH 15/22] Broaden JAX CP max logit tests

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_distributed_fused_attn.py | 106 +++++++----------------
 1 file changed, 32 insertions(+), 74 deletions(-)

diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
index b508120362..1e352c3c90 100644
--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -399,73 +399,10 @@ def test_softcap_score_mod_with_aux_params_backward(
     ),
 ]
 
-DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CASES = [
-    pytest.param(
-        QKVLayout.BSHD_BSHD_BSHD,
-        AttnMaskType.CAUSAL_MASK,
-        CPStrategy.ALL_GATHER,
-        (-1, -1),
-        None,
-        None,
-        False,
-        True,
-        id="AG-BSHD",
-    ),
-    pytest.param(
-        QKVLayout.THD_THD_THD,
-        AttnMaskType.PADDING_CAUSAL_MASK,
-        CPStrategy.ALL_GATHER,
-        (-1, -1),
-        64,
-        5,
-        False,
-        True,
-        id="AG-THD",
-    ),
-    pytest.param(
-        QKVLayout.BSHD_BSHD_BSHD,
-        AttnMaskType.CAUSAL_MASK,
-        CPStrategy.RING,
-        (-1, -1),
-        None,
-        None,
-        False,
-        True,
-        id="RING-BSHD-NO_SCAN",
-    ),
-    pytest.param(
-        QKVLayout.BSHD_BSHD_BSHD,
-        AttnMaskType.CAUSAL_MASK,
-        CPStrategy.RING,
-        (-1, -1),
-        None,
-        None,
-        True,
-        True,
-        id="RING-BSHD-SCAN",
-    ),
-    pytest.param(
-        QKVLayout.THD_THD_THD,
-        AttnMaskType.PADDING_CAUSAL_MASK,
-        CPStrategy.RING,
-        (-1, -1),
-        1,
-        5,
-        False,
-        False,
-        id="RING-THD-NO_SCAN",
-    ),
-    pytest.param(
-        QKVLayout.THD_THD_THD,
-        AttnMaskType.PADDING_CAUSAL_MASK,
-        CPStrategy.RING,
-        (-1, -1),
-        1,
-        5,
-        True,
-        False,
-        id="RING-THD-SCAN",
-    ),
+DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CP_MODES = [
+    pytest.param(CPStrategy.ALL_GATHER, False, id="AG"),
+    pytest.param(CPStrategy.RING, False, id="RING-NO_SCAN"),
+    pytest.param(CPStrategy.RING, True, id="RING-SCAN"),
 ]
 
 
@@ -595,12 +532,22 @@ def check_has_backend_for_mask(mask_type):
         generate_context_parallel_configs_for_attn(),
     )
     @pytest.mark.parametrize("data_shape", DISTRIBUTED_CONTEXT_SELF_ATTN_DATA_SHAPES[:1])
-    @pytest.mark.parametrize("kv_groups", [1])
+    @pytest.mark.parametrize("kv_groups", [1, 8])
     @pytest.mark.parametrize("dtype", [pytest.param(jnp.bfloat16, id="BF16")])
     @pytest.mark.parametrize(
-        "qkv_layout, attn_mask_type, cp_strategy, window_size, stripe_size,"
-        " num_segments_per_seq, use_scan_ring, check_forward_output",
-        DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CASES,
+        "qkv_layout, attn_mask_type",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_LAYOUTS_MASKS,
+    )
+    @pytest.mark.parametrize(
+        "cp_strategy, use_scan_ring",
+        DISTRIBUTED_CONTEXT_SELF_ATTN_MAX_LOGIT_CP_MODES,
+    )
+    @pytest.mark.parametrize(
+        "window_size",
+        [
+            pytest.param((-1, -1), id="NO_SWA"),
+            pytest.param((20, 0), id="SWA"),
+        ],
     )
     def test_context_parallel_return_max_logit(
         self,
@@ -615,12 +562,23 @@ def test_context_parallel_return_max_logit(
         attn_mask_type,
         cp_strategy,
         window_size,
-        stripe_size,
-        num_segments_per_seq,
         use_scan_ring,
-        check_forward_output,
     ):
         """Check CP fused attention returns global per-head max_logit."""
+        is_thd = qkv_layout.is_thd()
+        if window_size != (-1, -1) and not (
+            is_thd and cp_strategy == CPStrategy.RING and not use_scan_ring
+        ):
+            pytest.skip("SWA max-logit coverage is limited to THD Ring without scan.")
+        # TODO: Evaluate cuDNN Max mismatches observed for striped multi-segment THD Ring GQA.
+        if is_thd and cp_strategy == CPStrategy.RING and kv_groups > 1:
+            pytest.skip("THD Ring GQA Max mismatches require further evaluation.")
+
+        stripe_size = 64 if is_thd and cp_strategy == CPStrategy.ALL_GATHER else None
+        if is_thd and cp_strategy == CPStrategy.RING:
+            stripe_size = 1
+        num_segments_per_seq = 5 if is_thd else None
+        check_forward_output = not (is_thd and cp_strategy == CPStrategy.RING)
         self.impl_test_context_parallel_attn(
             device_count,
             mesh_shape,

From bb7ce9b1b1cc003ae6c2b94604219370c4c904d1 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Sun, 23 Aug 2026 19:26:06 -0700
Subject: [PATCH 16/22] Reduce JAX max logit across DP

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 .../jax/cpp_extensions/attention.py           | 40 ++++++++++++++-----
 1 file changed, 31 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 45ce08cb73..96d6f56c32 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -831,7 +831,21 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
-        impl = partial(FusedAttnFwdPrimitive.impl, config=config)
+        max_logit_reduce_axes = (
+            FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding)
+            if config.return_max_logit
+            else ()
+        )
+
+        def impl(*args):
+            output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl(
+                *args, config=config
+            )
+            max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
+                max_logit, mesh, max_logit_reduce_axes, config
+            )
+            return output, softmax_aux, rng_state, max_logit
+
         return mesh, impl, out_shardings, arg_shardings
 
     @staticmethod
@@ -1931,8 +1945,10 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[5] = seed_sharding
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
-        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
-            mesh, max_logit_sharding
+        max_logit_reduce_axes = (
+            FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding)
+            if config.return_max_logit
+            else ()
         )
 
         def impl(
@@ -2236,8 +2252,10 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[5] = seed_sharding
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
-        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
-            mesh, max_logit_sharding
+        max_logit_reduce_axes = (
+            FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding)
+            if config.return_max_logit
+            else ()
         )
 
         def impl(
@@ -2698,8 +2716,10 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
-        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
-            mesh, max_logit_sharding
+        max_logit_reduce_axes = (
+            FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding)
+            if config.return_max_logit
+            else ()
         )
 
         def ring_attn_fwd_impl(
@@ -3223,8 +3243,10 @@ def partition(config, mesh, arg_infos, result_infos):
         arg_shardings[-2] = arg_shardings[-4]
         arg_shardings = tuple(arg_shardings)
         out_shardings = (out_sharding, softmax_aux_sharding, rng_state_sharding, max_logit_sharding)
-        max_logit_reduce_axes = FusedAttnFwdPrimitive._max_logit_reduce_axes(
-            mesh, max_logit_sharding
+        max_logit_reduce_axes = (
+            FusedAttnFwdPrimitive._max_logit_reduce_axes(mesh, max_logit_sharding)
+            if config.return_max_logit
+            else ()
         )
 
         def fwd_impl(

From c09332983e0f0ea1e5cb54a2dd71d17c612ce2d0 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Sun, 23 Aug 2026 20:04:46 -0700
Subject: [PATCH 17/22] Simplify JAX max logit return

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_fused_attn.py               |  7 +++---
 transformer_engine/jax/attention.py        | 16 ++++++--------
 transformer_engine/jax/flax/transformer.py | 25 +++++++---------------
 3 files changed, 18 insertions(+), 30 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 19aea7f58b..190367db7e 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -352,8 +352,8 @@ def customcall_fused_dpa(
         qkv_args, bias, sequence_descriptor, dropout_rng, softmax_offset=softmax_offset, **kwargs
     )
     if isinstance(result, tuple):
-        output, aux = result
-        return output.astype(query.dtype), aux
+        output, max_logit = result
+        return output.astype(query.dtype), max_logit
     return result.astype(query.dtype)
 
 
@@ -1018,8 +1018,7 @@ def test_forward(self, return_max_logit=False, check_output=True):
         with self.mesh, autocast(mesh_resource=self.mesh_resource):
             primitive_out = customcall_fused_dpa_jit(*customcall_args)
             if return_max_logit:
-                primitive_out, primitive_aux = primitive_out
-                primitive_max_logit = primitive_aux["max_logit"]
+                primitive_out, primitive_max_logit = primitive_out
             primitive_out = self.cp_inverse_reorder_fn(primitive_out)
 
         if return_max_logit:
diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index 5cd90140df..ee0f2b322c 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -1090,13 +1090,12 @@ def _legacy_fused_attn(
         softmax_offset (Optional[jnp.ndarray]): An optional learnable softmax offset tensor with shape
             [1, num_heads, 1, 1]. Used when softmax_type is AttnSoftmaxType.LEARNABLE_SOFTMAX.
         return_max_logit (bool): If True, also return per-head maximum attention logits
-            in an auxiliary dictionary under ``"max_logit"``.
+            with shape ``[h]``.
     Returns:
         jnp.ndarray:
             Attention output when ``return_max_logit`` is False.
-        tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
-            ``(output, aux)`` when ``return_max_logit`` is True. ``aux`` contains
-            ``"max_logit"`` with shape ``[h]``.
+        tuple[jnp.ndarray, jnp.ndarray]:
+            ``(output, max_logit)`` when ``return_max_logit`` is True.
     """
     assert (
         not qkv_layout.is_thd()
@@ -1335,7 +1334,7 @@ def _fused_attn_fwd_rule(
     softmax_aux = checkpoint_name(softmax_aux, context_checkpoint_name)
     rng_state = checkpoint_name(rng_state, context_checkpoint_name)
     max_logit = checkpoint_name(max_logit, context_checkpoint_name)
-    attn_output = (output, {"max_logit": max_logit}) if return_max_logit else output
+    attn_output = (output, max_logit) if return_max_logit else output
     return attn_output, (
         qkv,
         bias,
@@ -1551,13 +1550,12 @@ def fused_attn(
         score_mod_bprop_tensors (Optional[Mapping[str, Any]]): Additional tensors or
             Python/NumPy scalars made available to `score_mod_bprop`.
         return_max_logit (bool): If True, also return per-head maximum attention logits
-            in an auxiliary dictionary under ``"max_logit"``.
+            with shape ``[h]``.
     Returns:
         jnp.ndarray:
             Attention output when ``return_max_logit`` is False.
-        tuple[jnp.ndarray, dict[str, jnp.ndarray]]:
-            ``(output, aux)`` when ``return_max_logit`` is True. ``aux`` contains
-            ``"max_logit"`` with shape ``[h]``.
+        tuple[jnp.ndarray, jnp.ndarray]:
+            ``(output, max_logit)`` when ``return_max_logit`` is True.
 
     Examples (non-THD, also known as non-packed):
         >>> #  q_segment_ids = [[1, 1, 1, 0], [1, 1, 0, 0]], 0 means padded tokens
diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py
index ce3dab4818..4d4f66afa5 100644
--- a/transformer_engine/jax/flax/transformer.py
+++ b/transformer_engine/jax/flax/transformer.py
@@ -437,7 +437,7 @@ def __call__(
             raise ValueError(f"Unsupported {self.qkv_layout=}.")
 
         if self.return_max_logit:
-            x, aux = x
+            x, max_logit = x
 
         if self.transpose_batch_sequence:
             x = x.transpose([1, 0, 2, 3])
@@ -446,7 +446,7 @@ def __call__(
             x.dtype == query.dtype
         ), f"output dtype {x.dtype} does not match query dtype {query.dtype}"
         if self.return_max_logit:
-            return x, aux
+            return x, max_logit
         return x
 
 
@@ -627,7 +627,7 @@ class DotProductAttention(nn.Module):  # pylint: disable=too-few-public-methods
     score_mod_bprop_tensors: Optional[Mapping[str, Any]], default = None
         Additional tensors or pass-by-value scalars for ``score_mod_bprop``.
     return_max_logit: bool, default = False
-        If True, return ``(output, aux)`` where ``aux["max_logit"]`` contains the per-head
+        If True, return ``(output, max_logit)`` where ``max_logit`` contains the per-head
         maximum attention logits with shape ``[h]``. This path requires fused attention.
 
     Optimization parameters
@@ -728,8 +728,8 @@ def __call__(
 
         Returns
         -------
-        outputs: jax.numpy.ndarray or tuple[jax.numpy.ndarray, dict[str, jax.numpy.ndarray]]
-            Output tensor, or ``(output, aux)`` when ``return_max_logit`` is enabled.
+        outputs: jax.numpy.ndarray or tuple[jax.numpy.ndarray, jax.numpy.ndarray]
+            Output tensor, or ``(output, max_logit)`` when ``return_max_logit`` is enabled.
         """
         input_dtype = query.dtype
 
@@ -790,9 +790,7 @@ def __call__(
         enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1"))
         if self.return_max_logit and not enable_fused_attn:
             raise ValueError(
-                "DotProductAttention(return_max_logit=True) requires fused attention, but "
-                "NVTE_FUSED_ATTN=0 disables it. Set NVTE_FUSED_ATTN=1 or unset the variable, "
-                "then ensure a cuDNN fused attention kernel is available for this configuration."
+                "return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0."
             )
 
         sequence_dim = 0 if self.transpose_batch_sequence else 1
@@ -840,15 +838,8 @@ def __call__(
             )
         if self.return_max_logit and not has_fused_attn_kernel:
             raise ValueError(
-                "DotProductAttention(return_max_logit=True) requires a cuDNN fused attention "
-                "kernel, but no compatible kernel is available for this configuration. "
-                "Set NVTE_FUSED_ATTN=1 or unset it, update cuDNN/Transformer Engine if needed, "
-                "and check the configuration: "
-                f"{qkv_layout=}, {attn_bias_type=}, {attn_mask_type=}, "
-                f"{softmax_type=}, attention_dropout={self.attention_dropout}, "
-                f"num_attention_heads={self.num_attention_heads}, "
-                f"num_gqa_groups={self.num_gqa_groups}, {seqlen_q=}, {seqlen_kv=}, "
-                f"{head_dim_qk=}, {head_dim_v=}, window_size={self.window_size}."
+                "return_max_logit requires fused attention, but no fused attention kernel is "
+                "available."
             )
 
         use_fused_attn = enable_fused_attn and has_fused_attn_kernel

From 0f7bb0bd5e22afb3be9dcbab6fd5d0b59562462e Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Mon, 24 Aug 2026 10:19:20 -0700
Subject: [PATCH 18/22] Document JAX max logit reductions

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 .../jax/cpp_extensions/attention.py             | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py
index 96d6f56c32..6ea54195ab 100644
--- a/transformer_engine/jax/cpp_extensions/attention.py
+++ b/transformer_engine/jax/cpp_extensions/attention.py
@@ -377,6 +377,9 @@ def abstract(
             raise ValueError(f"Unsupported {backend=}")
         softmax_aux_aval = q_aval.update(shape=softmax_shape, dtype=softmax_dtype)
         if config.return_max_logit:
+            # cuDNN Max is row-wise over S_kv. Dense and SM120 THD use
+            # [..., H, S_q, 1]; cuDNN >= 9.6 non-SM120 THD uses [..., S_q, H, 1].
+            # Both raw layouts are reduced to the public per-head [H] result below.
             if FusedAttnFwdPrimitive._uses_thd_ragged_max_tensor(config):
                 max_tensor_shape = (*batch_shape, q_max_seqlen, attn_heads, 1)
             else:
@@ -675,7 +678,13 @@ def convert_to_2d(offsets, batch, max_seqlen):
 
     @staticmethod
     def _reduce_max_logit(max_tensor, output, q_seqlen, q_seq_offsets, config):
-        """Reduce cuDNN's raw Max tensor to framework-compatible per-head max_logit."""
+        """Reduce cuDNN's row-wise Max tensor to the public per-head max_logit.
+
+        Dense and SM120 THD use ``[..., H, S_q, 1]``; cuDNN >= 9.6 non-SM120
+        THD uses ``[..., S_q, H, 1]``. A rank-3 THD result is ``[T_q, H, 1]``.
+        All layouts reduce to ``[H]``. Static THD buffers can contain invalid query
+        rows, so those rows are masked before reduction.
+        """
         if not config.return_max_logit:
             return jnp.zeros((0,), dtype=output.dtype)
 
@@ -841,6 +850,7 @@ def impl(*args):
             output, softmax_aux, rng_state, max_logit = FusedAttnFwdPrimitive.impl(
                 *args, config=config
             )
+            # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding.
             max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
                 max_logit, mesh, max_logit_reduce_axes, config
             )
@@ -2035,7 +2045,7 @@ def _cross_attn(idx, q, k, v, bias, softmax_offset, q_seqlen, kv_seqlen, seed):
             ]
 
             output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions)
-            # Reduce over non-head mesh axes to make [H] global over batch/sequence.
+            # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding.
             max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
                 max_logit, mesh, max_logit_reduce_axes, config
             )
@@ -2361,6 +2371,7 @@ def _cross_attn(
                 for _ in range(cp_size)
             ]
             output, softmax_aux, rng_state, max_logit = lax.switch(cp_rank, functions)
+            # Globalize the rank-local [H] max across DP/CP while preserving TP head sharding.
             max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
                 max_logit, mesh, max_logit_reduce_axes, config
             )
@@ -2911,6 +2922,7 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
             (kv, output, softmax_aux, max_logit) = carry
 
             output = output.astype(q.dtype)
+            # Globalize the rank-local running [H] max across DP/CP.
             max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
                 max_logit, mesh, max_logit_reduce_axes, config
             )
@@ -3380,6 +3392,7 @@ def correction(output, softmax_aux, output_per_step, softmax_aux_per_step):
             (_, _, _, output, softmax_aux, max_logit) = carry
 
             output = output.astype(q.dtype)
+            # Globalize the rank-local running [H] max across DP/CP.
             max_logit = FusedAttnFwdPrimitive._reduce_max_logit_across_mesh(
                 max_logit, mesh, max_logit_reduce_axes, config
             )

From 90a3eaf1bc2f9b2d7ac82ea8f27ddb4ea81032a7 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Mon, 24 Aug 2026 14:13:17 -0700
Subject: [PATCH 19/22] Refine JAX max logit tests

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 tests/jax/test_distributed_fused_attn.py | 15 ++++++++++-----
 tests/jax/test_fused_attn.py             |  9 ---------
 2 files changed, 10 insertions(+), 14 deletions(-)

diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
index 1e352c3c90..fd6d83835b 100644
--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -522,7 +522,10 @@ def check_has_backend_for_mask(mask_type):
             pytest.skip(f"Skipping {kv_groups=} not multiple of {data_shape=} or {tp_size=}")
 
         if return_max_logit:
-            runner.test_forward_with_max_logit(check_output=check_forward_output)
+            runner.test_forward(
+                return_max_logit=True,
+                check_output=check_forward_output,
+            )
         else:
             runner.test_backward()
         del os.environ["NVTE_FUSED_RING_ATTENTION_USE_SCAN"]
@@ -566,10 +569,12 @@ def test_context_parallel_return_max_logit(
     ):
         """Check CP fused attention returns global per-head max_logit."""
         is_thd = qkv_layout.is_thd()
-        if window_size != (-1, -1) and not (
-            is_thd and cp_strategy == CPStrategy.RING and not use_scan_ring
-        ):
-            pytest.skip("SWA max-logit coverage is limited to THD Ring without scan.")
+        supports_swa = is_thd and (
+            cp_strategy == CPStrategy.ALL_GATHER
+            or (cp_strategy == CPStrategy.RING and not use_scan_ring)
+        )
+        if window_size != (-1, -1) and not supports_swa:
+            pytest.skip("CP SWA requires THD All-Gather or unrolled THD Ring.")
         # TODO: Evaluate cuDNN Max mismatches observed for striped multi-segment THD Ring GQA.
         if is_thd and cp_strategy == CPStrategy.RING and kv_groups > 1:
             pytest.skip("THD Ring GQA Max mismatches require further evaluation.")
diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 190367db7e..730b02c4db 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1303,15 +1303,6 @@ def check_dqkv(primitive, reference, pad, idx):
                 target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text()
             assert_equal_collectives(target_hlo, self.coll_count_ref)
 
-    def test_forward_with_max_logit(self, check_output=True):
-        """Test forward output and returned max_logit."""
-        self.test_forward(return_max_logit=True, check_output=check_output)
-
-    def test_backward_with_max_logit(self):
-        """Ensure aux-return cotangents do not break the fused attention backward path."""
-        self.test_backward(return_max_logit=True)
-
-
 FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS = [
     pytest.param(
         QKVLayout.BSHD_BSHD_BSHD,

From 704957afba4d7e432d6332fe0c7e2f4b5bc382fd Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Mon, 24 Aug 2026 14:56:39 -0700
Subject: [PATCH 20/22] Rename JAX max logit buffer

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 transformer_engine/jax/csrc/extensions/attention.cpp | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp
index 9728f2cbcf..47dcafbed6 100644
--- a/transformer_engine/jax/csrc/extensions/attention.cpp
+++ b/transformer_engine/jax/csrc/extensions/attention.cpp
@@ -48,7 +48,7 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
                                        const size_t bias_heads, const size_t q_max_seqlen,
                                        const size_t kv_max_seqlen, DType dtype,
                                        NVTE_Bias_Type bias_type, NVTE_Fused_Attn_Backend backend,
-                                       void *softmax_buf, void *max_buf = nullptr,
+                                       void *softmax_buf, void *max_logits_buf = nullptr,
                                        void *rng_state_buf = nullptr, void *bias_buf = nullptr,
                                        void *softmax_offset_buf = nullptr) {
   // all backends need softmax but expect different shapes/dtypes
@@ -67,10 +67,10 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
   if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) {
     int size = 1;  // Start after softmax.
 
-    if (max_buf != nullptr) {
+    if (max_logits_buf != nullptr) {
       NVTETensor &max_aux = tensor_pack->tensors[size++];
       NVTEBasicTensor max_aux_data;
-      max_aux_data.data_ptr = max_buf;
+      max_aux_data.data_ptr = max_logits_buf;
       max_aux_data.shape = {};
       max_aux_data.shape.ndim = 4;
       max_aux_data.shape.data[0] = input_batch;

From 2e71bc8aefe7e3fa2ec02ae7609b075f0f93fe53 Mon Sep 17 00:00:00 2001
From: Kshitij Lakhani <klakhani@nvidia.com>
Date: Mon, 24 Aug 2026 14:57:05 -0700
Subject: [PATCH 21/22] Check JAX attention tensor pack capacity

Signed-off-by: Kshitij Lakhani <klakhani@nvidia.com>
---
 .../jax/csrc/extensions/attention.cpp             | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp
index 47dcafbed6..41f87ae71b 100644
--- a/transformer_engine/jax/csrc/extensions/attention.cpp
+++ b/transformer_engine/jax/csrc/extensions/attention.cpp
@@ -66,9 +66,14 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
   // arbitrary sequence length backend needs the RNG state and a different shape/dtype softmax
   if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) {
     int size = 1;  // Start after softmax.
+    auto next_aux_tensor = [&]() -> NVTETensor & {
+      NVTE_CHECK(size < NVTETensorPack::MAX_SIZE,
+                 "Fused attention auxiliary tensor pack capacity exceeded.");
+      return tensor_pack->tensors[size++];
+    };
 
     if (max_logits_buf != nullptr) {
-      NVTETensor &max_aux = tensor_pack->tensors[size++];
+      NVTETensor &max_aux = next_aux_tensor();
       NVTEBasicTensor max_aux_data;
       max_aux_data.data_ptr = max_logits_buf;
       max_aux_data.shape = {};
@@ -81,7 +86,7 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
       nvte_set_tensor_param(&max_aux, kNVTERowwiseData, &max_aux_data);
     }
 
-    NVTETensor &rng_state_aux = tensor_pack->tensors[size++];
+    NVTETensor &rng_state_aux = next_aux_tensor();
     NVTEBasicTensor rng_state_aux_data;
     rng_state_aux_data.data_ptr = rng_state_buf;
     rng_state_aux_data.shape = {};
@@ -94,8 +99,7 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
 
     // include bias if enabled
     if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS && bias_type != NVTE_Bias_Type::NVTE_ALIBI) {
-      NVTETensor &bias_aux = tensor_pack->tensors[size];
-      size++;
+      NVTETensor &bias_aux = next_aux_tensor();
       NVTEBasicTensor bias_aux_data;
       bias_aux_data.data_ptr = bias_buf;
       bias_aux_data.shape.ndim = 4;
@@ -109,8 +113,7 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t
 
     // include softmax_offset if provided
     if (softmax_offset_buf != nullptr) {
-      NVTETensor &softmax_offset_aux = tensor_pack->tensors[size];
-      size++;
+      NVTETensor &softmax_offset_aux = next_aux_tensor();
       NVTEBasicTensor softmax_offset_aux_data;
       softmax_offset_aux_data.data_ptr = softmax_offset_buf;
       softmax_offset_aux_data.shape.ndim = 4;

From d7bf4cf799108de8c77997283726030ad8e2bb18 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 24 Aug 2026 22:00:38 +0000
Subject: [PATCH 22/22] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/jax/test_fused_attn.py               | 5 ++---
 transformer_engine/jax/attention.py        | 1 +
 transformer_engine/jax/flax/transformer.py | 8 ++++----
 3 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index 730b02c4db..d68e409331 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -1022,9 +1022,7 @@ def test_forward(self, return_max_logit=False, check_output=True):
             primitive_out = self.cp_inverse_reorder_fn(primitive_out)
 
         if return_max_logit:
-            reference_max_logit = jax_dpa(
-                *args, is_max_logit_enabled=True, **reference_kwargs
-            )
+            reference_max_logit = jax_dpa(*args, is_max_logit_enabled=True, **reference_kwargs)
 
         if check_output and not (self.is_training and self.dropout_prob > 0.0):
             reference_out = jax_dpa(*args, **reference_kwargs)
@@ -1303,6 +1301,7 @@ def check_dqkv(primitive, reference, pad, idx):
                 target_hlo = jitted_primitive.lower(*customcall_args).compile().as_text()
             assert_equal_collectives(target_hlo, self.coll_count_ref)
 
+
 FUSED_ATTN_MAX_LOGIT_QKV_LAYOUTS = [
     pytest.param(
         QKVLayout.BSHD_BSHD_BSHD,
diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py
index ee0f2b322c..6d7c823e12 100644
--- a/transformer_engine/jax/attention.py
+++ b/transformer_engine/jax/attention.py
@@ -1412,6 +1412,7 @@ def _fused_attn_bwd_rule(
         None,
     )
 
+
 _fused_attn.defvjp(_fused_attn_fwd_rule, _fused_attn_bwd_rule)
 
 
diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py
index 4d4f66afa5..4b497826cc 100644
--- a/transformer_engine/jax/flax/transformer.py
+++ b/transformer_engine/jax/flax/transformer.py
@@ -789,9 +789,7 @@ def __call__(
         # Use fused attn (if kernel check below passes) by default
         enable_fused_attn = int(os.getenv("NVTE_FUSED_ATTN", "1"))
         if self.return_max_logit and not enable_fused_attn:
-            raise ValueError(
-                "return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0."
-            )
+            raise ValueError("return_max_logit requires fused attention, but NVTE_FUSED_ATTN=0.")
 
         sequence_dim = 0 if self.transpose_batch_sequence else 1
         seqlen_q = query.shape[sequence_dim]
@@ -950,7 +948,9 @@ def __call__(
                 score_mod_bprop_tensors=score_mod_bprop_tensors,
             )
         output = x[0] if self.return_max_logit else x
-        assert output.dtype == input_dtype, f"output_dtype={output.dtype}, input_dtype={input_dtype}"
+        assert (
+            output.dtype == input_dtype
+        ), f"output_dtype={output.dtype}, input_dtype={input_dtype}"
         return x