From c650199a3521370d5d83555708604d7d4384d51e Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:49 +0000
Subject: [PATCH 01/19] Training engine: do not differentiate parameters
 trainable_parameters_mask freezes

The 397B RL trainer freezes the MoE router gate with trainable_parameters_mask,
but fwd_bwd still differentiated every parameter and _update_kernel zeroed the
frozen gradients afterwards. So every micro-batch computed each layer's gate
weight gradient, an f32 [4096, 512] HIGHEST-precision matmul, and
reduce-scattered it 256-way on the SparseCore, only to discard it.

Wrap the frozen leaves in stop_gradient inside diff_wrapper. Autodiff then
emits neither the weight-gradient matmuls nor their reduce-scatters. Their
gradients come out zero, as _update_kernel made them before (ahead of the grad
norm and clipping), so the update and grad_norm are unchanged.
---
 src/maxtext/training_engine/maxtext_engine.py   | 7 ++++++-
 tests/post_training/unit/maxtext_engine_test.py | 2 ++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 0f1d0ddc1b..e0fc9e6f18 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -1546,11 +1546,16 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None):
       keep the parameters' dtype; and the loss denominator to normalize them by. Under
       `retry_when_tokens_dropped`, followed by `has_moe_overflow`: whether any MoE layer's ragged
       buffer dropped tokens, in which case the gradients and denominator are zero and `fwd_bwd`
-      reruns the micro-batch once the flag reaches the host.
+      reruns the micro-batch once the flag reaches the host. Frozen parameters' gradients are zero.
     """
     loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn
 
     def diff_wrapper(p, r, b):
+      if self._freeze_mask_fn is not None:
+        # Parameters `trainable_parameters_mask` freezes enter as constants. `_update_kernel` zeroes their
+        # gradients anyway, so the update is unchanged, but the backward pass no longer computes their weight
+        # gradients or reduce-scatters them (the MoE router gate's, every layer, every micro-batch).
+        p = jax.tree.map(lambda x, frozen: jax.lax.stop_gradient(x) if frozen else x, p, self._freeze_mask_fn(p))
       mdl = nnx.merge(graphdef if graphdef is not None else self._model_graphdef, p, r, copy=True)
       if self._gen_model_input_fn is not None:
         # A gen_model_input_fn maps a payload to the loss fn's *keyword arguments* -- see
diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py
index 674b7dc280..b562444b17 100644
--- a/tests/post_training/unit/maxtext_engine_test.py
+++ b/tests/post_training/unit/maxtext_engine_test.py
@@ -336,6 +336,8 @@ def __init__(self):
       payload = DummyPayload()
       t.compile(payload)
       t.fwd_bwd(payload)
+      # Not differentiated at all, rather than differentiated and zeroed in `update()`.
+      np.testing.assert_array_equal(np.asarray(t._accumulated_grads["router_gate"]), [0.0, 0.0])
       t.update()
 
       np.testing.assert_allclose(np.asarray(t.model.router_gate.get_value()), [10.0, 10.0])

From d5b1be0df368b3fc099eddfa8a8d5040a385be53 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:49 +0000
Subject: [PATCH 02/19] GDN forward: skip the zero fill of the prefill kernel's
 output buffer

_run_local_gdn_decoupled_fwd called fused_conv1d_gdn with the default
zero_initialize_out=True, which aliases a zero buffer into the kernel output.
XLA hoists the zeros out of the layer scan into a loop-carried buffer and
copies it before every kernel call, and the GDN recompute in the backward
broadcasts it again. The prefill kernel never reads that buffer and writes
every row (seq_len is a multiple of chunk_size, which both functions already
check), so the fill is dead work.

On the 397B v7x bench (16,384 tokens and 64 value heads per device) this
removes six 256 MiB copies per forward layer-block iteration (about 1 ms per
iteration) and three zero broadcasts per backward iteration, and frees the
256 MiB loop-carried buffer. Outputs are bit-identical.
---
 src/maxtext/kernels/gdn/gdn_bwd/api.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index 4ea358228b..7f7b2b824f 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -267,6 +267,9 @@ def _run_local_gdn_decoupled_fwd(
       d_k=head_k_dim,
       d_v=head_v_dim,
       kernel_size=conv_kernel_size,
+      # seq_len is a multiple of chunk_size (checked above), so the prefill kernel writes every
+      # output row and the output buffer needs no zero fill.
+      zero_initialize_out=False,
       compute_precision=jnp.dtype(compute_dtype),
       mixed_tile_size=chunk_size,
       is_prefill_only=True,

From 0b7cdb7f028729e7b6610b084927ac839739e218 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:49 +0000
Subject: [PATCH 03/19] MoE: concatenate the prefused gate/up gradient instead
 of padding and adding its halves

With prefuse_moe_weights the gate and up projections are one GEMM, whose
[tokens, 2 * mlp] output gmm_up split into two halves. Autodiff transposes
each half's slice into a zero-padded [tokens, 2 * mlp] gradient and adds the
two, so XLA writes both half gradients and reads them back in a pad-add
fusion over the whole ragged buffer (bf16[225280,2048], 8 per micro-batch on
the 397B bench). The chunk pipeline's tie 4 held the first half, so XLA also
copied that half out of the GEMM output (bf16[225280,1024] slice) to feed
the barrier.

gmm_up now returns the fused output whole when there are no MoE biases, and
_down_projection applies fused_glu, a custom_vjp around the split plus the
activation. Its backward computes the two half gradients with jax.vjp of the
same apply_ffn_activation and concatenates them; jnp.where(d == 0, 0, d)
keeps the -0 -> +0 of the zero-padded add, so the gradient is bit-identical.
The halves keep their remat names (mlpwi_0, moe_mlpwi_0, moe_mlpwi_1) inside
the custom_vjp, so remat policies save the same residuals as before. Tie 4
now holds the GEMM output itself.

Checks: the tiny Qwen3.5 A/B on v6e (prefused fake-quantized wi, ring of
experts, chunk pipeline, CP/EP/FSDP=2) is bit-identical to the base in loss,
logits and all 66 gradient leaves. FusedGluTest compares the forward and the
gradient bit for bit against autodiff of the split, with zero cotangent rows,
plain and under jax.checkpoint.
---
 src/maxtext/layers/moe.py | 42 ++++++++++++++++++++++++++++++++++++++-
 tests/unit/moe_test.py    | 23 +++++++++++++++++++++
 2 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index dda2f599eb..2c8c24d6d7 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -301,6 +301,39 @@ def forward_row_major(x: jax.Array) -> jax.Array:
   return _forward_layout_constraint(x, tuple(range(x.ndim)))
 
 
+def _split_gate_up(fused):
+  """The gate and up halves of the prefused gate/up GEMM output, with the remat names of `gmm_up`."""
+  n = fused.shape[-1] // 2
+  gate = adc.checkpoint_name(adc.checkpoint_name(fused[..., :n], "mlpwi_0"), "moe_mlpwi_0")
+  up = adc.checkpoint_name(fused[..., n:], "moe_mlpwi_1")
+  return gate, up
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
+def fused_glu(fused, glu):
+  """`glu(gate, up)` of the two halves of `fused`, the [..., 2 * mlp] output of the prefused gate/up GEMM.
+
+  Autodiff transposes each half's slice into a zero-padded [..., 2 * mlp] gradient and adds the two, one more
+  full pass over the gradient. This backward computes the same two half gradients, with `jax.vjp` of `glu`, and
+  concatenates them, which XLA writes straight from the activation's backward fusion.
+  """
+  return glu(*_split_gate_up(fused))
+
+
+def _fused_glu_fwd(fused, glu):
+  gate, up = _split_gate_up(fused)
+  return glu(gate, up), (gate, up)
+
+
+def _fused_glu_bwd(glu, residuals, cotangent):
+  _, glu_vjp = jax.vjp(glu, *residuals)
+  # The zero-padded add turns -0 into +0; so does this, so the gradient is bit-identical to autodiff's.
+  return (jnp.concatenate([jnp.where(d == 0, 0, d) for d in glu_vjp(cotangent)], axis=-1),)
+
+
+fused_glu.defvjp(_fused_glu_fwd, _fused_glu_bwd)
+
+
 def get_batchsplit_init_kernel_axes():
   return (
       ("expert_only", "embed_moe", None),
@@ -2678,6 +2711,10 @@ def gmm_up(
         # Weights are stored as (G,K,2N); w0/w1 are adjacent slices so XLA elides this concat.
         w_fused = jnp.concatenate([w0, w1], axis=-1)
         out = gmm_fn(x, w_fused, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes)
+        if not self.config.mlp_bias:
+          # Returned whole, as (out, None): `_down_projection` splits it in `fused_glu`, whose backward concatenates
+          # the halves' gradients, and the chunk pipeline's tie 4 holds the GEMM output, not a copy of its first half.
+          return out, None
         n = out.shape[-1] // 2
         layer_w0, layer_w1 = out[:, :n], out[:, n:]
         if self.config.mlp_bias and w0_bias is not None and w1_bias is not None:
@@ -2931,7 +2968,10 @@ def _up_projection(x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias
 
     def _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata):
       """The activation and the down projection, per routed token."""
-      intermediate_layer = self.apply_ffn_activation(output0, output1)
+      if output1 is None:  # output0 is the prefused gate/up GEMM output (see `gmm_up`)
+        intermediate_layer = fused_glu(output0, self.apply_ffn_activation)
+      else:
+        intermediate_layer = self.apply_ffn_activation(output0, output1)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
       intermediate_output = gmm_fn(
           intermediate_layer,
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
index fa4025072d..c2e0946e9a 100644
--- a/tests/unit/moe_test.py
+++ b/tests/unit/moe_test.py
@@ -3507,5 +3507,28 @@ def test_fp8_moe_with_weight_quant_config(self):
     self.assertEqual(model.wi_scale.shape, expected_wi_shape)
 
 
+class FusedGluTest(unittest.TestCase):
+  """`fused_glu` matches autodiff of the split prefused gate/up output, bit for bit."""
+
+  def test_matches_autodiff_of_split_bitwise(self):
+    def glu(gate, up):
+      return (jax.nn.silu(gate) * up).astype(jnp.bfloat16)
+
+    def split_then_glu(fused):
+      n = fused.shape[-1] // 2
+      return glu(fused[:, :n], fused[:, n:])
+
+    k_fused, k_cotangent = jax.random.split(jax.random.key(0))
+    fused = (3 * jax.random.normal(k_fused, (256, 2 * 128))).astype(jnp.bfloat16)
+    # Zero cotangent rows give -0 half gradients, which autodiff's zero-padded add turns into +0.
+    cotangent = jax.random.normal(k_cotangent, (256, 128)).astype(jnp.bfloat16).at[64:128].set(0)
+    bits = lambda a: np.asarray(a).view(np.uint16)
+    for wrap in (lambda f: f, jax.checkpoint):
+      want, want_vjp = jax.vjp(jax.jit(wrap(split_then_glu)), fused)
+      got, got_vjp = jax.vjp(jax.jit(wrap(lambda f: moe.fused_glu(f, glu))), fused)
+      np.testing.assert_array_equal(bits(got), bits(want))
+      np.testing.assert_array_equal(bits(got_vjp(cotangent)[0]), bits(want_vjp(cotangent)[0]))
+
+
 if __name__ == "__main__":
   absltest.main()

From 906e157145b30c5e6aa18baa8ec43259c09f6de7 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:49 +0000
Subject: [PATCH 04/19] GDN seq-CP pass 1: states-only forward kernel, no
 output or chunk states

Pass 1 of sequence-sharded context parallelism (gdn_cp_mode=seq) only needs the local
final states and t_inv, but it ran the full fused Conv1D-GDN prefill kernel and dropped
the output and the chunk states. It now runs the kernel in a new GDNMode.PER_SEQ_STATES:
- chunked_gdn skips the output path (q.state, q.k^T and (q.k^T * gating).(u - ws)) and
  computes ws = w.state_prev on its own instead of as the first half of the merged
  [w; q].state_prev matmul. Each row's dot product is unchanged, so the states and
  t_inv are bit-identical.
- No output or chunk-states buffers are allocated, so the kernel writes neither to HBM.
  Per 397B call that drops the bf16[16384,64,128] output, its zero init, and the
  f32[256,64,128,128] chunk states (about 1.34 GB).
Fixed behaviour, no new flag: fused_conv1d_gdn and the PER_SEQ kernel are unchanged.
Pass 1 calls the new wrapper.fused_conv1d_gdn_states through
api._run_local_gdn_states_fwd, which reuses the pure-JAX forward on CPU. The kernel
input layout that both callers share moves to api._fused_kernel_inputs.

Verified locally:
- v6e model A/B (tiny Qwen3.5 with gdn_cp_mode=seq, gdn_cp_state=device, context=2)
  against the previous code: loss, logits and all 66 gradient leaves bit-identical.
- The kernel on the local v6e, base tree against this tree, at two geometries
  including the 397B per-device heads (16 k / 64 v heads, dim 12288), with bf16 and
  f32 activations, with and without packed segments and initial states: the PER_SEQ
  outputs (80 arrays) and the states-only conv state, recurrent state and t_inv
  (48 arrays) are bit-identical.
- The same comparison in Pallas TPU interpret mode on CPU with the v7x chip
  description (64 MiB VMEM): 40 + 24 arrays bit-identical, no unwritten (NaN) rows.
- v7x AOT of the 397B bench program (tpu7x-256): compiles. fwd_bwd temp 89.26 GiB and
  dropless temp 69.08 GiB, both unchanged. In fwd_bwd, 3 of the 9 PER_SEQ calls become
  fused_conv1d_gdn_per_seq_states_b1_c64 (outputs: conv state, recurrent state,
  t_inv); in the dropless program 6 of 12. Zero-init output copies go from 6 to 3
  in both.
- The new test test_run_local_gdn_states_fwd_matches_full_forward passes on CPU
  (pure-JAX path).
---
 src/maxtext/kernels/gdn/compute_gdn.py |  60 +++---
 src/maxtext/kernels/gdn/config.py      |   2 +
 src/maxtext/kernels/gdn/gdn_bwd/api.py | 175 +++++++++++++----
 src/maxtext/kernels/gdn/memory_ref.py  |  25 +--
 src/maxtext/kernels/gdn/wrapper.py     | 254 +++++++++++++++++++------
 tests/unit/gdn_bwd_pallas_test.py      |  51 +++++
 6 files changed, 438 insertions(+), 129 deletions(-)

diff --git a/src/maxtext/kernels/gdn/compute_gdn.py b/src/maxtext/kernels/gdn/compute_gdn.py
index 36d8c7bf62..45a6ab7a60 100644
--- a/src/maxtext/kernels/gdn/compute_gdn.py
+++ b/src/maxtext/kernels/gdn/compute_gdn.py
@@ -114,8 +114,8 @@ def chunked_gdn_per_seq(
     cfg: config.GDNConfig,
     seg_c_col: jax.Array | None = None,  # [1, chunk, 1]
     seg_aux_col: jax.Array | None = None,  # [1, chunk, 1]
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Perform chunked GDN over input [num_heads, chunk, head_dim]."""
+) -> tuple[jax.Array | None, jax.Array, jax.Array]:
+  """Perform chunked GDN over input [num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode."""
   # bf16 activations keep single-pass MXU matmuls even with an f32 compute dtype
   # (HIGHEST would triple the forward MXU cost); HIGHEST needs f32 activations and compute.
   dot_prec = (
@@ -236,23 +236,36 @@ def chunked_gdn_per_seq(
   # [num_v_heads, chunk, v_head_dim]
   u, w = jnp.split(merged_uw, [cfg.v_head_dim], axis=-1)
 
-  # [num_v_heads, chunk, kq_head_dim]
-  q_large_gating = q_repeat * gating_forward
-  # NOTE: Concatenate lhs with same rhs to leverage weight
-  # stationary architecture.
-  # [num_v_heads, 2 * chunk, kq_head_dim]
-  merged_w_q = jnp.concat([w, q_large_gating], axis=1)
-  # [num_v_heads, 2 * chunk, v_head_dim]
-  merged_ws_out_updated = jax.lax.dot(
-      merged_w_q,
-      state_prev,
-      dimension_numbers=(((2,), (1,)), ((0,), (0,))),
-      precision=dot_prec,
-      preferred_element_type=jnp.float32,
-  )
+  states_only = cfg.mode == config.GDNMode.PER_SEQ_STATES
+  if states_only:
+    # No output: only the w rows of the merged matmul below. Each row's dot product is the same,
+    # so ws is bit-identical.
+    # [num_v_heads, chunk, v_head_dim]
+    ws = jax.lax.dot(
+        w,
+        state_prev,
+        dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+        precision=dot_prec,
+        preferred_element_type=jnp.float32,
+    )
+  else:
+    # [num_v_heads, chunk, kq_head_dim]
+    q_large_gating = q_repeat * gating_forward
+    # NOTE: Concatenate lhs with same rhs to leverage weight
+    # stationary architecture.
+    # [num_v_heads, 2 * chunk, kq_head_dim]
+    merged_w_q = jnp.concat([w, q_large_gating], axis=1)
+    # [num_v_heads, 2 * chunk, v_head_dim]
+    merged_ws_out_updated = jax.lax.dot(
+        merged_w_q,
+        state_prev,
+        dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+        precision=dot_prec,
+        preferred_element_type=jnp.float32,
+    )
 
-  # NOTE: Splitting along non sublane/lane dim is free.
-  ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1)
+    # NOTE: Splitting along non sublane/lane dim is free.
+    ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1)
   ws = ws.astype(cfg.dtypes.compute)
 
   # [num_v_heads, chunk, v_head_dim]
@@ -273,6 +286,8 @@ def chunked_gdn_per_seq(
   # [num_v_heads, kq_head_dim, v_head_dim]
   state_updated = state_prev * gating_last
   state = state_updated + state_new
+  if states_only:
+    return None, state, t_inv
 
   # [num_kq_heads, chunk, chunk]
   out_qk = jax.lax.dot(
@@ -314,8 +329,8 @@ def chunked_gdn(
     a_log: jax.Array,
     dt_bias: jax.Array,
     cfg: config.GDNConfig,
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]."""
+) -> tuple[jax.Array | None, jax.Array, jax.Array]:
+  """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode."""
 
   mask_dtype = get_mask_dtype(cfg.dtypes.compute)
   iota = jax.lax.broadcasted_iota(mask_dtype, (cfg.seq_tile_size, 1, cfg.chunk_size, 1), 2)
@@ -373,10 +388,11 @@ def chunked_gdn(
         seg_c_col=seg_c_per_seq[idx] if cfg.has_seg_ids else None,
         seg_aux_col=seg_aux_per_seq[idx] if cfg.has_seg_ids else None,
     )
-    out_list.append(out.swapaxes(0, 1))
+    if out is not None:
+      out_list.append(out.swapaxes(0, 1))
     state_list.append(state)
     t_inv_list.append(t_inv)
-  out = jnp.stack(out_list, axis=0)
+  out = jnp.stack(out_list, axis=0) if out_list else None
   state = jnp.stack(state_list, axis=0)
   t_inv = jnp.stack(t_inv_list, axis=0)
   return out, state, t_inv
diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py
index 8ad3933295..bf0c005510 100644
--- a/src/maxtext/kernels/gdn/config.py
+++ b/src/maxtext/kernels/gdn/config.py
@@ -32,6 +32,8 @@ class GDNMode(enum.StrEnum):
 
   BATCHED = enum.auto()
   PER_SEQ = enum.auto()
+  # PER_SEQ that only carries the conv/recurrent states and t_inv: no output, no chunk states.
+  PER_SEQ_STATES = enum.auto()
 
   def get_seq_tile_size(self, tile_size: int) -> int:
     if self == GDNMode.BATCHED:
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index 7f7b2b824f..014ef8f98e 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -130,6 +130,50 @@ def decoupled_conv1d_gdn_bwd_kernel(
   )
 
 
+def _fused_kernel_inputs(
+    qkv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    conv_weight: jax.Array,
+    conv_state: Optional[jax.Array],
+    recurrent_state: Optional[jax.Array],
+    num_v_heads: int,
+    head_k_dim: int,
+    head_v_dim: int,
+    conv_kernel_size: int,
+) -> dict[str, jax.Array]:
+  """Lays out [batch, seq, ...] inputs for the fused Conv1D-GDN kernel: one equal-length prefill sequence per row."""
+  num_seqs, seq_len, dim_size = qkv.shape
+  has_init = (conv_state is not None) or (recurrent_state is not None)
+
+  if conv_state is None:
+    tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype)
+  elif conv_state.shape[0] == num_seqs:
+    tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0)))
+  else:
+    tokamax_conv_state = conv_state
+
+  if recurrent_state is None:
+    tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32)
+  elif recurrent_state.shape[0] == num_seqs:
+    tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0)))
+  else:
+    tokamax_recurrent_state = recurrent_state.astype(jnp.float32)
+
+  return {
+      "qkv": qkv.reshape(-1, dim_size),
+      "b": b.reshape(-1, b.shape[-1]),
+      "a": a.reshape(-1, a.shape[-1]),
+      "conv_state": tokamax_conv_state,
+      "recurrent_state": tokamax_recurrent_state,
+      "conv_weight": jnp.swapaxes(conv_weight, 0, 2),
+      "query_start_loc": jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32),
+      "state_indices": jnp.arange(1, num_seqs + 1, dtype=jnp.int32),
+      "distribution": jnp.array([0, 0, num_seqs], dtype=jnp.int32),
+      "seq_lens": jnp.full((num_seqs,), seq_len + (1 if has_init else 0), dtype=jnp.int32),
+  }
+
+
 def _run_local_gdn_decoupled_fwd(
     qkv: jax.Array,
     b: jax.Array,
@@ -158,7 +202,7 @@ def _run_local_gdn_decoupled_fwd(
     Optional[jax.Array],
 ]:
   """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU."""
-  batch_size, seq_len, dim_size = qkv.shape
+  batch_size, seq_len, _ = qkv.shape
   if seq_len % chunk_size != 0:
     raise ValueError(
         f"GDN forward kernel requires the local sequence length ({seq_len}) to be a multiple of"
@@ -215,53 +259,19 @@ def _run_local_gdn_decoupled_fwd(
     )
     return (out, states), t_inv, chunk_states
 
-  num_seqs = batch_size
   num_chunks = seq_len // chunk_size
-
-  qkv_flat = qkv.reshape(-1, dim_size)
-  b_flat = b.reshape(-1, b.shape[-1])
-  a_flat = a.reshape(-1, a.shape[-1])
-  tokamax_conv_weight = jnp.swapaxes(conv_weight, 0, 2)
-
-  has_init = (conv_state is not None) or (recurrent_state is not None)
-  query_start_loc = jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32)
-  state_indices = jnp.arange(1, num_seqs + 1, dtype=jnp.int32)
-  seq_lens = jnp.full((num_seqs,), seq_len + (1 if has_init else 0), dtype=jnp.int32)
-  distribution = jnp.array([0, 0, num_seqs], dtype=jnp.int32)
-
-  if conv_state is None:
-    tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype)
-  elif conv_state.shape[0] == num_seqs:
-    tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0)))
-  else:
-    tokamax_conv_state = conv_state
-
-  if recurrent_state is None:
-    tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32)
-  elif recurrent_state.shape[0] == num_seqs:
-    tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0)))
-  else:
-    tokamax_recurrent_state = recurrent_state.astype(jnp.float32)
-
   (
       core_attn_out_flat,
       (new_conv_state, new_recurrent_state),
       t_inv_raw,
       chunk_states_raw,
   ) = local_gdn_wrapper.fused_conv1d_gdn(
-      qkv=qkv_flat,
-      b=b_flat,
-      a=a_flat,
-      conv_state=tokamax_conv_state,
-      recurrent_state=tokamax_recurrent_state,
-      conv_weight=tokamax_conv_weight,
+      **_fused_kernel_inputs(
+          qkv, b, a, conv_weight, conv_state, recurrent_state, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size
+      ),
       conv_bias=conv_bias,
       a_log=a_log,
       dt_bias=dt_bias,
-      query_start_loc=query_start_loc,
-      state_indices=state_indices,
-      distribution=distribution,
-      seq_lens=seq_lens,
       n_kq=num_k_heads,
       n_v=num_v_heads,
       d_k=head_k_dim,
@@ -295,6 +305,90 @@ def _run_local_gdn_decoupled_fwd(
   )
 
 
+def _run_local_gdn_states_fwd(
+    qkv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    conv_weight: jax.Array,
+    conv_bias: Optional[jax.Array],
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    conv_state: Optional[jax.Array],
+    recurrent_state: Optional[jax.Array],
+    *,
+    num_k_heads: int,
+    num_v_heads: int,
+    head_k_dim: int,
+    head_v_dim: int,
+    conv_kernel_size: int,
+    chunk_size: int,
+    use_qk_norm_in_gdn: bool,
+    compute_dtype: jnp.dtype,
+    segment_ids: Optional[jax.Array] = None,
+    conv_halo_seg: Optional[jax.Array] = None,
+    init_seg: Optional[jax.Array] = None,
+) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]:
+  """_run_local_gdn_decoupled_fwd for callers that need only (states, t_inv), bit-identical to it.
+
+  On TPU the kernel skips the output and the chunk states; on CPU XLA drops the unused pure-JAX outputs.
+  """
+  if target_platform(qkv) == "cpu":
+    (_, states), t_inv, _ = _run_local_gdn_decoupled_fwd(
+        qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        num_k_heads=num_k_heads,
+        num_v_heads=num_v_heads,
+        head_k_dim=head_k_dim,
+        head_v_dim=head_v_dim,
+        conv_kernel_size=conv_kernel_size,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        compute_dtype=compute_dtype,
+        segment_ids=segment_ids,
+        conv_halo_seg=conv_halo_seg,
+        init_seg=init_seg,
+    )
+    return states, t_inv
+  batch_size, seq_len, _ = qkv.shape
+  if seq_len % chunk_size != 0:
+    raise ValueError(
+        f"GDN forward kernel requires the local sequence length ({seq_len}) to be a multiple of"
+        f" chunk_size ({chunk_size}); with sequence-sharded context parallelism this is seq_len / cp."
+    )
+  fallback_reason = pallas_unsupported_reason(head_k_dim=head_k_dim, head_v_dim=head_v_dim, chunk_size=chunk_size)
+  if fallback_reason is not None:
+    raise ValueError(f"GDN Pallas TPU forward kernel does not support {fallback_reason}.")
+
+  (new_conv_state, new_recurrent_state), t_inv_raw = local_gdn_wrapper.fused_conv1d_gdn_states(
+      **_fused_kernel_inputs(
+          qkv, b, a, conv_weight, conv_state, recurrent_state, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size
+      ),
+      conv_bias=conv_bias,
+      a_log=a_log,
+      dt_bias=dt_bias,
+      n_kq=num_k_heads,
+      n_v=num_v_heads,
+      d_k=head_k_dim,
+      d_v=head_v_dim,
+      kernel_size=conv_kernel_size,
+      compute_precision=jnp.dtype(compute_dtype),
+      mixed_tile_size=chunk_size,
+      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+      init_seg=init_seg,
+  )
+  t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, seq_len // chunk_size, num_v_heads, chunk_size, chunk_size)
+  return (new_conv_state[1:].astype(qkv.dtype), new_recurrent_state[1:].astype(jnp.float32)), t_inv
+
+
 def _is_cp_active(cp_axis_name: str | tuple[str, ...] | None) -> bool:
   if cp_axis_name is None:
     return False
@@ -380,8 +474,9 @@ def _run_cp_gdn_decoupled_fwd_impl(
   conv_halo = checkpoint_name(conv_halo, "gdn_cp_conv_halo")
   zero_rs = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32)
 
-  # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local
-  (_, states_pass1), t_inv, _ = _run_local_gdn_decoupled_fwd(
+  # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local. Its output
+  # is not needed, so the kernel neither computes it nor writes it or the chunk states.
+  states_pass1, t_inv = _run_local_gdn_states_fwd(
       qkv,
       b,
       a,
diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py
index f29d79d81d..481ca62066 100644
--- a/src/maxtext/kernels/gdn/memory_ref.py
+++ b/src/maxtext/kernels/gdn/memory_ref.py
@@ -496,14 +496,14 @@ def create_allocs(
     qkv_ref: jax.Array,
     b_ref: jax.Array,
     a_ref: jax.Array,
-    out_ref: jax.Array,
+    out_ref: jax.Array | None,
     conv_state_ref: jax.Array,
     recurrent_state_ref: jax.Array,
     cfg: config.GDNConfig,
     t_inv_ref: jax.Array | None = None,
     chunk_states_ref: jax.Array | None = None,
 ) -> tuple[Any, ...]:
-  """Creates all Pallas double-buffered allocations for GDN kernel."""
+  """Creates all Pallas double-buffered allocations for GDN kernel (no output buffers without out_ref)."""
   qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size)
   ba_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.aligned_num_v_heads)
 
@@ -543,15 +543,6 @@ def create_allocs(
   b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref)
   a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref)
 
-  out_alloc = OutBufferedRef.output(
-      spec=block_spec_partial(block_shape=out_shape),
-      dtype_or_type=out_ref,
-      buffer_count=pipeline_mode.buffer_count,
-      use_lookahead=pipeline_mode.use_lookahead,
-      cfg=cfg,
-      metadata_ref=metadata_ref,
-  )
-
   conv_spec = block_spec_partial(block_shape=conv_shape)
   recurrent_spec = block_spec_partial(block_shape=recurrent_shape)
   state_buffered_partial = functools.partial(
@@ -570,9 +561,19 @@ def create_allocs(
       a_alloc,
       conv_alloc,
       recurrent_alloc,
-      out_alloc,
   ]
 
+  if out_ref is not None:
+    out_alloc = OutBufferedRef.output(
+        spec=block_spec_partial(block_shape=out_shape),
+        dtype_or_type=out_ref,
+        buffer_count=pipeline_mode.buffer_count,
+        use_lookahead=pipeline_mode.use_lookahead,
+        cfg=cfg,
+        metadata_ref=metadata_ref,
+    )
+    allocs.append(out_alloc)
+
   if t_inv_ref is not None:
     t_inv_shape = (
         cfg.seq_tile_size,
diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py
index 7b951c7cee..f81d4a9c82 100644
--- a/src/maxtext/kernels/gdn/wrapper.py
+++ b/src/maxtext/kernels/gdn/wrapper.py
@@ -100,9 +100,9 @@ def inner_kernel(
     a_slot_ref: jax.Ref,  # [seq, chunk, 1, num_v_heads]
     conv_state_slot_ref: jax.Ref,  # [seq, prev_kernel_size, 1, dim_size]
     recurrent_slot_ref: jax.Ref,  # [seq, num_v_heads, kq_head, v_head]
-    # Outputs.
-    out_slot_ref: jax.Array,  # [seq * chunk, num_v_heads, v_head]
-    t_inv_slot_ref: jax.Array,  # [seq, num_v_heads, chunk, chunk]
+    # Outputs, then scratches:
+    # out [seq * chunk, num_v_heads, v_head] (not in PER_SEQ_STATES), t_inv [seq, num_v_heads, chunk, chunk],
+    # chunk_states [seq, num_v_heads, kq_head, v_head] (PER_SEQ only).
     *args,
     cfg: config.GDNConfig,
     **kwargs,
@@ -114,18 +114,18 @@ def inner_kernel(
   kernel invokes vmem_ldst to pre-processes data needed for compute and
   invokes compute_conv1d and compute_gdn for actual compute.
   """
-  if cfg.mode == config.GDNMode.PER_SEQ:
-    chunk_states_slot_ref = args[0]
-    metadata_ref = args[1]
-    weights_ref = args[2]
-    carry_conv_scratch_ref = args[3] if len(args) > 3 else None
-    carry_recurrent_scratch_ref = args[4] if len(args) > 4 else None
+  if cfg.mode == config.GDNMode.PER_SEQ_STATES:
+    out_slot_ref = chunk_states_slot_ref = None
+    t_inv_slot_ref, *args = args
+  elif cfg.mode == config.GDNMode.PER_SEQ:
+    out_slot_ref, t_inv_slot_ref, chunk_states_slot_ref, *args = args
   else:
+    out_slot_ref, t_inv_slot_ref, *args = args
     chunk_states_slot_ref = None
-    metadata_ref = args[0]
-    weights_ref = args[1]
-    carry_conv_scratch_ref = args[2] if len(args) > 2 else None
-    carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None
+  metadata_ref = args[0]
+  weights_ref = args[1]
+  carry_conv_scratch_ref = args[2] if len(args) > 2 else None
+  carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None
 
   p_id = pl.program_id(0)
 
@@ -216,7 +216,8 @@ def inner_kernel(
     )
 
   # Store output, recurrent, and t_inv to vmem.
-  out_slot_ref[...] = out.astype(out_slot_ref.dtype)
+  if out_slot_ref is not None:
+    out_slot_ref[...] = out.astype(out_slot_ref.dtype)
   recurrent_slot_ref[...] = new_recurrent_state.astype(recurrent_slot_ref.dtype)
   t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype)
   if chunk_states_slot_ref is not None:
@@ -236,11 +237,8 @@ def outer_kernel(
     recurrent_state_ref: jax.Array,
     _: jax.Array,
     weights_ref: memory_ref.WeightRefs,
-    # Outputs.
-    out_ref: jax.Array,
-    conv_state_out_ref: jax.Array,
-    recurrent_state_out_ref: jax.Array,
-    t_inv_ref: jax.Array,
+    # Outputs: out (not in PER_SEQ_STATES), conv_state, recurrent_state, t_inv, chunk_states (PER_SEQ only).
+    # The state outputs alias the state inputs, which the pipeline updates in place.
     *args,
     carry_conv_scratch_ref: jax.Array | None = None,
     carry_recurrent_scratch_ref: jax.Array | None = None,
@@ -248,7 +246,11 @@ def outer_kernel(
     **kwargs,
 ) -> None:
   """Setup memory allocations and emit pipeline for running inner_kernel."""
-  del conv_state_out_ref, recurrent_state_out_ref
+  if cfg.mode == config.GDNMode.PER_SEQ_STATES:
+    out_ref = None
+    _, _, t_inv_ref, *args = args
+  else:
+    out_ref, _, _, t_inv_ref, *args = args
 
   chunk_states_ref = args[0] if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) else None
 
@@ -269,16 +271,11 @@ def outer_kernel(
   a_alloc = allocs[2]
   conv_alloc = allocs[3]
   recurrent_alloc = allocs[4]
-  out_alloc = allocs[5]
-  t_inv_alloc = allocs[6]
-  chunk_states_alloc = allocs[7] if len(allocs) > 7 else None
+  # [out,] t_inv[, chunk_states], in the order of the pipelined output refs below.
+  out_specs = [alloc.spec for alloc in allocs[5:]]
 
   num_tiles = metadata_ref.num_tiles[...]
 
-  out_specs = [out_alloc.spec, t_inv_alloc.spec]
-  if chunk_states_alloc is not None:
-    out_specs.append(chunk_states_alloc.spec)
-
   pipeline_func = pltpu.emit_pipeline(
       body=functools.partial(
           inner_kernel,
@@ -297,9 +294,7 @@ def outer_kernel(
 
   @pl.with_scoped(allocations=allocs)
   def _run(allocations):
-    out_args = [out_ref, t_inv_ref]
-    if chunk_states_ref is not None:
-      out_args.append(chunk_states_ref)
+    out_args = [ref for ref in (out_ref, t_inv_ref, chunk_states_ref) if ref is not None]
     pipeline_func(
         qkv_ref,
         b_ref,
@@ -320,23 +315,7 @@ def _run(allocations):
   _run()
 
 
-@jax.jit(
-    donate_argnames=("conv_state", "recurrent_state"),
-    static_argnames=(
-        "n_kq",
-        "n_v",
-        "d_k",
-        "d_v",
-        "kernel_size",
-        "decode_tile_size",
-        "mixed_tile_size",
-        "zero_initialize_out",
-        "compute_precision",
-        "is_prefill_only",
-        "use_qk_norm_in_gdn",
-    ),
-)
-def fused_conv1d_gdn(
+def _fused_conv1d_gdn(
     qkv: jax.Array,  # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size]
     b: jax.Array,  # [batch_size, n_v]
     a: jax.Array,  # [batch_size, n_v]
@@ -365,8 +344,9 @@ def fused_conv1d_gdn(
     segment_ids: jax.Array | None = None,
     conv_halo_seg: jax.Array | None = None,
     init_seg: jax.Array | None = None,
-) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]:
-  """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states)."""
+    prefill_mode: config.GDNMode = config.GDNMode.PER_SEQ,
+) -> tuple[jax.Array | None, tuple[jax.Array, jax.Array], jax.Array, jax.Array | None]:
+  """Implements fused_conv1d_gdn; prefill_mode=PER_SEQ_STATES returns no out and no chunk_states."""
   act_in_dtype = qkv.dtype
   act_out_dtype = qkv.dtype
   conv_out_dtype = conv_state.dtype
@@ -534,7 +514,7 @@ def call_kernel(
     input_output_aliases = {len(metadata_obj) + 3: 1, len(metadata_obj) + 4: 2}
     out_shape = cfg.get_out_shape()
 
-    if in_act is None and zero_initialize_out:
+    if in_act is None and zero_initialize_out and mode != config.GDNMode.PER_SEQ_STATES:
       in_act = jnp.zeros_like(out_shape)
     if in_act is not None:
       out_shape = in_act
@@ -547,7 +527,12 @@ def call_kernel(
         cfg.dtypes.compute,
     )
 
-    if mode == config.GDNMode.PER_SEQ:
+    if mode == config.GDNMode.PER_SEQ_STATES:
+      # No output, so the states move up to outputs 0 and 1.
+      input_output_aliases = {len(metadata_obj) + 3: 0, len(metadata_obj) + 4: 1}
+      out_shape_tuple = (in_conv_state, in_recurrent_state, t_inv_shape)
+      out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec)
+    elif mode == config.GDNMode.PER_SEQ:
       chunk_states_shape = jax.ShapeDtypeStruct(
           (num_chunks, cfg.num_v_heads, cfg.kq_head_dim, cfg.v_head_dim),
           cfg.dtypes.compute,
@@ -611,11 +596,15 @@ def call_kernel(
     out_conv_state = conv_state
     out_recurrent_state = recurrent_state
 
-  out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel(
-      out_conv_state, out_recurrent_state, out_act, config.GDNMode.PER_SEQ
-  )
+  if prefill_mode == config.GDNMode.PER_SEQ_STATES:
+    out_conv_state, out_recurrent_state, t_inv = call_kernel(out_conv_state, out_recurrent_state, None, prefill_mode)
+    out_act = chunk_states = None
+  else:
+    out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel(
+        out_conv_state, out_recurrent_state, out_act, prefill_mode
+    )
+    out_act = out_act.reshape(padded_batch_size, -1)[:batch_size]
 
-  out_act = out_act.reshape(padded_batch_size, -1)[:batch_size]
   out_conv_state = out_conv_state.astype(conv_out_dtype)
   out_conv_state = out_conv_state.reshape(conv_state_shape)
   if has_seg_ids and segment_ids is not None:
@@ -630,3 +619,158 @@ def call_kernel(
   out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype)
 
   return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states
+
+
+@jax.jit(
+    donate_argnames=("conv_state", "recurrent_state"),
+    static_argnames=(
+        "n_kq",
+        "n_v",
+        "d_k",
+        "d_v",
+        "kernel_size",
+        "decode_tile_size",
+        "mixed_tile_size",
+        "zero_initialize_out",
+        "compute_precision",
+        "is_prefill_only",
+        "use_qk_norm_in_gdn",
+    ),
+)
+def fused_conv1d_gdn(
+    qkv: jax.Array,  # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size]
+    b: jax.Array,  # [batch_size, n_v]
+    a: jax.Array,  # [batch_size, n_v]
+    conv_state: jax.Array,  # [num_seqs + 1, kernel_size - 1, dim_size]
+    recurrent_state: jax.Array,  # [num_seqs + 1, nv, dk, dv]
+    conv_weight: jax.Array,  # [kernel_size - 1, dim_size]
+    conv_bias: jax.Array | None,  # [dim_size]
+    a_log: jax.Array,  # [n_v]
+    dt_bias: jax.Array,  # [n_v]
+    query_start_loc: jax.Array,  # [num_seqs + 1]
+    state_indices: jax.Array,  # [num_seqs]
+    distribution: jax.Array,  # [3]
+    seq_lens: jax.Array,  # [num_seqs]
+    *,
+    n_kq: int,
+    n_v: int,
+    d_k: int,
+    d_v: int,
+    kernel_size: int,
+    zero_initialize_out: bool = True,
+    compute_precision: jnp.dtype = jnp.float32.dtype,
+    decode_tile_size: int | None = None,
+    mixed_tile_size: int | None = None,
+    is_prefill_only: bool = False,
+    use_qk_norm_in_gdn: bool = True,
+    segment_ids: jax.Array | None = None,
+    conv_halo_seg: jax.Array | None = None,
+    init_seg: jax.Array | None = None,
+) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]:
+  """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states)."""
+  return _fused_conv1d_gdn(
+      qkv,
+      b,
+      a,
+      conv_state,
+      recurrent_state,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      query_start_loc,
+      state_indices,
+      distribution,
+      seq_lens,
+      n_kq=n_kq,
+      n_v=n_v,
+      d_k=d_k,
+      d_v=d_v,
+      kernel_size=kernel_size,
+      zero_initialize_out=zero_initialize_out,
+      compute_precision=compute_precision,
+      decode_tile_size=decode_tile_size,
+      mixed_tile_size=mixed_tile_size,
+      is_prefill_only=is_prefill_only,
+      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+      init_seg=init_seg,
+  )
+
+
+@jax.jit(
+    donate_argnames=("conv_state", "recurrent_state"),
+    static_argnames=(
+        "n_kq",
+        "n_v",
+        "d_k",
+        "d_v",
+        "kernel_size",
+        "mixed_tile_size",
+        "compute_precision",
+        "use_qk_norm_in_gdn",
+    ),
+)
+def fused_conv1d_gdn_states(
+    qkv: jax.Array,  # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size]
+    b: jax.Array,  # [batch_size, n_v]
+    a: jax.Array,  # [batch_size, n_v]
+    conv_state: jax.Array,  # [num_seqs + 1, kernel_size - 1, dim_size]
+    recurrent_state: jax.Array,  # [num_seqs + 1, nv, dk, dv]
+    conv_weight: jax.Array,  # [kernel_size - 1, dim_size]
+    conv_bias: jax.Array | None,  # [dim_size]
+    a_log: jax.Array,  # [n_v]
+    dt_bias: jax.Array,  # [n_v]
+    query_start_loc: jax.Array,  # [num_seqs + 1]
+    state_indices: jax.Array,  # [num_seqs]
+    distribution: jax.Array,  # [3]
+    seq_lens: jax.Array,  # [num_seqs]
+    *,
+    n_kq: int,
+    n_v: int,
+    d_k: int,
+    d_v: int,
+    kernel_size: int,
+    mixed_tile_size: int,
+    compute_precision: jnp.dtype = jnp.float32.dtype,
+    use_qk_norm_in_gdn: bool = True,
+    segment_ids: jax.Array | None = None,
+    conv_halo_seg: jax.Array | None = None,
+    init_seg: jax.Array | None = None,
+) -> tuple[tuple[jax.Array, jax.Array], jax.Array]:
+  """Prefill-only fused_conv1d_gdn that skips the output and chunk states, returning (states, t_inv).
+
+  For callers that only carry the state forward, such as pass 1 of sequence-sharded context
+  parallelism: the kernel neither computes the output nor writes it or the chunk states to HBM. The
+  states and t_inv are bit-identical to fused_conv1d_gdn's.
+  """
+  _, states, t_inv, _ = _fused_conv1d_gdn(
+      qkv,
+      b,
+      a,
+      conv_state,
+      recurrent_state,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      query_start_loc,
+      state_indices,
+      distribution,
+      seq_lens,
+      n_kq=n_kq,
+      n_v=n_v,
+      d_k=d_k,
+      d_v=d_v,
+      kernel_size=kernel_size,
+      compute_precision=compute_precision,
+      mixed_tile_size=mixed_tile_size,
+      is_prefill_only=True,
+      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+      init_seg=init_seg,
+      prefill_mode=config.GDNMode.PER_SEQ_STATES,
+  )
+  return states, t_inv
diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py
index ee91c2b33e..8e94bb2396 100644
--- a/tests/unit/gdn_bwd_pallas_test.py
+++ b/tests/unit/gdn_bwd_pallas_test.py
@@ -27,6 +27,7 @@
 
 from maxtext.configs import types as config_types
 from maxtext.kernels.gdn import gdn_bwd_pallas
+from maxtext.kernels.gdn.gdn_bwd import api as gdn_api
 from maxtext.models import qwen3
 from maxtext.utils import maxtext_utils
 
@@ -726,6 +727,56 @@ def test_run_local_gdn_decoupled_fwd_returns_cached_chunk_states(self):
     np.testing.assert_allclose(chunk_states, exp_chunk_states, rtol=tol, atol=tol)
     np.testing.assert_allclose(t_inv, exp_t_inv, rtol=tol, atol=tol)
 
+  def test_run_local_gdn_states_fwd_matches_full_forward(self):
+    """The states-only forward (CP pass 1) returns the full forward's states and t_inv bit for bit.
+
+    On TPU this compares the PER_SEQ_STATES kernel, which skips the output, with the PER_SEQ kernel.
+    """
+    is_tpu = jax.default_backend() == "tpu"
+    batch_size = 2
+    chunk_size = 64 if is_tpu else 32
+    seq_len = 4 * chunk_size
+    num_k_heads = 2
+    num_v_heads = 4
+    head_k_dim = 128 if is_tpu else 64
+    head_v_dim = 128 if is_tpu else 64
+    conv_kernel_size = 4
+    dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim
+
+    qkv, b, a, conv_weight, conv_bias, a_log, dt_bias = _init_fwd_inputs(
+        jax.random.PRNGKey(7), batch_size, seq_len, dim_size, num_v_heads, conv_kernel_size
+    )
+    qkv, b, a = (x.astype(jnp.bfloat16) for x in (qkv, b, a))
+    # Pass-1 inputs: the conv halo, a zero recurrent state, and packed rows with trailing padding.
+    conv_state = jax.random.normal(jax.random.PRNGKey(8), (batch_size, conv_kernel_size - 1, dim_size), jnp.bfloat16)
+    recurrent_state = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), jnp.float32)
+    seg = np.zeros((batch_size, seq_len), np.int32)
+    seg[0, : seq_len // 3] = 1
+    seg[0, seq_len // 3 : seq_len - 10] = 2
+    seg[1, : seq_len - 40] = 1
+    segment_ids, conv_halo_seg, init_seg = gdn_api._local_segment_metadata(jnp.asarray(seg), True, conv_kernel_size)
+    args = (qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state)
+    kwargs = {
+        "num_k_heads": num_k_heads,
+        "num_v_heads": num_v_heads,
+        "head_k_dim": head_k_dim,
+        "head_v_dim": head_v_dim,
+        "conv_kernel_size": conv_kernel_size,
+        "chunk_size": chunk_size,
+        "use_qk_norm_in_gdn": True,
+        "compute_dtype": jnp.float32,
+        "segment_ids": segment_ids,
+        "conv_halo_seg": conv_halo_seg,
+        "init_seg": init_seg,
+    }
+
+    (_, exp_states), exp_t_inv, _ = gdn_api._run_local_gdn_decoupled_fwd(*args, **kwargs)
+    states, t_inv = gdn_api._run_local_gdn_states_fwd(*args, **kwargs)
+
+    for got, exp in zip((*states, t_inv), (*exp_states, exp_t_inv)):
+      self.assertEqual(got.dtype, exp.dtype)
+      np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(exp, np.float32))
+
   def test_decoupled_conv1d_gdn_kernel_gradient_with_initial_states(self):
     """Verifies custom VJP gradients when initial conv_state and recurrent_state are provided."""
     is_tpu = jax.default_backend() == "tpu"

From 464767398e62b1a616f157866b8458036e4417c6 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 05/19] moe: run the router's top_k only on tokens that need it
 under forced routing

get_topk sorted every token's router probabilities for auto_indices, but
under router replay only tokens that are neither forced nor padding read
them (the rest are replaced by the forced ids or masked). When at most 256
tokens need them, gather those rows (jnp.nonzero(size=256)), run top_k on
the [256, num_experts] block and scatter the ids back; a lax.cond falls
back to the full top_k when more need it. Blocks of 256 tokens or fewer
keep the plain top_k.

top_k is per row, so every id that is read is unchanged: outputs and
gradients are bit-identical. Checked against the previous get_topk (weights,
ids, jax.grad) on CPU and on v6e at 8192x512 top-10, both branches and the
256/257 boundary, and with the tiny Qwen3.5 A/B run.
---
 src/maxtext/layers/moe.py | 38 ++++++++++++++++++++++++++++++++++++--
 tests/unit/moe_test.py    | 22 ++++++++++++++++++++++
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 2c8c24d6d7..8bffab758c 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -417,6 +417,37 @@ def fold_padding_into_forced_routing(forced_routed_experts, token_mask, num_expe
   return jnp.where(is_pad, jnp.asarray(PAD_TOKEN_ROUTED_EXPERT, base.dtype), base)
 
 
+def _top_k_indices_where_needed(probs, needed, k):
+  """`jax.lax.top_k(probs, k)[1]` on the tokens where `needed` is set; other tokens may get any ids.
+
+  Under router replay almost every token is forced or padding. When at most 256 tokens need the
+  router's own choice, top_k sorts only those rows (gathered, then scattered back) instead of every
+  row; otherwise it sorts every row. top_k is per row, so the needed rows are exact either way.
+
+  Args:
+    probs: `[..., num_experts]` router probabilities.
+    needed: `[..., 1]` True for tokens whose ids are read.
+    k: top_k.
+  """
+  max_rows = 256
+  flat = probs.reshape(-1, probs.shape[-1])
+  if flat.shape[0] <= max_rows:
+    return jax.lax.top_k(probs, k)[1]
+  needed = needed.reshape(-1)
+
+  def compact(flat, needed):
+    # Fill entries repeat row 0, which only writes row 0's own top_k to row 0.
+    rows = jnp.nonzero(needed, size=max_rows)[0]
+    return jnp.zeros((flat.shape[0], k), jnp.int32).at[rows].set(jax.lax.top_k(flat[rows], k)[1])
+
+  def full(flat, needed):
+    del needed
+    return jax.lax.top_k(flat, k)[1]
+
+  indices = jax.lax.cond(jnp.sum(needed) <= max_rows, compact, full, flat, needed)
+  return indices.reshape(probs.shape[:-1] + (k,))
+
+
 def _top_2_in_group_sum(scores_grouped: jax.Array) -> jax.Array:
   """Sums the two largest scores along the last axis, accumulating in float32.
 
@@ -1169,7 +1200,6 @@ def get_topk(
     valid_token_mask = None
     if forced_routed_experts is not None:
       router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1)
-      _, auto_indices = jax.lax.top_k(router_probs, self.num_experts_per_tok)
       slot_valid = valid_expert_mask(forced_routed_experts, self.num_experts)
       # A token has real forced routing only if its slots are within [0, num_experts).
       # Tokens with UNSET_ROUTED_EXPERT (-1) slots must fall back to auto_indices
@@ -1177,7 +1207,6 @@ def get_topk(
       sorted_experts = jnp.sort(forced_routed_experts, axis=-1)
       has_duplicate = jnp.any(sorted_experts[..., 1:] == sorted_experts[..., :-1], axis=-1, keepdims=True)
       token_is_forced = jnp.logical_and(jnp.all(slot_valid, axis=-1, keepdims=True), jnp.logical_not(has_duplicate))
-      top_k_indices = jnp.where(token_is_forced, forced_routed_experts, auto_indices)
       if self.config.moe_exclude_padding_from_routing:
         # A padding token (every slot PAD_TOKEN_ROUTED_EXPERT) is never forced (its slots are not
         # valid ids) and gets no valid slot, so its weights are zeroed below and its ids become -1.
@@ -1185,6 +1214,11 @@ def get_topk(
         valid_token_mask = jnp.where(token_is_forced, slot_valid, jnp.logical_not(pad_token))
       else:
         valid_token_mask = jnp.where(token_is_forced, slot_valid, jnp.ones_like(slot_valid, dtype=jnp.bool_))
+      # Only tokens that are neither forced nor padding read auto_indices: every other id is replaced
+      # or masked below.
+      uses_auto = jnp.logical_and(jnp.logical_not(token_is_forced), jnp.any(valid_token_mask, axis=-1, keepdims=True))
+      auto_indices = _top_k_indices_where_needed(router_probs, uses_auto, self.num_experts_per_tok)
+      top_k_indices = jnp.where(token_is_forced, forced_routed_experts, auto_indices)
       gather_indices = jnp.where(valid_token_mask, top_k_indices, 0)
       if self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4:
         top_k_weights = jnp.take_along_axis(router_probs, gather_indices, axis=-1).astype(self.dtype)
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
index c2e0946e9a..67fb571690 100644
--- a/tests/unit/moe_test.py
+++ b/tests/unit/moe_test.py
@@ -267,6 +267,28 @@ def test_init(self):
     self.model.init({"params": self.rng, "dropout": self.rng}, x)
 
 
+class TopKIndicesWhereNeededTest(unittest.TestCase):
+
+  def test_matches_top_k_on_needed_tokens(self):
+    """Both branches give jax.lax.top_k's ids on the needed tokens: 256 compact, 257 sorts every row."""
+    # Coarse logits make ties common; top_k breaks them by the lower index.
+    logits = jnp.round(jax.random.normal(jax.random.PRNGKey(0), (2, 512, 32)) * 4) / 4
+    probs = jax.nn.softmax(logits, axis=-1)
+    expected = np.asarray(jax.lax.top_k(probs, 4)[1])
+    top_k_where_needed = jax.jit(moe._top_k_indices_where_needed, static_argnums=2)  # pylint: disable=protected-access
+    order = jax.random.permutation(jax.random.PRNGKey(1), 2 * 512)
+    for num_needed in (256, 257):
+      needed = jnp.zeros(2 * 512, jnp.bool_).at[order[:num_needed]].set(True)
+      got = np.asarray(top_k_where_needed(probs, needed.reshape(2, 512, 1), 4)).reshape(-1, 4)
+      needed = np.asarray(needed)
+      np.testing.assert_array_equal(got[needed], expected.reshape(-1, 4)[needed])
+      if num_needed <= 256:
+        # The compact branch leaves the other tokens (but row 0) at 0.
+        np.testing.assert_array_equal(got[1:][~needed[1:]], 0)
+      else:
+        np.testing.assert_array_equal(got, expected.reshape(-1, 4))
+
+
 class DeepSeekRoutingTest(unittest.TestCase):
 
   def setUp(self):

From cd9ce6f582dbe1d0cd65d752d586bb198d4e1ad4 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 06/19] GDN backward kernel: read and write head groups in
 place, no group-major copies

pallas_gdn_bwd_kernel copied qkv_conv, do, chunk_states and t_inv into
group-major (B, G, C, ...) layouts before the kernel and transposed dy_conv
back after it. The kernel now reads head group g from the token-major arrays:
q, k and v as column blocks of qkv_conv (v is sliced out only when q + k is
not a whole number of v blocks), do / chunk_states / t_inv through
group-indexed reshape views. It writes dq, dk and dv as column blocks of three
token-major outputs. The kernel math is unchanged; head_tile stays 32 (64
would need ~86 MiB VMEM on v7x).

397B bench program (tpu7x-256 AOT; per device 64 v heads in 2 groups), per
GDN layer backward, 3 -> 0 each: copies to f32[1,256,2,32,128,128]
(chunk_states), f32[1,256,2,32,64,64] (t_inv), bf16[1,256,64,2,32,128] (do),
the qkv concat pad_maximum_fusion.*.refs_0_, and the dy_conv transposes back
(bf16[1,256,64,2,4096] / [1,256,64,2,1024]). jit_fwd_bwd copies 718 -> 694.

Verified locally:
- Kernel outputs bit-identical to the previous kernel: local v6e, 34/34
  outputs over 6 geometries (397B layer with 2 groups, packed segments and
  dht/dh0; no segments; head_tile=64; f32 with 4 groups; head_tile=4; 32
  heads); Pallas interpret on CPU with the v7x chip description (64 MiB
  VMEM tiling), 23/23 outputs over 4 geometries.
- New tpu_only test_gdn_kernel_bwd_group_blocks_match_per_group_calls
  (2-group call vs per-group calls, bitwise) passes on the v6e; the CPU
  interpret tests around the kernel pass.
- 397B AOT for tpu7x-256 in the bench image: compiles (Mosaic for v7x);
  same 316 collectives (types and shapes) as the base; fwd_bwd temp
  89.26 -> 90.07 GiB (HBM 77.73 -> 78.13 GiB), dropless temp
  69.08 -> 68.24 GiB (HBM 65.48 -> 64.56 GiB).
- Tiny Qwen3.5 A/B on the v6e (gdn_cp_mode=seq): loss, logits and the
  embedding, norm, A_log, dt_bias, conv1d, router and routed-expert
  gradients are bit-identical, but 13 dense weight-gradient leaves differ
  by bf16-ulp amounts (max abs 3.1e-5; both programs deterministic). One
  of them is the last layer's attention out kernel, which no GDN backward
  feeds. Cause (HLO dumps of both): XLA's all-reduce combiner groups the
  tiny model's bf16 gradient all-reduces (cast_grads_after_all_reduce)
  differently, moving one bf16[2,256,8] in_proj_ba gradient into the other
  combined all-reduce, which changes the cross-device summation order of
  the leaves after it. The 397B program's collectives are unchanged (above).
---
 .../kernels/gdn/gdn_bwd/bwd_memory_ref.py     |  80 +++++++----
 .../gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py      | 126 ++++++++---------
 tests/unit/gdn_bwd_pallas_test.py             | 130 ++++++++++++++++--
 3 files changed, 236 insertions(+), 100 deletions(-)

diff --git a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
index 3d379e5cf7..f52565aa50 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
@@ -22,15 +22,23 @@
 def make_bwd_block_specs(
     num_chunks: int,
     chunk_size: int,
-    dim_size: int,
+    qk_tile_size: int,
+    v_tile_size: int,
     num_v_heads: int,
     kq_head_dim: int,
     v_head_dim: int,
+    k_block_offset: int,
+    v_block_offset: int,
     padded_num_v_heads: int | None = None,
     has_dht: bool = False,
     has_dh0: bool = False,
 ) -> Tuple[list[pl.BlockSpec], list[pl.BlockSpec], int, int]:
-  """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward."""
+  """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward.
+
+  Head group g of q/k/v and dq/dk/dv is a column block of a token-major array (the k and v blocks of qkv_conv start
+  at block k_block_offset / v_block_offset), and do / chunk_states / t_inv are indexed by group, so the kernel reads
+  and writes them in place instead of through group-major copies.
+  """
   if padded_num_v_heads is None:
     padded_num_v_heads = ((num_v_heads + 127) // 128) * 128
 
@@ -38,47 +46,57 @@ def rc(c):
     return num_chunks - 1 - c
 
   in_specs = [
-      # 0: qkv_conv [B, G, C, chunk_size, dim_size]
+      # 0: q [B, C, chunk_size, dim_size], columns of group g
       pl.BlockSpec(
-          (None, None, None, chunk_size, dim_size),
-          lambda b, g, c: (b, g, rc(c), 0, 0),
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
+      ),
+      # 1: k [B, C, chunk_size, dim_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, k_block_offset + g),
+      ),
+      # 2: v [B, C, chunk_size, dim_size or v_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, v_tile_size),
+          lambda b, g, c: (b, rc(c), 0, v_block_offset + g),
       ),
-      # 1: b [B, G, C, chunk_size, padded_num_v_heads]
+      # 3: b [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 2: a [B, G, C, chunk_size, padded_num_v_heads]
+      # 4: a [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 3: do [B, G, C, chunk_size, num_v_heads, v_head_dim]
+      # 5: do [B, C, chunk_size, G, num_v_heads, v_head_dim]
       pl.BlockSpec(
-          (None, None, None, chunk_size, num_v_heads, v_head_dim),
-          lambda b, g, c: (b, g, rc(c), 0, 0, 0),
+          (None, None, chunk_size, None, num_v_heads, v_head_dim),
+          lambda b, g, c: (b, rc(c), 0, g, 0, 0),
       ),
-      # 4: chunk_states [B, G, C, num_v_heads, kq_head_dim, v_head_dim]
+      # 6: chunk_states [B, C, G, num_v_heads, kq_head_dim, v_head_dim]
       pl.BlockSpec(
           (None, None, None, num_v_heads, kq_head_dim, v_head_dim),
-          lambda b, g, c: (b, g, rc(c), 0, 0, 0),
+          lambda b, g, c: (b, rc(c), g, 0, 0, 0),
       ),
-      # 5: t_inv [B, G, C, num_v_heads, chunk_size, chunk_size]
+      # 7: t_inv [B, C, G, num_v_heads, chunk_size, chunk_size]
       pl.BlockSpec(
           (None, None, None, num_v_heads, chunk_size, chunk_size),
-          lambda b, g, c: (b, g, rc(c), 0, 0, 0),
+          lambda b, g, c: (b, rc(c), g, 0, 0, 0),
       ),
-      # 6: a_log [B, G, 1, padded_num_v_heads]
+      # 8: a_log [B, G, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, 0, 0),
       ),
-      # 7: dt_bias [B, G, 1, padded_num_v_heads]
+      # 9: dt_bias [B, G, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, 0, 0),
       ),
-      # 8: reset [B, G, C, 1, 128]
+      # 10: reset [B, G, C, 1, 128]
       pl.BlockSpec(
           (None, None, None, 1, 128),
           lambda b, g, c: (b, g, rc(c), 0, 0),
@@ -86,7 +104,7 @@ def rc(c):
   ]
   if has_dht:
     in_specs.append(
-        # 9: dht [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
+        # 11: dht [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
         pl.BlockSpec(
             (None, None, 1, num_v_heads, kq_head_dim, v_head_dim),
             lambda b, g, c: (b, g, 0, 0, 0, 0),
@@ -94,27 +112,37 @@ def rc(c):
     )
 
   out_specs = [
-      # 0: dy_conv [B, G, C, chunk_size, dim_size]
+      # 0: dq [B, C, chunk_size, q_size], columns of group g
       pl.BlockSpec(
-          (None, None, None, chunk_size, dim_size),
-          lambda b, g, c: (b, g, rc(c), 0, 0),
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
+      ),
+      # 1: dk [B, C, chunk_size, k_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
+      ),
+      # 2: dv [B, C, chunk_size, v_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, v_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
       ),
-      # 1: db [B, G, C, chunk_size, padded_num_v_heads]
+      # 3: db [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 2: da [B, G, C, chunk_size, padded_num_v_heads]
+      # 4: da [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 3: dal [B, G, C, 1, padded_num_v_heads]
+      # 5: dal [B, G, C, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 4: ddt [B, G, C, 1, padded_num_v_heads]
+      # 6: ddt [B, G, C, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
@@ -122,7 +150,7 @@ def rc(c):
   ]
   if has_dh0:
     out_specs.append(
-        # 5: dh0 [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
+        # 7: dh0 [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
         pl.BlockSpec(
             (None, None, 1, num_v_heads, kq_head_dim, v_head_dim),
             lambda b, g, c: (b, g, 0, 0, 0, 0),
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py
index f4c490af8b..d57de46af2 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py
@@ -139,13 +139,15 @@ def _bwd_gdn_pipeline_body(
     has_dh0: bool = False,
 ) -> None:
   """Inner kernel executed per (batch, group, chunk) by emit_pipeline with manual GDN backward."""
-  if len(refs) == 17:
+  if len(refs) == 21:
     has_dht = True
     has_dh0 = True
   idx = 0
   # pylint: disable=unbalanced-tuple-unpacking
   (
-      qkv_conv_ref,
+      q_ref,
+      k_ref,
+      v_ref,
       b_ref,
       a_ref,
       do_ref,
@@ -154,22 +156,24 @@ def _bwd_gdn_pipeline_body(
       a_log_ref,
       dt_bias_ref,
       reset_ref,
-  ) = refs[idx : idx + 9]
-  idx += 9
+  ) = refs[idx : idx + 11]
+  idx += 11
   if has_dht:
     dht_ref = refs[idx]
     idx += 1
   else:
     dht_ref = None
   (
-      dy_conv_ref,
+      dq_ref,
+      dk_ref,
+      dv_ref,
       d_b_ref,
       d_a_ref,
       d_a_log_ref,
       d_dt_bias_ref,
-  ) = refs[idx : idx + 5]
+  ) = refs[idx : idx + 7]
   # pylint: enable=unbalanced-tuple-unpacking
-  idx += 5
+  idx += 7
   if has_dh0:
     dh0_ref = refs[idx]
     idx += 1
@@ -198,17 +202,12 @@ def _init():
 
   is_reset = reset_ref[...][0, 0] > 0.5
   d_state = jnp.where(is_reset, 0.0, d_state_scr[...])
-  compute_dtype = qkv_conv_ref.dtype
-  y_c = qkv_conv_ref[...]
-
-  # Slice chunk inputs for this head group
-  q_orig = y_c[:, :q_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
-  k_orig = y_c[:, q_size : q_size + k_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
-  v = (
-      y_c[:, q_size + k_size : q_size + k_size + v_size]
-      .reshape((chunk_size, num_v_heads, v_head_dim))
-      .astype(jnp.float32)
-  )
+  compute_dtype = q_ref.dtype
+
+  # Chunk inputs for this head group
+  q_orig = q_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  v = v_ref[...].reshape((chunk_size, num_v_heads, v_head_dim)).astype(jnp.float32)
 
   b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32)
   a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32)
@@ -446,11 +445,10 @@ def _init():
     d_q = jnp.where(valid_c[:, :, None], d_q, 0.0)
     d_k = jnp.where(valid_c[:, :, None], d_k, 0.0)
 
-  # Flatten gradients to chunk_size x dim_size and write to refs
-  d_q_flat = d_q.reshape(chunk_size, q_size).astype(dy_conv_ref.dtype)
-  d_k_flat = d_k.reshape(chunk_size, k_size).astype(dy_conv_ref.dtype)
-  d_v_flat = d_v_val.reshape(chunk_size, v_size).astype(dy_conv_ref.dtype)
-  dy_conv_ref[...] = jnp.concatenate([d_q_flat, d_k_flat, d_v_flat], axis=-1)
+  # Flatten gradients to chunk_size x q/k/v size and write to refs
+  dq_ref[...] = d_q.reshape(chunk_size, q_size).astype(dq_ref.dtype)
+  dk_ref[...] = d_k.reshape(chunk_size, k_size).astype(dk_ref.dtype)
+  dv_ref[...] = d_v_val.reshape(chunk_size, v_size).astype(dv_ref.dtype)
 
   if padded_num_v_heads > num_v_heads:
     d_b_ref[...] = jnp.pad(d_b_val, ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype(d_b_ref.dtype)
@@ -619,14 +617,15 @@ def pallas_gdn_bwd_kernel(
 
   padded_tile_v_heads = cfg.padded_num_v_heads
 
-  # 1. Prepare qkv_conv into (B, G, C, chunk_size, group_dim_size)
-  q = qkv_conv[:, :, :q_size]
-  k = qkv_conv[:, :, q_size : q_size + k_size]
-  v = qkv_conv[:, :, q_size + k_size :]
-  q_5d = q.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_q_size).transpose((0, 3, 1, 2, 4))
-  k_5d = k.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_k_size).transpose((0, 3, 1, 2, 4))
-  v_5d = v.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_size).transpose((0, 3, 1, 2, 4))
-  qkv_conv_5d = jnp.concatenate([q_5d, k_5d, v_5d], axis=-1)
+  # 1. View qkv_conv as (B, C, chunk_size, dim_size); the kernel reads group g's q, k and v as column blocks
+  # (k after the num_groups q blocks, v after q + k). If q + k is not a whole number of v blocks, v is sliced out.
+  qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, dim_size)
+  if (q_size + k_size) % tile_v_size == 0:
+    v_4d = qkv_conv_4d
+    v_block_offset = (q_size + k_size) // tile_v_size
+  else:
+    v_4d = qkv_conv_4d[..., q_size + k_size :]
+    v_block_offset = 0
 
   # 2. Prepare b and a into (B, G, C, chunk_size, padded_tile_v_heads)
   # With sequence packing, the signed segment IDs are packed into the spare lane
@@ -659,7 +658,7 @@ def pallas_gdn_bwd_kernel(
         ),
     )
 
-  # 3. Prepare do into (B, G, C, chunk_size, tile_v_heads, v_head_dim)
+  # 3. View do as (B, C, chunk_size, G, tile_v_heads, v_head_dim); the kernel reads group g in place
   do_6d = do.reshape(
       batch_size,
       num_chunks,
@@ -667,9 +666,9 @@ def pallas_gdn_bwd_kernel(
       num_groups,
       tile_v_heads,
       cfg.v_head_dim,
-  ).transpose((0, 3, 1, 2, 4, 5))
+  )
 
-  # 4. Prepare chunk_states into (B, G, C, tile_v_heads, kq_head_dim, v_head_dim)
+  # 4. View chunk_states as (B, C, G, tile_v_heads, kq_head_dim, v_head_dim)
   chunk_states_6d = chunk_states.reshape(
       batch_size,
       num_chunks,
@@ -677,22 +676,18 @@ def pallas_gdn_bwd_kernel(
       tile_v_heads,
       cfg.kq_head_dim,
       cfg.v_head_dim,
-  ).swapaxes(1, 2)
-
-  # 5. Prepare t_inv into (B, G, C, tile_v_heads, chunk_size, chunk_size)
-  t_inv_6d = (
-      t_inv.reshape(
-          batch_size,
-          num_chunks,
-          num_groups,
-          tile_v_heads,
-          chunk_size,
-          chunk_size,
-      )
-      .swapaxes(1, 2)
-      .astype(jnp.float32)
   )
 
+  # 5. View t_inv as (B, C, G, tile_v_heads, chunk_size, chunk_size)
+  t_inv_6d = t_inv.reshape(
+      batch_size,
+      num_chunks,
+      num_groups,
+      tile_v_heads,
+      chunk_size,
+      chunk_size,
+  ).astype(jnp.float32)
+
   # 6. Prepare a_log and dt_bias into (B, G, 1, padded_tile_v_heads)
   a_log_2d = a_log.reshape(num_groups, tile_v_heads)
   if padded_tile_v_heads > tile_v_heads:
@@ -739,7 +734,9 @@ def pallas_gdn_bwd_kernel(
   reset_hbm_5d = jnp.broadcast_to(reset_hbm[:, None, :, :, :], (batch_size, num_groups, num_chunks, 1, 128))
 
   in_args = [
-      qkv_conv_5d,
+      qkv_conv_4d,
+      qkv_conv_4d,
+      v_4d,
       b_5d,
       a_5d,
       do_6d,
@@ -767,20 +764,22 @@ def pallas_gdn_bwd_kernel(
   in_specs, out_specs, num_in, num_out = make_bwd_block_specs(
       num_chunks=num_chunks,
       chunk_size=chunk_size,
-      dim_size=group_dim_size,
+      qk_tile_size=tile_q_size,
+      v_tile_size=tile_v_size,
       num_v_heads=tile_v_heads,
       kq_head_dim=cfg.kq_head_dim,
       v_head_dim=cfg.v_head_dim,
+      k_block_offset=q_size // tile_k_size,
+      v_block_offset=v_block_offset,
       padded_num_v_heads=padded_tile_v_heads,
       has_dht=has_dht,
       has_dh0=has_dh0,
   )
 
   out_shapes_list = [
-      jax.ShapeDtypeStruct(
-          (batch_size, num_groups, num_chunks, chunk_size, group_dim_size),
-          qkv_conv.dtype,
-      ),
+      jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, q_size), qkv_conv.dtype),
+      jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, k_size), qkv_conv.dtype),
+      jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, v_size), qkv_conv.dtype),
       jax.ShapeDtypeStruct(b_5d.shape, b_5d.dtype),
       jax.ShapeDtypeStruct(a_5d.shape, a_5d.dtype),
       jax.ShapeDtypeStruct(
@@ -849,7 +848,9 @@ def outer(*refs):
 
   if has_dh0:
     (
-        dy_conv_chunks,
+        dq_4d,
+        dk_4d,
+        dv_4d,
         d_b_chunks,
         d_a_chunks,
         d_a_log_chunks,
@@ -858,7 +859,9 @@ def outer(*refs):
     ) = pallas_out
   else:
     (
-        dy_conv_chunks,
+        dq_4d,
+        dk_4d,
+        dv_4d,
         d_b_chunks,
         d_a_chunks,
         d_a_log_chunks,
@@ -866,13 +869,10 @@ def outer(*refs):
     ) = pallas_out
     dh0_6d = None
 
-  # Reconstruct dy_conv: (B, G, C, S, group_dim_size) -> (B, seq_len, dim_size)
-  dq_5d = dy_conv_chunks[..., :tile_q_size]
-  dk_5d = dy_conv_chunks[..., tile_q_size : tile_q_size + tile_k_size]
-  dv_5d = dy_conv_chunks[..., tile_q_size + tile_k_size :]
-  dq = dq_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, q_size)
-  dk = dk_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, k_size)
-  dv = dv_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, v_size)
+  # Reconstruct dy_conv: (B, C, S, q/k/v size) -> (B, seq_len, dim_size)
+  dq = dq_4d.reshape(batch_size, seq_len, q_size)
+  dk = dk_4d.reshape(batch_size, seq_len, k_size)
+  dv = dv_4d.reshape(batch_size, seq_len, v_size)
   dy_conv_flat = jnp.concatenate([dq, dk, dv], axis=-1).astype(qkv_conv.dtype)
 
   # Reconstruct d_b and d_a: (B, G, C, S, padded_tile_v_heads) -> (B, seq_len, num_v_heads)
diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py
index 8e94bb2396..c9a195b3e0 100644
--- a/tests/unit/gdn_bwd_pallas_test.py
+++ b/tests/unit/gdn_bwd_pallas_test.py
@@ -24,6 +24,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest
 
 from maxtext.configs import types as config_types
 from maxtext.kernels.gdn import gdn_bwd_pallas
@@ -930,6 +931,101 @@ def test_gdn_kernel_bwd_multi_group_head_parallel(self):
     np.testing.assert_allclose(dal1, dal2, rtol=1e-3, atol=1e-3)
     np.testing.assert_allclose(ddt1, ddt2, rtol=1e-3, atol=1e-3)
 
+  @pytest.mark.tpu_only
+  def test_gdn_kernel_bwd_group_blocks_match_per_group_calls(self):
+    """Verifies the in-place group reads (q/k/v column blocks, do, chunk_states, t_inv) and dq/dk/dv writes.
+
+    Every head group runs the same kernel body as a single-group call on that group's heads, so the 2-group call at
+    the Qwen3.5-397B per-device layer geometry (64 v / 16 kq heads, bf16) must match per-group calls bit for bit.
+    """
+    batch_size = 1
+    chunk_size = 64
+    num_chunks = 8
+    seq_len = num_chunks * chunk_size
+    num_k_heads = 16
+    num_v_heads = 64
+    head_tile = 32
+    head_dim = 128
+    q_size = num_k_heads * head_dim
+    dim_size = 2 * q_size + num_v_heads * head_dim
+    tile_q = q_size * head_tile // num_v_heads
+    tile_v = head_tile * head_dim
+
+    k1, k2, k3, k4, k5, k6, k7, k8, k9 = jax.random.split(jax.random.PRNGKey(2468), 9)
+    qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32).astype(jnp.bfloat16)
+    b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32).astype(jnp.bfloat16)
+    a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32).astype(jnp.bfloat16)
+    a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32)
+    dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32)
+    do = jax.random.normal(k6, (batch_size, seq_len, num_v_heads, head_dim), dtype=jnp.float32).astype(jnp.bfloat16)
+    chunk_states = 0.1 * jax.random.normal(k7, (batch_size, num_chunks, num_v_heads, head_dim, head_dim))
+    t_inv = 0.1 * jax.random.normal(k8, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size))
+    dht = 0.1 * jax.random.normal(k9, (batch_size, num_v_heads, head_dim, head_dim))
+    # Two packed documents, the second starting mid-chunk, then padding.
+    segment_ids = jnp.zeros((batch_size, seq_len), jnp.int32)
+    segment_ids = segment_ids.at[:, : seq_len // 2 + 5].set(1).at[:, seq_len // 2 + 5 : seq_len - 16].set(2)
+
+    run = functools.partial(
+        gdn_bwd_pallas.pallas_gdn_bwd_kernel,
+        kq_head_dim=head_dim,
+        v_head_dim=head_dim,
+        chunk_size=chunk_size,
+        head_tile=head_tile,
+        segment_ids=segment_ids,
+        return_dh0=True,
+    )
+    full = run(
+        qkv_conv=qkv,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        do=do,
+        chunk_states=chunk_states,
+        t_inv=t_inv,
+        num_v_heads=num_v_heads,
+        d_recurrent_state=dht,
+    )
+
+    per_group = []
+    for g in range(num_v_heads // head_tile):
+      heads = slice(g * head_tile, (g + 1) * head_tile)
+      qkv_g = jnp.concatenate(
+          [
+              qkv[..., g * tile_q : (g + 1) * tile_q],
+              qkv[..., q_size + g * tile_q : q_size + (g + 1) * tile_q],
+              qkv[..., 2 * q_size + g * tile_v : 2 * q_size + (g + 1) * tile_v],
+          ],
+          axis=-1,
+      )
+      per_group.append(
+          run(
+              qkv_conv=qkv_g,
+              b=b[..., heads],
+              a=a[..., heads],
+              a_log=a_log[heads],
+              dt_bias=dt_bias[heads],
+              do=do[:, :, heads],
+              chunk_states=chunk_states[:, :, heads],
+              t_inv=t_inv[:, :, heads],
+              num_v_heads=head_tile,
+              d_recurrent_state=dht[:, heads],
+          )
+      )
+
+    dy = [out[0] for out in per_group]
+    ref_dy = jnp.concatenate(
+        [d[..., :tile_q] for d in dy] + [d[..., tile_q : 2 * tile_q] for d in dy] + [d[..., 2 * tile_q :] for d in dy],
+        axis=-1,
+    )
+    np.testing.assert_array_equal(full[0], ref_dy)
+    np.testing.assert_array_equal(full[1], jnp.concatenate([out[1] for out in per_group], axis=-1))
+    np.testing.assert_array_equal(full[2], jnp.concatenate([out[2] for out in per_group], axis=-1))
+    np.testing.assert_array_equal(full[5], jnp.concatenate([out[5] for out in per_group], axis=1))
+    # d_a_log / d_dt_bias also sum per-chunk partials in XLA, whose order may depend on the number of groups.
+    np.testing.assert_allclose(full[3], jnp.concatenate([out[3] for out in per_group]), rtol=1e-6, atol=0)
+    np.testing.assert_allclose(full[4], jnp.concatenate([out[4] for out in per_group]), rtol=1e-6, atol=0)
+
   def test_gdn_kernel_bwd_variable_length_segment_ids_reset(self):
     """Verifies segment_ids document boundaries reset carried state gradient to prevent leakage."""
     is_tpu = jax.default_backend() == "tpu"
@@ -1562,53 +1658,65 @@ def test_bwd_block_specs_and_symbolic_zero_dht_dh0_omission(self):
     in_00, out_00, n_in_00, n_out_00 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=False,
         has_dh0=False,
     )
-    self.assertEqual(n_in_00, 9)
-    self.assertEqual(n_out_00, 5)
-    self.assertLen(in_00, 9)
-    self.assertLen(out_00, 5)
+    self.assertEqual(n_in_00, 11)
+    self.assertEqual(n_out_00, 7)
+    self.assertLen(in_00, 11)
+    self.assertLen(out_00, 7)
 
     _, _, n_in_10, n_out_10 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=True,
         has_dh0=False,
     )
-    self.assertEqual((n_in_10, n_out_10), (10, 5))
+    self.assertEqual((n_in_10, n_out_10), (12, 7))
 
     _, _, n_in_01, n_out_01 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=False,
         has_dh0=True,
     )
-    self.assertEqual((n_in_01, n_out_01), (9, 6))
+    self.assertEqual((n_in_01, n_out_01), (11, 8))
 
     _, _, n_in_11, n_out_11 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=True,
         has_dh0=True,
     )
-    self.assertEqual((n_in_11, n_out_11), (10, 6))
+    self.assertEqual((n_in_11, n_out_11), (12, 8))
 
     is_tpu = jax.default_backend() == "tpu"
     batch_size = 1

From 9c48e653a3d419b8330e1cdb794ade0e110fa2fa Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 07/19] GDN forward kernel: DMA bf16 qkv chunks and upcast to
 fp32 in VMEM

Prefill-only calls (the training path through gdn_bwd/api.py) no longer
materialise qkv.astype(f32).reshape(N, 1, D) in HBM before the fused
Conv1D-GDN kernel. The wrapper reshapes the act_in-dtype qkv to
[num_chunks, chunk, D] (a bitcast), ChunkInBufferedRef DMAs one whole chunk
per tile, and the kernel upcasts it into a compact fp32 VMEM scratch with one
strided store per 128-lane column (the inverse of load_as_qkv_large). The
scratch then serves as the qkv slot, so Conv1D and the q/k/v loads see the
same fp32 values in the same layout as before. bf16 -> fp32 is exact, so the
outputs are bit-identical. Both prefill kernels share this input path through
_fused_conv1d_gdn: the full PER_SEQ kernel and the states-only PER_SEQ_STATES
kernel of CP pass 1. The general (decode/mixed) path keeps its HBM upcast and
row DMAs.

VMEM per tile for bf16 qkv: a 2 x 1.5 MiB bf16 window plus a 3 MiB fp32
scratch, the same 6 MiB as the old 2 x 3 MiB fp32 window. No new flag: the
existing is_prefill_only argument selects the path.

Verified locally against the previous kernel:
- Tiny Qwen3.5 A/B on v6e (GDN kernel, gdn_cp_mode=seq, gdn_cp_state=device,
  context=2, so CP pass 1 runs the states-only kernel): loss, logits and all
  66 gradient leaves bit-identical.
- Old vs new kernels on v6e: packed segment_ids, initial conv/recurrent
  states and fp32 qkv at 2 kq / 8 v heads (2 x 256 tokens); segment_ids and
  initial states at the 397B per-device heads (16 kq / 64 v heads, dim 12288,
  256 tokens); each through fused_conv1d_gdn and fused_conv1d_gdn_states;
  plus the general path (3 decode + 2 prefill sequences). All 53 arrays
  bit-identical, and the new states-only states and t_inv equal the new full
  kernel's.
- The same cases in Pallas TPU interpret mode on CPU with the v7x chip
  description (64 MiB VMEM): all bit-identical. The interpreter needed two
  patches in the check script: a patched store discharge rule that can write
  through a ref reshape (jax 0.11 raises NotImplementedError for the strided
  stores into the reshaped fp32 scratch), and static-size row and state DMAs
  (it cannot slice with traced sizes).
- New test test_prefill_bf16_qkv_matches_fp32_qkv_path (tpu_only) passes on
  v6e: the full and the states-only prefill kernels match the HBM-upcast
  general path bitwise.
- gdn_bwd_pallas_test.py on CPU: 38 passed, 3 skipped (two need 4 devices,
  one is the new tpu_only test).
---
 src/maxtext/kernels/gdn/config.py     |  8 +++++
 src/maxtext/kernels/gdn/memory_ref.py | 47 ++++++++++++++++++++++--
 src/maxtext/kernels/gdn/vmem_ldst.py  | 22 ++++++++++++
 src/maxtext/kernels/gdn/wrapper.py    | 21 ++++++++---
 tests/unit/gdn_bwd_pallas_test.py     | 51 +++++++++++++++++++++++++++
 5 files changed, 143 insertions(+), 6 deletions(-)

diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py
index bf0c005510..fbf30d3baf 100644
--- a/src/maxtext/kernels/gdn/config.py
+++ b/src/maxtext/kernels/gdn/config.py
@@ -84,6 +84,8 @@ class GDNConfig:
   num_buffers: int = 2
   has_seg_ids: bool = False
   use_qk_norm_in_gdn: bool = True
+  # Prefill-only tiles are whole chunks: qkv is DMA'd in act_in dtype and upcast in VMEM.
+  is_prefill_only: bool = False
 
   @property
   def chunk_size(self) -> int:
@@ -152,7 +154,13 @@ def get_scratch_shape_dict(self) -> dict[str, Any]:
       carry_conv_scratch = pltpu.VMEM(conv_shape, jnp.float32)
       carry_recurrent_scratch = pltpu.VMEM(recurrent_shape, jnp.float32)
 
+    # Compact fp32 copy of a prefill qkv chunk, read by Conv1D and the q/k/v loads.
+    qkv_scratch = None
+    if self.is_prefill_only:
+      qkv_scratch = pltpu.VMEM((self.seq_tile_size, self.chunk_size, 1, self.dim_size), jnp.float32)
+
     return {
         "carry_conv_scratch_ref": carry_conv_scratch,
         "carry_recurrent_scratch_ref": carry_recurrent_scratch,
+        "qkv_scratch_ref": qkv_scratch,
     }
diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py
index 481ca62066..24a7d0bcf8 100644
--- a/src/maxtext/kernels/gdn/memory_ref.py
+++ b/src/maxtext/kernels/gdn/memory_ref.py
@@ -281,6 +281,36 @@ def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
     ).wait()
 
 
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ChunkInBufferedRef(BaseBufferedRef):
+  """Input DMA reference that copies whole chunks of a [num_chunks, chunk, dim] array."""
+
+  def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
+    """Initiates asynchronous copy of each tile's chunk into current input VMEM slot."""
+    assert self.sem_recvs is not None
+    assert self.window_ref is not None
+    slot = self.current_copy_in_slot
+    sem = self.sem_recvs.at[slot]
+    vmem_ref = self.window_ref.at[slot]
+    p_id = grid_indices[0]
+
+    for idx in range(self.cfg.seq_tile_size):
+      chunk_idx = self.metadata_ref.get_record(p_id, idx).r_base // self.cfg.chunk_size
+      pltpu.make_async_copy(src_ref.at[chunk_idx], vmem_ref.at[idx], sem).start()
+
+  def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
+    """Waits for asynchronous chunk copy to complete."""
+    assert self.sem_recvs is not None
+    assert self.window_ref is not None
+    slot = self.current_wait_in_slot
+    sem = self.sem_recvs.at[slot]
+    vmem_ref = self.window_ref.at[slot]
+
+    for idx in range(self.cfg.seq_tile_size):
+      pltpu.make_async_copy(vmem_ref.at[idx], vmem_ref.at[idx], sem).wait()
+
+
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class OutBufferedRef(BaseBufferedRef):
@@ -504,7 +534,10 @@ def create_allocs(
     chunk_states_ref: jax.Array | None = None,
 ) -> tuple[Any, ...]:
   """Creates all Pallas double-buffered allocations for GDN kernel (no output buffers without out_ref)."""
-  qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size)
+  if cfg.is_prefill_only:
+    qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, cfg.dim_size)
+  else:
+    qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size)
   ba_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.aligned_num_v_heads)
 
   out_shape = (
@@ -539,7 +572,17 @@ def create_allocs(
       cfg=cfg,
       metadata_ref=metadata_ref,
   )
-  qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref)
+  if cfg.is_prefill_only:
+    qkv_alloc = ChunkInBufferedRef.input(
+        spec=qkv_spec,
+        dtype_or_type=qkv_ref,
+        buffer_count=pipeline_mode.buffer_count,
+        use_lookahead=pipeline_mode.use_lookahead,
+        cfg=cfg,
+        metadata_ref=metadata_ref,
+    )
+  else:
+    qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref)
   b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref)
   a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref)
 
diff --git a/src/maxtext/kernels/gdn/vmem_ldst.py b/src/maxtext/kernels/gdn/vmem_ldst.py
index 7d604a69b1..3cc5225c48 100644
--- a/src/maxtext/kernels/gdn/vmem_ldst.py
+++ b/src/maxtext/kernels/gdn/vmem_ldst.py
@@ -134,6 +134,28 @@ def load_compact_to_large(vmem_ref: jax.Ref) -> jax.Array:
   return jnp.concat(vreg_list, axis=-1).reshape(new_shape)
 
 
+def upcast_large_to_compact(large_ref: jax.Ref, compact_ref: jax.Ref) -> None:
+  """Upcast activations from large layout into a compact fp32 reference.
+
+  The inverse of load_as_qkv_large: one strided store per 128-lane column puts
+  each token's lanes into its own row of the compact reference.
+
+  Args:
+    large_ref: Activations in VMEM of shape [seq_tile_size, chunk_size,
+      dim_size].
+    compact_ref: fp32 reference in VMEM of shape [seq_tile_size, chunk_size, 1,
+      dim_size].
+  """
+
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  lanes_per_col = compact_ref.shape[-1] // num_lanes
+  for idx in range(compact_ref.shape[0]):
+    compact_flat_ref = compact_ref.at[idx].reshape(-1, num_lanes)  # pyrefly: ignore[missing-attribute]
+    for lane in range(lanes_per_col):
+      col = large_ref[idx, :, lane * num_lanes : (lane + 1) * num_lanes]
+      compact_flat_ref[lane::lanes_per_col] = col.astype(compact_ref.dtype)
+
+
 def load_and_select_states(
     metadata_ref: memory_ref.MetadataRef,
     p_id: jax.Array,
diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py
index f81d4a9c82..65b52a8a89 100644
--- a/src/maxtext/kernels/gdn/wrapper.py
+++ b/src/maxtext/kernels/gdn/wrapper.py
@@ -95,7 +95,7 @@ def _pack_fwd_segment_metadata(
 
 def inner_kernel(
     # Inputs.
-    qkv_slot_ref: jax.Ref,  # [seq, chunk, 1, dim_size]
+    qkv_slot_ref: jax.Ref,  # [seq, chunk, 1, dim_size] ([seq, chunk, dim_size] for prefill-only)
     b_slot_ref: jax.Ref,  # [seq, chunk, 1, num_v_heads]
     a_slot_ref: jax.Ref,  # [seq, chunk, 1, num_v_heads]
     conv_state_slot_ref: jax.Ref,  # [seq, prev_kernel_size, 1, dim_size]
@@ -126,9 +126,16 @@ def inner_kernel(
   weights_ref = args[1]
   carry_conv_scratch_ref = args[2] if len(args) > 2 else None
   carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None
+  qkv_scratch_ref = args[4] if len(args) > 4 else None
 
   p_id = pl.program_id(0)
 
+  # Prefill-only qkv arrives in act_in dtype and large layout. Upcast it into the
+  # compact fp32 scratch, which then serves as the qkv slot below.
+  if qkv_scratch_ref is not None:
+    vmem_ldst.upcast_large_to_compact(qkv_slot_ref, qkv_scratch_ref)
+    qkv_slot_ref = qkv_scratch_ref
+
   # Prepare states.
   real_sizes, prev_conv, prev_recurrent = vmem_ldst.load_and_select_states(
       metadata_ref=metadata_ref,
@@ -242,6 +249,7 @@ def outer_kernel(
     *args,
     carry_conv_scratch_ref: jax.Array | None = None,
     carry_recurrent_scratch_ref: jax.Array | None = None,
+    qkv_scratch_ref: jax.Array | None = None,
     cfg: config.GDNConfig,
     **kwargs,
 ) -> None:
@@ -307,6 +315,7 @@ def _run(allocations):
             weights_ref,
             carry_conv_scratch_ref,
             carry_recurrent_scratch_ref,
+            qkv_scratch_ref,
         ),
         allocations=allocations,
     )
@@ -353,7 +362,6 @@ def _fused_conv1d_gdn(
   recurrent_out_dtype = recurrent_state.dtype
   assert a.dtype == b.dtype == qkv.dtype == act_in_dtype
 
-  qkv = qkv.astype(jnp.float32)
   b = b.astype(jnp.float32)
   a = a.astype(jnp.float32)
   conv_state = conv_state.astype(jnp.float32)
@@ -421,7 +429,6 @@ def _fused_conv1d_gdn(
   batch_padding_size = padded_batch_size - batch_size
   aligned_num_v_heads = tiling.align_to(n_v + extra_lanes, num_lanes)
   num_v_padding_size = aligned_num_v_heads - n_v
-  qkv = jnp.pad(qkv, ((0, batch_padding_size), (0, 0)))
   b = jnp.pad(b, ((0, batch_padding_size), (0, num_v_padding_size)))
   a = jnp.pad(a, ((0, batch_padding_size), (0, num_v_padding_size)))
   if has_seg_ids:
@@ -436,7 +443,12 @@ def _fused_conv1d_gdn(
     b = b.at[:batch_size, n_v].set(s_enc_flat)
     b = b.at[:batch_size, n_v + 1].set(seg_aux_flat)
 
-  qkv = qkv.reshape(padded_batch_size, 1, -1)
+  if is_prefill_only:
+    # Whole chunks, DMA'd in act_in dtype and upcast to fp32 in VMEM by the kernel.
+    qkv = qkv.reshape(batch_size // mixed_tile_size, mixed_tile_size, dim)
+  else:
+    qkv = jnp.pad(qkv.astype(jnp.float32), ((0, batch_padding_size), (0, 0)))
+    qkv = qkv.reshape(padded_batch_size, 1, -1)
   b = b.reshape(padded_batch_size, 1, -1)
   a = a.reshape(padded_batch_size, 1, -1)
 
@@ -480,6 +492,7 @@ def call_kernel(
         v_head_dim=d_v,
         has_seg_ids=has_seg_ids,
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        is_prefill_only=is_prefill_only,
         dtypes=config.Dtypes(
             act_in=act_in_dtype,
             act_out=act_out_dtype,
diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py
index c9a195b3e0..3016785d5f 100644
--- a/tests/unit/gdn_bwd_pallas_test.py
+++ b/tests/unit/gdn_bwd_pallas_test.py
@@ -778,6 +778,57 @@ def test_run_local_gdn_states_fwd_matches_full_forward(self):
       self.assertEqual(got.dtype, exp.dtype)
       np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(exp, np.float32))
 
+  @pytest.mark.tpu_only
+  def test_prefill_bf16_qkv_matches_fp32_qkv_path(self):
+    """Prefill-only kernels upcast bf16 qkv chunks in VMEM; the general path upcasts in HBM. Same bits.
+
+    Covers the full prefill kernel (PER_SEQ) and the states-only kernel (PER_SEQ_STATES).
+    """
+    from maxtext.kernels.gdn import wrapper  # pylint: disable=import-outside-toplevel
+
+    batch_size, chunk_size, num_k_heads, num_v_heads, head_dim, conv_kernel_size = 2, 64, 2, 8, 128, 4
+    seq_len = 4 * chunk_size
+    dim_size = num_k_heads * head_dim * 2 + num_v_heads * head_dim
+    qkv, b, a, conv_weight, conv_bias, a_log, dt_bias = _init_fwd_inputs(
+        jax.random.PRNGKey(5), batch_size, seq_len, dim_size, num_v_heads, conv_kernel_size
+    )
+
+    def run(fn, **kwargs):
+      # Fresh states for every call: the wrapper donates them.
+      k1, k2 = jax.random.split(jax.random.PRNGKey(6))
+      conv_state = jax.random.normal(k1, (batch_size + 1, conv_kernel_size - 1, dim_size), jnp.bfloat16)
+      recurrent_state = 0.1 * jax.random.normal(k2, (batch_size + 1, num_v_heads, head_dim, head_dim), jnp.float32)
+      return fn(
+          qkv=qkv.astype(jnp.bfloat16).reshape(-1, dim_size),
+          b=b.astype(jnp.bfloat16).reshape(-1, num_v_heads),
+          a=a.astype(jnp.bfloat16).reshape(-1, num_v_heads),
+          conv_state=conv_state,
+          recurrent_state=recurrent_state,
+          conv_weight=jnp.swapaxes(conv_weight, 0, 2),
+          conv_bias=conv_bias,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          query_start_loc=jnp.arange(0, (batch_size + 1) * seq_len, seq_len, dtype=jnp.int32),
+          state_indices=jnp.arange(1, batch_size + 1, dtype=jnp.int32),
+          distribution=jnp.array([0, 0, batch_size], dtype=jnp.int32),
+          seq_lens=jnp.full((batch_size,), seq_len + 1, dtype=jnp.int32),
+          n_kq=num_k_heads,
+          n_v=num_v_heads,
+          d_k=head_dim,
+          d_v=head_dim,
+          kernel_size=conv_kernel_size,
+          mixed_tile_size=chunk_size,
+          **kwargs,
+      )
+
+    out, (conv_out, rec_out), t_inv, chunk_states = run(wrapper.fused_conv1d_gdn, is_prefill_only=True)
+    (states_conv, states_rec), states_t_inv = run(wrapper.fused_conv1d_gdn_states)
+    ref_out, (ref_conv, ref_rec), ref_t_inv, ref_chunk_states = run(wrapper.fused_conv1d_gdn, is_prefill_only=False)
+    got = (out, conv_out, rec_out, t_inv, chunk_states, states_conv, states_rec, states_t_inv)
+    exp = (ref_out, ref_conv, ref_rec, ref_t_inv, ref_chunk_states, ref_conv, ref_rec, ref_t_inv)
+    for g, e in zip(got, exp):
+      np.testing.assert_array_equal(np.asarray(g), np.asarray(e))
+
   def test_decoupled_conv1d_gdn_kernel_gradient_with_initial_states(self):
     """Verifies custom VJP gradients when initial conv_state and recurrent_state are provided."""
     is_tpu = jax.default_backend() == "tpu"

From bb836e7e63dc34f59a614b9ea19a7337cc8974f4 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 08/19] ring_ragged_unsort backward: reuse the sort permutation
 instead of argsorting

The unsort backward argsorted the revert indices to rebuild their inverse,
which is the forward sort's permutation. ring_ragged_sort now also returns that
permutation and RouteOutput carries it to ring_ragged_unsort, so the backward
reads it instead of sorting again. The permutation is named for the `moe_route`
remat key in place of the sorted token indices (same size), which are now the
permutation // topk (flat slot j belongs to token j // topk) instead of an s32
gather of the repeated token ids.

Exact by construction: the argsort of a permutation is its inverse (distinct
keys, so stability does not matter) and arange(n).repeat(topk)[p] == p // topk.
Per MoE layer and token chunk this drops one argsort from the backward and the
s32 token-index gather from the forward and its replay.

The buffered branch's routing-weight gradient stays a scatter into zeros. A
gather over the revert indices would be equivalent, but on v7x XLA compiles it
to a TensorCore gather fusion over the f32[655360] slots, about 30 ms per call
against about 2 ms for the scatter (measured: unpermute went from 589 ms to
4,085 ms per micro-batch on the 128-chip bench with the gather).

Test: RingRaggedUnsortGradTest compares the custom backward with autodiff of a
plain-gather forward, bitwise, unbuffered and with a buffer that drops slots,
and checks that the gradient contains no sort.
---
 src/maxtext/kernels/ragged/ragged_sort.py | 28 +++++++----
 src/maxtext/layers/moe.py                 | 13 +++++-
 tests/unit/moe_test.py                    | 57 +++++++++++++++++++++++
 3 files changed, 88 insertions(+), 10 deletions(-)

diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index dc563cef68..35ac56d09c 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -72,6 +72,8 @@ def ring_ragged_sort(
         only the tokens destined for local experts, padded with zeros elsewhere.
       - 1D tensor ``group_sizes_local`` tracking expert token counts.
       - 1D tensor ``topk_argsort_revert_indices`` for inverse routing.
+      - 1D tensor ``topk_argsort_indices``, the sort permutation (the inverse of
+        ``topk_argsort_revert_indices``), for the backward of :func:`ring_ragged_unsort`.
   """
 
   @jax.custom_vjp
@@ -88,17 +90,16 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
     topk_indices_flat = topk_indices_local.flatten()  # num_tokens_local x topk
     topk_argsort_indices = jnp.argsort(topk_indices_flat)  # num_tokens_local x topk
 
-    token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk)  # num_tokens_local x topk
-    token_indices_sorted = token_indices[topk_argsort_indices]  # num_tokens_local x topk
-
     group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0)  # GLOBAL_NUM_EXPERTS
 
     topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices)  # num_tokens_local x topk
     # Named for the `moe_route` remat key. Saving these integer results lets the backward replay
     # the token gather, whose output the expert GMMs need again, without replaying both argsorts.
-    token_indices_sorted = checkpoint_name(token_indices_sorted, "moe_route")
+    topk_argsort_indices = checkpoint_name(topk_argsort_indices, "moe_route")
     group_sizes_local = checkpoint_name(group_sizes_local, "moe_route")
     topk_argsort_revert_indices = checkpoint_name(topk_argsort_revert_indices, "moe_route")
+    # Flat slot j belongs to token j // topk.
+    token_indices_sorted = topk_argsort_indices // topk  # num_tokens_local x topk
     shard_idx = jax.lax.axis_index(ep_name)
 
     local_num_experts = num_experts // ep_size
@@ -145,7 +146,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
           use_single_sparsecore=use_single_sparsecore,
       )
 
-    out = (x, group_sizes_local, topk_argsort_revert_indices)
+    out = (x, group_sizes_local, topk_argsort_revert_indices, topk_argsort_indices)
 
     res = (
         topk_argsort_revert_indices,
@@ -173,7 +174,7 @@ def _ring_ragged_sort_bwd(res, g_out):
         local_buffer_size,
         _,
     ) = res
-    g_x, _, _ = g_out
+    g_x, _, _, _ = g_out
     # Restrict to the [start, end) source range via a validity bitmask. The
     # ragged kernel packs valid rows to the front of each row-partition and
     # only iterates over the populated prefix, so we hand it the mask directly
@@ -236,6 +237,7 @@ def ring_ragged_unsort(
     sorted_tokens_local,
     group_sizes_local,
     topk_argsort_revert_indices,
+    topk_argsort_indices,
     topk,
     local_num_experts,
     ep_name,
@@ -268,6 +270,8 @@ def ring_ragged_unsort(
     sorted_tokens_local: 2D ``[buffer_size, hidden]`` output tensor from local experts.
     group_sizes_local: 1D tensor tracking the token loads per expert.
     topk_argsort_revert_indices: 1D permutation restoring flat token positions.
+    topk_argsort_indices: 1D sort permutation from :func:`ring_ragged_sort`, the inverse of
+      ``topk_argsort_revert_indices``. Read by the backward only.
     topk: scalar ``int`` representing the routing top-k factor.
     local_num_experts: scalar ``int`` representing the count of experts hosted on this shard.
     ep_name: ``str`` identifying the expert parallel axis name.
@@ -284,6 +288,7 @@ def _ring_ragged_unsort(
       sorted_tokens_local,
       group_sizes_local,
       topk_argsort_revert_indices,
+      topk_argsort_indices,
       topk_weights_flat,
   ):
     """Unsort and scatter activations."""
@@ -291,6 +296,7 @@ def _ring_ragged_unsort(
         sorted_tokens_local,
         group_sizes_local,
         topk_argsort_revert_indices,
+        topk_argsort_indices,
         topk_weights_flat,
     )[0]
 
@@ -299,6 +305,7 @@ def _ring_ragged_unsort_fwd(
       sorted_tokens_local,
       group_sizes_local,
       topk_argsort_revert_indices,
+      topk_argsort_indices,
       topk_weights_flat,
   ):
     """Executes unsorting sending tokens back."""
@@ -361,6 +368,7 @@ def _ring_ragged_unsort_fwd(
     res = (
         sorted_tokens_local,
         topk_argsort_revert_indices,
+        topk_argsort_indices,
         topk_weights_flat,
         shard_output_start,
         shard_output_end,
@@ -388,6 +396,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
     (
         sorted_tokens_local,
         topk_argsort_revert_indices,
+        topk_argsort_indices,
         topk_weights_flat,
         shard_output_start,
         shard_output_end,
@@ -396,9 +405,9 @@ def _ring_ragged_unsort_bwd(res, g_out):
     g_hidden_states_local = g_out
 
     n = topk_argsort_revert_indices.shape[0]
-    # Build the inverse permutation idx_inv such that idx_inv[j] = i
+    # The forward sort's permutation is the inverse permutation: idx_inv[j] = i
     # where revert[i] = j.
-    idx_inv = jnp.argsort(topk_argsort_revert_indices)
+    idx_inv = topk_argsort_indices
 
     # Handle the same two buffering modes for backward pass.
     # ragged_gather does the fan-out, by indexing into the un-expanded
@@ -451,7 +460,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
       dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1)
       slots = jnp.where(mask, sliced_idx_inv, n)
       grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop")
-    return grad_sorted_tokens, None, None, grad_topk_weights
+    return grad_sorted_tokens, None, None, None, grad_topk_weights
 
   _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd)
 
@@ -462,6 +471,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
       sorted_tokens_local,
       group_sizes_local,
       topk_argsort_revert_indices,
+      topk_argsort_indices,
       topk_weights_flat,
   )
 
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 8bffab758c..8899fc12e9 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -109,6 +109,9 @@ class RouteOutput:
   # log_moe_drops=True on the ring-of-experts ragged path only (None otherwise): the drop audit
   # `(dropped_assignments, max_load_ratio)`, see `RoutedMoE.permute`. Logging only.
   moe_audit: Optional[tuple[jax.Array, jax.Array]] = None
+  # Ring-of-experts ragged path only (None otherwise): the sort permutation, the inverse of
+  # `sorted_selected_experts` there. The unsort backward reads it instead of argsorting that.
+  topk_argsort_indices: Optional[jax.Array] = None
 
 
 def combine_moe_audits(audits):
@@ -1487,7 +1490,7 @@ def permute(
       else:
         buffer_size = None
 
-      sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort(
+      sorted_inputs, group_size, sorted_selected_experts, topk_argsort_indices = ring_ragged_sort(
           inputs_2d,
           topk_indices_2d,
           self.config.num_experts,
@@ -1504,6 +1507,7 @@ def permute(
           use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
       )
     else:
+      topk_argsort_indices = None
       flatten_selected_experts = jnp.ravel(selected_experts)
 
       # Must precede roll_to_expert_id: `(-1 - roll) % num_experts` wraps padding
@@ -1609,6 +1613,7 @@ def permute(
         local_group_size,
         has_overflow,
         moe_audit,
+        topk_argsort_indices,
     )
 
   def unpermute(
@@ -1620,6 +1625,7 @@ def unpermute(
       sequence_length,
       use_custom_sort_vjp=True,
       group_sizes=None,
+      topk_argsort_indices=None,
   ):
     """Unpermute tokens to original order and combine weights."""
 
@@ -1633,6 +1639,7 @@ def unpermute(
           intermediate,
           group_sizes,
           sorted_selected_experts,
+          topk_argsort_indices,
           self.num_experts_per_tok,
           local_num_experts,
           self._expert_parallelism_name,
@@ -2481,6 +2488,7 @@ def _ep_all_gather(z):
           local_group_sizes,
           has_overflow,
           moe_audit,
+          topk_argsort_indices,
       ) = self.permute(
           x,
           logits,
@@ -2506,6 +2514,7 @@ def _ep_all_gather(z):
               local_group_sizes=local_group_sizes,
               has_overflow=has_overflow,
               moe_audit=moe_audit,
+              topk_argsort_indices=topk_argsort_indices,
           ),
           RouteMetadata(
               expert_shard_id=expert_shard_id,
@@ -2539,6 +2548,7 @@ def ra2a_and_route(
           local_group_sizes,
           _,
           _,
+          _,
       ) = self.permute(
           x,
           logits,
@@ -3037,6 +3047,7 @@ def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length):
           sequence_length=sequence_length,
           use_custom_sort_vjp=self.config.use_custom_sort_vjp,
           group_sizes=routing.group_sizes,
+          topk_argsort_indices=routing.topk_argsort_indices,
       )
       return jnp.reshape(
           output,
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
index 67fb571690..fa8f72cadc 100644
--- a/tests/unit/moe_test.py
+++ b/tests/unit/moe_test.py
@@ -14,6 +14,7 @@
 """Mixture of Experts (MoE) tests."""
 
 import functools
+import re
 from types import SimpleNamespace
 import unittest
 from unittest import mock
@@ -30,6 +31,7 @@
 from jax.sharding import Mesh, PartitionSpec as P
 from maxtext.common.common_types import Config, DType
 from maxtext.configs import pyconfig
+from maxtext.kernels.ragged import ragged_sort
 from maxtext.layers import linears
 from maxtext.layers import moe
 from maxtext.layers import nnx_wrappers
@@ -3552,5 +3554,60 @@ def split_then_glu(fused):
       np.testing.assert_array_equal(bits(got_vjp(cotangent)[0]), bits(want_vjp(cotangent)[0]))
 
 
+class RingRaggedUnsortGradTest(parameterized.TestCase):
+  """ring_ragged_unsort's custom backward against autodiff of the same forward written with plain gathers."""
+
+  @parameterized.named_parameters(("unbuffered", None), ("buffered", 8))
+  def test_grad_matches_autodiff_without_argsort(self, buffer_size):
+    ep, num_tokens, hidden, topk, num_experts = 2, 8, 128, 2, 4
+    local_num_experts = num_experts // ep
+    rng = np.random.default_rng(0)
+    # Small integers keep every product and sum exact, so the comparison is bitwise.
+    hidden_states = jnp.asarray(rng.integers(-4, 5, (ep, num_tokens, hidden)), jnp.bfloat16)
+    experts = jnp.asarray(
+        [[rng.choice(num_experts, topk, replace=False) for _ in range(num_tokens)] for _ in range(ep)], jnp.int32
+    )
+    weights = jnp.asarray(rng.integers(1, 4, (ep, num_tokens * topk)), jnp.float32)
+    cotangent = jnp.asarray(rng.integers(-4, 5, (ep, num_tokens, hidden)), jnp.float32)
+    fallbacks = {"enforce_gather_fallback": True, "enforce_gather_reduce_fallback": True}
+
+    def sort(x, e):
+      return ragged_sort.ring_ragged_sort(x, e, num_experts, topk, "expert", ep, buffer_size=buffer_size, **fallbacks)
+
+    # Each vmap lane is one expert shard. buffer_size=8 is the balanced load: the shards receive 9 and 6 slots,
+    # so one drops a slot and the other leaves buffer rows unused.
+    tokens, group_sizes, revert, perm = jax.vmap(sort, axis_name="expert")(hidden_states, experts)
+
+    def unsort(tokens, weights, group_sizes, revert, perm):
+      return ragged_sort.ring_ragged_unsort(
+          tokens, group_sizes, revert, perm, topk, local_num_experts, "expert", weights, **fallbacks
+      )
+
+    def reference(tokens, weights, group_sizes, revert, perm):
+      del perm
+      offsets = jnp.cumulative_sum(group_sizes, include_initial=True)
+      shard = jax.lax.axis_index("expert")
+      start, end = offsets[shard * local_num_experts], offsets[(shard + 1) * local_num_experts]
+      if tokens.shape[0] < revert.shape[0]:  # packed buffer: this shard's rows start at 0
+        pos, valid = revert - start, (revert >= start) & (revert - start < jnp.minimum(end - start, tokens.shape[0]))
+      else:
+        pos, valid = revert, (revert >= start) & (revert < end)
+      rows = jnp.where(valid[:, None], tokens[jnp.where(valid, pos, 0)].astype(jnp.float32) * weights[:, None], 0.0)
+      return rows.reshape(num_tokens, topk, hidden).sum(axis=1)
+
+    def grads(fn):
+      def loss(tokens, weights):
+        out = jax.vmap(fn, axis_name="expert")(tokens, weights, group_sizes, revert, perm)
+        return jnp.sum(out.astype(jnp.float32) * cotangent)
+
+      return jax.grad(loss, argnums=(0, 1))
+
+    for actual, expected in zip(grads(unsort)(tokens, weights), grads(reference)(tokens, weights)):
+      np.testing.assert_array_equal(actual, expected)
+    # The backward reuses the sort's permutation instead of argsorting the revert indices.
+    jaxpr = str(jax.make_jaxpr(grads(unsort))(tokens, weights))
+    self.assertIsNone(re.search(r"\bsort\[", jaxpr), "the unsort backward argsorts")
+
+
 if __name__ == "__main__":
   absltest.main()

From e88a5b980ef64151192a050bfaf7e932882e9834 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 09/19] GDN forward kernels: right-looking substitution in the
 64x64 triangular inverse

compute_gdn.invert_triangular_matrix solved each 16-row diagonal block
left-looking: every row re-stacked the rows solved so far and moved its row of
t from lanes to sublanes to form sum_j t[i, j] * x_j. The right-looking form
keeps the block in place as two sublane-aligned 8-row halves and, once row j is
final, subtracts t[i, j] * x_j from every later row. Each row subtracts the
same products, one at a time in column order, so only the summation order
changes. The inter-block matmuls are unchanged. Both prefill kernels (the full
PER_SEQ kernel and the CP pass-1 states-only kernel) use this function, as does
the pure-JAX path.

tpu7x LLO at the 397B per-device shape (16,384 tokens, 16 kq / 64 v heads,
chunk 64), loop bundles per chunk:
- full kernel 43,448 -> 39,389 (-9.3%); XLU ops 29,633 -> 18,370, VALU ops
  88,676 -> 67,243.
- states-only kernel 40,049 -> 36,330 (-9.3%); XLU ops 26,945 -> 15,669.

The unrolled right-looking loop needs more VMEM for spills: the full kernel
would need 62.2 MiB, over the 57.6 MiB limit (0.90 x 64 MiB). Two changes make
room:
- create_allocs single-buffers the conv and recurrent state refs of a prefill
  call with one sequence (the training path: one row per device). Its states
  are copied in only before the first tile and out only after the last, so no
  copy is in flight while a tile uses the slot. This saves 4.1 MiB, down to
  58.05 MiB.
- DEFAULT_VMEM_LIMIT_FACTOR goes from 0.90 to 0.95 (60.8 MiB on TPU7x).

Verified against the previous kernels:
- v6e at the 397B per-device geometry (16,384 tokens, bf16 qkv/b/a, a packed
  row of 7 documents plus padding and a one-document row, with and without
  caller states). These are bit-identical: the full kernel's out, final states
  and chunk_states, the states-only kernel's states, and the loss. t_inv from
  both kernels differs by at most 1.5e-8 (|t_inv| <= 1). Gradients through
  gdn_decoupled_conv1d, whose Pallas backward reads the saved t_inv: d_b,
  d_conv_state and d_recurrent_state are bit-identical. Max |new - old| /
  max |old| is 6.2e-6 for qkv, 6.5e-6 for the conv weight, 3.7e-5 for a and
  4.6e-7 for A_log and dt_bias. Against an f32 HIGHEST pure-JAX reference, old
  and new have the same error to 4 digits (2e-3 to 6e-3 of max).
- Pallas interpret mode on CPU with the TPU7x chip description:
  - chunked_gdn inside a Pallas kernel, full and states-only modes, with and
    without segments: out and state bit-identical, t_inv within 1.2e-10.
  - The full and states-only kernels, 1 and 2 sequences per call, with the
    interpreter patches from the bf16-qkv change (static-size DMAs, stores
    through a ref reshape): final states bit-identical, t_inv within 3.7e-8,
    chunk states within 8e-8 of max (f32 matmuls on CPU carry the t_inv
    rounding into them), out within one bf16 ulp. The new states-only
    kernel's states and t_inv equal the new full kernel's bit for bit.
- Tiny Qwen3.5 A/B on v6e (ab_model.py: gdn_cp_mode=seq, gdn_cp_state=device,
  context 2, packed rows): loss and logits bit-identical. 64 of 66 gradient
  leaves are bit-identical; the other two (one GDN layer's A_log and conv
  kernel) differ by at most 8.3e-8 of max. Gradient norm rel diff 1.4e-14.
- 397B v7x AOT (wi_tile_dlhs 2048/2048, moe_route=device) compiles. fwd_bwd
  temp is 93.31 GiB and dropless temp 67.94 GiB, both the same as the base
  compiled with the same flags. Outside the 9 GDN kernel bodies and their
  scoped VMEM limit (60,397,977 -> 63,753,420 bytes), the optimized HLO is
  identical to the base's.
- New test test_invert_triangular_matrix_matches_float64_inverse checks the
  inverse against float64 at chunk 8 to 64 (atol 1e-6; dropping one column's
  update gives 0.2). gdn_bwd_pallas_test.py: 39 passed and 4 skipped on CPU;
  the 2 tpu_only tests pass on v6e.
---
 src/maxtext/kernels/gdn/compute_gdn.py | 28 ++++++++++++++++----------
 src/maxtext/kernels/gdn/config.py      |  3 ++-
 src/maxtext/kernels/gdn/memory_ref.py  | 13 ++++++++----
 tests/unit/gdn_bwd_pallas_test.py      | 23 +++++++++++++++++++++
 4 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/src/maxtext/kernels/gdn/compute_gdn.py b/src/maxtext/kernels/gdn/compute_gdn.py
index 45a6ab7a60..45a09d5eff 100644
--- a/src/maxtext/kernels/gdn/compute_gdn.py
+++ b/src/maxtext/kernels/gdn/compute_gdn.py
@@ -51,18 +51,24 @@ def invert_triangular_matrix(
   num_blocks = chunk // block_size
 
   def local_forward_sub(t_mat: jax.Array, b_mat: jax.Array) -> jax.Array:
-    x_list = []
-    for i in range(block_size):
-      b_i = b_mat[:, i, :]
-      if i == 0:
-        x_i = b_i
+    # Right-looking substitution: once row j is final, subtract t[i, j] * x_j from every later row i.
+    # The rows stay in place as two sublane-aligned halves, so no step re-stacks the solved rows or
+    # moves a row of t from lanes to sublanes. Each row subtracts the same products as a left-looking
+    # solve, one at a time in column order.
+    half = block_size // 2
+    x_lo, x_hi = b_mat[:, :half, :], b_mat[:, half:, :]
+    t_lo, t_hi = t_mat[:, :half, :], t_mat[:, half:, :]
+    rows_lo = jax.lax.broadcasted_iota(jnp.int32, (1, half, 1), 1)
+    rows_hi = jax.lax.broadcasted_iota(jnp.int32, (1, block_size - half, 1), 1)
+    for j in range(block_size - 1):
+      if j < half:
+        x_j = x_lo[:, j : j + 1, :]
+        x_lo = jnp.where(rows_lo > j, x_lo - t_lo[:, :, j : j + 1] * x_j, x_lo)
+        x_hi = x_hi - t_hi[:, :, j : j + 1] * x_j
       else:
-        stacked_x = jnp.stack(x_list, axis=1)
-        all_prev_t = t_mat[:, i, :i]
-        prev_sum = jnp.sum(all_prev_t[..., None] * stacked_x, axis=1)
-        x_i = b_i - prev_sum
-      x_list.append(x_i)
-    return jnp.stack(x_list, axis=1)
+        x_j = x_hi[:, j - half : j - half + 1, :]
+        x_hi = jnp.where(rows_hi > j - half, x_hi - t_hi[:, :, j : j + 1] * x_j, x_hi)
+    return jnp.concatenate([x_lo, x_hi], axis=1)
 
   x_blocks = []
   iota_r = jax.lax.broadcasted_iota(jnp.int32, t.shape, 1)
diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py
index fbf30d3baf..b06926d24c 100644
--- a/src/maxtext/kernels/gdn/config.py
+++ b/src/maxtext/kernels/gdn/config.py
@@ -24,7 +24,8 @@
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 
-DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90
+# The 64-v-head prefill kernel needs about 58 MiB of the 64 MiB VMEM on TPU7x.
+DEFAULT_VMEM_LIMIT_FACTOR: float = 0.95
 
 
 class GDNMode(enum.StrEnum):
diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py
index 24a7d0bcf8..a1c4545ef5 100644
--- a/src/maxtext/kernels/gdn/memory_ref.py
+++ b/src/maxtext/kernels/gdn/memory_ref.py
@@ -586,12 +586,17 @@ def create_allocs(
   b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref)
   a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref)
 
-  conv_spec = block_spec_partial(block_shape=conv_shape)
-  recurrent_spec = block_spec_partial(block_shape=recurrent_shape)
+  # A prefill call with one sequence copies the states in only before its first tile and out only after
+  # its last, so no copy is in flight while a tile uses the slot: one buffer suffices and frees VMEM.
+  state_pipeline_mode = pipeline_mode
+  if cfg.is_prefill_only and metadata_ref.s_idx_has_initial_state.shape[0] == 1:
+    state_pipeline_mode = pl.Buffered(buffer_count=1, use_lookahead=False)
+  conv_spec = block_spec_partial(block_shape=conv_shape, pipeline_mode=state_pipeline_mode)
+  recurrent_spec = block_spec_partial(block_shape=recurrent_shape, pipeline_mode=state_pipeline_mode)
   state_buffered_partial = functools.partial(
       StateBufferedRef.input_output,
-      buffer_count=pipeline_mode.buffer_count,
-      use_lookahead=pipeline_mode.use_lookahead,
+      buffer_count=state_pipeline_mode.buffer_count,
+      use_lookahead=state_pipeline_mode.use_lookahead,
       cfg=cfg,
       metadata_ref=metadata_ref,
   )
diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py
index 3016785d5f..d44e8f44b7 100644
--- a/tests/unit/gdn_bwd_pallas_test.py
+++ b/tests/unit/gdn_bwd_pallas_test.py
@@ -27,6 +27,7 @@
 import pytest
 
 from maxtext.configs import types as config_types
+from maxtext.kernels.gdn import compute_gdn
 from maxtext.kernels.gdn import gdn_bwd_pallas
 from maxtext.kernels.gdn.gdn_bwd import api as gdn_api
 from maxtext.models import qwen3
@@ -140,6 +141,28 @@ def test_chunk_forward_matches_jax_reference(self):
     np.testing.assert_allclose(out_emit, expected_out[0], rtol=1e-3, atol=1e-3)
     np.testing.assert_allclose(state_emit, expected_state[0], rtol=5e-3, atol=5e-3)
 
+  def test_invert_triangular_matrix_matches_float64_inverse(self):
+    """The forward kernels' blocked inverse of T = I + strictly-lower(decay * beta k k^T) is exact to f32 rounding.
+
+    Chunks of 16 and less are one substitution block; 32 and 64 also substitute into rows that the
+    inter-block matmul has already updated.
+    """
+    rng = np.random.default_rng(0)
+    num_heads, head_dim = 4, 128
+    for chunk_size in (8, 16, 32, 64):
+      k = rng.standard_normal((num_heads, chunk_size, head_dim))
+      k /= np.linalg.norm(k, axis=-1, keepdims=True)
+      beta = 1 / (1 + np.exp(-rng.standard_normal((num_heads, chunk_size, 1))))
+      log_g = np.cumsum(-0.05 * np.abs(rng.standard_normal((num_heads, chunk_size))), axis=1)
+      decay = np.exp(log_g[:, :, None] - log_g[:, None, :]) * np.tril(np.ones((chunk_size, chunk_size)), -1)
+      t = (np.eye(chunk_size) + decay * np.einsum("hik,hjk->hij", k * beta, k)).astype(np.float32)
+
+      t_inv = compute_gdn.invert_triangular_matrix(jnp.asarray(t), precision=jax.lax.Precision.HIGHEST)
+
+      np.testing.assert_allclose(
+          np.asarray(t_inv, np.float64), np.linalg.inv(t.astype(np.float64)), rtol=0, atol=1e-6, err_msg=f"{chunk_size=}"
+      )
+
   def test_compute_forward_conv_and_states(self):
     batch_size = 1
     chunk_size = 16

From 92ad4cccaa83f6f1917f907bff584aa578910d03 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 10/19] GDN CP forward: fold the rank's chunk transitions in a
 Pallas scan

With gdn_cp_matmul_precision=high on TPU, the sequence-sharded CP forward
computes m_local (the product of the rank's 256 chunk transitions
M_c = g_last I - k_bwd^T (t_inv k_beta_g)) with gdn_cp_m_scan, a Pallas
emit_pipeline over (batch, head group, chunk) that keeps the running product
in VMEM, instead of compose_local_from_t_inv's balanced tree of HBM-bound f32
[chunks, heads, 128, 128] XLA fusions. s_ext_local is pass 1's end state, as
the tree returns it. CPU and 'highest' keep the XLA tree.

Same terms in sequential instead of tree order, every matmul at 3-pass bf16
like the tree's Precision.HIGH, the segment-masked gating cumsum at HIGHEST.
The scan reads k / t_inv / b / a with pallas_gdn_bwd_kernel's block layout and
runs all 64 heads of the 397B layer in one group (head_tile 64).

Numerics (v6e, 397B per-device geometry: 16 kq / 64 v heads, dims 128,
16384 tokens, chunk 64, real pass-1 t_inv): against the XLA tree at HIGHEST,
the scan's m_local is within 1.7e-6 (relative to max) and the old HIGH tree
within 9.9e-4 (long-memory gates; with model-like gates m_local is exactly 0
in all three). The whole CP layer on 4 devices (S = 65536, packed and single
long documents) gives a bit-identical output.

v6e per call: 3.3 ms against 8.3 ms for the XLA tree. On the v7x heavy
device the tree costs 3.2 ms per GDN layer (144 ms per micro-batch).
397B v7x AOT: fwd_bwd temp 93.31 GiB (93.37 before); the forward f32
[256,1,64,128,128] tree fusions are gone.

Test: gdn_cp_dh_scan_test compares the interpreted scan with the tree
(no segments, one long document, packed documents; one and two head groups),
with chunk-order and segment-start negative controls.
---
 src/maxtext/configs/types.py               |   6 +-
 src/maxtext/kernels/gdn/gdn_bwd/api.py     |  56 ++--
 src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py | 337 +++++++++++++++++++++
 tests/unit/gdn_cp_dh_scan_test.py          | 122 ++++++++
 4 files changed, 501 insertions(+), 20 deletions(-)
 create mode 100644 src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
 create mode 100644 tests/unit/gdn_cp_dh_scan_test.py

diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index ba872c9a4b..04ddb4dbec 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1523,8 +1523,10 @@ class Qwen3Next(BaseModel):
       "highest",
       description=(
           "Precision of the f32 matmuls that compose the GDN recurrent state across sequence-sharded context"
-          " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3."
-          " Only the XLA composition math changes; the Pallas kernels keep their own precision. It applies only"
+          " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3;"
+          " with 'high' the rank-local forward transition product runs as a Pallas scan (kernels/gdn/gdn_bwd/"
+          "dh_scan.py) at the same 3 passes. Only the composition math changes; the Pallas GDN kernels keep their"
+          " own precision. It applies only"
           " with use_gdn_kernel=True and sequence-sharded GDN context parallelism: gdn_cp_mode='seq', or 'auto',"
           " which shards the sequence unless the GDN context parallelism is 2 and divides the number of key"
           " heads. It has no effect with use_gdn_kernel=False (that path's context parallelism always uses"
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index 014ef8f98e..2c4adf3c80 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -25,6 +25,7 @@
 from .. import compute_conv1d as local_compute_conv1d
 from .. import wrapper as local_gdn_wrapper
 from . import cp_gdn
+from . import dh_scan
 from .compute_conv1d_bwd import conv1d_silu_bwd
 from .compute_conv1d_bwd import conv1d_silu_fwd
 from .jax_compute_gdn_states import _compute_forward_conv_and_states
@@ -512,24 +513,43 @@ def _run_cp_gdn_decoupled_fwd_impl(
       conv_halo_seg=conv_halo_seg,
       output_dtype=jnp.float32,
   )
-  m_local, s_ext_local = cp_gdn.compose_local_from_t_inv(
-      qkv_conv=k_conv,
-      b=b,
-      a=a,
-      a_log=a_log,
-      dt_bias=dt_bias,
-      t_inv=t_inv,
-      num_k_heads=num_k_heads,
-      num_v_heads=num_v_heads,
-      head_k_dim=head_k_dim,
-      head_v_dim=head_v_dim,
-      chunk_size=chunk_size,
-      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
-      s_ext_pass1=states_pass1[1],
-      segment_ids=s_enc_local,
-      init_seg=init_seg,
-      precision=cp_matmul_precision,
-  )
+  if target_platform(qkv) == "tpu" and cp_matmul_precision == jax.lax.Precision.HIGH:
+    # The same chunk transitions as compose_local_from_t_inv at the same 3-pass precision, multiplied in
+    # sequence in a Pallas scan instead of XLA's HBM-bound tree of f32 [chunks, heads, d_k, d_k] products.
+    m_local = dh_scan.gdn_cp_m_scan(
+        k_conv,
+        b,
+        a,
+        a_log,
+        dt_bias,
+        t_inv,
+        num_v_heads=num_v_heads,
+        kq_head_dim=head_k_dim,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        segment_ids=s_enc_local,
+        init_seg=init_seg,
+    )
+    s_ext_local = states_pass1[1].astype(jnp.float32)
+  else:
+    m_local, s_ext_local = cp_gdn.compose_local_from_t_inv(
+        qkv_conv=k_conv,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        t_inv=t_inv,
+        num_k_heads=num_k_heads,
+        num_v_heads=num_v_heads,
+        head_k_dim=head_k_dim,
+        head_v_dim=head_v_dim,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        s_ext_pass1=states_pass1[1],
+        segment_ids=s_enc_local,
+        init_seg=init_seg,
+        precision=cp_matmul_precision,
+    )
   m_local = checkpoint_name(m_local, "gdn_cp_m_local")
 
   h_init = recurrent_state.astype(jnp.float32) if recurrent_state is not None else zero_rs
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
new file mode 100644
index 0000000000..aa1b1a7a55
--- /dev/null
+++ b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
@@ -0,0 +1,337 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""Pallas scans that fold one sequence-sharded CP rank's GDN chunk transitions on Mosaic TPU.
+
+Per chunk $c$ of the rank, with $w_c = T_c^{-1} k^{\beta g}_c$, the forward state map is
+$h \mapsto M_c h + S_c$ with $M_c = g^{last}_c I - (k^{bwd}_c)^\top w_c$. `gdn_cp_m_scan` returns
+$M_{N-1} \cdots M_0$, the `m_local` of `cp_gdn.compose_local_from_t_inv` (its `s_ext_pass1` branch): the same
+terms as its balanced tree, multiplied in sequence, with every matmul at 3-pass bf16 (`Precision.HIGH`, as
+`gdn_cp_matmul_precision=high` gives the tree) and the segment-masked gating cumsum at `HIGHEST`. It reads the head
+groups of k / t_inv in place, with the block layout of `pallas_gdn_bwd_kernel`, and keeps the running product in
+VMEM, so no [chunks, heads, d_k, d_k] f32 intermediate reaches HBM.
+"""
+
+import functools
+from typing import Optional
+
+import jax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+from .. import compute_conv1d as local_compute_conv1d
+from .pallas_mosaic_tpu_bwd import _gdn_matmul
+from .runtime_utils import ensure_cpu_interpret_registered
+from .runtime_utils import target_platform
+
+# Value heads per grid step: all 64 heads of the Qwen3.5-397B layer in one group.
+_M_SCAN_HEAD_TILE = 64
+
+_precise_matmul = functools.partial(_gdn_matmul, compute_dtype=jnp.bfloat16, precise=True)
+
+
+def _pick_head_tile(num_v_heads: int, repeats: int, target: int) -> int:
+  """Largest divisor of num_v_heads up to target that holds whole GVA groups, as pallas_gdn_bwd_kernel picks it."""
+  for candidate in range(min(num_v_heads, target), 0, -1):
+    if num_v_heads % candidate == 0 and candidate % repeats == 0:
+      return candidate
+  return repeats if num_v_heads % repeats == 0 else num_v_heads
+
+
+def _grouped_gate_inputs(
+    b: jax.Array,
+    a: jax.Array,
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    segment_ids: Optional[jax.Array],
+    init_seg: Optional[jax.Array],
+    *,
+    num_chunks: int,
+    chunk_size: int,
+    num_groups: int,
+    tile_v_heads: int,
+):
+  """Lays out b, a, a_log, dt_bias and the segment metadata per head group, as pallas_gdn_bwd_kernel does.
+
+  Returns b / a as [B, G, C, chunk_size, padded], a_log / dt_bias as [B, G, 1, padded] and, per chunk, the active
+  segment entering it (lane 0 of [B, G, C, 1, 128]). With segment IDs, their signed encoding rides in b's spare lane
+  `tile_v_heads`, and b stays f32 then: bf16 rounds integers above 256, so distinct documents would share an ID.
+  """
+  batch_size = b.shape[0]
+  seq_len = num_chunks * chunk_size
+  has_seg_ids = segment_ids is not None
+  padded = ((tile_v_heads + int(has_seg_ids) + 127) // 128) * 128
+  lane_pad = ((0, 0),) * 4 + ((0, padded - tile_v_heads),)
+
+  def per_group(x):
+    x = x.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_heads).transpose((0, 3, 1, 2, 4))
+    return jnp.pad(x, lane_pad)
+
+  def per_group_param(x):
+    x = jnp.pad(x.reshape(num_groups, tile_v_heads), ((0, 0), (0, padded - tile_v_heads)))
+    return jnp.broadcast_to(x[None, :, None, :], (batch_size, num_groups, 1, padded))
+
+  b_5d = per_group(b.astype(jnp.float32) if has_seg_ids else b)
+  a_5d = per_group(a)
+  if has_seg_ids:
+    s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch_size, seq_len), init_seg=init_seg)
+    s_enc_3d = s_enc.reshape(batch_size, num_chunks, chunk_size)
+    b_5d = b_5d.at[:, :, :, :, tile_v_heads].set(s_enc_3d[:, None, :, :].astype(b_5d.dtype))
+    if init_seg is not None:
+      init_active = jnp.abs(init_seg.astype(jnp.float32).reshape(batch_size, 1))
+    else:
+      init_active = jnp.zeros((batch_size, 1), dtype=jnp.float32)
+    seg_prev = jnp.concatenate([init_active, jnp.abs(s_enc_3d[:, :-1, -1])], axis=1)
+    seg_prev_hbm = jnp.pad(seg_prev[:, :, None, None], ((0, 0), (0, 0), (0, 0), (0, 127)))
+  else:
+    seg_prev_hbm = jnp.zeros((batch_size, num_chunks, 1, 128), dtype=jnp.float32)
+  seg_prev_5d = jnp.broadcast_to(seg_prev_hbm[:, None], (batch_size, num_groups, num_chunks, 1, 128))
+  return b_5d, a_5d, per_group_param(a_log), per_group_param(dt_bias), seg_prev_5d, padded
+
+
+def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_size, num_v_heads, has_seg_ids):
+  """Segment masks and decay gates of one chunk and head group, as cp_gdn computes them.
+
+  Returns (valid [C, 1] or None, beta_h [H, C], gating_forward [H, C, 1], gating_last [H, 1, 1],
+  gating_backward [H, C, 1]).
+  """
+  b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32)
+  a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32)
+  a_log_val = a_log_ref[...][0, :num_v_heads].astype(jnp.float32)
+  dt_bias_val = dt_bias_ref[...][0, :num_v_heads].astype(jnp.float32)
+  beta = jax.nn.sigmoid(b_val)
+  log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val)
+  mask_tril = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32))
+
+  valid_c = None
+  mask_cumsum = mask_tril
+  if has_seg_ids:
+    seg_c = b_ref[...][:, num_v_heads : num_v_heads + 1].astype(jnp.float32)
+    seg_prev = seg_prev_ref[...][0, 0].astype(jnp.float32)
+    valid_c = seg_c > 0.5
+    active_c = jnp.abs(seg_c)
+    beta = jnp.where(valid_c, beta, 0.0)
+    log_g = jnp.where(valid_c, log_g, 0.0)
+    same_active = (jnp.abs(active_c - active_c.T) < 0.5) & (active_c > 0.5)
+    mask_cumsum = mask_tril * same_active.astype(jnp.float32)
+
+  # The cumsum feeds exp(); like cp_gdn it stays at HIGHEST (0/1 mask times f32: exact up to summation order).
+  cumsum_h = jnp.transpose(jnp.dot(mask_cumsum, log_g, precision=jax.lax.Precision.HIGHEST), (1, 0))
+  gating_forward = jnp.exp(cumsum_h)[:, :, None]
+  gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None]
+  if has_seg_ids:
+    m_in = ((jnp.abs(seg_c - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(jnp.float32)
+    m_out = ((jnp.abs(seg_c - active_c[-1:]) < 0.5) & (active_c[-1:] > 0.5)).astype(jnp.float32)
+    m_keep = ((jnp.abs(active_c[-1:] - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(jnp.float32)
+    gating_forward = gating_forward * m_in[None, :, :]
+    gating_last = gating_last * m_keep[None, :, :]
+    bwd_diff = jnp.where(m_out.T > 0.5, cumsum_h[:, -1:] - cumsum_h, -1e4)
+    gating_backward = jnp.exp(bwd_diff)[:, :, None] * m_out[None, :, :]
+  else:
+    gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None]
+  beta_h = jnp.transpose(beta, (1, 0))
+  return valid_c, beta_h, gating_forward, gating_last, gating_backward
+
+
+def _unit_rows(x: jax.Array, use_qk_norm: bool) -> jax.Array:
+  """The q / k l2 normalization of cp_gdn: x * rsqrt(sum(x**2) + 1e-6) over the head dim."""
+  if not use_qk_norm:
+    return x
+  return x * jax.lax.rsqrt(jnp.sum(x**2, axis=-1, keepdims=True) + 1e-6)
+
+
+def _m_scan_body(
+    k_ref,
+    b_ref,
+    a_ref,
+    t_inv_ref,
+    a_log_ref,
+    dt_bias_ref,
+    seg_prev_ref,
+    m_ref,
+    acc_scr,
+    *,
+    chunk_size: int,
+    num_kq_heads: int,
+    num_v_heads: int,
+    kq_head_dim: int,
+    num_chunks: int,
+    has_seg_ids: bool,
+    use_qk_norm: bool,
+):
+  """Forward step per (batch, group, chunk): acc <- M_c acc = g_last acc - k_bwd^T (w (acc)), from acc = I."""
+  c = pl.program_id(2)
+
+  @pl.when(c == 0)
+  def _init():
+    rows = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 0)
+    cols = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 1)
+    acc_scr[...] = jnp.broadcast_to((rows == cols).astype(jnp.float32)[None], acc_scr.shape)
+
+  valid_c, beta_h, gating_forward, gating_last, gating_backward = _chunk_gates(
+      b_ref,
+      a_ref,
+      a_log_ref,
+      dt_bias_ref,
+      seg_prev_ref,
+      chunk_size=chunk_size,
+      num_v_heads=num_v_heads,
+      has_seg_ids=has_seg_ids,
+  )
+  k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  if valid_c is not None:
+    k_orig = jnp.where(valid_c[:, :, None], k_orig, 0.0)
+  k_h = jnp.transpose(jnp.repeat(_unit_rows(k_orig, use_qk_norm), num_v_heads // num_kq_heads, axis=1), (1, 0, 2))
+  k_beta_g = k_h * beta_h[:, :, None] * gating_forward
+  k_bwd = k_h * gating_backward
+
+  acc = acc_scr[...]
+  w = _precise_matmul(t_inv_ref[...].astype(jnp.float32), k_beta_g)
+  acc_new = acc * gating_last - _precise_matmul(jnp.swapaxes(k_bwd, -1, -2), _precise_matmul(w, acc))
+  acc_scr[...] = acc_new
+
+  @pl.when(c == num_chunks - 1)
+  def _store():
+    m_ref[0, ...] = acc_new
+
+
+def _run_scan(body, args, in_specs, out_shape, out_spec, acc_shape, *, grid, interpret, name):
+  """One pallas_call that pipelines `body` over grid=(B, G, C) with the running product in a VMEM scratch."""
+
+  def outer(*refs):
+    pltpu.emit_pipeline(body, grid=grid, in_specs=in_specs, out_specs=[out_spec])(
+        *refs[: len(args) + 1], scratches=tuple(refs[len(args) + 1 :])
+    )
+
+  hbm = pltpu.MemorySpace.HBM
+  return pl.pallas_call(
+      outer,
+      grid=(),
+      out_shape=out_shape,
+      in_specs=[pl.BlockSpec(memory_space=hbm)] * len(args),
+      out_specs=pl.BlockSpec(memory_space=hbm),
+      scratch_shapes=[pltpu.VMEM(acc_shape, jnp.float32)],
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=int(0.9 * pltpu.get_tpu_info().vmem_capacity_bytes),
+          disable_bounds_checks=True,
+      ),
+      interpret=interpret,
+      name=name,
+  )(*args)
+
+
+def _resolve_interpret(x: jax.Array, interpret: bool | pltpu.InterpretParams | None):
+  if interpret is None:
+    interpret = target_platform(x) == "cpu"
+  if interpret:
+    ensure_cpu_interpret_registered()
+  return interpret
+
+
+def gdn_cp_m_scan(
+    k_conv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    t_inv: jax.Array,
+    *,
+    num_v_heads: int,
+    kq_head_dim: int,
+    chunk_size: int = 64,
+    use_qk_norm_in_gdn: bool = False,
+    segment_ids: Optional[jax.Array] = None,
+    init_seg: Optional[jax.Array] = None,
+    interpret: bool | pltpu.InterpretParams | None = None,
+) -> jax.Array:
+  """M_local of `cp_gdn.compose_local_from_t_inv` (s_ext_pass1 branch), [B, num_v_heads, d_k, d_k] f32.
+
+  Args:
+    k_conv: [B, S, num_kq_heads * d_k] conv output of the K channels, as passed to compose_local_from_t_inv.
+    b, a: [B, S, num_v_heads] beta and decay pre-activations.
+    a_log, dt_bias: [num_v_heads] decay parameters.
+    t_inv: [B, S / chunk_size, num_v_heads, chunk_size, chunk_size] per-chunk T^{-1} of the forward kernel.
+    num_v_heads, kq_head_dim, chunk_size, use_qk_norm_in_gdn: GDN layer geometry.
+    segment_ids, init_seg: Rank-local segment metadata, as passed to compose_local_from_t_inv.
+    interpret: Pallas interpret mode; None interprets on CPU.
+  """
+  interpret = _resolve_interpret(k_conv, interpret)
+  batch_size, seq_len, k_size = k_conv.shape
+  num_chunks = seq_len // chunk_size
+  num_kq_heads = k_size // kq_head_dim
+  repeats = num_v_heads // num_kq_heads
+  tile_v_heads = _pick_head_tile(num_v_heads, repeats, _M_SCAN_HEAD_TILE)
+  num_groups = num_v_heads // tile_v_heads
+  tile_k_size = (tile_v_heads // repeats) * kq_head_dim
+  has_seg_ids = segment_ids is not None
+
+  b_5d, a_5d, a_log_4d, dt_bias_4d, seg_prev_5d, padded = _grouped_gate_inputs(
+      b,
+      a,
+      a_log,
+      dt_bias,
+      segment_ids,
+      init_seg,
+      num_chunks=num_chunks,
+      chunk_size=chunk_size,
+      num_groups=num_groups,
+      tile_v_heads=tile_v_heads,
+  )
+  args = [
+      k_conv.reshape(batch_size, num_chunks, chunk_size, k_size),
+      b_5d,
+      a_5d,
+      t_inv.reshape(batch_size, num_chunks, num_groups, tile_v_heads, chunk_size, chunk_size).astype(jnp.float32),
+      a_log_4d,
+      dt_bias_4d,
+      seg_prev_5d,
+  ]
+  in_specs = [
+      # k [B, C, chunk_size, k_size], columns of group g
+      pl.BlockSpec((None, None, chunk_size, tile_k_size), lambda b_, g, c: (b_, c, 0, g)),
+      # b, a [B, G, C, chunk_size, padded]
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, c, 0, 0)),
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, c, 0, 0)),
+      # t_inv [B, C, G, tile_v_heads, chunk_size, chunk_size]
+      pl.BlockSpec((None, None, None, tile_v_heads, chunk_size, chunk_size), lambda b_, g, c: (b_, c, g, 0, 0, 0)),
+      # a_log, dt_bias [B, G, 1, padded]
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      # segment entering the chunk [B, G, C, 1, 128]
+      pl.BlockSpec((None, None, None, 1, 128), lambda b_, g, c: (b_, g, c, 0, 0)),
+  ]
+  acc_shape = (tile_v_heads, kq_head_dim, kq_head_dim)
+  body = functools.partial(
+      _m_scan_body,
+      chunk_size=chunk_size,
+      num_kq_heads=tile_v_heads // repeats,
+      num_v_heads=tile_v_heads,
+      kq_head_dim=kq_head_dim,
+      num_chunks=num_chunks,
+      has_seg_ids=has_seg_ids,
+      use_qk_norm=use_qk_norm_in_gdn,
+  )
+  m_6d = _run_scan(
+      body,
+      args,
+      in_specs,
+      jax.ShapeDtypeStruct((batch_size, num_groups, 1) + acc_shape, jnp.float32),
+      pl.BlockSpec((None, None, 1) + acc_shape, lambda b_, g, c: (b_, g, 0, 0, 0, 0)),
+      acc_shape,
+      grid=(batch_size, num_groups, num_chunks),
+      interpret=interpret,
+      name="gdn_cp_m_scan",
+  )
+  return m_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, kq_head_dim)
diff --git a/tests/unit/gdn_cp_dh_scan_test.py b/tests/unit/gdn_cp_dh_scan_test.py
new file mode 100644
index 0000000000..6655fd43f4
--- /dev/null
+++ b/tests/unit/gdn_cp_dh_scan_test.py
@@ -0,0 +1,122 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The Pallas CP composition scan (`dh_scan`) against the XLA `cp_gdn` tree it replaces on TPU.
+
+Runs the kernel in Pallas interpret mode. XLA:CPU computes the tree's f32 matmuls in plain f32, so the comparison
+bounds the scan's 3-pass bf16 matmul error and the reduction order, not TPU numerics.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.kernels.gdn.gdn_bwd import cp_gdn
+from maxtext.kernels.gdn.gdn_bwd import dh_scan
+
+_NK, _NV, _D, _C, _NC = 2, 8, 128, 64, 4
+
+
+def _inputs(seed=0):
+  """One rank's K-slice conv output, gates and per-chunk T^{-1}, with weak decay."""
+  keys = jax.random.split(jax.random.PRNGKey(seed), 6)
+  s = _NC * _C
+  k_conv = jax.random.normal(keys[0], (1, s, _NK * _D), jnp.float32)
+  b = jax.random.normal(keys[1], (1, s, _NV), jnp.float32).astype(jnp.bfloat16)
+  # Weak decay keeps the product of the 4 chunk transitions far from 0.
+  a = (jax.random.normal(keys[2], (1, s, _NV), jnp.float32) - 6.0).astype(jnp.bfloat16)
+  a_log = jax.random.normal(keys[3], (_NV,), jnp.float32) * 0.3
+  dt_bias = jax.random.normal(keys[4], (_NV,), jnp.float32) * 0.3
+  t_inv = jnp.tril(jax.random.normal(keys[5], (1, _NC, _NV, _C, _C), jnp.float32) * 0.05, -1) + jnp.eye(_C)
+  return k_conv, b, a, a_log, dt_bias, t_inv
+
+
+def _tree_m(k_conv, b, a, a_log, dt_bias, t_inv, segment_ids, init_seg):
+  """m_local of the XLA tree (f32 on CPU)."""
+  m, _ = cp_gdn.compose_local_from_t_inv(
+      k_conv,
+      b,
+      a,
+      a_log,
+      dt_bias,
+      t_inv,
+      num_k_heads=_NK,
+      num_v_heads=_NV,
+      head_k_dim=_D,
+      head_v_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      s_ext_pass1=jnp.zeros((1, _NV, _D, _D), jnp.float32),
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+  )
+  return m
+
+
+def _scan_m(k_conv, b, a, a_log, dt_bias, t_inv, segment_ids, init_seg):
+  """m_local of the Pallas scan, interpreted."""
+  return dh_scan.gdn_cp_m_scan(
+      k_conv,
+      b,
+      a,
+      a_log,
+      dt_bias,
+      t_inv,
+      num_v_heads=_NV,
+      kq_head_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+      interpret=True,
+  )
+
+
+def _rel_err(got, ref):
+  got, ref = np.asarray(got, np.float64), np.asarray(ref, np.float64)
+  return np.abs(got - ref).max() / np.abs(ref).max()
+
+
+# A document continuing from the previous rank and spanning the whole rank (the long single-document case), and a
+# packed row: that document ends inside chunk 1, two more follow with boundaries inside chunks, then padding.
+_SINGLE = np.ones((1, _NC * _C), np.int32)
+_PACKED = np.zeros((1, _NC * _C), np.int32)
+_PACKED[0, :70], _PACKED[0, 70:200], _PACKED[0, 200:240] = 1, 2, 3
+
+
+@pytest.mark.parametrize("head_tile", [64, 4])  # one head group, two head groups
+@pytest.mark.parametrize("layout", ["none", "single", "packed"])
+def test_m_scan_matches_compose_local_tree(layout, head_tile, monkeypatch):
+  monkeypatch.setattr(dh_scan, "_M_SCAN_HEAD_TILE", head_tile)
+  segment_ids = {"none": None, "single": jnp.asarray(_SINGLE), "packed": jnp.asarray(_PACKED)}[layout]
+  init_seg = None if segment_ids is None else jnp.ones((1,), jnp.float32)
+  args = _inputs()
+  ref = _tree_m(*args, segment_ids, init_seg)
+  got = _scan_m(*args, segment_ids, init_seg)
+  assert float(jnp.abs(ref).max()) > 1e-4, "degenerate reference: the transitions vanish"
+  assert _rel_err(got, ref) < 5e-5
+
+
+def test_m_scan_sees_chunk_order_and_segment_starts():
+  """Negative controls: the comparison above can fail."""
+  args = _inputs()
+  segment_ids = jnp.asarray(_SINGLE)
+  ref = _tree_m(*args, segment_ids, jnp.ones((1,), jnp.float32))
+  k_conv, b, a, a_log, dt_bias, t_inv = args
+  reversed_chunks = _scan_m(k_conv, b, a, a_log, dt_bias, t_inv[:, ::-1], segment_ids, jnp.ones((1,), jnp.float32))
+  assert _rel_err(reversed_chunks, ref) > 1e-3
+  # Without the incoming document, the rank's first token starts a new one, and nothing from the incoming state
+  # reaches the rank's end state.
+  fresh = _scan_m(*args, segment_ids, jnp.zeros((1,), jnp.float32))
+  assert float(jnp.abs(fresh).max()) == 0.0

From 4026f507997db825e2a635f014e8eb6b46b39fda Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:11:50 +0000
Subject: [PATCH 11/19] GDN CP backward: dS_ext_local as a reverse Pallas scan

With gdn_cp_matmul_precision=high on TPU, the sequence-sharded CP backward
computes ds_ext_local with gdn_cp_dh_scan, a reverse Pallas emit_pipeline over
(batch, 2 groups of 32 heads, chunk) in the backward kernel's form of the state
recurrence (acc <- g_last acc + q_g^T do - k_beta_g^T t_inv^T (attn^T do +
k_bwd acc)), instead of compose_bwd_local_from_t_inv's balanced tree of
HBM-bound f32 [chunks, heads, 128, 128] XLA fusions. dm_local is the transposed
forward m_local, as the tree returns it with m_local_cached. CPU and 'highest'
keep the XLA tree.

Numerics (v6e, 397B per-device geometry, real pass-1 t_inv): against the XLA
tree at HIGHEST, the scan's ds_ext_local is within 3.1e-6 (relative to max),
the old HIGH tree within 1.7e-5 to 4.8e-4. Whole CP layer on 4 devices
(S = 65536): output bit-identical; bf16 input grads within 1 bf16 ulp (as HIGH
vs HIGHEST); f32 param grads differ by 1e-4 to 2e-3 relative, the size of the
old path's own HIGH-vs-HIGHEST difference.

v6e per call: 4.8 ms against 24 ms for the XLA tree (v7x heavy device: tree
8.95 ms per GDN layer, 403 ms per micro-batch).

Memory: 397B v7x AOT fwd_bwd temp +1.3 GiB. In the backward loop body XLA now
materializes the f32 conv recompute (about 2.25 GiB) before the second GDN
layer's gdn_bwd_kernel instead of after it, so it overlaps that layer's
replayed chunk_states. The bench config still fits and runs on 128 v7x chips.
---
 src/maxtext/configs/types.py               |  13 +-
 src/maxtext/kernels/gdn/gdn_bwd/api.py     |  60 ++++--
 src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py | 228 ++++++++++++++++++++-
 tests/unit/gdn_cp_dh_scan_test.py          |  71 ++++++-
 4 files changed, 333 insertions(+), 39 deletions(-)

diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 04ddb4dbec..fca14624ed 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1524,13 +1524,12 @@ class Qwen3Next(BaseModel):
       description=(
           "Precision of the f32 matmuls that compose the GDN recurrent state across sequence-sharded context"
           " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3;"
-          " with 'high' the rank-local forward transition product runs as a Pallas scan (kernels/gdn/gdn_bwd/"
-          "dh_scan.py) at the same 3 passes. Only the composition math changes; the Pallas GDN kernels keep their"
-          " own precision. It applies only"
-          " with use_gdn_kernel=True and sequence-sharded GDN context parallelism: gdn_cp_mode='seq', or 'auto',"
-          " which shards the sequence unless the GDN context parallelism is 2 and divides the number of key"
-          " heads. It has no effect with use_gdn_kernel=False (that path's context parallelism always uses"
-          " HIGHEST) or with head-sharded context parallelism."
+          " with 'high' the rank-local forward transition product and backward state gradient run as Pallas scans"
+          " (kernels/gdn/gdn_bwd/dh_scan.py) at the same 3 passes. Only the composition math changes; the Pallas GDN"
+          " kernels keep their own precision. It applies only with use_gdn_kernel=True and sequence-sharded GDN"
+          " context parallelism: gdn_cp_mode='seq', or 'auto', which shards the sequence unless the GDN context"
+          " parallelism is 2 and divides the number of key heads. It has no effect with use_gdn_kernel=False (that"
+          " path's context parallelism always uses HIGHEST) or with head-sharded context parallelism."
       ),
   )
 
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index 2c4adf3c80..e56aa53cd8 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -1027,25 +1027,47 @@ def _gdn_decoupled_conv1d_bwd(
         if d_conv_state is not None
         else None
     )
-    dm_local, ds_ext_local = cp_gdn.compose_bwd_local_from_t_inv(
-        qkv_conv=qkv_conv,
-        b=b,
-        a=a,
-        a_log=a_log,
-        dt_bias=dt_bias,
-        do=d_out,
-        t_inv=t_inv,
-        num_k_heads=num_k_heads,
-        num_v_heads=num_v_heads,
-        head_k_dim=head_k_dim,
-        head_v_dim=head_v_dim,
-        chunk_size=chunk_size,
-        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
-        m_local_cached=m_local_fwd,
-        segment_ids=segment_ids,
-        init_seg=init_seg,
-        precision=cp_matmul_precision,
-    )
+    if target_platform(qkv_conv) == "tpu" and m_local_fwd is not None and cp_matmul_precision == jax.lax.Precision.HIGH:
+      # ds_ext_local of compose_bwd_local_from_t_inv at the same 3-pass precision, as a reverse Pallas scan
+      # (Horner order) instead of XLA's HBM-bound tree. dm_local is the transposed forward transition, which
+      # compose_bwd_local_from_t_inv also returns when given m_local_cached.
+      ds_ext_local = dh_scan.gdn_cp_dh_scan(
+          qkv_conv,
+          b,
+          a,
+          a_log,
+          dt_bias,
+          d_out,
+          t_inv,
+          num_v_heads=num_v_heads,
+          kq_head_dim=head_k_dim,
+          v_head_dim=head_v_dim,
+          chunk_size=chunk_size,
+          use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+          segment_ids=segment_ids,
+          init_seg=init_seg,
+      )
+      dm_local = jnp.swapaxes(m_local_fwd.astype(jnp.float32), -1, -2)
+    else:
+      dm_local, ds_ext_local = cp_gdn.compose_bwd_local_from_t_inv(
+          qkv_conv=qkv_conv,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          do=d_out,
+          t_inv=t_inv,
+          num_k_heads=num_k_heads,
+          num_v_heads=num_v_heads,
+          head_k_dim=head_k_dim,
+          head_v_dim=head_v_dim,
+          chunk_size=chunk_size,
+          use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+          m_local_cached=m_local_fwd,
+          segment_ids=segment_ids,
+          init_seg=init_seg,
+          precision=cp_matmul_precision,
+      )
     dht_local, _ = cp_gdn.incoming_grad_state(
         dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision
     )
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
index aa1b1a7a55..b4dba19cf4 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
@@ -15,12 +15,20 @@
 r"""Pallas scans that fold one sequence-sharded CP rank's GDN chunk transitions on Mosaic TPU.
 
 Per chunk $c$ of the rank, with $w_c = T_c^{-1} k^{\beta g}_c$, the forward state map is
-$h \mapsto M_c h + S_c$ with $M_c = g^{last}_c I - (k^{bwd}_c)^\top w_c$. `gdn_cp_m_scan` returns
-$M_{N-1} \cdots M_0$, the `m_local` of `cp_gdn.compose_local_from_t_inv` (its `s_ext_pass1` branch): the same
-terms as its balanced tree, multiplied in sequence, with every matmul at 3-pass bf16 (`Precision.HIGH`, as
-`gdn_cp_matmul_precision=high` gives the tree) and the segment-masked gating cumsum at `HIGHEST`. It reads the head
-groups of k / t_inv in place, with the block layout of `pallas_gdn_bwd_kernel`, and keeps the running product in
-VMEM, so no [chunks, heads, d_k, d_k] f32 intermediate reaches HBM.
+$h \mapsto M_c h + S_c$ with $M_c = g^{last}_c I - (k^{bwd}_c)^\top w_c$, and the backward one is
+$\mathrm{d}h \mapsto M_c^\top \mathrm{d}h + \mathrm{d}S_c$ with
+$\mathrm{d}S_c = q_g^\top \mathrm{d}o - (k^{\beta g}_c)^\top T_c^{-\top} (\mathrm{attn}_c^\top \mathrm{d}o)$.
+
+- `gdn_cp_m_scan` returns $M_{N-1} \cdots M_0$, the `m_local` of `cp_gdn.compose_local_from_t_inv`
+  (its `s_ext_pass1` branch).
+- `gdn_cp_dh_scan` returns $\mathrm{d}S_0 + M_0^\top(\mathrm{d}S_1 + M_1^\top(\cdots))$, the `ds_ext_local` of
+  `cp_gdn.compose_bwd_local_from_t_inv`, as a reverse scan in the form of the backward kernel's state recurrence.
+
+Both are the same terms as the `cp_gdn` balanced trees in sequential order, with every matmul at 3-pass bf16
+(`Precision.HIGH`, as `gdn_cp_matmul_precision=high` gives the trees) and the segment-masked gating cumsum at
+`HIGHEST`. They read the head groups of qkv / do / t_inv in place, with the block layout of
+`pallas_gdn_bwd_kernel`, and keep the running product in VMEM, so no [chunks, heads, d_k, d_k] f32 intermediate
+reaches HBM.
 """
 
 import functools
@@ -36,7 +44,9 @@
 from .runtime_utils import ensure_cpu_interpret_registered
 from .runtime_utils import target_platform
 
-# Value heads per grid step: all 64 heads of the Qwen3.5-397B layer in one group.
+# Value heads per grid step. The dh-scan's per-chunk working set grows with the head tile; the M-scan is lighter
+# and runs all 64 heads of the Qwen3.5-397B layer in one group.
+_DH_SCAN_HEAD_TILE = 32
 _M_SCAN_HEAD_TILE = 64
 
 _precise_matmul = functools.partial(_gdn_matmul, compute_dtype=jnp.bfloat16, precise=True)
@@ -104,8 +114,8 @@ def per_group_param(x):
 def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_size, num_v_heads, has_seg_ids):
   """Segment masks and decay gates of one chunk and head group, as cp_gdn computes them.
 
-  Returns (valid [C, 1] or None, beta_h [H, C], gating_forward [H, C, 1], gating_last [H, 1, 1],
-  gating_backward [H, C, 1]).
+  Returns (valid [C, 1] or None, beta_h [H, C], cumsum_h [H, C], gating_forward [H, C, 1], gating_last [H, 1, 1],
+  gating_backward [H, C, 1], mask_causal [C, C]).
   """
   b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32)
   a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32)
@@ -117,6 +127,7 @@ def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_si
 
   valid_c = None
   mask_cumsum = mask_tril
+  mask_causal = mask_tril
   if has_seg_ids:
     seg_c = b_ref[...][:, num_v_heads : num_v_heads + 1].astype(jnp.float32)
     seg_prev = seg_prev_ref[...][0, 0].astype(jnp.float32)
@@ -125,7 +136,9 @@ def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_si
     beta = jnp.where(valid_c, beta, 0.0)
     log_g = jnp.where(valid_c, log_g, 0.0)
     same_active = (jnp.abs(active_c - active_c.T) < 0.5) & (active_c > 0.5)
+    same_valid = (jnp.abs(seg_c - seg_c.T) < 0.5) & valid_c
     mask_cumsum = mask_tril * same_active.astype(jnp.float32)
+    mask_causal = mask_tril * same_valid.astype(jnp.float32)
 
   # The cumsum feeds exp(); like cp_gdn it stays at HIGHEST (0/1 mask times f32: exact up to summation order).
   cumsum_h = jnp.transpose(jnp.dot(mask_cumsum, log_g, precision=jax.lax.Precision.HIGHEST), (1, 0))
@@ -142,7 +155,7 @@ def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_si
   else:
     gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None]
   beta_h = jnp.transpose(beta, (1, 0))
-  return valid_c, beta_h, gating_forward, gating_last, gating_backward
+  return valid_c, beta_h, cumsum_h, gating_forward, gating_last, gating_backward, mask_causal
 
 
 def _unit_rows(x: jax.Array, use_qk_norm: bool) -> jax.Array:
@@ -180,7 +193,7 @@ def _init():
     cols = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 1)
     acc_scr[...] = jnp.broadcast_to((rows == cols).astype(jnp.float32)[None], acc_scr.shape)
 
-  valid_c, beta_h, gating_forward, gating_last, gating_backward = _chunk_gates(
+  valid_c, beta_h, _, gating_forward, gating_last, gating_backward, _ = _chunk_gates(
       b_ref,
       a_ref,
       a_log_ref,
@@ -207,6 +220,88 @@ def _store():
     m_ref[0, ...] = acc_new
 
 
+def _dh_scan_body(
+    q_ref,
+    k_ref,
+    b_ref,
+    a_ref,
+    do_ref,
+    t_inv_ref,
+    a_log_ref,
+    dt_bias_ref,
+    seg_prev_ref,
+    ds_ref,
+    acc_scr,
+    *,
+    chunk_size: int,
+    num_kq_heads: int,
+    num_v_heads: int,
+    kq_head_dim: int,
+    num_chunks: int,
+    has_seg_ids: bool,
+    use_qk_norm: bool,
+):
+  """Reverse step per (batch, group, chunk): acc <- dS_c + M_c^T acc, from acc = 0.
+
+  M_c^T acc + dS_c is evaluated as the backward kernel's state recurrence does it:
+  g_last acc + q_g^T do - k_beta_g^T (t_inv^T (attn^T do + k_bwd acc)).
+  """
+  c = pl.program_id(2)
+  repeats = num_v_heads // num_kq_heads
+
+  @pl.when(c == 0)
+  def _init():
+    acc_scr[...] = jnp.zeros_like(acc_scr)
+
+  valid_c, beta_h, cumsum_h, gating_forward, gating_last, gating_backward, mask_causal = _chunk_gates(
+      b_ref,
+      a_ref,
+      a_log_ref,
+      dt_bias_ref,
+      seg_prev_ref,
+      chunk_size=chunk_size,
+      num_v_heads=num_v_heads,
+      has_seg_ids=has_seg_ids,
+  )
+  q_orig = q_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  do_val = do_ref[...].astype(jnp.float32)
+  if valid_c is not None:
+    q_orig = jnp.where(valid_c[:, :, None], q_orig, 0.0)
+    k_orig = jnp.where(valid_c[:, :, None], k_orig, 0.0)
+    do_val = jnp.where(valid_c[:, :, None], do_val, 0.0)
+  q_scaled = _unit_rows(q_orig, use_qk_norm) * (1.0 / jnp.sqrt(kq_head_dim))
+  k_scaled = _unit_rows(k_orig, use_qk_norm)
+
+  q_h = jnp.transpose(jnp.repeat(q_scaled, repeats, axis=1), (1, 0, 2))
+  k_h = jnp.transpose(jnp.repeat(k_scaled, repeats, axis=1), (1, 0, 2))
+  do_h = jnp.transpose(do_val, (1, 0, 2))
+  diff = cumsum_h[:, :, None] - cumsum_h[:, None, :]
+  g_mat_causal = jnp.exp(jnp.where(mask_causal[None] > 0.5, diff, -1e4)) * mask_causal[None]
+  q_g = q_h * gating_forward
+  k_beta_g = k_h * beta_h[:, :, None] * gating_forward
+  k_bwd = k_h * gating_backward
+
+  # q k^T once per kq head, then repeated to its value heads: the same dot products as with repeated q and k.
+  qk = _precise_matmul(jnp.transpose(q_scaled, (1, 0, 2)), jnp.swapaxes(jnp.transpose(k_scaled, (1, 0, 2)), -1, -2))
+  attn = jnp.repeat(qk, repeats, axis=0) * g_mat_causal
+  dv_attn = _precise_matmul(jnp.swapaxes(attn, -1, -2), do_h)
+
+  acc = acc_scr[...]
+  t_inv_t = jnp.swapaxes(t_inv_ref[...].astype(jnp.float32), -1, -2)
+  d_v_beta = _precise_matmul(t_inv_t, dv_attn + _precise_matmul(k_bwd, acc))
+  # q_g^T do - k_beta_g^T d_v_beta as one matmul over the stacked chunk rows.
+  acc_new = acc * gating_last + _precise_matmul(
+      jnp.swapaxes(jnp.concatenate([q_g, -k_beta_g], axis=1), -1, -2),
+      jnp.concatenate([do_h, d_v_beta], axis=1),
+  )
+  acc_scr[...] = acc_new
+
+  @pl.when(c == num_chunks - 1)
+  def _store():
+    ds_ref[0, ...] = acc_new
+
+
 def _run_scan(body, args, in_specs, out_shape, out_spec, acc_shape, *, grid, interpret, name):
   """One pallas_call that pipelines `body` over grid=(B, G, C) with the running product in a VMEM scratch."""
 
@@ -335,3 +430,114 @@ def gdn_cp_m_scan(
       name="gdn_cp_m_scan",
   )
   return m_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, kq_head_dim)
+
+
+def gdn_cp_dh_scan(
+    qkv_conv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    do: jax.Array,
+    t_inv: jax.Array,
+    *,
+    num_v_heads: int,
+    kq_head_dim: int,
+    v_head_dim: int,
+    chunk_size: int = 64,
+    use_qk_norm_in_gdn: bool = False,
+    segment_ids: Optional[jax.Array] = None,
+    init_seg: Optional[jax.Array] = None,
+    interpret: bool | pltpu.InterpretParams | None = None,
+) -> jax.Array:
+  """dS_ext_local of `cp_gdn.compose_bwd_local_from_t_inv`, [B, num_v_heads, d_k, d_v] f32.
+
+  Args:
+    qkv_conv: [B, S, 2 * num_kq_heads * d_k + num_v_heads * d_v] conv output (q, k, v channels).
+    b, a: [B, S, num_v_heads] beta and decay pre-activations.
+    a_log, dt_bias: [num_v_heads] decay parameters.
+    do: [B, S, num_v_heads, d_v] cotangent of the GDN output.
+    t_inv: [B, S / chunk_size, num_v_heads, chunk_size, chunk_size] per-chunk T^{-1} of the forward kernel.
+    num_v_heads, kq_head_dim, v_head_dim, chunk_size, use_qk_norm_in_gdn: GDN layer geometry.
+    segment_ids, init_seg: Rank-local segment metadata, as passed to compose_bwd_local_from_t_inv.
+    interpret: Pallas interpret mode; None interprets on CPU.
+  """
+  interpret = _resolve_interpret(qkv_conv, interpret)
+  batch_size, seq_len, dim_size = qkv_conv.shape
+  num_chunks = seq_len // chunk_size
+  num_kq_heads = (dim_size - num_v_heads * v_head_dim) // (kq_head_dim * 2)
+  repeats = num_v_heads // num_kq_heads
+  tile_v_heads = _pick_head_tile(num_v_heads, repeats, _DH_SCAN_HEAD_TILE)
+  num_groups = num_v_heads // tile_v_heads
+  tile_q_size = (tile_v_heads // repeats) * kq_head_dim
+  k_block_offset = num_kq_heads * kq_head_dim // tile_q_size
+  has_seg_ids = segment_ids is not None
+
+  b_5d, a_5d, a_log_4d, dt_bias_4d, seg_prev_5d, padded = _grouped_gate_inputs(
+      b,
+      a,
+      a_log,
+      dt_bias,
+      segment_ids,
+      init_seg,
+      num_chunks=num_chunks,
+      chunk_size=chunk_size,
+      num_groups=num_groups,
+      tile_v_heads=tile_v_heads,
+  )
+  qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, dim_size)
+  args = [
+      qkv_conv_4d,
+      qkv_conv_4d,
+      b_5d,
+      a_5d,
+      do.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_heads, v_head_dim),
+      t_inv.reshape(batch_size, num_chunks, num_groups, tile_v_heads, chunk_size, chunk_size).astype(jnp.float32),
+      a_log_4d,
+      dt_bias_4d,
+      seg_prev_5d,
+  ]
+
+  def rc(c):
+    return num_chunks - 1 - c
+
+  in_specs = [
+      # q, k [B, C, chunk_size, dim_size], columns of group g (k after the q blocks)
+      pl.BlockSpec((None, None, chunk_size, tile_q_size), lambda b_, g, c: (b_, rc(c), 0, g)),
+      pl.BlockSpec((None, None, chunk_size, tile_q_size), lambda b_, g, c: (b_, rc(c), 0, k_block_offset + g)),
+      # b, a [B, G, C, chunk_size, padded]
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, rc(c), 0, 0)),
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, rc(c), 0, 0)),
+      # do [B, C, chunk_size, G, tile_v_heads, d_v]
+      pl.BlockSpec((None, None, chunk_size, None, tile_v_heads, v_head_dim), lambda b_, g, c: (b_, rc(c), 0, g, 0, 0)),
+      # t_inv [B, C, G, tile_v_heads, chunk_size, chunk_size]
+      pl.BlockSpec((None, None, None, tile_v_heads, chunk_size, chunk_size), lambda b_, g, c: (b_, rc(c), g, 0, 0, 0)),
+      # a_log, dt_bias [B, G, 1, padded]
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      # segment entering the chunk [B, G, C, 1, 128]
+      pl.BlockSpec((None, None, None, 1, 128), lambda b_, g, c: (b_, g, rc(c), 0, 0)),
+  ]
+  acc_shape = (tile_v_heads, kq_head_dim, v_head_dim)
+  body = functools.partial(
+      _dh_scan_body,
+      chunk_size=chunk_size,
+      num_kq_heads=tile_v_heads // repeats,
+      num_v_heads=tile_v_heads,
+      kq_head_dim=kq_head_dim,
+      num_chunks=num_chunks,
+      has_seg_ids=has_seg_ids,
+      use_qk_norm=use_qk_norm_in_gdn,
+  )
+  ds_6d = _run_scan(
+      body,
+      args,
+      in_specs,
+      jax.ShapeDtypeStruct((batch_size, num_groups, 1) + acc_shape, jnp.float32),
+      pl.BlockSpec((None, None, 1) + acc_shape, lambda b_, g, c: (b_, g, 0, 0, 0, 0)),
+      acc_shape,
+      grid=(batch_size, num_groups, num_chunks),
+      interpret=interpret,
+      name="gdn_cp_dh_scan",
+  )
+  return ds_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, v_head_dim)
diff --git a/tests/unit/gdn_cp_dh_scan_test.py b/tests/unit/gdn_cp_dh_scan_test.py
index 6655fd43f4..28857fa05b 100644
--- a/tests/unit/gdn_cp_dh_scan_test.py
+++ b/tests/unit/gdn_cp_dh_scan_test.py
@@ -11,10 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""The Pallas CP composition scan (`dh_scan`) against the XLA `cp_gdn` tree it replaces on TPU.
+"""The Pallas CP composition scans (`dh_scan`) against the XLA `cp_gdn` trees they replace on TPU.
 
 Runs the kernel in Pallas interpret mode. XLA:CPU computes the tree's f32 matmuls in plain f32, so the comparison
-bounds the scan's 3-pass bf16 matmul error and the reduction order, not TPU numerics.
+bounds the scans' 3-pass bf16 matmul error and the reduction order, not TPU numerics.
 """
 
 import jax
@@ -120,3 +120,70 @@ def test_m_scan_sees_chunk_order_and_segment_starts():
   # reaches the rank's end state.
   fresh = _scan_m(*args, segment_ids, jnp.zeros((1,), jnp.float32))
   assert float(jnp.abs(fresh).max()) == 0.0
+
+
+def _bwd_inputs(seed=1):
+  """One rank's full conv output (q, k, v), gates, output cotangent and per-chunk T^{-1}."""
+  keys = jax.random.split(jax.random.PRNGKey(seed), 7)
+  s = _NC * _C
+  qkv_conv = jax.random.normal(keys[0], (1, s, 2 * _NK * _D + _NV * _D), jnp.float32).astype(jnp.bfloat16)
+  b = jax.random.normal(keys[1], (1, s, _NV), jnp.float32).astype(jnp.bfloat16)
+  a = (jax.random.normal(keys[2], (1, s, _NV), jnp.float32) - 2.0).astype(jnp.bfloat16)
+  a_log = jax.random.normal(keys[3], (_NV,), jnp.float32) * 0.3
+  dt_bias = jax.random.normal(keys[4], (_NV,), jnp.float32) * 0.3
+  do = jax.random.normal(keys[5], (1, s, _NV, _D), jnp.float32).astype(jnp.bfloat16)
+  t_inv = jnp.tril(jax.random.normal(keys[6], (1, _NC, _NV, _C, _C), jnp.float32) * 0.05, -1) + jnp.eye(_C)
+  return qkv_conv, b, a, a_log, dt_bias, do, t_inv
+
+
+@pytest.mark.parametrize("head_tile", [32, 4])  # one head group, two head groups
+@pytest.mark.parametrize("layout", ["none", "single", "packed"])
+def test_dh_scan_matches_compose_bwd_local_tree(layout, head_tile, monkeypatch):
+  monkeypatch.setattr(dh_scan, "_DH_SCAN_HEAD_TILE", head_tile)
+  segment_ids = {"none": None, "single": jnp.asarray(_SINGLE), "packed": jnp.asarray(_PACKED)}[layout]
+  init_seg = None if segment_ids is None else jnp.ones((1,), jnp.float32)
+  args = _bwd_inputs()
+  _, ref = cp_gdn.compose_bwd_local_from_t_inv(
+      *args,
+      num_k_heads=_NK,
+      num_v_heads=_NV,
+      head_k_dim=_D,
+      head_v_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+  )
+  got = dh_scan.gdn_cp_dh_scan(
+      *args,
+      num_v_heads=_NV,
+      kq_head_dim=_D,
+      v_head_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+      interpret=True,
+  )
+  assert float(jnp.abs(ref).max()) > 1e-4, "degenerate reference"
+  assert _rel_err(got, ref) < 5e-5
+  # Negative control: a perturbed t_inv must show.
+  qkv_conv, b, a, a_log, dt_bias, do, t_inv = args
+  bad = dh_scan.gdn_cp_dh_scan(
+      qkv_conv,
+      b,
+      a,
+      a_log,
+      dt_bias,
+      do,
+      t_inv * 1.01,
+      num_v_heads=_NV,
+      kq_head_dim=_D,
+      v_head_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+      interpret=True,
+  )
+  assert _rel_err(bad, ref) > 1e-3

From 498dea11241cd45929739b9ed5a528386b414f0a Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 14:26:41 +0000
Subject: [PATCH 12/19] MoE wo: FSDP-shard its mlp dimension under cp-as-ep so
 its gradient reduce-scatter scatters a middle dim

wo is [experts, mlp, embed]. Sharding its embed dimension over FSDP made XLA transpose each wo
gradient to [embed, experts*mlp] and reduce-scatter its major dimension on the SparseCore, about 3x
slower than wi's middle-dimension scatter for the same bytes (12.9 vs 4.2 ms per call on v7x), and
that chain of expert-weight-gradient reduce-scatters ends every backward layer iteration. Mesh rules
that define mlp_moe_out/embed_moe_out (cp-as-ep) now shard wo's mlp dimension over FSDP; rule sets
without those names keep the old layout. The forward all-gather becomes [64,32,4096]->[64,1024,4096].
---
 src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml | 3 +++
 src/maxtext/layers/moe.py                             | 9 ++++++++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml
index bd3828825d..bd2d4d1bf3 100644
--- a/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml
+++ b/src/maxtext/configs/custom_mesh_and_rule/cp-as-ep.yml
@@ -63,6 +63,9 @@ logical_axis_rules: [
                       # MoE Weights
                       ['exp', ['context', 'expert']],
                       ['embed_moe', ['fsdp']],
+                      # wo's dims: FSDP on its mlp dim (see RoutedMoE.wo_kernel_axes)
+                      ['mlp_moe_out', ['fsdp']],
+                      ['embed_moe_out', []],
                       # ==========================================
                       # Standard MLP / Dense Layers / Model Structure
                       # ==========================================
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 8899fc12e9..266bc82289 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -829,7 +829,14 @@ def __init__(
       self.wi_kernel_axes, self.wo_kernel_axes = get_batchsplit_init_kernel_axes()
     else:
       self.wi_kernel_axes = ("exp", "embed_moe", "mlp_moe")
-      self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe")
+      # wo is [experts, mlp, embed]. Mesh rules that define `mlp_moe_out`/`embed_moe_out` (cp-as-ep) shard its mlp
+      # dim over FSDP instead of its embed dim, so its gradient reduce-scatter scatters a middle dimension, like wi's,
+      # instead of being transposed to [embed, experts * mlp] and scattered on its major dimension (about 3x slower
+      # on the SparseCore for the same bytes).
+      if {"mlp_moe_out", "embed_moe_out"} <= {rule[0] for rule in self.config.logical_axis_rules}:
+        self.wo_kernel_axes = ("exp", "mlp_moe_out", "embed_moe_out")
+      else:
+        self.wo_kernel_axes = ("exp", "mlp_moe", "embed_moe")
 
     if self.config.attention in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"):
       # vLLM uses 'model' as the tensor parallelism axis name

From 6b58e2a4caa6d23a11119023e7c129ec8aabe470 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 15:40:50 +0000
Subject: [PATCH 13/19] MoE ring: unsort backward rows and the SiLU GLU visit
 only the packed buffer's valid rows

Only rows [0, limit) of the ring's packed [buffer_size, ...] expert buffer hold tokens (about 56% on
the bench's hot shard at ragged_buffer_factor 2.75). The grouped matmuls already skip the rest, but
four XLA fusion families per layer chunk still streamed the whole buffer on every device: the unsort
backward (f32[225280] + bf16[225280,4096]), the GLU forward and its backward recompute
(bf16[225280,1024]), and the GLU backward, a select fusion writing both half gradients plus a
separate pad/maximum pass that concatenated them (bf16[225280,2048]).

- kernels/ragged/unsort_bwd_rows.py: a Pallas TensorCore kernel for the buffered branch of
  _ring_ragged_unsort_bwd. It computes bf16(f32(g) * w) and the f32 routing-weight row dot in
  128x128 tiles over (1024, 2048) blocks, with `limit` as a prefetched scalar; blocks past the one
  holding row limit - 1 map to that block, so they are neither fetched nor written. The weights and
  the dot are lane-dense [rows / 128, 128] (free reshapes), turned into per-row columns with one
  128x128 transpose each. The scatter of the dot is unchanged; off a TPU the XLA code runs.
- kernels/ragged/glu_rows.py: forward and backward kernels for silu(gate) * up over the same rows,
  reading gate and up in place from the fused buffer. The backward writes [d_gate | d_up] in one
  pass, with the -0 -> +0 fix. They follow XLA's bf16 fusions exactly: the same op sequence
  (negate, exp, + 1, divide, multiply), in f32, rounding only the outputs (XLA keeps excess precision
  inside the fusion; op-by-op bf16 rounding differs from it in ~3% of elements).
- moe.py: fused_glu(fused, glu, limit=None). _down_projection passes valid_token_count as the limit
  when the buffer is the ring's packed one and apply_ffn_activation is the plain bf16 SiLU GLU; the
  kernel path keeps the whole buffer as its residual, under the remat names of both halves. The
  dropless replay program (unpacked buffer) and every other path keep the XLA code.

Rows left unwritten are never read: gmm_v2 reads lhs rows up to the sublane tile of row limit - 1,
tgmm_v2 masks rows outside the groups with a select, and the dot scatter drops rows >= limit.

Numerics: the weighted unsort gradient and both GLU results are bit-identical; the routing-weight
dot changes summation order only.

Checks:
- v6e, bench shape [225280, 4096], limits 0, 1, 1023, 125952, 126156, 225280: unsort gradient
  bitwise equal to XLA (rows up to the end of limit's block), dot max |d| 5.3e-5. Interpret mode
  with the TPU7x chip description and a partial last block: bitwise as well.
- v6e, GLU forward and backward against fused_glu's XLA path, every bf16 bit pattern as the gate
  (x64 random up/d, plus +-0, +-inf, NaN, denormals), 225280-row and partial-block buffers: 0
  mismatching bits.
- v6e timing at 126156 of 225280 valid rows (XLA -> kernel): unsort 4.71 -> 2.61 ms, GLU forward
  1.08 -> 0.61 ms, GLU backward 1.92 -> 1.02 ms; at a full buffer the kernels are within 1-9% of XLA.
- Tiny Qwen3.5 on v6e (ab_model.py, use_ragged_sort=true ragged_buffer_factor=2.75): loss, logits
  and all 66 gradient leaves bit-identical to the unmodified tree.
- 397B v7x AOT (wi dlhs tiles 2048, moe_route=device, recompute limit 100): fwd_bwd temp
  94.95 -> 92.25 GiB, dropless 56.69 GiB unchanged. module_0000 has 8 unsort_bwd_rows, 16
  glu_fwd_rows and 8 glu_bwd_rows calls in place of the 16 activation, 8 unsort, 8 select and
  8 concatenate fusions, with no new 225280-row copies.
---
 src/maxtext/kernels/ragged/glu_rows.py        | 167 ++++++++++++++++
 src/maxtext/kernels/ragged/ragged_sort.py     |   9 +-
 src/maxtext/kernels/ragged/unsort_bwd_rows.py | 179 ++++++++++++++++++
 src/maxtext/layers/moe.py                     |  53 +++++-
 tests/unit/moe_valid_rows_kernels_test.py     | 102 ++++++++++
 5 files changed, 501 insertions(+), 9 deletions(-)
 create mode 100644 src/maxtext/kernels/ragged/glu_rows.py
 create mode 100644 src/maxtext/kernels/ragged/unsort_bwd_rows.py
 create mode 100644 tests/unit/moe_valid_rows_kernels_test.py

diff --git a/src/maxtext/kernels/ragged/glu_rows.py b/src/maxtext/kernels/ragged/glu_rows.py
new file mode 100644
index 0000000000..b69d422a10
--- /dev/null
+++ b/src/maxtext/kernels/ragged/glu_rows.py
@@ -0,0 +1,167 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""SiLU-GLU of the prefused gate/up GEMM output, and its gradient, over the valid rows of the expert buffer only.
+
+Only rows ``[0, limit)`` of the packed ``[buffer_size, 2 * mlp]`` buffer hold tokens. XLA's fusions stream the
+whole buffer, and the backward writes the two half gradients and then concatenates them, one more pass. These
+kernels visit the row blocks up to the one holding row ``limit - 1`` (later blocks map to it: no fetch, no
+write, those output rows are left as they are), and the backward writes both halves of its block at once.
+Consumers never read the rows left unwritten: gmm_v2 reads rows up to the sublane tile of row ``limit - 1``,
+tgmm_v2 masks them.
+
+Both kernels run the op sequence of XLA's fusions for ``silu(gate) * up`` and its autodiff gradient (negate,
+exp, + 1, divide, multiply) the way XLA runs those bf16 fusions on a TPU: in f32, rounding only the outputs.
+"""
+
+import jax
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+# Rows per grid step. A (1024, 1024) bf16 block is 2 MiB. A partial last block is fine: Pallas drops its
+# out-of-range writes.
+_BLOCK_ROWS = 1024
+# Rows per inner-loop tile.
+_TILE_ROWS = 64
+
+
+def _sigmoid(x):
+  """XLA's expansion of logistic: 1 / (exp(-x) + 1)."""
+  return 1.0 / (jnp.exp(-x) + 1.0)
+
+
+def _glu_fwd_tile(gate, up):
+  """`silu(gate) * up` as XLA's fusion computes it: in f32, rounded once to the output dtype."""
+  g = gate.astype(jnp.float32)
+  return ((g * _sigmoid(g)) * up.astype(jnp.float32)).astype(gate.dtype)
+
+
+def _glu_bwd_tile(gate, up, d_out):
+  """The autodiff gradient of `_glu_fwd_tile` as XLA's fusion computes it (f32, rounded once), -0 as +0."""
+  g, u, d = gate.astype(jnp.float32), up.astype(jnp.float32), d_out.astype(jnp.float32)
+  sig = _sigmoid(g)
+  d_up_in = d * u
+  d_gate = d_up_in * sig + (g * d_up_in) * (sig * (1.0 - sig))
+  d_up = (g * sig) * d
+  d_gate = jnp.where(d_gate == 0, 0.0, d_gate)
+  d_up = jnp.where(d_up == 0, 0.0, d_up)
+  return d_gate.astype(gate.dtype), d_up.astype(gate.dtype)
+
+
+def kernel_supported(fused):
+  rows, width = fused.shape
+  return fused.dtype == jnp.bfloat16 and rows % 16 == 0 and width % 256 == 0
+
+
+def _last_block(limit):
+  """The row block holding row ``limit - 1`` (block 0 if there are no valid rows)."""
+  return jnp.maximum(pl.cdiv(limit, _BLOCK_ROWS) - 1, 0)
+
+
+def _row_map(column_block):
+  def index_map(i, limit_ref):
+    return jnp.minimum(i, _last_block(limit_ref[0])), column_block
+
+  return index_map
+
+
+def _run(kernel, operands, in_specs, out_spec, out_shape, *, limit, flops, transcendentals, interpret, name):
+  rows = out_shape.shape[0]
+  bytes_accessed = sum(x.size * x.dtype.itemsize for x in operands) + out_shape.size * out_shape.dtype.itemsize
+  return pl.pallas_call(
+      kernel,
+      grid_spec=pltpu.PrefetchScalarGridSpec(
+          num_scalar_prefetch=1,
+          grid=(pl.cdiv(rows, _BLOCK_ROWS),),
+          in_specs=in_specs,
+          out_specs=out_spec,
+      ),
+      out_shape=out_shape,
+      compiler_params=pltpu.CompilerParams(dimension_semantics=("arbitrary",), vmem_limit_bytes=48 * 1024 * 1024),
+      cost_estimate=pl.CostEstimate(flops=flops, transcendentals=transcendentals, bytes_accessed=bytes_accessed),
+      interpret=interpret,
+      name=name,
+  )(jnp.reshape(jnp.asarray(limit).astype(jnp.int32), (1,)), *operands)
+
+
+def _visit_rows(i, limit, tile_fn):
+  """Runs `tile_fn(rows)` over the row tiles of block i if the block holds valid rows (or is block 0)."""
+
+  @pl.when(i <= _last_block(limit))
+  def _():
+    def body(t, carry):
+      tile_fn(pl.ds(pl.multiple_of(t * _TILE_ROWS, _TILE_ROWS), _TILE_ROWS))
+      return carry
+
+    lax.fori_loop(0, _BLOCK_ROWS // _TILE_ROWS, body, 0)
+
+
+def _glu_fwd_kernel(limit_ref, gate_ref, up_ref, out_ref):
+  def tile(rows):
+    out_ref[rows, :] = _glu_fwd_tile(gate_ref[rows, :], up_ref[rows, :])
+
+  _visit_rows(pl.program_id(0), limit_ref[0], tile)
+
+
+def _glu_bwd_kernel(limit_ref, gate_ref, up_ref, d_ref, out_ref):
+  mlp = gate_ref.shape[1]
+
+  def tile(rows):
+    d_gate, d_up = _glu_bwd_tile(gate_ref[rows, :], up_ref[rows, :], d_ref[rows, :])
+    out_ref[rows, :mlp] = d_gate
+    out_ref[rows, mlp:] = d_up
+
+  _visit_rows(pl.program_id(0), limit_ref[0], tile)
+
+
+def glu_fwd_rows(fused, limit, *, interpret=False):
+  """``silu(gate) * up`` of rows ``[0, limit)`` of ``fused = [gate | up]``, ``[rows, 2 * mlp]`` -> ``[rows, mlp]``."""
+  rows, width = fused.shape
+  mlp = width // 2
+  return _run(
+      _glu_fwd_kernel,
+      (fused, fused),
+      [pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)), pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(1))],
+      pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)),
+      jax.ShapeDtypeStruct((rows, mlp), fused.dtype),
+      limit=limit,
+      flops=4 * rows * mlp,
+      transcendentals=2 * rows * mlp,
+      interpret=interpret,
+      name="glu_fwd_rows",
+  )
+
+
+def glu_bwd_rows(fused, d_out, limit, *, interpret=False):
+  """The ``[rows, 2 * mlp]`` gradient ``[d_gate | d_up]`` of `glu_fwd_rows` at rows ``[0, limit)``, -0 as +0."""
+  rows, width = fused.shape
+  mlp = width // 2
+  return _run(
+      _glu_bwd_kernel,
+      (fused, fused, d_out),
+      [
+          pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)),
+          pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(1)),
+          pl.BlockSpec((_BLOCK_ROWS, mlp), _row_map(0)),
+      ],
+      pl.BlockSpec((_BLOCK_ROWS, width), _row_map(0)),
+      jax.ShapeDtypeStruct((rows, width), fused.dtype),
+      limit=limit,
+      flops=12 * rows * mlp,
+      transcendentals=2 * rows * mlp,
+      interpret=interpret,
+      name="glu_bwd_rows",
+  )
diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index 35ac56d09c..45f2a34bde 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -19,6 +19,7 @@
 import jax.numpy as jnp
 from maxtext.kernels.ragged.ragged_gather import ragged_gather
 from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce
+from maxtext.kernels.ragged.unsort_bwd_rows import unsort_bwd_rows
 
 
 def ring_ragged_sort(
@@ -451,13 +452,11 @@ def _ring_ragged_unsort_bwd(res, g_out):
           bytes_accessed_override=gather_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
       )
-      # Mask out gradients for elements beyond the valid limit of the local buffer.
+      # Only rows [0, limit) of the local buffer hold tokens: weight their gradients and take their dots.
       limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size)
-      mask = jnp.arange(buffer_size) < limit
-      gathered = jnp.where(mask[:, None], gathered, 0.0)
-      grad_sorted_tokens = (gathered * sliced_weights[:, None]).astype(gathered.dtype)
+      grad_sorted_tokens, dot_local = unsort_bwd_rows(gathered, sorted_tokens_local, sliced_weights, limit)
       # Scatter the per-slot dot back to flat slot order; dropped slots stay zero.
-      dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1)
+      mask = jnp.arange(buffer_size) < limit
       slots = jnp.where(mask, sliced_idx_inv, n)
       grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop")
     return grad_sorted_tokens, None, None, None, grad_topk_weights
diff --git a/src/maxtext/kernels/ragged/unsort_bwd_rows.py b/src/maxtext/kernels/ragged/unsort_bwd_rows.py
new file mode 100644
index 0000000000..148870b9cf
--- /dev/null
+++ b/src/maxtext/kernels/ragged/unsort_bwd_rows.py
@@ -0,0 +1,179 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The row-wise part of the ring unsort backward, on the TensorCore, over the valid rows of the buffer only.
+
+Only rows ``[0, limit)`` of the packed ``[buffer_size, hidden]`` expert buffer hold tokens; the rest is
+ragged-buffer padding. XLA's fusion of the same three lines streams the whole buffer. This kernel visits the
+row blocks up to the one holding row ``limit - 1``: the later blocks map to that block, so they are neither
+fetched nor written, and those output rows are left as they are. No consumer reads them: the gmm_v2 /
+tgmm_v2 kernels read rows up to the sublane tile of row ``limit - 1`` (tgmm masks them), and the weight-dot
+scatter drops rows ``>= limit``.
+"""
+
+import functools
+
+import jax
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+# Rows per grid step: one (8, 128) tile of the lane-dense [rows / 128, 128] weights and dot. A partial last block
+# is fine: Pallas drops its out-of-range writes, and its rows are past `limit`.
+_BLOCK_ROWS = 1024
+# Hidden columns per grid step, at most. A (1024, 2048) bf16 block is 4 MiB; 3 double-buffered fit v7x VMEM.
+_BLOCK_COLS = 2048
+_LANES = 128
+
+
+def unsort_bwd_rows_reference(gathered, sorted_tokens, weights, limit):
+  """XLA version: the weighted gradient and the f32 row dot of every row, zeroing rows ``>= limit``."""
+  mask = jnp.arange(gathered.shape[0]) < limit
+  gathered = jnp.where(mask[:, None], gathered, 0.0)
+  grad = (gathered * weights[:, None]).astype(gathered.dtype)
+  dot = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens.astype(jnp.float32), axis=-1)
+  return grad, dot
+
+
+def _block_cols(hidden):
+  return hidden if hidden <= _BLOCK_COLS else _BLOCK_COLS
+
+
+def _kernel_supported(gathered, sorted_tokens, weights):
+  rows, hidden = gathered.shape
+  return (
+      gathered.dtype == jnp.bfloat16
+      and sorted_tokens.dtype == jnp.bfloat16
+      and sorted_tokens.shape == gathered.shape
+      and weights.shape == (rows,)
+      and rows % _LANES == 0
+      and hidden % _LANES == 0
+      and hidden % _block_cols(hidden) == 0
+  )
+
+
+def _last_block(limit):
+  """The row block holding row ``limit - 1`` (block 0 if there are no valid rows)."""
+  return jnp.maximum(pl.cdiv(limit, _BLOCK_ROWS) - 1, 0)
+
+
+def _unsort_bwd_rows_kernel(limit_ref, g_ref, y_ref, w_ref, grad_ref, dot_ref, *, num_col_blocks):  # pylint: disable=too-many-positional-arguments
+  """One (1024, cols) block: grad = bf16(f32(g) * w) and the partial f32 row dot, in 128 x 128 tiles."""
+  i, j = pl.program_id(0), pl.program_id(1)
+  limit = limit_ref[0]
+  row0 = i * _BLOCK_ROWS
+  num_col_tiles = g_ref.shape[1] // _LANES
+
+  def row_tile(masked, r, carry):
+    rows = pl.ds(pl.multiple_of(r * _LANES, _LANES), _LANES)
+    # Row a of the tile gets weight w[row0 + 128 r + a], broadcast along the lanes.
+    w_col = jnp.transpose(jnp.broadcast_to(w_ref[pl.ds(r, 1), :], (_LANES, _LANES)))
+    if masked:
+      row_ids = row0 + r * _LANES + lax.broadcasted_iota(jnp.int32, (_LANES, _LANES), 0)
+      valid = row_ids < limit
+    acc = jnp.zeros((_LANES, _LANES), jnp.float32)
+    for c in range(num_col_tiles):
+      cols = pl.ds(c * _LANES, _LANES)
+      g = g_ref[rows, cols].astype(jnp.float32)
+      if masked:
+        g = jnp.where(valid, g, 0.0)
+      grad_ref[rows, cols] = (g * w_col).astype(grad_ref.dtype)
+      acc = acc + g * y_ref[rows, cols].astype(jnp.float32)
+    # Lane-dense: entry a of the row sum over the transposed tile is row a's dot over this block's columns.
+    part = jnp.sum(jnp.transpose(acc), axis=0, keepdims=True)
+    if num_col_blocks > 1:
+      part = part + jnp.where(j == 0, 0.0, dot_ref[pl.ds(r, 1), :])
+    dot_ref[pl.ds(r, 1), :] = part
+    return carry
+
+  num_row_tiles = _BLOCK_ROWS // _LANES
+
+  @pl.when(row0 + _BLOCK_ROWS <= limit)
+  def _full():
+    lax.fori_loop(0, num_row_tiles, functools.partial(row_tile, False), 0)
+
+  @pl.when(jnp.logical_and(row0 + _BLOCK_ROWS > limit, i == _last_block(limit)))
+  def _partial():
+    lax.fori_loop(0, num_row_tiles, functools.partial(row_tile, True), 0)
+
+
+def _unsort_bwd_rows_pallas(gathered, sorted_tokens, weights, limit, *, interpret=False):
+  """The Pallas TensorCore kernel; see `unsort_bwd_rows`."""
+  rows, hidden = gathered.shape
+  block_cols = _block_cols(hidden)
+  num_row_blocks = pl.cdiv(rows, _BLOCK_ROWS)
+  num_col_blocks = hidden // block_cols
+
+  def data_map(i, j, limit_ref):
+    last = _last_block(limit_ref[0])
+    return jnp.minimum(i, last), jnp.where(i <= last, j, num_col_blocks - 1)
+
+  def row_map(i, j, limit_ref):
+    del j
+    return jnp.minimum(i, _last_block(limit_ref[0])), 0
+
+  data_spec = pl.BlockSpec((_BLOCK_ROWS, block_cols), data_map)
+  lane_dense_spec = pl.BlockSpec((_BLOCK_ROWS // _LANES, _LANES), row_map)
+  grad, dot = pl.pallas_call(
+      functools.partial(_unsort_bwd_rows_kernel, num_col_blocks=num_col_blocks),
+      grid_spec=pltpu.PrefetchScalarGridSpec(
+          num_scalar_prefetch=1,
+          grid=(num_row_blocks, num_col_blocks),
+          in_specs=[data_spec, data_spec, lane_dense_spec],
+          out_specs=[data_spec, lane_dense_spec],
+      ),
+      out_shape=[
+          jax.ShapeDtypeStruct((rows, hidden), gathered.dtype),
+          jax.ShapeDtypeStruct((rows // _LANES, _LANES), jnp.float32),
+      ],
+      compiler_params=pltpu.CompilerParams(
+          dimension_semantics=("arbitrary", "arbitrary"),
+          vmem_limit_bytes=48 * 1024 * 1024,
+      ),
+      cost_estimate=pl.CostEstimate(
+          flops=4 * rows * hidden,
+          transcendentals=0,
+          bytes_accessed=3 * rows * hidden * gathered.dtype.itemsize + 8 * rows,
+      ),
+      interpret=interpret,
+      name="unsort_bwd_rows",
+  )(
+      jnp.reshape(limit.astype(jnp.int32), (1,)),
+      gathered,
+      sorted_tokens,
+      jnp.reshape(weights.astype(jnp.float32), (rows // _LANES, _LANES)),
+  )
+  return grad, jnp.reshape(dot, (rows,))
+
+
+def unsort_bwd_rows(gathered, sorted_tokens, weights, limit, *, interpret=None):
+  """The weighted gradient and the routing-weight row dot of the valid rows ``[0, limit)`` of the expert buffer.
+
+  Args:
+    gathered: ``[rows, hidden]`` bf16 output gradient, gathered into the packed expert buffer order.
+    sorted_tokens: ``[rows, hidden]`` bf16 expert outputs in the same order.
+    weights: ``[rows]`` f32 routing weight of each buffer row.
+    limit: int32 scalar, the number of valid rows.
+    interpret: Pallas interpret mode (tests). None runs the kernel on a TPU target and the XLA version elsewhere.
+
+  Returns:
+    ``(grad, dot)``: for rows ``< limit``, ``grad = bf16(f32(gathered) * weights)`` (bit-identical to the XLA
+    version) and ``dot = sum(f32(gathered) * f32(sorted_tokens), -1)`` (the summation order differs). Rows from
+    ``limit`` to the end of its 1024-row block are computed as zero rows; later rows are not written.
+  """
+  use_kernel = interpret if interpret is not None else pltpu.is_tpu_device()
+  if not use_kernel or not _kernel_supported(gathered, sorted_tokens, weights):
+    return unsort_bwd_rows_reference(gathered, sorted_tokens, weights, limit)
+  return _unsort_bwd_rows_pallas(gathered, sorted_tokens, weights, jnp.asarray(limit), interpret=bool(interpret))
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 266bc82289..f8689e3e22 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -30,6 +30,7 @@
 from jax.experimental.compute_on import compute_on
 from jax.experimental import xla_metadata
 from jax.experimental.layout import Layout, with_layout_constraint
+from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 from jax.sharding import Mesh, NamedSharding
 from jax.sharding import PartitionSpec as P
@@ -37,6 +38,7 @@
 from maxtext.common.common_types import ShardMode
 from maxtext.kernels import megablox as mblx
 from maxtext.kernels import sort_activations
+from maxtext.kernels.ragged import glu_rows
 from maxtext.kernels.ragged.ragged_sort import a2a_ragged_sort
 from maxtext.kernels.ragged.ragged_sort import a2a_ragged_unsort
 from maxtext.kernels.ragged.ragged_sort import ring_ragged_sort
@@ -312,14 +314,25 @@ def _split_gate_up(fused):
   return gate, up
 
 
-@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
-def fused_glu(fused, glu):
+def fused_glu(fused, glu, limit=None):
   """`glu(gate, up)` of the two halves of `fused`, the [..., 2 * mlp] output of the prefused gate/up GEMM.
 
   Autodiff transposes each half's slice into a zero-padded [..., 2 * mlp] gradient and adds the two, one more
   full pass over the gradient. This backward computes the same two half gradients, with `jax.vjp` of `glu`, and
   concatenates them, which XLA writes straight from the activation's backward fusion.
+
+  `limit`, if given, says that only rows [0, limit) of the [rows, 2 * mlp] expert buffer `fused` hold tokens and
+  that `glu` is the bf16 SiLU GLU. On a TPU, Pallas kernels then compute the activation and its gradient on those
+  rows only, bit-identical to the XLA fusions, and write the gradient's two halves at once; the rows past the
+  1024-row block of row limit - 1 are left unwritten in both results.
   """
+  if limit is not None and pltpu.is_tpu_device() and glu_rows.kernel_supported(fused):
+    return _silu_glu_rows(fused, limit)
+  return _fused_glu(fused, glu)
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
+def _fused_glu(fused, glu):
   return glu(*_split_gate_up(fused))
 
 
@@ -334,7 +347,26 @@ def _fused_glu_bwd(glu, residuals, cotangent):
   return (jnp.concatenate([jnp.where(d == 0, 0, d) for d in glu_vjp(cotangent)], axis=-1),)
 
 
-fused_glu.defvjp(_fused_glu_fwd, _fused_glu_bwd)
+_fused_glu.defvjp(_fused_glu_fwd, _fused_glu_bwd)
+
+
+@jax.custom_vjp
+def _silu_glu_rows(fused, limit):
+  return glu_rows.glu_fwd_rows(fused, limit)
+
+
+def _silu_glu_rows_fwd(fused, limit):
+  # The kernels read both halves in place, so the residual is the whole buffer, under the names of both halves.
+  fused = adc.checkpoint_name(adc.checkpoint_name(adc.checkpoint_name(fused, "mlpwi_0"), "moe_mlpwi_0"), "moe_mlpwi_1")
+  return glu_rows.glu_fwd_rows(fused, limit), (fused, limit)
+
+
+def _silu_glu_rows_bwd(residuals, cotangent):
+  fused, limit = residuals
+  return glu_rows.glu_bwd_rows(fused, cotangent, limit), None
+
+
+_silu_glu_rows.defvjp(_silu_glu_rows_fwd, _silu_glu_rows_bwd)
 
 
 def get_batchsplit_init_kernel_axes():
@@ -3020,7 +3052,20 @@ def _up_projection(x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias
     def _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata):
       """The activation and the down projection, per routed token."""
       if output1 is None:  # output0 is the prefused gate/up GEMM output (see `gmm_up`)
-        intermediate_layer = fused_glu(output0, self.apply_ffn_activation)
+        # The ring's packed buffer holds its tokens in rows [0, valid_token_count); the SiLU GLU kernels skip the
+        # rest. apply_ffn_activation is the plain SiLU GLU unless the decoder block clips its inputs.
+        silu_glu = (
+            self.config.mlp_activations[0] == "silu"
+            and self.dtype == jnp.bfloat16
+            and self.config.decoder_block != ctypes.DecoderBlockType.GPT_OSS
+            and not (
+                self.config.decoder_block in (ctypes.DecoderBlockType.DEEPSEEK, ctypes.DecoderBlockType.DEEPSEEK4)
+                and self.config.mlp_activations_limit > 0.0
+            )
+        )
+        packed = self.config.use_ring_of_experts and output0.shape[0] < routing.sorted_selected_experts.shape[0]
+        limit = valid_token_count(output0, routing, route_metadata) if silu_glu and packed else None
+        intermediate_layer = fused_glu(output0, self.apply_ffn_activation, limit)
       else:
         intermediate_layer = self.apply_ffn_activation(output0, output1)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
diff --git a/tests/unit/moe_valid_rows_kernels_test.py b/tests/unit/moe_valid_rows_kernels_test.py
new file mode 100644
index 0000000000..2edc4891d8
--- /dev/null
+++ b/tests/unit/moe_valid_rows_kernels_test.py
@@ -0,0 +1,102 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The valid-rows MoE kernels (`unsort_bwd_rows`, `glu_rows`) against the XLA code they replace.
+
+On CPU the kernels run in Pallas interpret mode with the TPU7x chip description, over a buffer whose last
+1024-row block is partial. On a TPU the SiLU-GLU kernels must match XLA's fusions bit for bit for every bf16 gate.
+"""
+
+import jax
+from jax._src.pallas.mosaic import tpu_info
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.kernels.ragged import glu_rows
+from maxtext.kernels.ragged import unsort_bwd_rows as unsort
+from maxtext.layers import moe
+
+_ROWS, _HIDDEN, _MLP = 2560, 256, 128
+
+
+@pytest.fixture(autouse=True)
+def _tpu7x_interpret():
+  if jax.default_backend() == "cpu":
+    tpu_info.registry["cpu"] = lambda: tpu_info.get_tpu_info_for_chip(tpu_info.ChipVersion.TPU_7X, 1)
+    tpu_info.get_tpu_info.cache_clear()
+
+
+def _bf16(key, shape, scale=1.0):
+  return (jax.random.normal(key, shape, jnp.float32) * scale).astype(jnp.bfloat16)
+
+
+def _bits(x, rows):
+  return np.asarray(x)[:rows].view(np.uint16)
+
+
+def _silu_glu(gate, up):
+  return (jax.nn.silu(gate) * up).astype(jnp.bfloat16)
+
+
+@pytest.mark.parametrize("limit", [0, 1000, 1700, _ROWS])
+def test_unsort_bwd_rows(limit):
+  """Weighted gradient bit-identical and row dot to f32 summation order on rows < limit, zeros up to the block end."""
+  keys = jax.random.split(jax.random.PRNGKey(0), 3)
+  gathered, tokens = _bf16(keys[0], (_ROWS, _HIDDEN)), _bf16(keys[1], (_ROWS, _HIDDEN))
+  weights = jax.random.uniform(keys[2], (_ROWS,), jnp.float32, -1.0, 1.0)
+  interpret = True if jax.default_backend() == "cpu" else None
+  grad, dot = unsort.unsort_bwd_rows(gathered, tokens, weights, jnp.int32(limit), interpret=interpret)
+  ref_grad, ref_dot = unsort.unsort_bwd_rows_reference(gathered, tokens, weights, jnp.int32(limit))
+  block_end = min(_ROWS, max(-(-limit // 1024), 1) * 1024)
+  np.testing.assert_array_equal(_bits(grad, block_end), _bits(ref_grad, block_end))
+  np.testing.assert_allclose(np.asarray(dot)[:limit], np.asarray(ref_dot)[:limit], rtol=0, atol=1e-4)
+  np.testing.assert_array_equal(np.asarray(dot)[limit:block_end], 0.0)
+
+
+def _glu_inputs(rows, mlp):
+  """Every bf16 bit pattern as a gate value, against random up and output gradients."""
+  gate = np.resize(np.arange(1 << 16, dtype=np.uint16).view(jnp.bfloat16), (rows, mlp))
+  keys = jax.random.split(jax.random.PRNGKey(1), 2)
+  fused = jnp.concatenate([jnp.asarray(gate), _bf16(keys[0], (rows, mlp), 4.0)], axis=-1)
+  return fused, _bf16(keys[1], (rows, mlp), 4.0)
+
+
+def test_glu_rows_interpret_rows_and_blocks():
+  """CPU: the visited rows (partial last block) hold the per-tile formula; XLA:CPU may contract it differently."""
+  if jax.default_backend() != "cpu":
+    pytest.skip("interpret-mode check")
+  fused, d_out = _glu_inputs(_ROWS, _MLP)
+  gate, up = fused[:, :_MLP], fused[:, _MLP:]
+  act = glu_rows.glu_fwd_rows(fused, jnp.int32(2100), interpret=True)
+  d_fused = glu_rows.glu_bwd_rows(fused, d_out, jnp.int32(2100), interpret=True)
+  ref_act = glu_rows._glu_fwd_tile(gate, up)  # pylint: disable=protected-access
+  ref_d_fused = jnp.concatenate(glu_rows._glu_bwd_tile(gate, up, d_out), axis=-1)  # pylint: disable=protected-access
+  for got, want in ((act, ref_act), (d_fused, ref_d_fused)):
+    np.testing.assert_allclose(np.asarray(got, np.float32), np.asarray(want, np.float32), rtol=2**-6, atol=1e-37)
+
+
+@pytest.mark.tpu_only
+def test_glu_rows_match_xla_fusions():
+  """TPU: activation and gradient bit-identical to `fused_glu`'s XLA path, for every bf16 gate."""
+  rows = 8192
+  fused, d_out = _glu_inputs(rows, 1024)
+  limit = 6000
+  block_end = 6144
+  act = jax.jit(glu_rows.glu_fwd_rows)(fused, jnp.int32(limit))
+  d_fused = jax.jit(glu_rows.glu_bwd_rows)(fused, d_out, jnp.int32(limit))
+  # Jitted, so XLA fuses the reference as it does in the model (f32 inside the fusion, bf16 outputs).
+  ref_act = jax.jit(lambda f: moe.fused_glu(f, _silu_glu))(fused)
+  ref_d_fused = jax.jit(lambda f, d: jax.vjp(lambda x: moe.fused_glu(x, _silu_glu), f)[1](d)[0])(fused, d_out)
+  np.testing.assert_array_equal(_bits(act, block_end), _bits(ref_act, block_end))
+  np.testing.assert_array_equal(_bits(d_fused, block_end), _bits(ref_d_fused, block_end))

From 80c1495724b2764a243c45fb204fe326529fcb37 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:41:42 +0000
Subject: [PATCH 14/19] GDN forward out_proj: pin the gated-norm output
 token-minor in training

In training, layout assignment gives the forward out_proj's LHS the gated RMSNorm's head-tiled
layout, so the projection re-tiles it on every MXU fetch, while the backward's recomputed
projection reads it token-minor. Constrain the forward value to the token-minor layout with the
existing _forward_layout_constraint custom_vjp, which leaves the cotangent's layout to XLA.
---
 src/maxtext/models/qwen3.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/maxtext/models/qwen3.py b/src/maxtext/models/qwen3.py
index 84c53cf96a..90d06c693f 100644
--- a/src/maxtext/models/qwen3.py
+++ b/src/maxtext/models/qwen3.py
@@ -1473,6 +1473,10 @@ def shard_mapped_delta_rule(q, k, v, g_val, beta_val, init_h):
     # Reshape back to a single feature dimension for the final projection.
     # Shape from (B, S, H_v, D_v) -> (B, S, value_dim)
     gated_output = jnp.reshape(gated_output_reshaped, (batch, seq_len, -1), out_sharding=flat_sharding)
+    if model_mode == MODEL_MODE_TRAIN:
+      # Token-minor, as the backward's recomputed projection reads it: left to itself, layout assignment gives
+      # the forward's operand the norm's head-tiled layout and out_proj re-tiles it on every MXU fetch.
+      gated_output = moe._forward_layout_constraint(gated_output, (0, 2, 1))
 
     # Final output shape: (B, S, E)
     output = self.out_proj(gated_output, out_sharding=out_sharding)

From 75dae0de442cdf5821a56f721dfce7d2c4e8b82d Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Thu, 15 Oct 2026 16:41:42 +0000
Subject: [PATCH 15/19] MoE chunked ring: backward-only barrier ties the
 layer-input gradient to the expert-weight gradients

A custom_vjp that is the identity in the forward and an optimization barrier on the cotangents in
the backward, applied to the layer input and the expert weights. The input gradient is handed on
only once the expert-weight gradients exist, so each layer's weight-gradient grouped matmuls run
under that layer's token reduce-scatters instead of piling up at the end of the backward
iteration, where every deferred layer's sorted-token buffers would still be live.
---
 src/maxtext/layers/moe.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index f8689e3e22..857c1a98b5 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -250,6 +250,23 @@ def _forward_optimization_barrier_bwd(_, cotangents):
 _forward_optimization_barrier.defvjp(_forward_optimization_barrier_fwd, _forward_optimization_barrier_bwd)
 
 
+@jax.custom_vjp
+def _backward_optimization_barrier(operands):
+  """The identity on the forward values and `jax.lax.optimization_barrier` on their cotangents."""
+  return operands
+
+
+def _backward_optimization_barrier_fwd(operands):
+  return operands, None
+
+
+def _backward_optimization_barrier_bwd(_, cotangents):
+  return (jax.lax.optimization_barrier(cotangents),)
+
+
+_backward_optimization_barrier.defvjp(_backward_optimization_barrier_fwd, _backward_optimization_barrier_bwd)
+
+
 def barrier_tie(x, dep, backward: bool = False):
   """Returns `(x, dep)` unchanged, tied together by `jax.lax.optimization_barrier`.
 
@@ -3303,6 +3320,11 @@ def _pipelined_route_and_compute(force_dropless):
         reads those barrier outputs, and XLA removes the reduce-scatter and its
         gather-reduce. The compiled backward has the collectives of the plain chunk loop.
         """
+        nonlocal x, w0, w1, wo
+        # The backward hands the layer-input gradient on only once the expert-weight gradients exist, so each
+        # layer's tgmm runs under its own token reduce-scatters instead of at the end of the backward iteration,
+        # where every deferred layer's sorted-token buffers would still be live.
+        x, (w0, w1, wo) = _backward_optimization_barrier((x, (w0, w1, wo)))
         seq_len = x.shape[1]
         chunk = seq_len // n_chunks
 

From 06902ae25fd88f74d201559d26df1e057bbef2f5 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Wed, 14 Oct 2026 22:45:39 +0000
Subject: [PATCH 16/19] Gather MoE router kernels once per scanned layer stack,
 not per layer

The router kernel (embed x num_experts, f32 under float32_gate_logits) is
FSDP-sharded, so every layer all-gathered it on the critical path and the
backward gathered it again; on v7x those all-gathers run synchronously on the
TensorCore (~22 ms each on the 397B bench). Replicate the stacked kernels
before the scan (one all-gather), and give each layer's slice a backward that
constrains its gradient to the kernel's sharding, so gradients stay per-layer
reduce-scatters into a stacked shard rather than full-size all-reduces.

(cherry picked from commit 4f5093f3c3e677ac51118174e742c1ba221184ab)
---
 src/maxtext/layers/moe.py          | 70 ++++++++++++++++++++++++++++++
 src/maxtext/layers/nnx_decoders.py |  7 +++
 2 files changed, 77 insertions(+)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 857c1a98b5..91829464b3 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -24,6 +24,7 @@
 
 from aqt.jax.v2 import aqt_tensor as aqt
 from flax import nnx
+from flax.linen import partitioning as nn_partitioning
 from flax import struct
 import jax
 from jax import ad_checkpoint as adc
@@ -783,6 +784,75 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax.
     return output, pre_bias_logits
 
 
+def _gathered_kernels(layers: nnx.Module) -> dict:
+  """{kernel path within a scanned layer stack: its module}, for the kernels gathered once per stack."""
+  kernels = {}
+  for path, module in nnx.iter_modules(layers):
+    if isinstance(module, GateLogit):
+      kernels[path + ("kernel",)] = module
+  return kernels
+
+
+def gather_scanned_kernels_once(layers: nnx.Module, params, mesh: Mesh, shard_mode):
+  """Gathers small FSDP-sharded kernels of a scanned layer stack once, before the scan, not in every layer.
+
+  The MoE router kernel (embed x num_experts, f32 under float32_gate_logits) is small but FSDP-sharded,
+  so every layer all-gathers it on the critical path, synchronously on TPUs that run all-gathers on the
+  TensorCore, and the backward pass gathers it again. Replicating the stacked kernels up front is one
+  all-gather for all layers. `params` is the stack's Param state with the scan axis leading; pair with
+  `reduce_scatter_gathered_grads` inside the scan so each layer's gradient still lands sharded.
+  """
+  gathered = _gathered_kernels(layers)
+  replicated = NamedSharding(mesh, P())
+
+  def replicate(path, x):
+    if tuple(k.key for k in path if isinstance(k, jax.tree_util.DictKey)) not in gathered:
+      return x
+    if shard_mode == ShardMode.EXPLICIT:
+      return jax.sharding.reshard(x, replicated)
+    return jax.lax.with_sharding_constraint(x, replicated)
+
+  return jax.tree_util.tree_map_with_path(replicate, params)
+
+
+def gathered_kernel_grad_shardings(layers: nnx.Module, mesh: Mesh, rules) -> dict:
+  """{kernel path within one layer: the kernel's own sharding} for the kernels `gather_scanned_kernels_once` gathers."""
+  return {
+      path: NamedSharding(mesh, nn_partitioning.logical_to_mesh_axes(tuple(module.kernel_axes), rules))
+      for path, module in _gathered_kernels(layers).items()
+  }
+
+
+def _identity_with_grad_sharding(x, sharding):
+  @jax.custom_vjp
+  def ident(y):
+    return y
+
+  def ident_fwd(y):
+    return y, None
+
+  def ident_bwd(_, ct):
+    return (jax.lax.with_sharding_constraint(ct, sharding),)
+
+  ident.defvjp(ident_fwd, ident_bwd)
+  return ident(x)
+
+
+def reduce_scatter_gathered_grads(params, shardings: dict):
+  """Identity on one layer's params whose backward puts the gathered kernels' gradients back on their sharding.
+
+  A layer handed a slice of a replicated kernel otherwise gets a replicated gradient: an all-reduce per
+  layer, stacked at full size across the scan. Constraining it to the kernel's own sharding makes it a
+  reduce-scatter into a stacked shard, as when the layer gathered the kernel itself.
+  """
+
+  def apply(path, x):
+    keys = tuple(k.key for k in path if isinstance(k, jax.tree_util.DictKey))
+    return _identity_with_grad_sharding(x, shardings[keys]) if keys in shardings else x
+
+  return jax.tree_util.tree_map_with_path(apply, params)
+
+
 class RoutedMoE(nnx.Module):
   """Implements a routed MoE block."""
 
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 890bd01179..de3d9be6a9 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -1071,6 +1071,10 @@ def _apply_layers_sequentially(
     scan_axis = self.config.param_scan_axis
     if scan_axis != 0:
       params = jax.tree.map(lambda x: jnp.moveaxis(x, scan_axis, 0) if x.ndim > scan_axis else x, params)
+    gathered_grad_shardings = {}
+    if kv_caches_stacked is None and not self.config.parameter_memory_host_offload:
+      params = moe.gather_scanned_kernels_once(layers, params, self.mesh, self.config.shard_mode)
+      gathered_grad_shardings = moe.gathered_kernel_grad_shardings(layers, self.mesh, self.config.logical_axis_rules)
 
     layer_cls = layers.__class__
     sig = inspect.signature(layer_cls.__call__)
@@ -1105,6 +1109,9 @@ def layer_fn(carry, scanned_vars):
         kv_cache_layer = None
         forced_routed_experts_layer = None
 
+      if gathered_grad_shardings:
+        current_params = moe.reduce_scatter_gathered_grads(current_params, gathered_grad_shardings)
+
       if self.config.parameter_memory_host_offload:
         current_params = jax.tree.map(
             lambda x: jax.device_put(x, max_utils.device_space()),

From a0038f3f18f4ca41716513661af5165a6fe2a0d0 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Wed, 14 Oct 2026 22:45:39 +0000
Subject: [PATCH 17/19] Also gather the attention query kernel once per scanned
 layer stack

On the 397B bench (v7x) the query kernel's per-layer all-gathers are synchronous
and 256-way, and the backward gathers it twice per layer (rematerialized
projection and input gradient): ~0.54 s of exposed wait per micro-batch.

(cherry picked from commit 035ccaf65945b5ea7ce4c7baac84f5ae40091cfc)
---
 src/maxtext/layers/moe.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 91829464b3..333b2f5283 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -790,6 +790,10 @@ def _gathered_kernels(layers: nnx.Module) -> dict:
   for path, module in nnx.iter_modules(layers):
     if isinstance(module, GateLogit):
       kernels[path + ("kernel",)] = module
+    elif isinstance(module, attentions.Attention):
+      # The backward pass otherwise gathers the query kernel twice per layer: for the rematerialized
+      # projection and for the input gradient.
+      kernels[path + ("query", "kernel")] = module.query
   return kernels
 
 

From 9d7ccea6eb758fc1e3e889ac709549b7ca3e3f74 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Wed, 14 Oct 2026 22:53:16 +0000
Subject: [PATCH 18/19] Address review: module-level custom_vjp, MaxText
 logical_to_mesh_axes, list keys in param paths

- _identity_with_grad_sharding is a module-level custom_vjp with the sharding
  as a nondiff argument, as _forward_layout_constraint, instead of a new
  custom_vjp per call.
- Gradient shardings use maxtext.utils.sharding.logical_to_mesh_axes, which
  drops size-one mesh axes; the flax partitioning import goes.
- Param paths keep SequenceKey indices as well as DictKey keys (shared
  _param_path); a leaf's trailing GetAttrKey('value') is still dropped, since
  nnx.iter_modules paths do not have it.

(cherry picked from commit a0eaad4fc709a19a8e3961f6fbcee5aa188ce21c)
(cherry picked from commit ee1f10a91dc11c68ada50c3ee0407caac3285202)
---
 src/maxtext/layers/moe.py | 42 +++++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 13 deletions(-)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 333b2f5283..5c0116dd9b 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -24,7 +24,6 @@
 
 from aqt.jax.v2 import aqt_tensor as aqt
 from flax import nnx
-from flax.linen import partitioning as nn_partitioning
 from flax import struct
 import jax
 from jax import ad_checkpoint as adc
@@ -784,6 +783,20 @@ def __call__(self, inputs: jax.Array, _initializing: bool = False) -> Tuple[jax.
     return output, pre_bias_logits
 
 
+def _param_path(path) -> tuple:
+  """A params pytree path as the module path keys `nnx.iter_modules` yields, plus the param name.
+
+  Drops the trailing `GetAttrKey("value")` of each Variable leaf.
+  """
+  keys = []
+  for k in path:
+    if isinstance(k, jax.tree_util.DictKey):
+      keys.append(k.key)
+    elif isinstance(k, jax.tree_util.SequenceKey):
+      keys.append(k.idx)
+  return tuple(keys)
+
+
 def _gathered_kernels(layers: nnx.Module) -> dict:
   """{kernel path within a scanned layer stack: its module}, for the kernels gathered once per stack."""
   kernels = {}
@@ -810,7 +823,7 @@ def gather_scanned_kernels_once(layers: nnx.Module, params, mesh: Mesh, shard_mo
   replicated = NamedSharding(mesh, P())
 
   def replicate(path, x):
-    if tuple(k.key for k in path if isinstance(k, jax.tree_util.DictKey)) not in gathered:
+    if _param_path(path) not in gathered:
       return x
     if shard_mode == ShardMode.EXPLICIT:
       return jax.sharding.reshard(x, replicated)
@@ -822,24 +835,27 @@ def replicate(path, x):
 def gathered_kernel_grad_shardings(layers: nnx.Module, mesh: Mesh, rules) -> dict:
   """{kernel path within one layer: the kernel's own sharding} for the kernels `gather_scanned_kernels_once` gathers."""
   return {
-      path: NamedSharding(mesh, nn_partitioning.logical_to_mesh_axes(tuple(module.kernel_axes), rules))
+      path: NamedSharding(mesh, logical_to_mesh_axes(tuple(module.kernel_axes), mesh, rules))
       for path, module in _gathered_kernels(layers).items()
   }
 
 
+@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
 def _identity_with_grad_sharding(x, sharding):
-  @jax.custom_vjp
-  def ident(y):
-    return y
+  """The identity on `x`; its cotangent is constrained to `sharding`."""
+  return x
+
+
+def _identity_with_grad_sharding_fwd(x, sharding):
+  del sharding
+  return x, None
+
 
-  def ident_fwd(y):
-    return y, None
+def _identity_with_grad_sharding_bwd(sharding, _, cotangent):
+  return (jax.lax.with_sharding_constraint(cotangent, sharding),)
 
-  def ident_bwd(_, ct):
-    return (jax.lax.with_sharding_constraint(ct, sharding),)
 
-  ident.defvjp(ident_fwd, ident_bwd)
-  return ident(x)
+_identity_with_grad_sharding.defvjp(_identity_with_grad_sharding_fwd, _identity_with_grad_sharding_bwd)
 
 
 def reduce_scatter_gathered_grads(params, shardings: dict):
@@ -851,7 +867,7 @@ def reduce_scatter_gathered_grads(params, shardings: dict):
   """
 
   def apply(path, x):
-    keys = tuple(k.key for k in path if isinstance(k, jax.tree_util.DictKey))
+    keys = _param_path(path)
     return _identity_with_grad_sharding(x, shardings[keys]) if keys in shardings else x
 
   return jax.tree_util.tree_map_with_path(apply, params)

From dd546aa1f432789e55c423afc69a091124acb982 Mon Sep 17 00:00:00 2001
From: Mohit Khatwani <mohitkhatwani@google.com>
Date: Wed, 14 Oct 2026 22:53:27 +0000
Subject: [PATCH 19/19] Address review: skip the query kernel when Attention
 has none (fused_qkv)

Under fused_qkv the Attention module creates qkv_proj and no query, so the
gather-once selection must not touch module.query there.

(cherry picked from commit dcd4a5c7d3fe5ffd66bfe39343c915d9faceef29)
---
 src/maxtext/layers/moe.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 5c0116dd9b..cf2bafa5c8 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -803,9 +803,9 @@ def _gathered_kernels(layers: nnx.Module) -> dict:
   for path, module in nnx.iter_modules(layers):
     if isinstance(module, GateLogit):
       kernels[path + ("kernel",)] = module
-    elif isinstance(module, attentions.Attention):
+    elif isinstance(module, attentions.Attention) and hasattr(module, "query"):
       # The backward pass otherwise gathers the query kernel twice per layer: for the rematerialized
-      # projection and for the input gradient.
+      # projection and for the input gradient. Under fused_qkv there is no `query` (it is `qkv_proj`).
       kernels[path + ("query", "kernel")] = module.query
   return kernels