From dede90a9f44c70e56409a2a3145cbef6502df774 Mon Sep 17 00:00:00 2001
From: Rohan Bierneni <rbierneni@google.com>
Date: Mon, 12 Oct 2026 18:03:57 +0000
Subject: [PATCH 01/10] Optimize GDN triangular inverse with xor_dnc and
 support gdn_chunk_size=128

---
 src/maxtext/configs/types.py                  |   4 +-
 src/maxtext/kernels/gdn/compute_gdn.py        | 324 ++++++++++++---
 src/maxtext/kernels/gdn/config.py             |  24 +-
 src/maxtext/kernels/gdn/gdn_bwd/api.py        |  67 ++--
 .../kernels/gdn/gdn_bwd/bwd_memory_ref.py     |  21 +-
 src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py     |  70 +++-
 .../gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py      | 140 +++++--
 .../kernels/gdn/gdn_bwd/runtime_utils.py      |   9 +-
 src/maxtext/kernels/gdn/memory_ref.py         |  58 ++-
 src/maxtext/kernels/gdn/wrapper.py            |  93 ++++-
 tests/unit/gdn_bwd_pallas_test.py             | 379 +++++++++++++++++-
 tests/unit/gdn_cp_test.py                     | 236 ++++++++++-
 tests/unit/pyconfig_test.py                   |  27 +-
 13 files changed, 1277 insertions(+), 175 deletions(-)

diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 88ca7271ed..730217f5ef 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -4964,8 +4964,8 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None:
           f"'tensor' (got ici_tensor_parallelism={cfg.ici_tensor_parallelism}, "
           f"dcn_tensor_parallelism={cfg.dcn_tensor_parallelism})."
       )
-    if cfg.gdn_chunk_size != 64:
-      raise ValueError(f"use_gdn_kernel=True requires gdn_chunk_size=64, got {cfg.gdn_chunk_size}.")
+    if cfg.gdn_chunk_size not in (64, 128):
+      raise ValueError(f"use_gdn_kernel=True requires gdn_chunk_size in (64, 128), got {cfg.gdn_chunk_size}.")
     if cfg.gdn_key_head_dim % 128 or cfg.gdn_value_head_dim % 128:
       raise ValueError(
           "use_gdn_kernel=True requires gdn_key_head_dim and gdn_value_head_dim to be multiples of 128, "
diff --git a/src/maxtext/kernels/gdn/compute_gdn.py b/src/maxtext/kernels/gdn/compute_gdn.py
index 36d8c7bf62..a19b43790a 100644
--- a/src/maxtext/kernels/gdn/compute_gdn.py
+++ b/src/maxtext/kernels/gdn/compute_gdn.py
@@ -15,11 +15,58 @@
 
 """Core GDN forward computation with triangular inverse matrix caching."""
 
+import dataclasses
+
 import jax
 import jax.numpy as jnp
 
 from . import config
 
+# ==============================================================================
+# Triangular-inverse algorithm selection.
+#
+# `xor_dnc` is the default algorithm (pure-MXU divide-and-conquer with a
+# half-row final merge). `baseline_b16` is the legacy 16-row block forward
+# substitution algorithm. Both compute (I + L)^-1 for a batch of unit
+# lower-triangular matrices `t: [batch, chunk, chunk]`.
+# ==============================================================================
+
+DEFAULT_TRI_INV_ALGO = "xor_dnc"
+TRI_INV_ALGOS = ("xor_dnc", "baseline_b16")
+
+
+def validate_tri_inv_algo(algo: str) -> str:
+  if algo not in TRI_INV_ALGOS:
+    raise ValueError(f"Unknown triangular-inverse algorithm {algo!r}; expected one of {TRI_INV_ALGOS}.")
+  return algo
+
+
+def default_head_group_size(
+    chunk_size: int,
+    num_kq_heads: int,
+    num_v_heads: int,
+    act_dtype: jnp.dtype | None = None,
+) -> int | None:
+  """Returns the default KQ-head group size for `chunk_size`, or None for all heads at once.
+
+  At `chunk_size > 64`, processing all heads at once can exceed VMEM with `[num_v_heads,
+  chunk_size, chunk_size]` intermediates. Bounding each group to at most 32 value heads
+  for `bfloat16` activations (or 8 value heads for `float32` activations) at `chunk_size = 128`
+  keeps peak VMEM comfortably within budget while preserving full MXU utilization.
+  """
+  if chunk_size <= 64 or num_kq_heads <= 0 or num_v_heads <= 0:
+    return None
+  v_per_kq = max(1, num_v_heads // num_kq_heads)
+  base_v_heads = 8 if (act_dtype is not None and jnp.dtype(act_dtype) == jnp.float32) else 32
+  max_v_heads = max(v_per_kq, (base_v_heads * 128) // chunk_size)
+  max_kq = max(1, max_v_heads // v_per_kq)
+  if num_kq_heads <= max_kq:
+    return None
+  for g in range(max_kq, 0, -1):
+    if num_kq_heads % g == 0:
+      return g
+  return 1
+
 
 def l2_norm(x: jax.Array, eps: float = 1e-6) -> jax.Array:
   norm = jnp.sqrt(jnp.sum(x * x, axis=-1, keepdims=True, dtype=x.dtype) + eps)
@@ -37,12 +84,12 @@ def get_mask_dtype(dtype: jnp.dtype) -> jnp.dtype:
 
 
 # NOTE: Fork of recurrent_scan_v2.py but applied various optimizations.
-def invert_triangular_matrix(
+def _tri_inv_baseline_b16(
     t: jax.Array,
     block_size: int = 16,
     precision: jax.lax.Precision | None = None,
 ) -> jax.Array:
-  """Compute invert matrix of a given triauglar matrix."""
+  """Compute inverse matrix of a given triangular matrix using 16-row block forward substitution."""
 
   # NOTE: if chunk_size=1, compiler will perform DCE.
   out_dtype = t.dtype
@@ -94,6 +141,107 @@ def local_forward_sub(t_mat: jax.Array, b_mat: jax.Array) -> jax.Array:
   return jnp.concatenate(x_blocks, axis=1)
 
 
+def _bmm(lhs: jax.Array, rhs: jax.Array, precision: jax.lax.Precision | None) -> jax.Array:
+  """Batched [n, a, b] @ [n, b, c] -> [n, a, c] with f32 accumulation."""
+  return jax.lax.dot_general(
+      lhs,
+      rhs,
+      dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+      precision=precision,
+      preferred_element_type=jnp.float32,
+  )
+
+
+def _row_col_iotas(shape: tuple[int, ...]) -> tuple[jax.Array, jax.Array]:
+  """Broadcastable [1, R, C] row and column int32 iotas for the last two dims of `shape`."""
+  mask_shape = (1, shape[-2], shape[-1])
+  return (
+      jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1),
+      jax.lax.broadcasted_iota(jnp.int32, mask_shape, 2),
+  )
+
+
+def _dnc_merge_levels(
+    x: jax.Array,
+    lower_part: jax.Array,
+    first_block: int,
+    precision: jax.lax.Precision | None,
+    xor_rc: jax.Array | None = None,
+) -> jax.Array:
+  """Merges a block-diagonal inverse with `first_block`-sized blocks up to full size.
+
+  At level b every 2b-window [[A, 0], [B, C]] gets its lower-left block
+  -C^-1 B A^-1, i.e. X <- X - X ((M_b * L) X) at full width, where M_b selects the
+  lower-left bxb block of each 2b-window ((r ^ c) in [b, 2b) and r > c).
+
+  Args:
+    x: [n, c, c] f32 block-diagonal inverse with `first_block`-sized blocks.
+    lower_part: [n, c, c] f32 lower-triangular matrix (zeros above the diagonal)
+      whose strictly lower part is L of T = I + L.
+    first_block: Size of the diagonal blocks already inverted in `x`.
+    precision: Matmul precision.
+    xor_rc: Optional precomputed `bitwise_xor(iota_r, iota_c)` of `[1, c, c]`.
+
+  Returns:
+    [n, c, c] f32 inverse of T.
+  """
+  chunk = x.shape[-1]
+  if xor_rc is None:
+    iota_r, iota_c = _row_col_iotas(x.shape)
+    xor_rc = jnp.bitwise_xor(iota_r, iota_c)
+  b = first_block
+  while b < chunk:
+    if 2 * b == chunk and b >= 8:
+      lower_bot = jnp.where(xor_rc[:, b:, :] >= b, lower_part[:, b:, :], 0.0)
+      mid = _bmm(lower_bot, x, precision)
+      mid_padded = jnp.concatenate([jnp.zeros_like(mid), mid], axis=1)
+      bot = x[:, b:, :] - _bmm(x[:, b:, :], mid_padded, precision)
+      x = jnp.concatenate([x[:, :b, :], bot], axis=1)
+    else:
+      level_mask = (xor_rc >= b) & (xor_rc < 2 * b)
+      lower_b = jnp.where(level_mask, lower_part, 0.0)
+      x = x - _bmm(x, _bmm(lower_b, x, precision), precision)
+    b *= 2
+  return x
+
+
+def _tri_inv_xor_dnc(t: jax.Array, precision: jax.lax.Precision | None = None) -> jax.Array:
+  """Pure-MXU divide and conquer (gdm/universal_ai/clup/delta_rule.py), 2(log2 C - 1) matmuls."""
+  chunk = t.shape[-1]
+  if chunk <= 1 or (chunk & (chunk - 1)) != 0:
+    return _tri_inv_baseline_b16(t, precision=precision)
+  out_dtype = t.dtype
+  t32 = t.astype(jnp.float32)
+  iota_r, iota_c = _row_col_iotas(t.shape)
+  xor_rc = jnp.bitwise_xor(iota_r, iota_c)
+  x = jnp.where(xor_rc == 0, 1.0, jnp.where(xor_rc == 1, -t32, 0.0))
+  return _dnc_merge_levels(x, t32, 2, precision, xor_rc).astype(out_dtype)
+
+
+def invert_triangular_matrix(
+    t: jax.Array,
+    block_size: int = 16,
+    precision: jax.lax.Precision | None = None,
+    *,
+    algo: str = DEFAULT_TRI_INV_ALGO,
+) -> jax.Array:
+  """Computes (I + L)^-1 for unit lower-triangular `t: [batch, chunk, chunk]`.
+
+  Args:
+    t: Unit lower-triangular matrices (ones on the diagonal, zeros above it).
+    block_size: Block size of the `baseline_b16` algorithm.
+    precision: Matmul precision.
+    algo: One of `TRI_INV_ALGOS`; defaults to `"xor_dnc"`.
+
+  Returns:
+    The inverses, in `t.dtype`.
+  """
+  validate_tri_inv_algo(algo)
+  if algo == "xor_dnc":
+    return _tri_inv_xor_dnc(t, precision)
+  return _tri_inv_baseline_b16(t, block_size, precision)
+
+
 def fused_transpose_broadcast(x: jax.Array, src_dim: int, dst_dim: int) -> jax.Array:
   """Perform 1D transpose where results are broadcasted along src_dim."""
   assert x.shape[dst_dim] == 1
@@ -114,33 +262,76 @@ 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]."""
-  # 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 = (
-      jax.lax.Precision.HIGHEST
-      if jnp.dtype(cfg.dtypes.act_in) == jnp.float32 and jnp.dtype(cfg.dtypes.compute) == jnp.float32
-      else jax.lax.Precision.DEFAULT
-  )
+) -> tuple[jax.Array, jax.Array, jax.Array | None]:
+  """Perform chunked GDN over input [num_heads, chunk, head_dim].
+
+  With `cfg.head_group_size` set, kq heads (and their value heads) are processed
+  `head_group_size` at a time, which bounds the live [heads, chunk, chunk]
+  intermediates at large chunk sizes.
+  """
+  group = cfg.head_group_size
+  if group is None or group >= cfg.num_kq_heads:
+    return _chunked_gdn_per_seq_heads(
+        q_large, k_large, v_large, gating_log, beta, state_prev, cfg, seg_c_col=seg_c_col, seg_aux_col=seg_aux_col
+    )
+  if cfg.num_kq_heads % group != 0:
+    raise ValueError(f"head_group_size={group} must divide num_kq_heads={cfg.num_kq_heads}.")
+  group_v = group * cfg.v_per_kq_head
+  group_cfg = dataclasses.replace(cfg, num_kq_heads=group, num_v_heads=group_v, head_group_size=None)
+  # The decay cumsum is a chunk-long serial chain on one lane-dense vreg row that
+  # every head shares: compute it once instead of once per head group. The lane ->
+  # head transposes stay per group, so only that group's lane-padded
+  # [heads, chunk, 1] columns are live at a time.
+  valid_c, active_c, _ = _segment_columns(cfg, seg_c_col, seg_aux_col)
+  if valid_c is not None:
+    gating_log = jnp.where(valid_c, gating_log, 0.0)
+    beta = jnp.where(valid_c, beta, 0.0)
+  # [1, chunk, aligned_num_v_heads]
+  g_cum_sum_all = _decay_cumsum(gating_log, cfg, active_c)
+  outs, states, t_invs = [], [], []
+  for g in range(cfg.num_kq_heads // group):
+    kq_sl = slice(g * group, (g + 1) * group)
+    v_sl = slice(g * group_v, (g + 1) * group_v)
+    out, state, t_inv = _chunked_gdn_per_seq_heads(
+        q_large[kq_sl],
+        k_large[kq_sl],
+        v_large[v_sl],
+        None,
+        beta[:, :, v_sl],
+        state_prev[v_sl],
+        group_cfg,
+        seg_c_col=seg_c_col,
+        seg_aux_col=seg_aux_col,
+        decay_cumsum=g_cum_sum_all[:, :, v_sl],
+    )
+    outs.append(out)
+    states.append(state)
+    if cfg.emit_t_inv:
+      t_invs.append(t_inv)
+  t_inv_cat = jnp.concatenate(t_invs, axis=0) if cfg.emit_t_inv else None
+  return jnp.concatenate(outs, axis=0), jnp.concatenate(states, axis=0), t_inv_cat
 
-  valid_c = None
-  active_c = None
-  seg_prev = None
+
+def _segment_columns(
+    cfg: config.GDNConfig,
+    seg_c_col: jax.Array | None,
+    seg_aux_col: jax.Array | None,
+) -> tuple[jax.Array | None, jax.Array | None, jax.Array | None]:
+  """(valid_c, active_c, seg_prev) [1, chunk, 1] columns for packed sequences, else Nones."""
   if cfg.has_seg_ids and seg_c_col is not None and seg_aux_col is not None:
     valid_c = seg_c_col > 0.5
     active_c = jnp.abs(seg_c_col)
     seg_prev = jnp.abs(seg_aux_col[:, cfg.prev_kernel_size : cfg.prev_kernel_size + 1, :])
-    q_large = jnp.where(valid_c, q_large, 0.0)
-    k_large = jnp.where(valid_c, k_large, 0.0)
-    v_large = jnp.where(valid_c, v_large, 0.0)
-    gating_log = jnp.where(valid_c, gating_log, 0.0)
-    beta = jnp.where(valid_c, beta, 0.0)
+    return valid_c, active_c, seg_prev
+  return None, None, None
 
-  # NOTE: Repeat along non lane/sublane dim is free.
-  q_repeat = jnp.repeat(q_large, cfg.v_per_kq_head, axis=0)
-  k_repeat = jnp.repeat(k_large, cfg.v_per_kq_head, axis=0)
 
+def _decay_cumsum(
+    gating_log: jax.Array,  # [1, chunk, lanes], masked for packed sequences
+    cfg: config.GDNConfig,
+    active_c: jax.Array | None,
+) -> jax.Array:
+  """[1, chunk, lanes] f32 cumulative log-decay (restarting at document starts when packed)."""
   # Compute cumulative sum of decay in float32 to prevent bf16 cumsum drift.
   gating_log = gating_log.astype(jnp.float32)
   # [1, 1, num_v_heads]
@@ -153,7 +344,50 @@ def chunked_gdn_per_seq(
     else:
       g_cum_sum_list.append(g_cum_sum_list[-1] + gating_log[:, row : row + 1])
   # [1, chunk, num_v_heads]
-  g_cum_sum_log = jnp.concat(g_cum_sum_list, axis=1)
+  return jnp.concat(g_cum_sum_list, axis=1)
+
+
+def _chunked_gdn_per_seq_heads(
+    q_large: jax.Array,  # [num_kq_heads, chunk, kq_head_dim]
+    k_large: jax.Array,  # [num_kq_heads, chunk, kq_head_dim]
+    v_large: jax.Array,  # [num_v_heads, chunk, v_head_dim]
+    gating_log: jax.Array | None,  # [1, chunk, >= num_v_heads]; None with decay_cumsum
+    beta: jax.Array,  # [1, chunk, >= num_v_heads]; already masked with decay_cumsum
+    state_prev: jax.Array,  # [num_v_heads, kq_head_dim, v_head_dim]
+    cfg: config.GDNConfig,
+    seg_c_col: jax.Array | None = None,  # [1, chunk, 1]
+    seg_aux_col: jax.Array | None = None,  # [1, chunk, 1]
+    decay_cumsum: jax.Array | None = None,  # [1, chunk, >= num_v_heads]
+) -> tuple[jax.Array, jax.Array, jax.Array]:
+  """Perform chunked GDN over all heads of input [num_heads, chunk, head_dim].
+
+  `decay_cumsum` is the precomputed (`_decay_cumsum`) cumulative log-decay of these
+  heads; it replaces `gating_log`, and `beta` must then already be masked.
+  """
+  # 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 = (
+      jax.lax.Precision.HIGHEST
+      if jnp.dtype(cfg.dtypes.act_in) == jnp.float32 and jnp.dtype(cfg.dtypes.compute) == jnp.float32
+      else jax.lax.Precision.DEFAULT
+  )
+
+  valid_c, active_c, seg_prev = _segment_columns(cfg, seg_c_col, seg_aux_col)
+  if valid_c is not None:
+    q_large = jnp.where(valid_c, q_large, 0.0)
+    k_large = jnp.where(valid_c, k_large, 0.0)
+    v_large = jnp.where(valid_c, v_large, 0.0)
+    if decay_cumsum is None:
+      gating_log = jnp.where(valid_c, gating_log, 0.0)
+      beta = jnp.where(valid_c, beta, 0.0)
+
+  # NOTE: Repeat along non lane/sublane dim is free.
+  # [num_v_heads, chunk, kq_head_dim]
+  q_repeat = jnp.repeat(q_large, cfg.v_per_kq_head, axis=0)
+  k_repeat = jnp.repeat(k_large, cfg.v_per_kq_head, axis=0)
+
+  # [1, chunk, num_v_heads]
+  g_cum_sum_log = _decay_cumsum(gating_log, cfg, active_c) if decay_cumsum is None else decay_cumsum
 
   # [num_v_heads, chunk, 1]
   g_cum_sum_log = fused_transpose_broadcast(g_cum_sum_log, src_dim=2, dst_dim=0)
@@ -177,49 +411,44 @@ def chunked_gdn_per_seq(
     m_out = ((jnp.abs(active_c - active_last) < 0.5) & (active_last > 0.5)).astype(g_cum_sum_log.dtype)
     m_keep = ((jnp.abs(active_last - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(g_cum_sum_log.dtype)
 
-    bwd_diff_safe = jnp.where(m_out > 0.5, -g_cum_sum_diff_log[..., -1:], -1e4)
+    bwd_diff_safe = jnp.where(m_out > 0.5, g_cum_sum_log[:, -1:] - g_cum_sum_log, -1e4)
     gating_backward = (jnp.exp(bwd_diff_safe) * m_out).astype(cfg.dtypes.compute)
     gating_forward = (jnp.exp(g_cum_sum_log) * m_in).astype(cfg.dtypes.compute)
     gating_last = (jnp.exp(g_cum_sum_log[:, -1:]) * m_keep).astype(cfg.dtypes.compute)
   else:
     gating_map = jnp.exp(g_cum_sum_diff_log).astype(cfg.dtypes.compute)
     # [num_v_heads, chunk, 1]
-    gating_backward = jnp.exp(-g_cum_sum_diff_log[..., -1:]).astype(cfg.dtypes.compute)
+    gating_backward = jnp.exp(g_cum_sum_log[:, -1:] - g_cum_sum_log).astype(cfg.dtypes.compute)
     # [num_v_heads, chunk, 1]
     gating_forward = jnp.exp(g_cum_sum_log).astype(cfg.dtypes.compute)
     # [num_v_heads, 1, 1]
     gating_last = gating_forward[:, -1:]
 
   mask_dtype = get_mask_dtype(cfg.dtypes.compute)
-  iota_r = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 1)
-  iota_c = jax.lax.broadcasted_iota(mask_dtype, gating_map.shape, 2)
-  identity_mask = iota_r == iota_c
-  strictly_lower_mask = iota_r > iota_c
-  lower_mask = iota_r >= iota_c
-  # [num_v_heads, chunk, chunk]
-  gating_map_masked = jnp.where(strictly_lower_mask, gating_map, 0)
-
-  # [num_v_heads, chunk, kq_head_dim]
-  k_beta_repeat = k_repeat * beta_large
+  mask_shape = (1, cfg.chunk_size, cfg.chunk_size)
+  iota_r = jax.lax.broadcasted_iota(mask_dtype, mask_shape, 1)
+  iota_c = jax.lax.broadcasted_iota(mask_dtype, mask_shape, 2)
+  gating_map = jnp.where(iota_r >= iota_c, gating_map, 0)
 
-  # [num_v_heads, chunk, chunk]
-  beta_k_k_t = jax.lax.dot(
-      k_beta_repeat,
-      k_repeat,
+  # Compute k @ k^T at num_kq_heads (4x fewer matmuls) before repeating to num_v_heads.
+  k_k_t = jax.lax.dot(
+      k_large,
+      k_large,
       dimension_numbers=(((2,), (2,)), ((0,), (0,))),
       precision=dot_prec,
       preferred_element_type=jnp.float32,
   ).astype(cfg.dtypes.compute)
-  gating_beta_k_k_t = gating_map_masked * beta_k_k_t
-  t = jnp.where(identity_mask, 1, gating_beta_k_k_t)
+  beta_k_k_t = jnp.repeat(k_k_t, cfg.v_per_kq_head, axis=0) * beta_large
+  gating_beta_k_k_t = jnp.where(iota_r > iota_c, gating_map * beta_k_k_t, 0)
+  t = jnp.where(iota_r == iota_c, 1, gating_beta_k_k_t)
 
   # [num_v_heads, chunk, chunk]
-  t_inv = invert_triangular_matrix(t, precision=dot_prec)
+  t_inv = invert_triangular_matrix(t, precision=dot_prec, algo=cfg.tri_inv_algo)
 
   # [num_v_heads, chunk, v_head_dim]
   v_beta_large = v_large * beta_large
   # [num_v_heads, chunk, kv_head_dim]
-  k_beta_gating = k_beta_repeat * gating_forward
+  k_beta_gating = k_repeat * (beta_large * gating_forward)
   # NOTE: If v_head_dim < mxu size, concatenating them will help increase mxu
   # utilization. Also, if v_head_dim is multiple of lane size, concat / split
   # along lane dim is free - making this optimization strictly beneficial.
@@ -286,7 +515,6 @@ def chunked_gdn_per_seq(
   # [num_v_heads, chunk, chunk]
   out_qk = jnp.repeat(out_qk, cfg.v_per_kq_head, axis=0)
   out_qk *= gating_map
-  out_qk = jnp.where(lower_mask, out_qk, 0)
 
   # [num_v_heads, chunk, v_head_dim]
   out_new = jax.lax.dot(
@@ -314,7 +542,7 @@ def chunked_gdn(
     a_log: jax.Array,
     dt_bias: jax.Array,
     cfg: config.GDNConfig,
-) -> tuple[jax.Array, jax.Array, jax.Array]:
+) -> tuple[jax.Array, jax.Array, jax.Array | None]:
   """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]."""
 
   mask_dtype = get_mask_dtype(cfg.dtypes.compute)
@@ -342,6 +570,7 @@ def chunked_gdn(
   dt_bias_f32 = dt_bias.reshape(1, 1, 1, -1).astype(jnp.float32)
 
   # NOTE: Any element-wise computations should occur before repeat.
+  # [seqs, num_kq_heads, chunk, kq_head_dim]
   if cfg.use_qk_norm_in_gdn:
     q_large = l2_norm(q_large)
     k_large = l2_norm(k_large)
@@ -375,10 +604,11 @@ def chunked_gdn(
     )
     out_list.append(out.swapaxes(0, 1))
     state_list.append(state)
-    t_inv_list.append(t_inv)
+    if cfg.emit_t_inv:
+      t_inv_list.append(t_inv)
   out = jnp.stack(out_list, axis=0)
   state = jnp.stack(state_list, axis=0)
-  t_inv = jnp.stack(t_inv_list, axis=0)
+  t_inv = jnp.stack(t_inv_list, axis=0) if cfg.emit_t_inv else None
   return out, state, t_inv
 
 
diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py
index 8ad3933295..72c15644e7 100644
--- a/src/maxtext/kernels/gdn/config.py
+++ b/src/maxtext/kernels/gdn/config.py
@@ -26,6 +26,13 @@
 
 DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90
 
+# Optional GDNConfig fields and their defaults.
+_TRI_INV_FIELD_DEFAULTS = {
+    "tri_inv_algo": "xor_dnc",
+    "head_group_size": None,
+    "emit_t_inv": True,
+}
+
 
 class GDNMode(enum.StrEnum):
   """Execution mode for GDN kernel."""
@@ -82,6 +89,9 @@ class GDNConfig:
   num_buffers: int = 2
   has_seg_ids: bool = False
   use_qk_norm_in_gdn: bool = True
+  tri_inv_algo: str = "xor_dnc"
+  head_group_size: int | None = None
+  emit_t_inv: bool = True
 
   @property
   def chunk_size(self) -> int:
@@ -115,10 +125,22 @@ def aligned_num_v_heads(self) -> int:
     return pl.cdiv(self.num_v_heads + extra_lanes, num_lanes) * num_lanes
 
   def get_kernel_name(self) -> str:
-    return f"fused_conv1d_gdn_{self.mode.value}_b{self.seq_tile_size}" f"_c{self.chunk_size}"
+    """Pallas kernel name; non-default triangular-inverse options are appended as suffixes."""
+    name = f"fused_conv1d_gdn_{self.mode.value}_b{self.seq_tile_size}" f"_c{self.chunk_size}"
+    if self.tri_inv_algo != "xor_dnc":
+      name += f"_{self.tri_inv_algo}"
+    if self.head_group_size is not None:
+      name += f"_hg{self.head_group_size}"
+    if not self.emit_t_inv:
+      name += "_notinv"
+    return name
 
   def get_metadata(self) -> dict[str, str | int | float]:
+    """Flattened config fields for the kernel's metadata (knobs at their defaults are left out)."""
     cfgs_dict = dataclasses.asdict(self)
+    for field_name, default in _TRI_INV_FIELD_DEFAULTS.items():
+      if cfgs_dict.get(field_name) == default:
+        cfgs_dict.pop(field_name)
     ret = {}
     for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict):
       key = jax.tree_util.keystr(path, simple=True, separator=".")
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index 4ea358228b..f703c1dda2 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -23,6 +23,7 @@
 import jax.numpy as jnp
 
 from .. import compute_conv1d as local_compute_conv1d
+from .. import compute_gdn as local_compute_gdn
 from .. import wrapper as local_gdn_wrapper
 from . import cp_gdn
 from .compute_conv1d_bwd import conv1d_silu_bwd
@@ -61,15 +62,7 @@ def decoupled_conv1d_gdn_bwd_kernel(
     conv_halo_seg: Optional[jax.Array] = None,
     init_seg: Optional[jax.Array] = None,
     interpret: bool | pltpu.InterpretParams | None = None,
-) -> Tuple[
-    jax.Array,
-    jax.Array,
-    jax.Array,
-    jax.Array,
-    Optional[jax.Array],
-    jax.Array,
-    jax.Array,
-]:
+) -> Tuple[jax.Array, jax.Array, jax.Array, jax.Array, Optional[jax.Array], jax.Array, jax.Array,]:
   """Decoupled Conv1D + GDN backward combining Pallas GDN bwd and JAX Conv1D bwd."""
   del seq_lens, qkv
   if segment_ids is not None and conv_halo_seg is None and init_seg is None:
@@ -104,6 +97,7 @@ def decoupled_conv1d_gdn_bwd_kernel(
       segment_ids=segment_ids,
       init_seg=init_seg,
       interpret=interpret,
+      recompute_t_inv=t_inv is None,
   )
 
   d_pre_conv_qkv, d_conv_weight, d_conv_bias = conv1d_silu_bwd(
@@ -152,11 +146,7 @@ def _run_local_gdn_decoupled_fwd(
     segment_ids: Optional[jax.Array] = None,
     conv_halo_seg: Optional[jax.Array] = None,
     init_seg: Optional[jax.Array] = None,
-) -> Tuple[
-    Tuple[jax.Array, Tuple[jax.Array, jax.Array]],
-    Optional[jax.Array],
-    Optional[jax.Array],
-]:
+) -> Tuple[Tuple[jax.Array, Tuple[jax.Array, jax.Array]], Optional[jax.Array], 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
   if seq_len % chunk_size != 0:
@@ -168,6 +158,8 @@ def _run_local_gdn_decoupled_fwd(
   fallback_reason = pallas_unsupported_reason(head_k_dim=head_k_dim, head_v_dim=head_v_dim, chunk_size=chunk_size)
   if not on_cpu and fallback_reason is not None:
     raise ValueError(f"GDN Pallas TPU forward kernel does not support {fallback_reason}.")
+  emit_t_inv = chunk_size <= 64
+  head_group_size = local_compute_gdn.default_head_group_size(chunk_size, num_k_heads, num_v_heads, qkv.dtype)
   if on_cpu:
     out, states = pure_jax_decoupled_conv1d_gdn(
         qkv=qkv,
@@ -213,6 +205,8 @@ def _run_local_gdn_decoupled_fwd(
         conv_halo_seg=conv_halo_seg,
         init_seg=init_seg,
     )
+    if not emit_t_inv:
+      t_inv = None
     return (out, states), t_inv, chunk_states
 
   num_seqs = batch_size
@@ -274,10 +268,16 @@ def _run_local_gdn_decoupled_fwd(
       segment_ids=segment_ids,
       conv_halo_seg=conv_halo_seg,
       init_seg=init_seg,
+      tri_inv_algo=local_compute_gdn.DEFAULT_TRI_INV_ALGO,
+      head_group_size=head_group_size,
+      emit_t_inv=emit_t_inv,
   )
 
   core_attn_out = core_attn_out_flat.reshape(batch_size, seq_len, num_v_heads, head_v_dim)
-  t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, chunk_size, chunk_size)
+  if t_inv_raw is None:
+    t_inv = None
+  else:
+    t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, chunk_size, chunk_size)
   chunk_states = chunk_states_raw.astype(jnp.float32).reshape(batch_size, num_chunks, num_v_heads, head_k_dim, head_v_dim)
   return (
       (
@@ -377,7 +377,7 @@ 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
+  # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv (when chunk_size <= 64) and S_ext_local
   (_, states_pass1), t_inv, _ = _run_local_gdn_decoupled_fwd(
       qkv,
       b,
@@ -598,20 +598,18 @@ def _gdn_decoupled_conv1d_fwd(
       conv_state,
       recurrent_state,
       segment_ids,
-  ) = _unwrap_primal(
-      (
-          qkv,
-          b,
-          a,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          conv_state,
-          recurrent_state,
-          segment_ids,
-      )
-  )
+  ) = _unwrap_primal((
+      qkv,
+      b,
+      a,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      conv_state,
+      recurrent_state,
+      segment_ids,
+  ))
   qkv = checkpoint_name(qkv, "gdn_fwd_conv")
   qkv = checkpoint_name(qkv, "gdn_conv_out")
   has_initial_state = (conv_state is not None) or (recurrent_state is not None)
@@ -848,9 +846,12 @@ def _gdn_decoupled_conv1d_bwd(
         has_initial_state=bool(has_user_conv_state) or bool(has_user_recurrent_state),
     )
 
-  # Recompute forward chunk states and t_inv if not cached in residuals
+  # When chunk_states is cached in residuals and t_inv was omitted (chunk_size > 64),
+  # the Pallas backward kernel rebuilds T^-1 per chunk in VMEM using xor_dnc instead
+  # of reading it from HBM or recomputing it with XLA.
+  recompute_t_inv_in_kernel = chunk_states is not None and t_inv_fwd is None
   conv_out_cached = None
-  if chunk_states is None or t_inv_fwd is None:
+  if not recompute_t_inv_in_kernel and (chunk_states is None or t_inv_fwd is None):
     qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states(
         qkv=pre_conv_qkv,
         b=b,
@@ -949,6 +950,7 @@ def _gdn_decoupled_conv1d_bwd(
         init_seg=init_seg,
         d_recurrent_state=dht_local,
         return_dh0=need_dh0,
+        recompute_t_inv=recompute_t_inv_in_kernel,
     )
     if need_dh0:
       dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
@@ -1034,6 +1036,7 @@ def _gdn_decoupled_conv1d_bwd(
       init_seg=init_seg,
       d_recurrent_state=d_recurrent_state,
       return_dh0=need_dh0,
+      recompute_t_inv=recompute_t_inv_in_kernel,
   )
   # pylint: disable=unbalanced-tuple-unpacking
   if need_dh0:
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..ef29e24b9e 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
@@ -29,8 +29,12 @@ def make_bwd_block_specs(
     padded_num_v_heads: int | None = None,
     has_dht: bool = False,
     has_dh0: bool = False,
+    has_t_inv: bool = True,
 ) -> 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.
+
+  With `has_t_inv=False` (t_inv recomputed in-kernel) the t_inv input spec is omitted.
+  """
   if padded_num_v_heads is None:
     padded_num_v_heads = ((num_v_heads + 127) // 128) * 128
 
@@ -63,11 +67,16 @@ def rc(c):
           (None, None, None, num_v_heads, kq_head_dim, v_head_dim),
           lambda b, g, c: (b, g, rc(c), 0, 0, 0),
       ),
-      # 5: t_inv [B, G, C, 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),
-      ),
+  ]
+  if has_t_inv:
+    in_specs.append(
+        # 5: t_inv [B, G, C, 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),
+        )
+    )
+  in_specs += [
       # 6: a_log [B, G, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, 1, padded_num_v_heads),
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py b/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py
index 4958bfccde..95358a2da4 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/cp_gdn.py
@@ -36,12 +36,49 @@
 import jax.numpy as jnp
 
 from .. import compute_conv1d as local_compute_conv1d
+from .. import compute_gdn as local_compute_gdn
 
 # Precision of the f32 composition matmuls below: 6 bf16 passes on TPU. `gdn_cp_matmul_precision=high`
 # passes Precision.HIGH (3 passes) to them through their `precision` argument instead.
 DEFAULT_PRECISION = jax.lax.Precision.HIGHEST
 
 
+def _compute_t_inv_batched(
+    k_orig: jax.Array,  # [B, N_c, C, H_k, d_k] f32, already qk-normed if enabled
+    beta_h: jax.Array,  # [N_c, B, H_v, C] f32, already masked if segment_ids
+    cumsum_h: jax.Array,  # [N_c, B, H_v, C] f32
+    seg_c: Optional[jax.Array],  # [N_c, B, 1, C] f32 or None
+    valid_c: Optional[jax.Array],  # [N_c, B, 1, C] bool or None
+    *,
+    repeats: int,
+    precision: jax.lax.Precision = DEFAULT_PRECISION,
+) -> jax.Array:
+  """Computes `t_inv_c: [N_c, B, H_v, C, C]` across all local chunks in one batched call."""
+  num_chunks, batch, num_v_heads, chunk_size = beta_h.shape
+  k_kq = jnp.transpose(k_orig, (1, 0, 3, 2, 4))  # [N_c, B, H_k, C, d_k]
+  if valid_c is not None:
+    k_kq = jnp.where(valid_c[..., None], k_kq, 0.0)
+  kk_kq = jnp.matmul(k_kq, jnp.swapaxes(k_kq, -1, -2), precision=precision)  # [N_c, B, H_k, C, C]
+  kk_h = jnp.repeat(kk_kq, repeats, axis=2)  # [N_c, B, H_v, C, C]
+
+  mask_strict_base = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=-1)
+  if seg_c is not None and valid_c is not None:
+    same_valid = (jnp.abs(seg_c[..., :, None] - seg_c[..., None, :]) < 0.5) & valid_c[..., :, None]
+    mask_strict = mask_strict_base[None, None, None, :, :] * same_valid.astype(jnp.float32)
+  else:
+    mask_strict = mask_strict_base
+  diff = cumsum_h[..., :, None] - cumsum_h[..., None, :]
+  safe_diff = jnp.where(mask_strict > 0.5, diff, -1e4)
+  g_mat_strict = jnp.exp(safe_diff) * mask_strict
+
+  t_lower = kk_h * (beta_h[..., :, None] * g_mat_strict)
+  eye = jnp.eye(chunk_size, dtype=jnp.float32)
+  t_full = eye + t_lower
+  t_flat = t_full.reshape(num_chunks * batch * num_v_heads, chunk_size, chunk_size)
+  t_inv_flat = local_compute_gdn.invert_triangular_matrix(t_flat, precision=precision)
+  return t_inv_flat.reshape(num_chunks, batch, num_v_heads, chunk_size, chunk_size)
+
+
 def gather_cp_segment_metadata(
     segment_ids: Optional[jax.Array],
     cp_axis: str | tuple[str, ...],
@@ -279,7 +316,7 @@ def compose_local_from_t_inv(
     a: jax.Array,
     a_log: jax.Array,
     dt_bias: jax.Array,
-    t_inv: jax.Array,
+    t_inv: Optional[jax.Array] = None,
     *,
     num_k_heads: int,
     num_v_heads: int,
@@ -324,6 +361,7 @@ def compose_local_from_t_inv(
   dt_bias_f32 = dt_bias.astype(jnp.float32)[None, None, :, None]
   log_g = -jnp.exp(a_log_f32) * jax.nn.softplus(a_h + dt_bias_f32)
 
+  seg_c = None
   valid_c = None
   if segment_ids is not None:
     s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch, seq_len), init_seg=init_seg)
@@ -364,7 +402,18 @@ def compose_local_from_t_inv(
 
   k_beta_g = k_h * beta_h[..., None] * gating_forward
   k_scaled_bwd = k_h * gating_backward
-  t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4))  # [N_c, B, H_v, C, C]
+  if t_inv is None:
+    t_inv_c = _compute_t_inv_batched(
+        k_orig,
+        beta_h,
+        cumsum_h,
+        seg_c,
+        valid_c,
+        repeats=repeats,
+        precision=precision,
+    )
+  else:
+    t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4))  # [N_c, B, H_v, C, C]
 
   # Pre-batched GEMMs across all chunks outside any sequential loop
   w_all = jnp.matmul(t_inv_c, k_beta_g, precision=precision)  # [N_c, B, H_v, C, d_k]
@@ -422,7 +471,7 @@ def compose_bwd_local_from_t_inv(
     a_log: jax.Array,
     dt_bias: jax.Array,
     do: jax.Array,
-    t_inv: jax.Array,
+    t_inv: Optional[jax.Array] = None,
     *,
     num_k_heads: int,
     num_v_heads: int,
@@ -478,6 +527,8 @@ def compose_bwd_local_from_t_inv(
   dt_bias_f32 = dt_bias.astype(jnp.float32)[None, None, :, None]
   log_g = -jnp.exp(a_log_f32) * jax.nn.softplus(a_h + dt_bias_f32)
 
+  seg_c = None
+  valid_c = None
   mask_causal_base = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32), k=0)
   if segment_ids is not None:
     s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch, seq_len), init_seg=init_seg)
@@ -523,7 +574,18 @@ def compose_bwd_local_from_t_inv(
   q_g = q_h * gating_forward
   k_beta_g = k_h * beta_h[..., None] * gating_forward
   k_scaled_bwd = k_h * gating_backward
-  t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4))
+  if t_inv is None:
+    t_inv_c = _compute_t_inv_batched(
+        k_orig,
+        beta_h,
+        cumsum_h,
+        seg_c,
+        valid_c,
+        repeats=repeats,
+        precision=precision,
+    )
+  else:
+    t_inv_c = jnp.transpose(t_inv.astype(jnp.float32), (1, 0, 2, 3, 4))
 
   diff = cumsum_h[..., :, None] - cumsum_h[..., None, :]
   safe_diff_causal = jnp.where(mask_causal > 0.5, diff, -1e4)
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..b5b97e0454 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
@@ -24,6 +24,7 @@
 import jax.numpy as jnp
 
 from .. import compute_conv1d as local_compute_conv1d
+from .. import compute_gdn as local_compute_gdn
 from .bwd_memory_ref import make_bwd_block_specs
 from .runtime_utils import ensure_cpu_interpret_registered
 from .runtime_utils import pallas_unsupported_reason
@@ -121,6 +122,10 @@ class GDNBackwardConfig:
   vmem_limit_mb: Optional[int] = None
   use_qk_norm_in_gdn: bool = False
   has_seg_ids: bool = False
+  # When True, t_inv is not an input: each chunk rebuilds T and inverts it in VMEM
+  # with `tri_inv_algo`.
+  recompute_t_inv: bool = False
+  tri_inv_algo: str = local_compute_gdn.DEFAULT_TRI_INV_ALGO
 
   @property
   def repeats(self) -> int:
@@ -139,23 +144,37 @@ 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) == 17 and not cfg.recompute_t_inv:
     has_dht = True
     has_dh0 = True
   idx = 0
   # pylint: disable=unbalanced-tuple-unpacking
-  (
-      qkv_conv_ref,
-      b_ref,
-      a_ref,
-      do_ref,
-      chunk_states_ref,
-      t_inv_ref,
-      a_log_ref,
-      dt_bias_ref,
-      reset_ref,
-  ) = refs[idx : idx + 9]
-  idx += 9
+  if cfg.recompute_t_inv:
+    (
+        qkv_conv_ref,
+        b_ref,
+        a_ref,
+        do_ref,
+        chunk_states_ref,
+        a_log_ref,
+        dt_bias_ref,
+        reset_ref,
+    ) = refs[idx : idx + 8]
+    t_inv_ref = None
+    idx += 8
+  else:
+    (
+        qkv_conv_ref,
+        b_ref,
+        a_ref,
+        do_ref,
+        chunk_states_ref,
+        t_inv_ref,
+        a_log_ref,
+        dt_bias_ref,
+        reset_ref,
+    ) = refs[idx : idx + 9]
+    idx += 9
   if has_dht:
     dht_ref = refs[idx]
     idx += 1
@@ -214,7 +233,7 @@ def _init():
   a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32)
   do_val = do_ref[...].astype(jnp.float32)
   state_prev = chunk_states_ref[...].astype(jnp.float32)
-  t_inv_val = t_inv_ref[...].astype(jnp.float32)
+  t_inv_val = t_inv_ref[...].astype(jnp.float32) if t_inv_ref is not None else None
   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)
 
@@ -266,8 +285,10 @@ def _init():
     q_scaled = jnp.where(valid_c[:, :, None], q_scaled, 0.0)
     k_scaled_val = jnp.where(valid_c[:, :, None], k_scaled_val, 0.0)
 
-  q_rep = jnp.repeat(q_scaled, repeats, axis=1)
-  k_rep = jnp.repeat(k_scaled_val, repeats, axis=1)
+  q_kq = jnp.transpose(q_scaled, (1, 0, 2))
+  k_kq = jnp.transpose(k_scaled_val, (1, 0, 2))
+  q_h = jnp.repeat(q_kq, repeats, axis=0)
+  k_h = jnp.repeat(k_kq, repeats, axis=0)
   beta = jax.nn.sigmoid(b_val)
 
   sp_input = a_val + dt_bias_val
@@ -281,19 +302,15 @@ def _init():
 
   cumsum_log_g = jnp.dot(mask_cumsum, log_g, precision=jax.lax.Precision.HIGHEST)
 
-  q_h = jnp.transpose(q_rep, (1, 0, 2))
-  k_h = jnp.transpose(k_rep, (1, 0, 2))
   v_h = jnp.transpose(v, (1, 0, 2))
   beta_h = jnp.transpose(beta, (1, 0))
   cumsum_h = jnp.transpose(cumsum_log_g, (1, 0))
   do_h = jnp.transpose(do_val, (1, 0, 2))
 
   diff = cumsum_h[:, :, None] - cumsum_h[:, None, :]
-  safe_diff_strict = jnp.where(mask_strict[None, :, :] > 0.5, diff, -1e4)
-  g_mat_strict = jnp.exp(safe_diff_strict) * mask_strict[None, :, :]
-
   safe_diff_causal = jnp.where(mask_causal[None, :, :] > 0.5, diff, -1e4)
   g_mat_causal = jnp.exp(safe_diff_causal) * mask_causal[None, :, :]
+  g_mat_strict = g_mat_causal * mask_strict[None, :, :]
 
   gating_forward = jnp.exp(cumsum_h)[:, :, None]
   gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None]
@@ -306,9 +323,21 @@ def _init():
     gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None]
 
   k_beta = k_h * beta_h[:, :, None]
-  k_h_T = jnp.swapaxes(k_h, -1, -2)
-
-  A = t_inv_val
+  k_kq_T = jnp.swapaxes(k_kq, -1, -2)
+
+  if cfg.recompute_t_inv:
+    # Rebuild T = I + strictly_lower(beta_i k_i.k_j exp(g_i - g_j)) and invert it here
+    # instead of reading the forward kernel's t_inv from HBM. Compute both q @ k^T and
+    # k @ k^T at num_kq_heads (4x fewer 3-pass bf16 matmuls) before repeating to num_v_heads.
+    kk_kq = _gdn_matmul(k_kq, k_kq_T, compute_dtype, precise=True)
+    t_lower = jnp.repeat(kk_kq, repeats, axis=0) * (beta_h[:, :, None] * g_mat_strict)
+    t_rows = jax.lax.broadcasted_iota(jnp.int32, t_lower.shape, 1)
+    t_cols = jax.lax.broadcasted_iota(jnp.int32, t_lower.shape, 2)
+    t_full = jnp.where(t_rows == t_cols, 1.0, t_lower)
+    inv_prec = jax.lax.Precision.HIGHEST if jnp.dtype(compute_dtype) == jnp.float32 else jax.lax.Precision.DEFAULT
+    A = local_compute_gdn.invert_triangular_matrix(t_full, precision=inv_prec, algo=cfg.tri_inv_algo).astype(jnp.float32)
+  else:
+    A = t_inv_val
 
   v_beta = v_h * beta_h[:, :, None]
   k_beta_g = k_beta * gating_forward
@@ -331,7 +360,8 @@ def _init():
   )
 
   q_g = q_h * gating_forward
-  attn_unmasked = _gdn_matmul(q_h, k_h_T, compute_dtype, precise=True)
+  attn_kq = _gdn_matmul(q_kq, k_kq_T, compute_dtype, precise=True)
+  attn_unmasked = jnp.repeat(attn_kq, repeats, axis=0)
   attn = attn_unmasked * g_mat_causal
 
   k_scaled_bwd = k_h * gating_backward
@@ -484,7 +514,7 @@ def _pallas_gdn_bwd_kernel_single_group(
     dt_bias: jax.Array,
     do: jax.Array,
     chunk_states: jax.Array,
-    t_inv: jax.Array,
+    t_inv: Optional[jax.Array],
     *,
     cfg: GDNBackwardConfig,
     segment_ids: Optional[jax.Array] = None,
@@ -517,6 +547,8 @@ def _pallas_gdn_bwd_kernel_single_group(
       return_dh0=return_dh0,
       interpret=interpret,
       name=name,
+      recompute_t_inv=cfg.recompute_t_inv,
+      tri_inv_algo=cfg.tri_inv_algo,
   )
 
 
@@ -528,7 +560,7 @@ def pallas_gdn_bwd_kernel(
     dt_bias: jax.Array,
     do: jax.Array,
     chunk_states: jax.Array,
-    t_inv: jax.Array,
+    t_inv: Optional[jax.Array],
     *,
     num_v_heads: int,
     kq_head_dim: int,
@@ -543,12 +575,18 @@ def pallas_gdn_bwd_kernel(
     return_dh0: bool = False,
     interpret: bool | pltpu.InterpretParams | None = None,
     name: str = "gdn_bwd_kernel",
+    recompute_t_inv: bool = False,
+    tri_inv_algo: str = local_compute_gdn.DEFAULT_TRI_INV_ALGO,
 ):
   """Executes the Pallas reverse-chunk GDNv3 backward kernel using emit_pipeline.
 
   Dispatches a single kernel call with 3D grid=(batch_size, num_groups,
-  num_chunks).
+  num_chunks). With `recompute_t_inv=True`, `t_inv` is ignored (may be None) and
+  each chunk rebuilds T^-1 in VMEM with `tri_inv_algo`.
   """
+  if t_inv is None and not recompute_t_inv:
+    raise ValueError("pallas_gdn_bwd_kernel needs t_inv unless recompute_t_inv=True.")
+  local_compute_gdn.validate_tri_inv_algo(tri_inv_algo)
   if interpret is None:
     on_cpu = target_platform(qkv_conv) == "cpu"
     fallback_reason = pallas_unsupported_reason(head_k_dim=kq_head_dim, head_v_dim=v_head_dim, chunk_size=chunk_size)
@@ -579,10 +617,14 @@ def pallas_gdn_bwd_kernel(
       _, init_seg = local_compute_conv1d.initial_state_segment_metadata(segment_ids, 4)
   if head_tile is not None:
     target_tile = head_tile
-  elif jnp.dtype(qkv_conv.dtype) == jnp.float32:
-    target_tile = _F32_HEAD_TILE
   else:
-    target_tile = _BF16_HEAD_TILE
+    if jnp.dtype(qkv_conv.dtype) == jnp.float32:
+      target_tile = _F32_HEAD_TILE
+    else:
+      target_tile = _BF16_HEAD_TILE
+    if chunk_size > 64:
+      # ~10 [tile, C, C] f32 intermediates are live per step; keep them at the C=64 budget.
+      target_tile = max(1, target_tile * 64 // chunk_size)
 
   max_possible = min(num_v_heads, target_tile)
   tile_v_heads = None
@@ -615,6 +657,8 @@ def pallas_gdn_bwd_kernel(
       vmem_limit_mb=vmem_limit_mb,
       use_qk_norm_in_gdn=use_qk_norm_in_gdn,
       has_seg_ids=has_seg_ids,
+      recompute_t_inv=recompute_t_inv,
+      tri_inv_algo=tri_inv_algo,
   )
 
   padded_tile_v_heads = cfg.padded_num_v_heads
@@ -680,18 +724,20 @@ def pallas_gdn_bwd_kernel(
   ).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)
-  )
+  t_inv_6d = None
+  if not recompute_t_inv:
+    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)
+    )
 
   # 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)
@@ -744,7 +790,10 @@ def pallas_gdn_bwd_kernel(
       a_5d,
       do_6d,
       chunk_states_6d,
-      t_inv_6d,
+  ]
+  if t_inv_6d is not None:
+    in_args.append(t_inv_6d)
+  in_args += [
       a_log_4d,
       dt_bias_4d,
       reset_hbm_5d,
@@ -774,7 +823,10 @@ def pallas_gdn_bwd_kernel(
       padded_num_v_heads=padded_tile_v_heads,
       has_dht=has_dht,
       has_dh0=has_dh0,
+      has_t_inv=not recompute_t_inv,
   )
+  if recompute_t_inv and name == "gdn_bwd_kernel":
+    name = f"gdn_bwd_kernel_recompute_{tri_inv_algo}"
 
   out_shapes_list = [
       jax.ShapeDtypeStruct(
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py b/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py
index b70a0bf07f..6bdaaba0e1 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/runtime_utils.py
@@ -21,6 +21,9 @@
 
 from .. import compute_gdn as local_compute_gdn
 
+# Chunk sizes supported by the Pallas TPU GDN forward and backward kernels.
+_SUPPORTED_PALLAS_CHUNK_SIZES = (64, 128)
+
 
 def target_platform(*arrays: Any) -> str:
   """Returns the platform ('tpu', 'cpu', 'gpu') the computation is being compiled for.
@@ -43,8 +46,8 @@ def pallas_unsupported_reason(*, head_k_dim: int, head_v_dim: int, chunk_size: i
     return f"head_k_dim={head_k_dim} is not a multiple of 128"
   if head_v_dim % 128 != 0:
     return f"head_v_dim={head_v_dim} is not a multiple of 128"
-  if chunk_size != 64:
-    return f"chunk_size={chunk_size} != 64"
+  if chunk_size not in _SUPPORTED_PALLAS_CHUNK_SIZES:
+    return f"chunk_size={chunk_size} not in {_SUPPORTED_PALLAS_CHUNK_SIZES}"
   return None
 
 
@@ -67,7 +70,7 @@ def ensure_cpu_interpret_registered() -> None:
 
 @jax.custom_vjp
 def invert_triangular_matrix(t: jax.Array) -> jax.Array:
-  """Computes inverse of unit lower-triangular matrix using Tokamax block forward substitution."""
+  """Computes inverse of unit lower-triangular matrix using `compute_gdn.invert_triangular_matrix`."""
   return local_compute_gdn.invert_triangular_matrix(t, block_size=16, precision=jax.lax.Precision.HIGHEST)
 
 
diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py
index f29d79d81d..6d9ce6492e 100644
--- a/src/maxtext/kernels/gdn/memory_ref.py
+++ b/src/maxtext/kernels/gdn/memory_ref.py
@@ -366,6 +366,7 @@ class ChunkStatesBufferedRef(BaseBufferedRef):
   """DMA buffer for caching intermediate recurrent chunk states."""
 
   def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
+    """Initiates asynchronous chunk-state copy from current output VMEM slot."""
     assert self.sem_sends is not None
     assert self.window_ref is not None
     slot = self.current_copy_out_slot
@@ -379,8 +380,18 @@ def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
           dst_ref.at[p_id * self.cfg.seq_tile_size + idx],
           sem,
       ).start()
+    if self.buffer_count == 1:
+      for idx in range(self.cfg.seq_tile_size):
+        pltpu.make_async_copy(
+            vmem_ref.at[idx],
+            vmem_ref.at[idx],
+            sem,
+        ).wait()
 
   def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
+    """Waits for asynchronous chunk-state output copy to complete."""
+    if self.buffer_count == 1:
+      return
     assert self.sem_sends is not None
     assert self.window_ref is not None
     slot = self.current_wait_out_slot
@@ -457,21 +468,31 @@ def copy_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
     vmem_ref = self.window_ref.at[slot]
     p_id = grid_indices[0]
 
+    total_dma_size = 0
     for idx in range(self.cfg.seq_tile_size):
       record = self.metadata_ref.get_record(p_id, idx)
       is_last_tile = record.is_last_tile
       s_idx = record.s_idx
       state_idx = self.metadata_ref.s_idx_to_state_indices[s_idx]
       dma_size = jnp.where(is_last_tile, 1, 0)
+      total_dma_size += dma_size
 
       pltpu.make_async_copy(
           vmem_ref.at[pl.ds(idx, dma_size)],  # pyrefly: ignore[missing-attribute]
           dst_ref.at[pl.ds(state_idx, dma_size)],
           sem,
       ).start()
+    if self.buffer_count == 1:
+      pltpu.make_async_copy(
+          vmem_ref.at[pl.ds(0, total_dma_size)],  # pyrefly: ignore[missing-attribute]
+          vmem_ref.at[pl.ds(0, total_dma_size)],  # pyrefly: ignore[missing-attribute]
+          sem,
+      ).wait()
 
   def wait_out(self, dst_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
     """Waits for asynchronous state output copy to complete."""
+    if self.buffer_count == 1:
+      return
     assert self.sem_sends is not None
     assert self.window_ref is not None
     slot = self.current_wait_out_slot
@@ -530,7 +551,18 @@ def create_allocs(
       pipeline_mode=pipeline_mode,
   )
 
-  qkv_spec = block_spec_partial(block_shape=qkv_shape)
+  large_fp32_buffer_count = (
+      1
+      if (cfg.chunk_size > 64 and cfg.num_v_heads >= 64 and jnp.dtype(cfg.dtypes.act_in) == jnp.float32)
+      else pipeline_mode.buffer_count
+  )
+  large_fp32_spec_partial = functools.partial(
+      pl.BlockSpec,
+      memory_space=pltpu.VMEM,
+      index_map=lambda i: (i,),
+      pipeline_mode=pl.Buffered(buffer_count=large_fp32_buffer_count, use_lookahead=False),
+  )
+  qkv_spec = large_fp32_spec_partial(block_shape=qkv_shape)
   ba_spec = block_spec_partial(block_shape=ba_shape)
   in_buffered_partial = functools.partial(
       InBufferedRef.input,
@@ -539,7 +571,14 @@ def create_allocs(
       cfg=cfg,
       metadata_ref=metadata_ref,
   )
-  qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref)
+  qkv_alloc = InBufferedRef.input(
+      spec=qkv_spec,
+      dtype_or_type=qkv_ref,
+      buffer_count=large_fp32_buffer_count,
+      use_lookahead=pipeline_mode.use_lookahead,
+      cfg=cfg,
+      metadata_ref=metadata_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)
 
@@ -553,7 +592,7 @@ def create_allocs(
   )
 
   conv_spec = block_spec_partial(block_shape=conv_shape)
-  recurrent_spec = block_spec_partial(block_shape=recurrent_shape)
+  recurrent_spec = large_fp32_spec_partial(block_shape=recurrent_shape)
   state_buffered_partial = functools.partial(
       StateBufferedRef.input_output,
       buffer_count=pipeline_mode.buffer_count,
@@ -562,7 +601,14 @@ def create_allocs(
       metadata_ref=metadata_ref,
   )
   conv_alloc = state_buffered_partial(spec=conv_spec, dtype_or_type=conv_state_ref)
-  recurrent_alloc = state_buffered_partial(spec=recurrent_spec, dtype_or_type=recurrent_state_ref)
+  recurrent_alloc = StateBufferedRef.input_output(
+      spec=recurrent_spec,
+      dtype_or_type=recurrent_state_ref,
+      buffer_count=large_fp32_buffer_count,
+      use_lookahead=pipeline_mode.use_lookahead,
+      cfg=cfg,
+      metadata_ref=metadata_ref,
+  )
 
   allocs = [
       qkv_alloc,
@@ -598,9 +644,9 @@ def create_allocs(
         cfg.v_head_dim,
     )
     chunk_states_alloc = ChunkStatesBufferedRef.output(
-        spec=block_spec_partial(block_shape=chunk_states_shape),
+        spec=large_fp32_spec_partial(block_shape=chunk_states_shape),
         dtype_or_type=chunk_states_ref,
-        buffer_count=pipeline_mode.buffer_count,
+        buffer_count=large_fp32_buffer_count,
         use_lookahead=pipeline_mode.use_lookahead,
         cfg=cfg,
         metadata_ref=metadata_ref,
diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py
index 7b951c7cee..686eac6a4d 100644
--- a/src/maxtext/kernels/gdn/wrapper.py
+++ b/src/maxtext/kernels/gdn/wrapper.py
@@ -139,6 +139,8 @@ def inner_kernel(
       carry_recurrent_scratch_ref=carry_recurrent_scratch_ref,
       cfg=cfg,
   )
+  if chunk_states_slot_ref is not None:
+    chunk_states_slot_ref[...] = prev_recurrent.astype(chunk_states_slot_ref.dtype)
 
   # Step 1: Conv1D.
   qkv_in_compact = qkv_slot_ref[...].astype(jnp.float32)
@@ -218,14 +220,39 @@ def inner_kernel(
   # Store output, recurrent, and t_inv to vmem.
   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:
-    chunk_states_slot_ref[...] = prev_recurrent.astype(chunk_states_slot_ref.dtype)
+  if t_inv_slot_ref is not None:
+    t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype)
 
   if carry_recurrent_scratch_ref is not None:
     carry_recurrent_scratch_ref[...] = new_recurrent_state
 
 
+def inner_kernel_without_t_inv(
+    qkv_slot_ref: jax.Ref,
+    b_slot_ref: jax.Ref,
+    a_slot_ref: jax.Ref,
+    conv_state_slot_ref: jax.Ref,
+    recurrent_slot_ref: jax.Ref,
+    out_slot_ref: jax.Array,
+    *args,
+    cfg: config.GDNConfig,
+    **kwargs,
+) -> None:
+  """`inner_kernel` for pipelines without a t_inv output (cfg.emit_t_inv=False)."""
+  inner_kernel(
+      qkv_slot_ref,
+      b_slot_ref,
+      a_slot_ref,
+      conv_state_slot_ref,
+      recurrent_slot_ref,
+      out_slot_ref,
+      None,
+      *args,
+      cfg=cfg,
+      **kwargs,
+  )
+
+
 def outer_kernel(
     # Inputs.
     metadata_ref: memory_ref.MetadataRef,
@@ -251,6 +278,7 @@ def outer_kernel(
   del conv_state_out_ref, recurrent_state_out_ref
 
   chunk_states_ref = args[0] if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) else None
+  emit_t_inv = cfg.emit_t_inv
 
   allocs = memory_ref.create_allocs(
       metadata_ref=metadata_ref,
@@ -261,7 +289,7 @@ def outer_kernel(
       conv_state_ref=conv_state_ref,
       recurrent_state_ref=recurrent_state_ref,
       cfg=cfg,
-      t_inv_ref=t_inv_ref,
+      t_inv_ref=t_inv_ref if emit_t_inv else None,
       chunk_states_ref=chunk_states_ref,
   )
   qkv_alloc = allocs[0]
@@ -270,18 +298,24 @@ def outer_kernel(
   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
+  next_alloc = 6
+  t_inv_alloc = None
+  if emit_t_inv:
+    t_inv_alloc = allocs[next_alloc]
+    next_alloc += 1
+  chunk_states_alloc = allocs[next_alloc] if len(allocs) > next_alloc else None
 
   num_tiles = metadata_ref.num_tiles[...]
 
-  out_specs = [out_alloc.spec, t_inv_alloc.spec]
+  out_specs = [out_alloc.spec]
+  if t_inv_alloc is not None:
+    out_specs.append(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,
+          inner_kernel if emit_t_inv else inner_kernel_without_t_inv,
           cfg=cfg,
       ),
       grid=(num_tiles,),
@@ -297,7 +331,7 @@ def outer_kernel(
 
   @pl.with_scoped(allocations=allocs)
   def _run(allocations):
-    out_args = [out_ref, t_inv_ref]
+    out_args = [out_ref, t_inv_ref] if emit_t_inv else [out_ref]
     if chunk_states_ref is not None:
       out_args.append(chunk_states_ref)
     pipeline_func(
@@ -334,6 +368,9 @@ def _run(allocations):
         "compute_precision",
         "is_prefill_only",
         "use_qk_norm_in_gdn",
+        "tri_inv_algo",
+        "head_group_size",
+        "emit_t_inv",
     ),
 )
 def fused_conv1d_gdn(
@@ -365,8 +402,16 @@ 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)."""
+    tri_inv_algo: str = compute_gdn.DEFAULT_TRI_INV_ALGO,
+    head_group_size: int | None = None,
+    emit_t_inv: bool = True,
+) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array | None, jax.Array]:
+  """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states).
+
+  When `head_group_size` is None, `compute_gdn.default_head_group_size` selects an
+  appropriate KQ-head group size for `chunk_size > 64`. With `emit_t_inv=False`,
+  the returned `t_inv` is None.
+  """
   act_in_dtype = qkv.dtype
   act_out_dtype = qkv.dtype
   conv_out_dtype = conv_state.dtype
@@ -485,8 +530,19 @@ def call_kernel(
   ) -> tuple[jax.Array, ...]:
     if mode == config.GDNMode.BATCHED:
       tile_size = decode_tile_size
+      tri_inv_kwargs = {}
     else:
       tile_size = mixed_tile_size
+      resolved_head_group_size = (
+          compute_gdn.default_head_group_size(tile_size, n_kq, n_v, act_in_dtype)
+          if head_group_size is None
+          else head_group_size
+      )
+      tri_inv_kwargs = {
+          "tri_inv_algo": compute_gdn.validate_tri_inv_algo(tri_inv_algo),
+          "head_group_size": resolved_head_group_size,
+          "emit_t_inv": emit_t_inv,
+      }
 
     cfg = config.GDNConfig(
         mode=mode,
@@ -507,6 +563,7 @@ def call_kernel(
             recurrent_state=in_recurrent_state.dtype,
             conv_state=in_conv_state.dtype,
         ),
+        **tri_inv_kwargs,
     )
 
     if mode == config.GDNMode.BATCHED:
@@ -542,10 +599,14 @@ def call_kernel(
       input_output_aliases[len(metadata_obj) + 5] = 0
 
     num_chunks = cfg.batch_size // cfg.chunk_size
-    t_inv_shape = jax.ShapeDtypeStruct(
-        (num_chunks, cfg.num_v_heads, cfg.chunk_size, cfg.chunk_size),
-        cfg.dtypes.compute,
-    )
+    if cfg.emit_t_inv:
+      t_inv_shape = jax.ShapeDtypeStruct(
+          (num_chunks, cfg.num_v_heads, cfg.chunk_size, cfg.chunk_size),
+          cfg.dtypes.compute,
+      )
+    else:
+      # Never written; keeps the pallas_call output arity unchanged.
+      t_inv_shape = jax.ShapeDtypeStruct((1, 1, 8, 128), cfg.dtypes.compute)
 
     if mode == config.GDNMode.PER_SEQ:
       chunk_states_shape = jax.ShapeDtypeStruct(
@@ -628,5 +689,7 @@ def call_kernel(
     ).astype(conv_out_dtype)
     out_conv_state = out_conv_state.at[state_indices].set(masked_cs)
   out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype)
+  if not emit_t_inv:
+    t_inv = None
 
   return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states
diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py
index ee91c2b33e..6f201ff581 100644
--- a/tests/unit/gdn_bwd_pallas_test.py
+++ b/tests/unit/gdn_bwd_pallas_test.py
@@ -1582,12 +1582,10 @@ def test_bwd_block_specs_and_symbolic_zero_dht_dh0_omission(self):
     orig_kernel = bwd_api.pallas_gdn_bwd_kernel
 
     def wrapped_kernel(*args, **kwargs):
-      recorded_calls.append(
-          (
-              kwargs.get("d_recurrent_state") is not None,
-              bool(kwargs.get("return_dh0", False)),
-          )
-      )
+      recorded_calls.append((
+          kwargs.get("d_recurrent_state") is not None,
+          bool(kwargs.get("return_dh0", False)),
+      ))
       return orig_kernel(*args, **kwargs)
 
     with mock.patch.object(bwd_api, "pallas_gdn_bwd_kernel", side_effect=wrapped_kernel):
@@ -2783,6 +2781,375 @@ def operand_dtypes(eqn):
     self.assertLen(f32_dots, 17)
     self.assertEqual({e.params["precision"] for e in f32_dots}, {highest})
 
+  def test_chunk_size_128_fwd_bwd_and_remat_policies(self):
+    """Verifies chunk_size=128 forward + backward across bf16/fp32, packing, initial states, and remat."""
+    is_tpu = jax.default_backend() == "tpu"
+    chunk_size = 128
+    dims = _packing_dims()
+    batch_size, seq_len = (2 if is_tpu else 1), 2 * chunk_size
+    key = jax.random.PRNGKey(2026)
+    qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims)
+    cs, rs = _packing_states(jax.random.PRNGKey(2027), batch_size, dims)
+    seg = jnp.zeros((batch_size, seq_len), dtype=jnp.int32)
+    # Row 0: doc 1 across chunk boundary, doc 2 with trailing 0-pad
+    seg = seg.at[0, : chunk_size + 20].set(1)
+    seg = seg.at[0, chunk_size + 20 : 2 * chunk_size - 12].set(2)
+    if batch_size > 1:
+      # Row 1: leading 0-pad, doc 1, doc 2
+      seg = seg.at[1, 8 : chunk_size - 8].set(1)
+      seg = seg.at[1, chunk_size - 8 :].set(2)
+
+    def ref_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_):
+      o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size)
+      loss = jnp.sum(o * up) + 0.1 * jnp.sum(ncs) + 0.2 * jnp.sum(nrs)
+      return loss, (o, ncs, nrs)
+
+    ref_vg = jax.value_and_grad(ref_loss, argnums=tuple(range(9)), has_aux=True)
+    if is_tpu:
+      ref_vg = jax.jit(ref_vg)
+    (_, ref_aux), ref_grads = ref_vg(qkv, b, a, cw, cb, al, dt, cs, rs)
+
+    cp_state_policy = jax.checkpoint_policies.save_only_these_names("gdn_cp_state")
+    cases = (
+        [
+            (act_dtype, tol, policy)
+            for act_dtype, tol in ((jnp.float32, 1.5e-2), (jnp.bfloat16, 4e-2))
+            for policy in (None, jax.checkpoint_policies.nothing_saveable, cp_state_policy)
+        ]
+        if is_tpu
+        else [
+            (jnp.float32, 2e-3, None),
+            (jnp.bfloat16, 4e-2, cp_state_policy),
+        ]
+    )
+
+    for act_dtype, tol, policy in cases:
+      qkv_d = qkv.astype(act_dtype)
+      b_d = b.astype(act_dtype)
+      a_d = a.astype(act_dtype)
+      cw_d = cw.astype(act_dtype)
+      cb_d = cb.astype(act_dtype)
+      up_d = up.astype(act_dtype)
+
+      def kern_step(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, _act_dtype=act_dtype):
+        return gdn_bwd_pallas.gdn_decoupled_conv1d(
+            qkv_,
+            b_,
+            a_,
+            cw_,
+            cb_,
+            al_,
+            dt_,
+            cs_,
+            rs_,
+            dims["num_k_heads"],
+            dims["num_v_heads"],
+            dims["head_k_dim"],
+            dims["head_v_dim"],
+            dims["conv_kernel_size"],
+            chunk_size,
+            True,
+            _act_dtype,
+            None,
+            seg,
+        )
+
+      step_fn = jax.checkpoint(kern_step, policy=policy) if policy is not None else kern_step
+
+      def kern_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, _step=step_fn, _up_d=up_d):
+        o, (ncs, nrs) = _step(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_)
+        loss = (
+            jnp.sum(o.astype(jnp.float32) * _up_d.astype(jnp.float32))
+            + 0.1 * jnp.sum(ncs.astype(jnp.float32))
+            + 0.2 * jnp.sum(nrs)
+        )
+        return loss, (o, ncs, nrs)
+
+      kern_vg = jax.value_and_grad(kern_loss, argnums=tuple(range(9)), has_aux=True)
+      if is_tpu:
+        kern_vg = jax.jit(kern_vg)
+      (_, got_aux), got_grads = kern_vg(qkv_d, b_d, a_d, cw_d, cb_d, al, dt, cs, rs)
+      for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), got_aux, ref_aux):
+        np.testing.assert_allclose(
+            np.asarray(got, dtype=np.float32),
+            np.asarray(exp, dtype=np.float32),
+            rtol=tol,
+            atol=tol,
+            err_msg=f"chunk_size={chunk_size} dtype={act_dtype} policy={policy}: {name}",
+        )
+      for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), got_grads, ref_grads):
+        np.testing.assert_allclose(
+            np.asarray(got, dtype=np.float32),
+            np.asarray(exp, dtype=np.float32),
+            rtol=tol,
+            atol=tol,
+            err_msg=f"chunk_size={chunk_size} dtype={act_dtype} policy={policy}: {name}",
+        )
+
+  def test_chunk_size_128_head_group_size_64_v_heads(self):
+    """Verifies chunk_size=128 with num_k_heads=16, num_v_heads=64 for both bf16 and fp32 activations + packing."""
+    if jax.default_backend() != "tpu":
+      self.skipTest("TPU-specific VMEM head_group_size test at chunk_size=128")
+    from maxtext.kernels.gdn import compute_gdn as local_compute_gdn  # pylint: disable=import-outside-toplevel
+
+    self.assertIsNone(local_compute_gdn.default_head_group_size(64, 16, 64))
+    self.assertEqual(local_compute_gdn.default_head_group_size(128, 16, 64), 8)
+    self.assertEqual(local_compute_gdn.default_head_group_size(128, 16, 64, jnp.bfloat16), 8)
+    self.assertEqual(local_compute_gdn.default_head_group_size(128, 16, 64, jnp.float32), 2)
+    self.assertIsNone(local_compute_gdn.default_head_group_size(128, 8, 32, jnp.bfloat16))
+
+    chunk_size = 128
+    dims = {"num_k_heads": 16, "num_v_heads": 64, "head_k_dim": 128, "head_v_dim": 128, "conv_kernel_size": 4}
+    batch_size, seq_len = 1, 256
+    key = jax.random.PRNGKey(3030)
+    qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims)
+    cs, rs = _packing_states(jax.random.PRNGKey(3031), batch_size, dims)
+    seg = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+    seg = seg.at[0, :8].set(0)
+    seg = seg.at[0, chunk_size + 20 : 2 * chunk_size - 6].set(2)
+    seg = seg.at[0, 2 * chunk_size - 6 :].set(0)
+
+    def ref_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_):
+      o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size)
+      loss = jnp.sum(o * up) + 0.1 * jnp.sum(ncs) + 0.2 * jnp.sum(nrs)
+      return loss, (o, ncs, nrs)
+
+    (_, ref_aux), ref_grads = jax.jit(jax.value_and_grad(ref_loss, argnums=tuple(range(9)), has_aux=True))(
+        qkv, b, a, cw, cb, al, dt, cs, rs
+    )
+
+    for act_dtype, tol in ((jnp.bfloat16, 4e-2), (jnp.float32, 2e-2)):
+      qkv_d, b_d, a_d, cw_d, cb_d, up_d = (x.astype(act_dtype) for x in (qkv, b, a, cw, cb, up))
+
+      def kern_loss(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, _act_dtype=act_dtype, _up_d=up_d):
+        o, (ncs, nrs) = gdn_bwd_pallas.gdn_decoupled_conv1d(
+            qkv_,
+            b_,
+            a_,
+            cw_,
+            cb_,
+            al_,
+            dt_,
+            cs_,
+            rs_,
+            dims["num_k_heads"],
+            dims["num_v_heads"],
+            dims["head_k_dim"],
+            dims["head_v_dim"],
+            dims["conv_kernel_size"],
+            chunk_size,
+            True,
+            _act_dtype,
+            None,
+            seg,
+        )
+        loss = (
+            jnp.sum(o.astype(jnp.float32) * _up_d.astype(jnp.float32))
+            + 0.1 * jnp.sum(ncs.astype(jnp.float32))
+            + 0.2 * jnp.sum(nrs)
+        )
+        return loss, (o, ncs, nrs)
+
+      (_, got_aux), got_grads = jax.jit(jax.value_and_grad(kern_loss, argnums=tuple(range(9)), has_aux=True))(
+          qkv_d, b_d, a_d, cw_d, cb_d, al, dt, cs, rs
+      )
+      for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), got_aux, ref_aux):
+        np.testing.assert_allclose(
+            np.asarray(got, dtype=np.float32),
+            np.asarray(exp, dtype=np.float32),
+            rtol=tol,
+            atol=tol,
+            err_msg=f"64 V heads chunk_size=128 dtype={act_dtype}: {name}",
+        )
+      for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), got_grads, ref_grads):
+        np.testing.assert_allclose(
+            np.asarray(got, dtype=np.float32),
+            np.asarray(exp, dtype=np.float32),
+            rtol=tol,
+            atol=tol,
+            err_msg=f"64 V heads chunk_size=128 dtype={act_dtype}: {name}",
+        )
+
+  def test_chunk_size_128_head_cp_packed_with_initial_states(self):
+    """Verifies head-sharded CP (gdn_cp_mode='head') at chunk_size=128 with packed segment_ids and initial states."""
+    devices = jax.devices()
+    if len(devices) < 2:
+      self.skipTest("Requires at least 2 devices for head-CP")
+    is_tpu = jax.default_backend() == "tpu"
+    chunk_size = 128
+    dims = _packing_dims()
+    cp_size = 2
+    local_dims = {
+        **dims,
+        "num_k_heads": dims["num_k_heads"] // cp_size,
+        "num_v_heads": dims["num_v_heads"] // cp_size,
+    }
+    batch_size, seq_len = (2 if is_tpu else 1), 2 * chunk_size
+    key = jax.random.PRNGKey(5050)
+    qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims)
+    cs, rs = _packing_states(jax.random.PRNGKey(5051), batch_size, dims)
+    seg = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+    seg = seg.at[0, :10].set(0)
+    seg = seg.at[0, chunk_size + 14 : 2 * chunk_size - 6].set(2)
+    seg = seg.at[0, 2 * chunk_size - 6 :].set(0)
+    if batch_size > 1:
+      seg = seg.at[1, :chunk_size].set(0)
+      seg = seg.at[1, chunk_size + 8 :].set(2)
+
+    q_size = dims["num_k_heads"] * dims["head_k_dim"]
+    k_size = dims["num_k_heads"] * dims["head_k_dim"]
+    v_size = dims["num_v_heads"] * dims["head_v_dim"]
+
+    def split_qkv_heads(x, axis):
+      q, k, v = jnp.split(x, [q_size, q_size + k_size], axis=axis)
+      q_s = q.reshape(*q.shape[:axis], cp_size, q_size // cp_size, *q.shape[axis + 1 :])
+      k_s = k.reshape(*k.shape[:axis], cp_size, k_size // cp_size, *k.shape[axis + 1 :])
+      v_s = v.reshape(*v.shape[:axis], cp_size, v_size // cp_size, *v.shape[axis + 1 :])
+      return jnp.concatenate([q_s, k_s, v_s], axis=axis + 1)
+
+    def merge_qkv_heads(x, axis):
+      lq = q_size // cp_size
+      lk = k_size // cp_size
+      q_s, k_s, v_s = jnp.split(x, [lq, lq + lk], axis=axis + 1)
+      q = q_s.reshape(*q_s.shape[:axis], q_size, *q_s.shape[axis + 2 :])
+      k = k_s.reshape(*k_s.shape[:axis], k_size, *k_s.shape[axis + 2 :])
+      v = v_s.reshape(*v_s.shape[:axis], v_size, *v_s.shape[axis + 2 :])
+      return jnp.concatenate([q, k, v], axis=axis)
+
+    mesh = jax.sharding.Mesh(np.array(devices[:cp_size]), ("context",))
+    ps_h2 = jax.sharding.PartitionSpec(None, None, "context", None)
+    ps_cw = jax.sharding.PartitionSpec(None, None, "context", None)
+    ps_h0 = jax.sharding.PartitionSpec("context", None)
+    ps_h0_1d = jax.sharding.PartitionSpec("context", None)
+    ps_rs = jax.sharding.PartitionSpec(None, "context", None, None)
+    ps_rep = jax.sharding.PartitionSpec()
+
+    def loss_single(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_):
+      o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size)
+      return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs)
+
+    @functools.partial(
+        jax.shard_map,
+        mesh=mesh,
+        in_specs=(ps_h2, ps_h2, ps_h2, ps_cw, ps_h0, ps_h0_1d, ps_h0_1d, ps_h2, ps_rs, ps_rep),
+        out_specs=(ps_h2, ps_h2, ps_rs),
+        check_vma=False,
+    )
+    def run_head_cp(qkv_h, b_h, a_h, cw_h, cb_h, al_h, dt_h, cs_h, rs_h, seg_rep):
+      o_h, (ncs_h, nrs_h) = _call_kernel_api(
+          qkv_h.squeeze(2),
+          b_h.squeeze(2),
+          a_h.squeeze(2),
+          cw_h.squeeze(2),
+          cb_h.squeeze(0),
+          al_h.squeeze(0),
+          dt_h.squeeze(0),
+          cs_h.squeeze(2),
+          rs_h,
+          seg_rep,
+          local_dims,
+          chunk_size,
+          None,
+      )
+      return jnp.expand_dims(o_h, 2), jnp.expand_dims(ncs_h, 2), nrs_h
+
+    def loss_head_cp(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_):
+      nv_local = local_dims["num_v_heads"]
+      o_h, ncs_h, nrs = run_head_cp(
+          split_qkv_heads(qkv_, 2),
+          b_.reshape(batch_size, seq_len, cp_size, nv_local),
+          a_.reshape(batch_size, seq_len, cp_size, nv_local),
+          split_qkv_heads(cw_, 2),
+          split_qkv_heads(cb_, 0),
+          al_.reshape(cp_size, nv_local),
+          dt_.reshape(cp_size, nv_local),
+          split_qkv_heads(cs_, 2),
+          rs_,
+          seg,
+      )
+      o = o_h.reshape(batch_size, seq_len, dims["num_v_heads"], dims["head_v_dim"])
+      ncs = merge_qkv_heads(ncs_h, 2)
+      return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs)
+
+    args = (qkv, b, a, cw, cb, al, dt, cs, rs)
+    vg_s = jax.value_and_grad(loss_single, argnums=tuple(range(9)), has_aux=True)
+    vg_h = jax.value_and_grad(loss_head_cp, argnums=tuple(range(9)), has_aux=True)
+    if is_tpu:
+      vg_s = jax.jit(vg_s)
+      vg_h = jax.jit(vg_h)
+    (_, aux_s), grads_s = vg_s(*args)
+    (_, aux_h), grads_h = vg_h(*args)
+    tol = 2e-2 if is_tpu else 2e-3
+    for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), aux_h, aux_s):
+      np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"head-CP c=128: {name}")
+    for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), grads_h, grads_s):
+      np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"head-CP c=128: {name}")
+
+  def test_chunk_size_128_seq_cp_packed_with_initial_states(self):
+    """Verifies seq-CP (cp_axis_name='context') at chunk_size=128 with on-the-fly t_inv recomputation."""
+    devices = jax.devices()
+    if len(devices) < 2:
+      self.skipTest("Requires at least 2 devices for seq-CP")
+    is_tpu = jax.default_backend() == "tpu"
+    chunk_size = 128
+    dims = _packing_dims()
+    batch_size, seq_len = (2 if is_tpu else 1), 4 * chunk_size  # 2 chunks per rank across 2 CP ranks
+    key = jax.random.PRNGKey(4040)
+    qkv, b, a, cw, cb, al, dt, up = _packing_inputs(key, batch_size, seq_len, dims)
+    cs, rs = _packing_states(jax.random.PRNGKey(4041), batch_size, dims)
+    seg = jnp.ones((batch_size, seq_len), dtype=jnp.int32)
+    seg = seg.at[0, :12].set(0)
+    seg = seg.at[0, chunk_size + 16 : 3 * chunk_size - 8].set(2)
+    seg = seg.at[0, 3 * chunk_size - 8 :].set(0)
+    if batch_size > 1:
+      seg = seg.at[1, : 2 * chunk_size].set(0)
+      seg = seg.at[1, 2 * chunk_size + 10 :].set(2)
+
+    mesh = jax.sharding.Mesh(np.array(devices[:2]), ("context",))
+    ps_seq = jax.sharding.PartitionSpec(None, "context", None)
+    ps_seg = jax.sharding.PartitionSpec(None, "context")
+    ps_rep = jax.sharding.PartitionSpec()
+    ps_up = jax.sharding.PartitionSpec(None, "context", None, None)
+
+    def loss_single(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_):
+      o, (ncs, nrs) = _call_pure_jax(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, dims, chunk_size)
+      return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs)
+
+    @functools.partial(
+        jax.shard_map,
+        mesh=mesh,
+        in_specs=(ps_seq, ps_seq, ps_seq, ps_rep, ps_rep, ps_rep, ps_rep, ps_rep, ps_rep, ps_seg, ps_up),
+        out_specs=(ps_up, ps_rep, ps_rep),
+        check_vma=False,
+    )
+    def run_cp(qkv_s, b_s, a_s, cw_, cb_, al_, dt_, cs_, rs_, seg_s, up_s):
+      del up_s
+      o, (ncs, nrs) = _call_kernel_api(qkv_s, b_s, a_s, cw_, cb_, al_, dt_, cs_, rs_, seg_s, dims, chunk_size, "context")
+      return o, ncs, nrs
+
+    # Also exercise custom remat saving only gdn_cp_state on TPU to verify chunk_states != None and t_inv_fwd == None.
+    run_cp_fn = (
+        jax.checkpoint(run_cp, policy=jax.checkpoint_policies.save_only_these_names("gdn_cp_state")) if is_tpu else run_cp
+    )
+
+    def loss_cp(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_):
+      o, ncs, nrs = run_cp_fn(qkv_, b_, a_, cw_, cb_, al_, dt_, cs_, rs_, seg, up)
+      return jnp.sum(o * up) + 0.13 * jnp.sum(ncs) + 0.19 * jnp.sum(nrs), (o, ncs, nrs)
+
+    args = (qkv, b, a, cw, cb, al, dt, cs, rs)
+    vg_s = jax.value_and_grad(loss_single, argnums=tuple(range(9)), has_aux=True)
+    vg_c = jax.value_and_grad(loss_cp, argnums=tuple(range(9)), has_aux=True)
+    if is_tpu:
+      vg_s = jax.jit(vg_s)
+      vg_c = jax.jit(vg_c)
+    (_, aux_s), grads_s = vg_s(*args)
+    (_, aux_c), grads_c = vg_c(*args)
+    tol = 2e-2 if is_tpu else 2e-3
+    for name, got, exp in zip(("out", "next_conv_state", "next_recurrent_state"), aux_c, aux_s):
+      np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"CP c=128: {name}")
+    for name, got, exp in zip(_GRAD_NAMES + ("d_conv_state", "d_recurrent_state"), grads_c, grads_s):
+      np.testing.assert_allclose(np.asarray(got), np.asarray(exp), rtol=tol, atol=tol, err_msg=f"CP c=128: {name}")
+
 
 def _find_eqns(jaxpr, primitive_name):
   """Returns every `primitive_name` equation in `jaxpr`, including those in nested (tuples of) sub-jaxprs."""
diff --git a/tests/unit/gdn_cp_test.py b/tests/unit/gdn_cp_test.py
index 94394fcbd0..851578d63a 100644
--- a/tests/unit/gdn_cp_test.py
+++ b/tests/unit/gdn_cp_test.py
@@ -372,7 +372,7 @@ def test_cp_gdn_kernel_module_primitives(self):
       np.testing.assert_allclose(m_loc_konly, m_loc, rtol=1e-6, atol=1e-6)
 
       # Independently verify s_ext (initial_state=0) and m_loc @ h0 + s_ext (initial_state=h0)
-      # against pure_jax_decoupled_conv1d_gdn so neither m_loc nor s_ext is compared with itself.
+      # against pure_jax_decoupled_conv1d_gdn so neither m_loc nor s_ext is compared with itself (C6).
       _, (_, s_ext_ref) = gdn_bwd_pallas.pure_jax_decoupled_conv1d_gdn(
           qkv=qkv_raw,
           b=b,
@@ -448,6 +448,240 @@ def test_cp_gdn_kernel_module_primitives(self):
       np.testing.assert_allclose(dm_uncached, dm_cached, rtol=1e-5, atol=1e-5)
       np.testing.assert_allclose(ds_uncached, ds_cached, rtol=1e-5, atol=1e-5)
 
+      # Verify on-the-fly t_inv recomputation (t_inv=None) used at chunk_size=128 matches explicit t_inv.
+      m_loc_no_tinv, s_ext_no_tinv = cp_gdn.compose_local_from_t_inv(
+          qkv_conv=qkv_conv,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          t_inv=None,
+          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=True,
+      )
+      np.testing.assert_allclose(m_loc_no_tinv, m_loc, rtol=1e-5, atol=1e-5)
+      np.testing.assert_allclose(s_ext_no_tinv, s_ext, rtol=1e-5, atol=1e-5)
+
+      dm_no_tinv, ds_no_tinv = 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=do,
+          t_inv=None,
+          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=True,
+      )
+      np.testing.assert_allclose(dm_no_tinv, dm_uncached, rtol=1e-5, atol=1e-5)
+      np.testing.assert_allclose(ds_no_tinv, ds_uncached, rtol=1e-5, atol=1e-5)
+
+      # Verify on-the-fly t_inv recomputation with packed segment_ids (including intra-chunk reset + padding).
+      seg = jnp.ones((batch, seq_len), dtype=jnp.int32)
+      seg = seg.at[0, :3].set(0)
+      seg = seg.at[0, chunk_size + 5 : 2 * chunk_size + 4].set(2)
+      seg = seg.at[0, -4:].set(0)
+      init_seg = jnp.array([1], dtype=jnp.int32)
+      qkv_conv_seg, _, t_inv_seg = gdn_bwd_pallas._compute_forward_conv_and_states(
+          qkv=qkv_raw,
+          b=b,
+          a=a,
+          conv_weight=cw_id,
+          conv_bias=None,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          recurrent_state=None,
+          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=4,
+          chunk_size=chunk_size,
+          use_qk_norm_in_gdn=True,
+          segment_ids=seg,
+          init_seg=init_seg,
+      )
+      m_seg_ref, s_seg_ref = cp_gdn.compose_local_from_t_inv(
+          qkv_conv=qkv_conv_seg,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          t_inv=t_inv_seg,
+          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=True,
+          segment_ids=seg,
+          init_seg=init_seg,
+      )
+      m_seg_no_tinv, s_seg_no_tinv = cp_gdn.compose_local_from_t_inv(
+          qkv_conv=qkv_conv_seg,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          t_inv=None,
+          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=True,
+          segment_ids=seg,
+          init_seg=init_seg,
+      )
+      np.testing.assert_allclose(m_seg_no_tinv, m_seg_ref, rtol=1e-5, atol=1e-5)
+      np.testing.assert_allclose(s_seg_no_tinv, s_seg_ref, rtol=1e-5, atol=1e-5)
+
+      dm_seg_ref, ds_seg_ref = cp_gdn.compose_bwd_local_from_t_inv(
+          qkv_conv=qkv_conv_seg,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          do=do,
+          t_inv=t_inv_seg,
+          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=True,
+          segment_ids=seg,
+          init_seg=init_seg,
+      )
+      dm_seg_no_tinv, ds_seg_no_tinv = cp_gdn.compose_bwd_local_from_t_inv(
+          qkv_conv=qkv_conv_seg,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          do=do,
+          t_inv=None,
+          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=True,
+          segment_ids=seg,
+          init_seg=init_seg,
+      )
+      np.testing.assert_allclose(dm_seg_no_tinv, dm_seg_ref, rtol=1e-5, atol=1e-5)
+      np.testing.assert_allclose(ds_seg_no_tinv, ds_seg_ref, rtol=1e-5, atol=1e-5)
+
+    # 2. Verify chunk_size=128 with t_inv=None and packed segment_ids.
+    batch, chunk_size, num_chunks = 1, 128, 2
+    seq_len = num_chunks * chunk_size
+    num_k_heads, num_v_heads = 2, 4
+    head_k_dim, head_v_dim = 32, 32
+    dim_size = 2 * num_k_heads * head_k_dim + num_v_heads * head_v_dim
+    k0, k1, k2, k3, k4, k5 = jax.random.split(jax.random.PRNGKey(228), 6)
+    qkv_raw = jax.random.normal(k0, (batch, seq_len, dim_size), jnp.float32) * 0.1
+    b = jax.random.normal(k1, (batch, seq_len, num_v_heads), jnp.float32) * 0.2
+    a = jax.random.normal(k2, (batch, seq_len, num_v_heads), jnp.float32) * 0.2
+    a_log = jnp.log(jax.random.uniform(k3, (num_v_heads,), minval=0.1, maxval=1.0))
+    dt_bias = jax.random.normal(k4, (num_v_heads,), jnp.float32) * 0.1
+    do = jax.random.normal(k5, (batch, seq_len, num_v_heads, head_v_dim), jnp.float32) * 0.1
+    seg128 = jnp.ones((batch, seq_len), dtype=jnp.int32).at[0, chunk_size + 12 :].set(2)
+    cw_id = jnp.zeros((4, 1, dim_size), jnp.float32).at[3, 0, :].set(1.0)
+    qkv_conv128, _, t_inv128 = gdn_bwd_pallas._compute_forward_conv_and_states(
+        qkv=qkv_raw,
+        b=b,
+        a=a,
+        conv_weight=cw_id,
+        conv_bias=None,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        recurrent_state=None,
+        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=4,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=True,
+        segment_ids=seg128,
+    )
+    m128_ref, s128_ref = cp_gdn.compose_local_from_t_inv(
+        qkv_conv=qkv_conv128,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        t_inv=t_inv128,
+        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=True,
+        segment_ids=seg128,
+    )
+    m128_got, s128_got = cp_gdn.compose_local_from_t_inv(
+        qkv_conv=qkv_conv128,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        t_inv=None,
+        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=True,
+        segment_ids=seg128,
+    )
+    np.testing.assert_allclose(m128_got, m128_ref, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(s128_got, s128_ref, rtol=1e-5, atol=1e-5)
+
+    dm128_ref, ds128_ref = cp_gdn.compose_bwd_local_from_t_inv(
+        qkv_conv=qkv_conv128,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        do=do,
+        t_inv=t_inv128,
+        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=True,
+        segment_ids=seg128,
+    )
+    dm128_got, ds128_got = cp_gdn.compose_bwd_local_from_t_inv(
+        qkv_conv=qkv_conv128,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        do=do,
+        t_inv=None,
+        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=True,
+        segment_ids=seg128,
+    )
+    np.testing.assert_allclose(dm128_got, dm128_ref, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(ds128_got, ds128_ref, rtol=1e-5, atol=1e-5)
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/unit/pyconfig_test.py b/tests/unit/pyconfig_test.py
index 816a1c8b29..cd5fdb8bd0 100644
--- a/tests/unit/pyconfig_test.py
+++ b/tests/unit/pyconfig_test.py
@@ -308,14 +308,25 @@ def test_gdn_kernel_fast_fail_constraints(self):
     )
     self.assertEqual(cfg_cp4.max_target_length // cfg_cp4.ici_context_parallelism, cfg_cp4.gdn_chunk_size)
 
-    # gdn_chunk_size != 64 or head_dim % 128 != 0 raises ValueError when use_gdn_kernel=True.
-    with self.assertRaisesRegex(ValueError, "requires gdn_chunk_size=64"):
-      self._init_qwen3_next_packing(
-          packing=False,
-          use_gdn_kernel=True,
-          override_model_config=True,
-          gdn_chunk_size=32,
-      )
+    # gdn_chunk_size not in (64, 128) or head_dim % 128 != 0 raises ValueError when use_gdn_kernel=True.
+    for bad_chunk_size in (32, 256):
+      with self.assertRaisesRegex(ValueError, r"requires gdn_chunk_size in \(64, 128\)"):
+        self._init_qwen3_next_packing(
+            packing=False,
+            use_gdn_kernel=True,
+            override_model_config=True,
+            gdn_chunk_size=bad_chunk_size,
+        )
+    cfg_c128 = self._init_qwen3_next_packing(
+        packing=False,
+        use_gdn_kernel=True,
+        override_model_config=True,
+        gdn_chunk_size=128,
+        max_target_length=256,
+        ici_context_parallelism=2,
+        context_parallel_load_balance=False,
+    )
+    self.assertEqual(cfg_c128.gdn_chunk_size, 128)
     with self.assertRaisesRegex(ValueError, "multiples of 128"):
       self._init_qwen3_next_packing(
           packing=False,

From 7f5dcc7cedd06d33b6a68bee14ffecdf21d7471c Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Wed, 14 Oct 2026 23:28:06 -0700
Subject: [PATCH 02/10] [MaxTextEngine] Stop the gradient of frozen parameters
 inside fwd_bwd

fwd_bwd differentiated every parameter, and update() then zeroed the
gradients of the ones trainable_parameters_mask freezes before the norm,
clipping and the optimizer saw them. Every micro-batch computed,
reduce-scattered and accumulated those gradients only to throw them away.
diff_wrapper now applies jax.lax.stop_gradient to the frozen leaves,
inside the differentiated function: JAX emits no weight-gradient op for
them and fwd_bwd returns a constant zero in their place. The gradient
through a frozen weight into the activations still comes from the same
ops, so the trainable gradients, the norm and the update are the same
computation.

Unconditional whenever a mask is set, with no config key. Nothing in the
engine reads a frozen leaf's gradient: the norm, clipping, the NaN and
spike skips and the optimizer all come after update() zeroes it, and no
API returns gradients. update() still zeroes them, for an accumulator
restored from an intra-step checkpoint written by an older engine or
under another mask. The leaves keep their slots in the gradient tree, as
zeros, so kernel signatures, the accumulator and checkpoints are
unchanged. pre_train/train.py is untouched: it clips over the whole
gradient tree, frozen leaves included, so the same stop there would
change its updates (README, "Differences from pre_train/train.py").

Both kernels read the mask through one helper, _freeze_mask_for, so the
leaves fwd_bwd stops are exactly the leaves update() zeroes. It reads the
compile-time cache once and, when it is empty, computes the mask from the
tree, as _update_kernel already did: on the eager path, and for a
dropless or final kernel traced on a background thread after
_disable_pure_state emptied the cache, where checking the cache alone
would silently skip the stop.

The MLPerf RL recipes freeze the MoE router
(^(?!.*routed_experts/gate/kernel).*). Measured on Qwen3.5-397B-A17B,
128 TPU v7x, on the RL trainer benchmark (#5701 driver, mesh fsdp 32 x
context 4 x expert 2, 9 micro-batches per step), with the same stop
behind an environment switch:
- AOT with the TPU libtpu: the router's f32[4096,512] weight-gradient
  fusions and f32[16,512] reduce-scatters are gone from fwd_bwd (8 -> 0
  each in the program, one per MoE layer when run).
- Regular micro-batch 10.189 -> 10.131 s, both with
  moe_routing_maps=device (o275x-stk1a-a vs o275x-stk2, two jobs on one
  booking): -58 ms against a predicted -45 ms (-27 to -65).
- Numerics A/B (nxstk2 vs nxref, nxref2): 5 of 5 gates pass; loss
  4.6e-8 relative, grad_norm 4.8e-5 relative against a replicate spread
  of 1.9e-5.

Tests (maxtext_engine_frozen_grad_test.py, new), against the same engine
with the stop removed: on a tiny decoder with one MLP kernel frozen,
trained by AdamW with clipping on four CPU devices, fwd_bwd lowers to one
dot_general fewer, shaped like the kernel, and none added; it returns
exactly zero for the kernel; and every other gradient, the running sums,
the gradient norms and the trained weights are bitwise equal, under fsdp,
under the deferred all-reduce and on the eager path. On the one-layer
Qwen3.5 MoE rig, with the recipes' router mask and Tunix's router replay
payload and loss, the same holds. A stop that also hits one trainable
kernel is caught, and the failure names it. README: the config key, the
test and the difference from pre_train/train.py.
---
 src/maxtext/training_engine/README.md         |   6 +
 src/maxtext/training_engine/maxtext_engine.py |  56 ++-
 .../unit/maxtext_engine_frozen_grad_test.py   | 402 ++++++++++++++++++
 3 files changed, 460 insertions(+), 4 deletions(-)
 create mode 100644 tests/post_training/unit/maxtext_engine_frozen_grad_test.py

diff --git a/src/maxtext/training_engine/README.md b/src/maxtext/training_engine/README.md
index 58b2797dff..d3da9b81a9 100644
--- a/src/maxtext/training_engine/README.md
+++ b/src/maxtext/training_engine/README.md
@@ -344,6 +344,7 @@ sharding at construction (for example Tokamax ring attention), while under `jax.
 | `parameter_memory_host_offload` | not supported; raises |
 | `shard_optimizer_over_data` (Zero-1) | shards the optimizer state over the `data` axis inside `update`; requires `shard_mode=explicit` |
 | `skip_step_on_nan` (default true), `skip_step_on_spikes`, `max_grad_norm_spike` | a non-finite or spiking step leaves the state unchanged and records `step_skipped` |
+| `trainable_parameters_mask` | parameters it does not match are frozen. `fwd_bwd` stops their gradient inside the function it differentiates, so it returns zeros for them without computing a weight gradient, or reducing one inside `fwd_bwd`. The zeros keep their slots in the accumulator (and in `update`'s all-reduce, when that is deferred), so kernel signatures and checkpoints are unchanged. `update` zeroes them again before the gradient norm, clipping and the optimizer, so the norm and clipping cover the trainable parameters only. On CPU the trainable gradients, the norm and the trained weights are bitwise those of computing the frozen gradients and discarding them (`maxtext_engine_frozen_grad_test.py`) |
 | `retry_when_tokens_dropped`, `retry_num_moe_token_chunks` | a micro-batch whose MoE ragged buffer overflowed is accumulated as zero and rerun through `fwd_bwd_dropless`, built on a copy of the model with dropless MoE layers (that many token chunks, as a rematerialized scan) and full remat; compiled on a background thread from the first micro-batch. Its flag is read after the next micro-batch is dispatched, or in `update` for the last; `update` records `moe_dropless_replays`. `eval_step` and `model_scope`/`fwd_only` run the dropless graph too |
 | `enable_checkpointing` | `false` disables saving and restoring the train state; `load_parameters_path` still loads weights |
 | `gradient_accumulation_steps` | read only by the standalone script; Tunix chooses the number of micro-batches itself |
@@ -391,6 +392,7 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi
 | `maxtext_engine_model_build_test.py` | model construction with and without a mesh set by the caller |
 | `maxtext_engine_checkpoint_test.py` | resuming from a mid-step checkpoint through Orbax reproduces the uninterrupted step |
 | `maxtext_engine_grad_cast_test.py` | `cast_grads_after_all_reduce`: which gradients leave `fwd_bwd` uncast, and that moving their cast changes no number on CPU |
+| `maxtext_engine_frozen_grad_test.py` | `trainable_parameters_mask`: `fwd_bwd` has no weight-gradient matmul for a frozen kernel and returns zero for it, while the other gradients, the gradient norms and the trained weights stay bitwise those of computing it, under fsdp, under the deferred all-reduce and on the eager path; a stop on a trainable kernel is caught |
 | `maxtext_engine_prepared_input{,_dropless,_sharded}_test.py` | `prepare_fwd_bwd_input` and `wait_for_device`: a prepared micro-batch trains bit-identically on one device, through Tunix's `TrainerWorker` staging, under the dropless replay and final dropless programs, and on a four-device mesh with Tunix's payload; the wait leaves the in-flight queue and the metrics writes as they were |
 | `maxtext_engine_stage_before_update_test.py` | the order Tunix uses to stage before a deferred update (`fwd_bwd(last)`, `prepare_fwd_bwd_input(next)`, `wait_for_device()`, `update()`) under the dropless replay and final dropless programs, directly and through Tunix's `TrainerWorker`: bit-identical to the base order, including a rerun of the last micro-batch inside `update()` |
 | `maxtext_engine_{constructor,data_parallel,xaot,packing,profiling}_test.py`, `router_replay_engine_test.py` | construction, data parallelism, ahead-of-time compilation, sequence packing, profiling, router replay |
@@ -404,3 +406,7 @@ The engine's tests are in `tests/post_training/unit/` and run on CPU; they requi
   of 0, are unaffected.
 - `pre_train/train.py` sums micro-batch gradients in the dtype of the parameters it differentiates;
   the engine sums them in `grad_accumulation_dtype`, which defaults to `grad_dtype`.
+- With `trainable_parameters_mask`, `pre_train/train.py` computes the frozen parameters' gradients and
+  includes them in the gradient norm and clipping; only its optimizer zeroes their updates. The engine
+  leaves them out of the norm and clipping and does not compute them, which is why stopping their
+  gradient changes no number there.
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 0f1d0ddc1b..bee659577c 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -1530,6 +1530,38 @@ def _to_grad_dtype(self, x: Any) -> Any:
       return x
     return x.astype(self._config.grad_dtype)
 
+  def _freeze_mask_for(self, tree: Any) -> Any:
+    """Returns which leaves of `tree` `trainable_parameters_mask` freezes, or None if it freezes none.
+
+    A pytree of Python bools with `tree`'s structure, read off the leaves' paths alone, so `tree` can
+    be the parameters or a gradient tree shaped like them. `_compile_for_batch` caches the parameters'
+    mask in `_freeze_mask`; while that cache is empty -- on the eager path, or for a kernel traced
+    after `_invalidate_pure_state` -- the mask is computed from `tree` instead. `_fwd_bwd_kernel` and
+    `_update_kernel` both read it here, so the leaves whose gradient `fwd_bwd` stops are exactly the
+    leaves `update` zeroes.
+    """
+    # Read once: a background compile can trace a kernel while the main thread empties the cache
+    # (`_disable_pure_state`), and a second read could then return None and skip the stop.
+    freeze_mask = self._freeze_mask
+    if freeze_mask is not None:
+      return freeze_mask
+    if self._freeze_mask_fn is None:
+      return None
+    return self._freeze_mask_fn(tree)
+
+  def _stop_frozen_gradients(self, params: Any) -> Any:
+    """Returns `params` with `jax.lax.stop_gradient` on every leaf `trainable_parameters_mask` freezes.
+
+    Takes effect only inside the function being differentiated: `_fwd_bwd_kernel`'s `diff_wrapper`
+    calls it on the parameters it differentiates. See the comment there for why it changes no number.
+    The leaves stay in the gradient tree, as zeros, so every kernel's signature, the accumulator's
+    layout and the checkpoint format are unchanged.
+    """
+    freeze_mask = self._freeze_mask_for(params)
+    if freeze_mask is None:
+      return params
+    return jax.tree.map(lambda x, frozen: jax.lax.stop_gradient(x) if frozen else x, params, freeze_mask)
+
   def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None):
     """Executes a single forward and backward pass.
 
@@ -1543,7 +1575,8 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None):
     Returns:
       `(primary_loss, aux_metrics, new_rest, grads, denominator)`: this micro-batch's gradients,
       not yet normalized, in the accumulation dtype except the leaves in `_uncast_grad_mask`, which
-      keep the parameters' dtype; and the loss denominator to normalize them by. Under
+      keep the parameters' dtype, and exactly zero for the parameters `trainable_parameters_mask`
+      freezes; 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.
@@ -1551,6 +1584,19 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None):
     loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn
 
     def diff_wrapper(p, r, b):
+      # No gradient for frozen parameters. `update()` zeroes the gradient of every leaf that
+      # `trainable_parameters_mask` freezes before the norm, clipping and the optimizer see it, so
+      # nothing uses that gradient; stopped here, inside the differentiated function, it is not computed
+      # either (on `value_and_grad`'s arguments the stop would change nothing). JAX emits no
+      # weight-gradient op for those leaves -- a custom-VJP kernel still runs its whole backward, and
+      # XLA drops the cotangent nobody reads where it can -- and `fwd_bwd` returns a constant zero for
+      # them, with nothing to reduce-scatter or all-reduce inside it. The gradient through a frozen
+      # weight into the activations, and on to the trainable parameters, comes from the same ops as
+      # before, so the rest is the same computation: bitwise equal on CPU
+      # (`maxtext_engine_frozen_grad_test.py`), though on TPU XLA may fuse the smaller program
+      # differently. The MLPerf RL recipes freeze the MoE router (`routed_experts/gate/kernel`), so this
+      # drops the router's weight gradient and its reduce-scatter from every MoE layer.
+      p = self._stop_frozen_gradients(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
@@ -1789,10 +1835,12 @@ def normalize(g):
       # Cast to `grad_dtype` once, after the sum and the division, as `train.py` does with its
       # accumulated `raw_grads`. A no-op unless `grad_accumulation_dtype` differs from it.
       grads = jax.tree.map(self._to_grad_dtype, grads)
-      freeze_mask = self._freeze_mask
-      if freeze_mask is None and self._freeze_mask_fn is not None:
-        freeze_mask = self._freeze_mask_fn(accumulated_grads)
+      freeze_mask = self._freeze_mask_for(accumulated_grads)
       if freeze_mask is not None:
+        # `fwd_bwd` already returns zeros for these leaves (`_stop_frozen_gradients`). Zeroing them again
+        # costs nothing and still matters: an accumulator restored from an intra-step checkpoint written
+        # by an older engine, or under another mask, can hold non-zero gradients for them, and it keeps
+        # frozen parameters out of the norm and clipping below by construction.
 
         def _apply_freeze(g, is_frozen):
           return jnp.zeros_like(g) if is_frozen else g
diff --git a/tests/post_training/unit/maxtext_engine_frozen_grad_test.py b/tests/post_training/unit/maxtext_engine_frozen_grad_test.py
new file mode 100644
index 0000000000..5ede9104bb
--- /dev/null
+++ b/tests/post_training/unit/maxtext_engine_frozen_grad_test.py
@@ -0,0 +1,402 @@
+# 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.
+
+"""`trainable_parameters_mask`: `fwd_bwd` computes no gradient for a frozen parameter, and no number changes.
+
+`update()` zeroes the gradient of every parameter the mask freezes before the norm, clipping and the
+optimizer see it. `fwd_bwd` also stops that gradient inside the function it differentiates
+(`_stop_frozen_gradients`), so it is never computed. Each test compares the engine with the same engine
+whose `_stop_frozen_gradients` is the identity: the program that computes the frozen gradient and zeroes
+it in `update()` alone. Both rigs clip every update, so a change to any gradient or to the norm would
+reach the trained weights.
+
+On a tiny real decoder with one MLP kernel frozen, trained by AdamW:
+- the lowered `fwd_bwd` has one `dot_general` fewer, shaped like the frozen kernel, and none added;
+- `fwd_bwd` returns exactly zero for the frozen kernel, where the reference computes a non-zero gradient;
+- every other gradient, the running sums, each update's gradient norm and the trained weights are
+  bitwise equal: under fsdp, under the deferred data-parallel all-reduce, and on the eager path, which
+  does not cache the mask;
+- control: a stop that also hits one trainable kernel is caught, and the comparison names that kernel.
+These run in a subprocess with four CPU devices, for the reason given in
+`maxtext_engine_data_parallel_test.py`: a sharded mesh needs more than one device, and the device count is
+fixed before this file is imported.
+
+On the one-layer Qwen3.5 MoE rig of the dropless tests, with the MLPerf RL recipes' own mask, which freezes
+the router: the same comparison through Tunix's router replay payload and loss. It runs on one device, in
+the pytest process, as that rig's own tests do.
+"""
+
+import collections
+import contextlib
+import functools
+import os
+import re
+import subprocess
+import sys
+import tempfile
+from typing import Any
+import unittest
+from unittest import mock
+
+from absl.testing import absltest
+from absl.testing import parameterized
+from flax import nnx
+import jax
+from maxtext.configs import pyconfig
+from maxtext.models import models
+from maxtext.training_engine import maxtext_engine
+from maxtext.utils import maxtext_utils
+import numpy as np
+import pytest
+
+from tests.post_training.unit.maxtext_engine_dropless_replay_test import _cfg as _moe_rig_cfg
+from tests.post_training.unit.maxtext_engine_dropless_replay_test import _payload as _moe_rig_payload
+from tests.utils.test_helpers import get_test_config_path
+
+# training_engine imports tunix, so these tests need the post-training dependency bundle.
+pytestmark = [pytest.mark.post_training, pytest.mark.cpu_only]
+
+_REQUIRED_DEVICES = 4
+_SENTINEL = "MAXTEXT_ENGINE_FROZEN_GRAD_TESTS_PASSED"
+_RAN = re.compile(rf"{_SENTINEL} ran=(\d+)")
+
+# The kernel the mask freezes, and the trainable kernel the control stops by mistake. Paths are written
+# the way `trainable_parameters_mask` matches them (`optimizers.get_path_mask_fn`).
+_FROZEN = "layers_0/mlp/wi_0/kernel"
+_WRONGLY_STOPPED = "layers_1/mlp/wo/kernel"
+# The mask lists what trains, so freezing one kernel takes a negative lookahead.
+_MASK = f"^(?!.*{_FROZEN}).*"
+# The MLPerf RL recipes' mask, verbatim (`TRAINABLE_PARAMETERS_MASK` in Tunix's `mlperf_base.sh`).
+_ROUTER = "routed_experts/gate/kernel"
+_ROUTER_MASK = "^(?!.*routed_experts/gate/kernel).*"
+# Below every update's gradient norm (checked), so clipping scales each update by the norm and a change
+# to the norm would change the trained weights.
+_CLIP = 0.5
+
+# fsdp is the kind of mesh the RL recipes train on. The deferred all-reduce returns the gradients
+# `unreduced`, so a stopped leaf's zero has to come out with that tag too.
+_FSDP = {"ici_fsdp_parallelism": _REQUIRED_DEVICES, "ici_data_parallelism": 1}
+_DEFERRED = {"ici_fsdp_parallelism": 1, "ici_data_parallelism": _REQUIRED_DEVICES, "shard_mode": "explicit"}
+
+# Every engine's `base_output_directory`, removed when the process exits.
+_OUTPUT_DIR = tempfile.TemporaryDirectory(prefix="engine_frozen_grad_")  # pylint: disable=consider-using-with
+
+
+@pytest.mark.post_training
+@pytest.mark.cpu_only
+def test_frozen_gradients_on_a_four_device_cpu_mesh():
+  """Runs `FrozenGradientTest` in a child process with four CPU devices."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = f"{env.get('XLA_FLAGS', '')} --xla_force_host_platform_device_count={_REQUIRED_DEVICES}".strip()
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+
+  report = f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert result.returncode == 0, report
+  ran = _RAN.search(result.stdout)
+  assert ran, f"the child did not report a completed run\n{report}"
+  assert int(ran.group(1)) > 0, f"every test in the child skipped\n{report}"
+
+
+def _config(**overrides) -> pyconfig.HyperParameters:
+  """A tiny real decoder, float32 throughout, trained by AdamW with clipping, with `_FROZEN` frozen."""
+  fields = {
+      "model_name": "default",
+      "run_name": "engine_frozen_grad_test",
+      "base_output_directory": _OUTPUT_DIR.name,
+      "enable_checkpointing": False,
+      "convert_checkpoint_if_possible": False,
+      "skip_jax_distributed_system": True,
+      "enable_tensorboard": False,
+      "record_internal_nn_metrics": False,
+      "enable_dropout": False,
+      "init_weights_seed": 0,
+      "dtype": "float32",
+      "weight_dtype": "float32",
+      "grad_dtype": "float32",
+      "remat_policy": "none",
+      "scan_layers": False,
+      "attention": "dot_product",
+      "ici_tensor_parallelism": 1,
+      "per_device_batch_size": 1,
+      "vocab_size": 128,
+      "base_emb_dim": 64,
+      "base_mlp_dim": 128,
+      "base_num_decoder_layers": 2,
+      "base_num_query_heads": 4,
+      "base_num_kv_heads": 4,
+      "head_dim": 16,
+      "max_target_length": 32,
+      "opt_type": "adamw",
+      "learning_rate": 0.01,
+      "gradient_clipping_threshold": _CLIP,
+      "warmup_steps_fraction": 0.0,
+      "learning_rate_final_fraction": 1.0,
+      "profiler_steps": 0,
+      "trainable_parameters_mask": _MASK,
+      **_FSDP,
+  }
+  fields.update(overrides)
+  argv = ["maxtext_engine_frozen_grad_test.py", get_test_config_path("base.yml")]
+  return pyconfig.initialize(argv + [f"{k}={v}" for k, v in fields.items()])
+
+
+def _batch(cfg: pyconfig.HyperParameters, seed: int) -> dict[str, np.ndarray]:
+  """A `train.py` loss batch as host arrays, which the engine places itself."""
+  rows, seq = int(cfg.micro_batch_size_to_train_on), cfg.max_target_length
+  tokens = np.random.default_rng(seed).integers(1, cfg.vocab_size, size=(rows, seq)).astype(np.int32)
+  ones = np.ones((rows, seq), dtype=np.int32)
+  return {
+      "inputs": tokens,
+      "targets": np.roll(tokens, -1, axis=-1),
+      "inputs_position": np.tile(np.arange(seq, dtype=np.int32), (rows, 1)),
+      "inputs_segmentation": ones,
+      "targets_segmentation": ones,
+  }
+
+
+def _path(path) -> str:
+  """A leaf's path the way `trainable_parameters_mask` matches it, e.g. `decoder/layers_0/mlp/wi_0/kernel/value`."""
+  return jax.tree_util.keystr(path, simple=True, separator="/")
+
+
+def _bits(tree) -> dict[str, np.ndarray]:
+  """Every leaf's raw bits by path, so comparing two trees is bitwise."""
+  out = {}
+  for path, leaf in jax.tree_util.tree_leaves_with_path(tree):
+    value = np.asarray(jax.device_get(leaf))
+    out[_path(path)] = value.view(np.uint16 if value.dtype.itemsize == 2 else np.uint32)
+  return out
+
+
+def _no_stop(engine, params):
+  """The reference `_stop_frozen_gradients`: stops nothing, so `fwd_bwd` computes the frozen gradient too."""
+  del engine
+  return params
+
+
+_REAL_STOP = maxtext_engine.MaxTextTrainingEngine._stop_frozen_gradients  # pylint: disable=protected-access
+
+
+def _also_stopping(pattern: str):
+  """A deliberately wrong `_stop_frozen_gradients`: the real one, plus a stop on every leaf whose path has `pattern`."""
+
+  def stop(engine, params):
+    return jax.tree_util.tree_map_with_path(
+        lambda path, x: jax.lax.stop_gradient(x) if pattern in _path(path) else x, _REAL_STOP(engine, params)
+    )
+
+  return stop
+
+
+def _with_stop(stop):
+  """Replaces `MaxTextTrainingEngine._stop_frozen_gradients` by `stop` while active; None leaves the engine as it is."""
+  if stop is None:
+    return contextlib.nullcontext()
+  return mock.patch.object(maxtext_engine.MaxTextTrainingEngine, "_stop_frozen_gradients", stop)
+
+
+def _run(engine: maxtext_engine.MaxTextTrainingEngine, steps: list[list[Any]]) -> dict[str, Any]:
+  """Trains one update per entry of `steps`, over that entry's micro-batches, and closes the engine.
+
+  Returns, as raw bits by leaf path: `first`, the first micro-batch's gradients as `fwd_bwd` returned
+  them; `sums`, the running sum before each update; `norms`, each update's gradient norm; `weights`, the
+  trained parameters; `initial`, the parameters before training. Also `frozen`, the paths the mask froze;
+  `clip`, the clipping threshold; `deferred`, whether the all-reduce was deferred; and `mask_cached`,
+  whether the engine had cached the freeze mask.
+  """
+  params = nnx.state(engine.state.model, nnx.Param)
+  initial = _bits(params)
+  mask = engine._freeze_mask_for(params)  # pylint: disable=protected-access
+  frozen = sorted(_path(path) for path, is_frozen in jax.tree_util.tree_leaves_with_path(mask) if is_frozen)
+  first, sums, norms = None, [], []
+  for micro_batches in steps:
+    for micro_batch in micro_batches:
+      engine.fwd_bwd(micro_batch)
+      if first is None:
+        first = _bits(engine._reduced_accumulated_grads())  # pylint: disable=protected-access
+    sums.append(_bits(engine._reduced_accumulated_grads()))  # pylint: disable=protected-access
+    engine.update()
+    norms.append(np.asarray(engine.get_metrics(clear_cache=True).scalar_metrics["gradient_norm"]).reshape(-1))
+  run = {
+      "first": first,
+      "sums": sums,
+      "norms": norms,
+      "weights": _bits(nnx.state(engine.state.model, nnx.Param)),
+      "initial": initial,
+      "frozen": frozen,
+      "clip": engine._config.gradient_clipping_threshold,  # pylint: disable=protected-access
+      "deferred": engine._unreduced_grad_shardings is not None,  # pylint: disable=protected-access
+      "mask_cached": engine._freeze_mask is not None,  # pylint: disable=protected-access
+  }
+  engine.close()
+  return run
+
+
+def _train(stop=None, *, eager: bool = False, steps: int = 2, micro_batches: int = 2, **overrides) -> dict[str, Any]:
+  """`_run` on the tiny decoder, with `_stop_frozen_gradients` replaced by `stop` for the whole run, tracing included."""
+  with _with_stop(stop):
+    cfg = _config(**overrides)
+    engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=maxtext_utils.get_mesh_from_config(cfg))
+    if not eager:
+      engine.compile(_batch(cfg, 0))
+    seeds = np.arange(steps * micro_batches).reshape(steps, micro_batches) + 1
+    return _run(engine, [[_batch(cfg, int(seed)) for seed in row] for row in seeds])
+
+
+@functools.cache
+def _reference(eager: bool = False, steps: int = 2, mesh: str = "fsdp") -> dict[str, Any]:
+  """`_train` with nothing stopped, the program before `_stop_frozen_gradients`; cached, as several tests share it."""
+  return _train(_no_stop, eager=eager, steps=steps, **(_DEFERRED if mesh == "deferred" else _FSDP))
+
+
+def _fwd_bwd_dots(stop=None) -> collections.Counter[str]:
+  """The `dot_general` ops in `fwd_bwd` as lowered, before XLA could remove any, counted by result type."""
+  with _with_stop(stop):
+    cfg = _config()
+    engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=maxtext_utils.get_mesh_from_config(cfg))
+    text = engine._lower_kernels(_batch(cfg, 0))["fwd_bwd"].as_text()  # pylint: disable=protected-access
+    engine.close()
+  dots = [line for line in text.splitlines() if "stablehlo.dot_general" in line]
+  result_types = collections.Counter(re.search(r"->\s*(tensor<[^>]*>)", line).group(1) for line in dots)
+  assert sum(result_types.values()) == len(dots), "a dot_general line without a result type"
+  return result_types
+
+
+def _train_moe_rig(stop=None) -> dict[str, Any]:
+  """`_run` on the dropless tests' one-layer Qwen3.5 MoE rig, with the recipes' router mask.
+
+  Through Tunix's router replay payload and loss, and `compile(None)` as Tunix's `TrainerWorker` calls
+  it, so the kernels compile on the first `fwd_bwd`. Two updates of two micro-batches.
+  """
+  with _with_stop(stop):
+    cfg = _moe_rig_cfg(
+        trainable_parameters_mask=_ROUTER_MASK,
+        gradient_clipping_threshold=_CLIP,
+        base_output_directory=_OUTPUT_DIR.name,
+    )
+    mesh = jax.sharding.Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+    model = models.Transformer(config=cfg, mesh=mesh, quant=None, model_mode="train", rngs=nnx.Rngs(0))
+    with mock.patch.object(maxtext_engine.model_creation_utils, "from_pretrained", return_value=model):
+      engine = maxtext_engine.MaxTextTrainingEngine(cfg, mesh=mesh)
+    engine.with_gen_model_input_fn(maxtext_engine.router_replay_gen_model_input_fn)
+    engine.with_loss_fn(maxtext_engine.make_router_replay_loss_fn(cfg), has_aux=True)
+    engine.compile(None)
+    return _run(engine, [[_moe_rig_payload(10), _moe_rig_payload(20)], [_moe_rig_payload(30), _moe_rig_payload(40)]])
+
+
+class _FrozenGradientChecks:
+  """The comparison every test here makes, for a `unittest.TestCase`."""
+
+  def _differing(self, got: dict[str, np.ndarray], want: dict[str, np.ndarray]) -> list[str]:
+    """The paths whose bits differ between two `_bits` dicts over the same leaves."""
+    self.assertEqual(set(got), set(want))
+    return sorted(path for path in want if not np.array_equal(got[path], want[path]))
+
+  def _assert_frozen_gradient_skipped_and_nothing_else_changed(
+      self, run: dict[str, Any], reference: dict[str, Any], frozen_path: str
+  ):
+    """`run` returned zero for the gradient of the leaf at `frozen_path`, and matches `reference` bitwise elsewhere."""
+    # The mask froze exactly that one leaf, the same in both runs.
+    self.assertLen(run["frozen"], 1)
+    (frozen,) = run["frozen"]
+    self.assertIn(frozen_path, frozen)
+    self.assertEqual(reference["frozen"], run["frozen"])
+    # Not vacuous: the reference computes a non-zero gradient for the frozen leaf ...
+    self.assertTrue(np.any(reference["first"][frozen]), "the frozen leaf has no gradient, so this checks nothing")
+    # ... where the engine returns exactly zero (every bit 0, i.e. +0.0).
+    self.assertFalse(np.any(run["first"][frozen]), "fwd_bwd computed a gradient for the frozen leaf")
+    # Every other gradient is bitwise unchanged, in the first micro-batch and in the running sums.
+    self.assertEqual(self._differing(run["first"], reference["first"]), [frozen], "first micro-batch's gradients")
+    for step, (got, want) in enumerate(zip(run["sums"], reference["sums"], strict=True)):
+      self.assertEqual(self._differing(got, want), [frozen], f"running sum before update {step}")
+    # So is everything the update computes from them: the norm, which clips every update, ...
+    for step, (got, want) in enumerate(zip(run["norms"], reference["norms"], strict=True)):
+      self.assertGreater(float(want[0]), run["clip"], f"update {step} was not clipped, so its norm reached no weight")
+      np.testing.assert_array_equal(got.view(np.uint32), want.view(np.uint32), err_msg=f"gradient norm, update {step}")
+    # ... and every trained weight. The frozen leaf did not move.
+    self.assertEqual(self._differing(run["weights"], reference["weights"]), [], "trained weights")
+    np.testing.assert_array_equal(run["weights"][frozen], run["initial"][frozen], err_msg="the frozen leaf moved")
+
+
+class FrozenGradientTest(_FrozenGradientChecks, parameterized.TestCase):
+  """The tiny decoder with one kernel frozen, against the same engine computing its gradient."""
+
+  __test__ = False  # collected only via the subprocess entry point at the top of this file.
+
+  def test_fwd_bwd_drops_the_frozen_kernels_weight_gradient(self):
+    """One `dot_general` fewer, shaped like the frozen kernel, and none added: its weight gradient is never computed.
+
+    Its gradient through to the activations is still computed, since layer 0's input depends on trainable
+    parameters, so of the kernel's three matmuls (forward, input gradient, weight gradient) one goes.
+    """
+    reference, engine = _fwd_bwd_dots(_no_stop), _fwd_bwd_dots()
+    removed, added = reference - engine, engine - reference
+    self.assertEqual(dict(added), {}, f"dot_general ops the stop added, by result type; reference {dict(reference)}")
+    self.assertEqual(sum(removed.values()), 1, f"dot_general ops the stop removed, by result type: {dict(removed)}")
+    # The [emb, mlp] kernel's weight gradient, in either layout (JAX's transpose rule may emit it [mlp, emb]).
+    cfg = _config()
+    kernel_shaped = {f"tensor<{cfg.emb_dim}x{cfg.mlp_dim}xf32>", f"tensor<{cfg.mlp_dim}x{cfg.emb_dim}xf32>"}
+    self.assertIn(next(iter(removed)), kernel_shaped, "the dot_general the stop removed is not shaped like the kernel")
+
+  @parameterized.named_parameters(("fsdp", "fsdp"), ("deferred_all_reduce", "deferred"))
+  def test_frozen_gradient_is_skipped_and_nothing_else_changes(self, mesh):
+    run = _train(**(_DEFERRED if mesh == "deferred" else _FSDP))
+    self.assertTrue(run["mask_cached"], "the compiled path did not cache the mask")
+    self.assertEqual(run["deferred"], mesh == "deferred", "the deferral is not what this case is about")
+    self._assert_frozen_gradient_skipped_and_nothing_else_changed(run, _reference(mesh=mesh), _FROZEN)
+
+  def test_eager_path_computes_the_mask_and_stops_the_same_gradient(self):
+    """Without `compile()` the mask is never cached; `_freeze_mask_for` computes it from the parameters."""
+    run = _train(eager=True, steps=1)
+    self.assertFalse(run["mask_cached"], "the mask was cached, so the eager fallback did not run")
+    self._assert_frozen_gradient_skipped_and_nothing_else_changed(run, _reference(eager=True, steps=1), _FROZEN)
+
+  def test_control_a_stop_on_a_trainable_kernel_is_caught(self):
+    """Control: the same comparison, against a stop that also hits one trainable kernel, fails and names it."""
+    reference = _reference(mesh="fsdp")
+    wrong = _train(_also_stopping(_WRONGLY_STOPPED))
+    (frozen,) = wrong["frozen"]
+    (wrongly_stopped,) = [path for path in reference["first"] if _WRONGLY_STOPPED in path]
+    # The wrongly stopped kernel's gradient is zero, and it is the only other one that changed ...
+    self.assertFalse(np.any(wrong["first"][wrongly_stopped]))
+    self.assertEqual(self._differing(wrong["first"], reference["first"]), sorted([frozen, wrongly_stopped]))
+    # ... so the norm differs, and with it every update.
+    self.assertFalse(np.array_equal(wrong["norms"][0], reference["norms"][0]))
+    with self.assertRaisesRegex(AssertionError, re.escape(wrongly_stopped)):
+      self._assert_frozen_gradient_skipped_and_nothing_else_changed(wrong, reference, _FROZEN)
+
+
+class FrozenRouterTest(_FrozenGradientChecks, absltest.TestCase):
+  """The MLPerf RL recipes' router freeze on a one-layer Qwen3.5 MoE, against the same engine computing it."""
+
+  def test_the_router_gradient_is_skipped_and_nothing_else_changes(self):
+    run, reference = _train_moe_rig(), _train_moe_rig(_no_stop)
+    self.assertTrue(run["mask_cached"], "the compiled path did not cache the mask")
+    self._assert_frozen_gradient_skipped_and_nothing_else_changed(run, reference, _ROUTER)
+
+
+if __name__ == "__main__":
+  if jax.device_count() < _REQUIRED_DEVICES:
+    raise SystemExit(
+        f"needs {_REQUIRED_DEVICES} devices, got {jax.device_count()}; run this through pytest, which sets "
+        f"XLA_FLAGS=--xla_force_host_platform_device_count={_REQUIRED_DEVICES}"
+    )
+  _result = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(FrozenGradientTest))
+  if not _result.wasSuccessful():
+    sys.exit(1)
+  print(f"{_SENTINEL} ran={_result.testsRun - len(_result.skipped)}")

From add4a1e7457792cab0a52717b98af34dfd7d8544 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 02:08:04 -0700
Subject: [PATCH 03/10] MoE: compute the ring-of-experts routing index maps
 without sorts (moe_routing_index_math)

The ring-of-experts ragged sort derives its routing index maps from the
top-k expert ids with sorts: a stable argsort for the expert permutation,
a gather of a repeated iota for the sorted token indices, a second argsort
for the inverse permutation, a stable argsort in the SparseCore
gather-reduce's validity preprocessing, and, for tokens without forced ids
(every token under moe_exclude_padding_from_routing), jax.lax.top_k on the
router probabilities. XLA lowers each stable argsort in the forward pass
with an extra, duplicated iota operand.

moe_routing_index_math="count" (default "sort", unchanged) computes the
same maps with less work:
- the permutation by one single-operand sort of the packed uint32 key
  (expert << pos_bits | slot), whose values are distinct, so the unstable
  sort has exactly the stable order (packed_stable_sort);
- the sorted token indices as slot // topk, without a gather;
- the inverse permutation by counting, on this shard's slots only, -1
  elsewhere (ragged_sort._local_slot_positions): a slot's rank is the
  slots of lower local experts, plus the slots of its expert in earlier
  tokens (a blocked prefix count as a triangular bf16 matmul, exact), plus
  earlier slots of the same token with the same expert. Every reader of
  the inverse masks it to this shard's window, so nothing reads the -1s,
  and the unsort backward gets the forward permutation instead of
  argsorting the partial inverse back (it raises without it);
- the router-weight gradient back to slot order by a masked gather instead
  of a TensorCore scatter (the table is padded to the slot count so XLA
  offloads the gather to the SparseCore);
- the gather-reduce validity partition as one packed uint32 sort
  (_preprocess(packed=True));
- the top-k indices of non-forced tokens by k argmax passes
  (topk_indices_by_argmax), which break ties the way top_k does. That
  these are cheaper than top_k's sort is expected, not measured: they were
  timed only as part of this key.

Exact: the same permutation everywhere, so the outputs and gradients are
the same bits. The validator rejects "count" where no ring-of-experts
ragged sort runs (it needs sparse_matmul, te_moe_block=False,
use_ring_of_experts and use_ragged_sort). It starts the validator form the
later keys of this series share: _reject_unless names every unmet
requirement, and the fields of the other config mixins are read directly
rather than through getattr with a default. With the key off the lowered
TPU program of the layer is the same as before the change (checked on the
forward-plus-backward lowering of a shrunk Qwen3.5 MoE layer for v7x).

Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), with this
change switched by an environment variable in a build that a TPU lowering
check shows compiling to the same program as this commit: regular
micro-batch 10.022 -> 9.875 s (-0.147 s; arm o275x-mf-cnt-prof against
o275x-mf-ctl, the flags-off control in the same pod), 99.25 -> 98.05 s per
step, 0 dropped tokens in both. The same control measured in different
pods spreads by +/-0.02-0.03 s per micro-batch. The timing arms run with
the benchmark's log-prob gates on (loss and gradient norm read 0), so they
say nothing about numerics.

TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks with this key in both measured
stacks: with moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp,
moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run
nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm
difference 7.0e-6 against a replicate spread of 1.9e-5), and with the
fused SiLU GMM in place of the fused activation (run nxs6silu:
gradient-norm difference 2.5e-6). Caveats: that is the stack, not this key
alone; and it ran on the environment-switch build, not on this branch's
binary.

Tests (moe_routing_index_math_test.py, new), on CPU:
- _local_slot_positions equals argsort(argsort(ids)) minus the window
  start on every shard's slots and -1 elsewhere, across partial blocks,
  the padding id and repeated experts; topk_indices_by_argmax equals
  top_k's indices with ties; _preprocess(packed=True) equals the argsort
  version. The SparseCore gather-reduce is TPU-only (CPU takes its
  reference fallback), so its preprocessing is tested directly.
- On the shrunk Qwen3.5 routed MoE on 8 CPU devices (cp-as-ep, local
  routing, two chunks), "count" is bit-identical to "sort" in the output,
  the loss and every gradient: plain and pipelined chunk loop, padding
  excluded from routing, forced ids with a buffer that drops tokens,
  dropless, saved routing maps, bf16 router, and the benchmark's MoE
  settings (expert-weight prefetch, row-major reduce-scatter, routing maps
  on device, ragged_buffer_factor 2.75, fp8 fake-quant).
- Structure, two-sided: "count" traces fewer sorts than "sort", none of
  them stable or with more than one operand (so a partial fallback to
  jnp.argsort fails), and no top_k under padding exclusion; every
  ragged_gather_reduce call gets packed_preprocess=True under "count" and
  False under "sort" (a spy, since only the TPU kernel reads it). A
  counted inverse shifted by one is caught, and "count" without the
  forward permutation raises.
- The validator rejects "count" for each unmet requirement, naming it.
---
 src/maxtext/configs/base.yml                  |   5 +
 src/maxtext/configs/types.py                  |  51 ++
 .../kernels/ragged/ragged_gather_reduce_v2.py |  27 +-
 src/maxtext/kernels/ragged/ragged_sort.py     | 141 ++++-
 src/maxtext/layers/moe.py                     |  34 +-
 tests/unit/moe_routing_index_math_test.py     | 502 ++++++++++++++++++
 6 files changed, 745 insertions(+), 15 deletions(-)
 create mode 100644 tests/unit/moe_routing_index_math_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index aa5d701a73..98d0514258 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -355,6 +355,11 @@ moe_dispatch_no_expert_sharding: false
 use_ragged_sort: false # whether to use the Pallas ragged-sort kernels in the MoE permute path; valid both with and
                        # without `use_ring_of_experts` (with EP > 1). When `use_ring_of_experts=True` the kernels run
                        # inside `permute`/`unpermute`; otherwise they run inside `local_permute`/local-unpermute.
+# Ring-of-experts ragged sort routing index maps: 'sort' = two stable argsorts (default); 'count' = one packed
+# single-operand sort plus counting for the inverse (and argmax passes for the top-k of non-forced tokens under
+# router replay / padding exclusion). Same permutation, same values.
+# Needs sparse_matmul, use_ring_of_experts and use_ragged_sort (and te_moe_block=False).
+moe_routing_index_math: 'sort'
 moe_tc_ragged_sort: false # ring-of-experts ragged sort/unsort on TensorCore (Pallas DMA kernels) instead of the SparseCore
                           # kernels; requires use_ring_of_experts, use_ragged_sort and ragged_buffer_factor > 0.
 moe_tc_ragged_gather_block_size: 1024 # moe_tc_ragged_sort: rows per grid step of the TC ragged gather.
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 730217f5ef..e8aed40890 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1101,6 +1101,18 @@ class SplashAttention(BaseModel):
   )
 
 
+def _reject_unless(setting: str, requirements: list[tuple[bool, str]]) -> None:
+  """Raises a ValueError naming every requirement of `setting` that is not met.
+
+  Args:
+    setting: the setting being validated, as the user wrote it (e.g. "moe_unsplit_prefused_wi=True").
+    requirements: (met, requirement) pairs, `requirement` describing what `setting` needs.
+  """
+  unmet = [requirement for met, requirement in requirements if not met]
+  if unmet:
+    raise ValueError(f"{setting} requires: {'; '.join(unmet)}.")
+
+
 class MoEGeneral(BaseModel):
   """General configuration for Mixture of Experts (MoE) layers."""
 
@@ -1421,6 +1433,22 @@ class MoEGeneral(BaseModel):
       False,
       description="Whether to use ragged kernel for sorting, improve performance when EP is enabled.",
   )
+  moe_routing_index_math: Literal["sort", "count"] = Field(
+      "sort",
+      description=(
+          "How the ring-of-experts ragged sort computes its routing index maps. 'sort' (default): the expert"
+          " permutation and its inverse by two stable argsorts. 'count': the permutation by one single-operand"
+          " packed uint32 sort (XLA lowers the stable argsort with an extra, duplicated iota operand); the sorted"
+          " token indices by division instead of a gather; the inverse permutation by counting, on this shard's"
+          " slots only, instead of a second argsort; the router-weight gradient back to slot order by a gather"
+          " instead of a scatter; the gather-reduce validity sort packed into one uint32 operand; and, for tokens"
+          " without forced ids under router replay or moe_exclude_padding_from_routing, the router's top-k indices"
+          " by k argmax passes instead of a sort. Exact: the same permutation, so the same outputs and gradients,"
+          " dropped tokens included. With moe_tc_ragged_sort, which replaces the ragged sort of the dropping program,"
+          " it applies to the dropless program's ragged sort only. Requires sparse_matmul=True, te_moe_block=False,"
+          " use_ring_of_experts=True and use_ragged_sort=True."
+      ),
+  )
   moe_quantize_combine_bwd_method: str | None = Field(
       "",
       description=(
@@ -1691,6 +1719,29 @@ def validate_moe_sharding_strategy(self) -> "MoEGeneral":
       )
     return self
 
+  # The validators of the MoE performance keys below share one form: each lists (met, requirement) pairs and
+  # _reject_unless raises naming every unmet one. They read the fields of the other config mixins (sparse_matmul,
+  # mlp_bias, remat_policy, the parallelism degrees, ...) directly: they run on the combined configs (MaxTextConfig
+  # and RLConfig), which have all of them, so a renamed field fails loudly instead of passing the check.
+
+  def _ring_of_experts_ragged_sort_requirements(self) -> list[tuple[bool, str]]:
+    """What the MoE layer needs to run the ring-of-experts ragged sort."""
+    return [
+        (self.sparse_matmul, "sparse_matmul=True"),
+        (not self.te_moe_block, "te_moe_block=False"),
+        (self.use_ring_of_experts, "use_ring_of_experts=True"),
+        (self.use_ragged_sort, "use_ragged_sort=True"),
+    ]
+
+  @model_validator(mode="after")
+  def validate_moe_routing_index_math(self) -> "MoEGeneral":
+    """moe_routing_index_math='count' needs the ring-of-experts ragged sort, the only user of these index maps."""
+    if self.moe_routing_index_math != "sort":
+      _reject_unless(
+          f"moe_routing_index_math={self.moe_routing_index_math!r}", self._ring_of_experts_ragged_sort_requirements()
+      )
+    return self
+
 
 class MoEKernels(BaseModel):
   """Configuration for MoE-specific kernels like Megablox."""
diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
index 571200ffb7..3e94b5b08b 100644
--- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
+++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
@@ -269,9 +269,13 @@ def _preprocess(
     num_row_partitions: int,
     num_simd_lanes: int,
     row_chunk_size: int,
+    packed: bool = False,
 ) -> tuple[jax.Array, jax.Array, jax.Array]:
   """Sorts valid source rows to the front of each row partition.
 
+  With ``packed`` (``moe_routing_index_math="count"``) the stable partition is one single-operand sort of
+  packed uint32 keys instead of a stable argsort; the permutation is the same.
+
   Returns:
     sorted_by_validity: original row index of each slot after the stable
       sort, flattened across partitions and padded to ``row_chunk_size``.
@@ -282,9 +286,20 @@ def _preprocess(
   row_partition_size = valid_rows_mask.shape[0] // num_row_partitions
   valid_rows_mask_2d = valid_rows_mask.reshape(num_row_partitions, -1)
 
-  # Stable sort of a boolean key is a stable partition: valid rows keep their
-  # relative order and move ahead of the invalid ones.
-  sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, descending=False, stable=True, axis=-1)
+  if packed:
+    # Key (invalid << idx_bits) | index: valid rows sort first, each group in index order. The keys are all
+    # distinct, so an unstable sort has exactly one result and it is the stable partition below. It sorts ONE
+    # uint32 operand, where the stable boolean argsort sorts the key with an iota tie-break (XLA's stable-sort
+    # expansion adds a third operand in the forward pass).
+    idx_bits = max(1, (row_partition_size - 1).bit_length())
+    key = (valid_rows_mask_2d == 0).astype(jnp.uint32) << jnp.uint32(idx_bits)
+    key = key | jnp.arange(row_partition_size, dtype=jnp.uint32)[None, :]
+    key = jax.lax.sort(key, dimension=1, is_stable=False)
+    sorted_by_validity = (key & jnp.uint32((1 << idx_bits) - 1)).astype(jnp.int32)
+  else:
+    # Stable sort of a boolean key is a stable partition: valid rows keep their
+    # relative order and move ahead of the invalid ones.
+    sorted_by_validity = jnp.argsort(~valid_rows_mask_2d, descending=False, stable=True, axis=-1)
   sorted_by_validity += jnp.arange(num_row_partitions)[:, None] * row_partition_size
 
   pad_to = _align_to(row_partition_size, row_chunk_size)
@@ -636,6 +651,7 @@ def col_loop(col_compute_offset):
         "flops_override",
         "bytes_accessed_override",
         "use_single_sparsecore",
+        "packed_preprocess",
     ),
 )
 def ragged_gather_reduce(
@@ -648,6 +664,7 @@ def ragged_gather_reduce(
     flops_override: int = -1,
     bytes_accessed_override: int = -1,
     use_single_sparsecore: bool = False,
+    packed_preprocess: bool = False,
 ) -> jax.Array:
   """Gathers ``x`` by ``indices``, weights and masks, then reduces by group.
 
@@ -657,6 +674,9 @@ def ragged_gather_reduce(
     topk_weights: 1-D per-row weights, ``(input_size,)``.
     valid_rows_mask: 1-D bool mask of valid gathered rows, ``(input_size,)``.
     reduce_group_size: number of consecutive rows summed into one output row.
+    packed_preprocess: sort the valid rows to the front with one packed uint32 sort (``_preprocess(packed=True)``,
+      ``moe_routing_index_math="count"``). Same permutation, so the same output. Only the SparseCore path sorts;
+      the reference fallback (any non-TPU target) ignores it.
 
   Returns:
     Reduced output, ``(input_size // reduce_group_size, hidden_size)``.
@@ -720,6 +740,7 @@ def ragged_gather_reduce(
       num_row_partitions,
       num_simd_lanes,
       row_chunk_size,
+      packed=packed_preprocess,
   )
 
   # Step 4: Launch the SparseCore kernel.
diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index cb8041c2c3..9c57255ba0 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -18,6 +18,70 @@
 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.ragged_gather_tc import packed_stable_sort
+
+_INDEX_MATH_MODES = ("sort", "count")
+
+
+def _check_index_math(index_math):
+  if index_math not in _INDEX_MATH_MODES:
+    raise ValueError(f"index_math must be one of {_INDEX_MATH_MODES}, got {index_math!r}")
+
+
+def _local_slot_positions(topk_indices_local, num_experts, shard_idx, ep_size, block=128):
+  """Position of every flat (token, k) slot routed to this shard, inside this shard's window of the sorted slots.
+
+  Equals ``argsort(argsort(flat_ids)) - shard_output_start`` on the slots of this shard's experts (``argsort``
+  stable, by expert id) and is -1 on every other slot. It needs no sort and no scatter: in the stable order by
+  expert id, slot (t, k) with local expert e comes after
+    - every slot of a lower local expert (an exclusive prefix sum of the local expert counts),
+    - every slot of expert e in an earlier token (an exclusive prefix count over tokens), and
+    - every earlier slot of token t that also picked e (0 unless a token repeats an expert).
+  The prefix count over tokens runs per block of ``block`` tokens as a strictly-lower-triangular matmul (its
+  inputs are counts <= topk, exact in bfloat16, and it accumulates in float32, exact below 2**24), plus an int32
+  prefix sum over the block totals.
+
+  Args:
+    topk_indices_local: ``[num_tokens, topk]`` expert ids. Ids outside this shard's experts (including the
+      ``num_experts`` padding id) are never local.
+    num_experts: total number of experts.
+    shard_idx: this shard's index on the expert-parallel axis (a traced scalar is fine).
+    ep_size: number of expert shards.
+    block: tokens per block of the prefix count.
+
+  Returns:
+    ``[num_tokens * topk]`` int32 positions in ``[0, local slots)``, or -1.
+  """
+  num_tokens, topk = topk_indices_local.shape
+  local_num_experts = num_experts // ep_size
+  pad = (-num_tokens) % block
+  ids = topk_indices_local.astype(jnp.int32)
+  if pad:
+    ids = jnp.pad(ids, ((0, pad), (0, 0)), constant_values=num_experts)  # the padding id: never local
+  t = ids.shape[0]
+  le = ids - shard_idx * local_num_experts  # local expert id, in [0, local_num_experts) iff the slot is local
+  is_local = (le >= 0) & (le < local_num_experts)
+  hit = le[:, :, None] == jnp.arange(local_num_experts, dtype=jnp.int32)[None, None, :]  # [t, topk, E_local]
+  counts_per_token = jnp.sum(hit, axis=1, dtype=jnp.int32)  # [t, E_local]
+  kk = jnp.arange(topk, dtype=jnp.int32)
+  dup_before = jnp.sum(
+      (le[:, :, None] == le[:, None, :]) & (kk[None, None, :] < kk[None, :, None]), axis=-1, dtype=jnp.int32
+  )  # [t, topk]: earlier slots of the same token with the same expert
+  blocks = counts_per_token.reshape(t // block, block, local_num_experts)
+  strict_lower = jnp.tril(jnp.ones((block, block), jnp.bfloat16), k=-1)
+  # Slots of each local expert in the earlier tokens of the same block.
+  within = jnp.einsum(
+      "ij,bjk->bik", strict_lower, blocks.astype(jnp.bfloat16), preferred_element_type=jnp.float32
+  ).astype(jnp.int32)
+  block_totals = jnp.sum(blocks, axis=1, dtype=jnp.int32)  # [t // block, E_local]
+  before_block = jnp.cumsum(block_totals, axis=0) - block_totals
+  earlier = (within + before_block[:, None, :]).reshape(t, local_num_experts)  # [t, E_local]
+  local_counts = jnp.sum(block_totals, axis=0)
+  local_offsets = jnp.cumsum(local_counts) - local_counts
+  # One select-reduce over the local experts picks each slot's own expert's terms (no gather).
+  pos = jnp.sum(jnp.where(hit, earlier[:, None, :] + local_offsets[None, None, :], 0), axis=-1) + dup_before
+  pos = jnp.where(is_local, pos, -1)[:num_tokens]
+  return pos.reshape(-1)
 
 
 def ring_ragged_sort(
@@ -36,6 +100,7 @@ def ring_ragged_sort(
     gather_reduce_bytes_accessed_override=-1,
     use_single_sparsecore=False,
     tag_routing_fn=None,
+    index_math="sort",
 ):
   """Ragged-gather variant for AG-RS Expert Parallelism token routing.
 
@@ -66,6 +131,13 @@ def ring_ragged_sort(
     buffer_size: optional scalar ``int`` representing the size of the local buffer.
     tag_routing_fn: optional function applied to each integer routing array (the sort permutation, its inverse,
       the sorted token indices and the group sizes).
+    index_math: how the routing index maps are computed (``moe_routing_index_math``).
+      "sort" (default): the permutation and its inverse by two stable argsorts.
+      "count": the permutation by one single-operand packed uint32 sort, the sorted token indices by division,
+      and the inverse permutation by counting, on this shard's slots only (-1 on every other slot); the
+      gather-reduce of the backward packs its validity sort too. Every reader of the inverse masks it to this
+      shard's window, so the forward output and the gradients are the same. The unsort of these rows must then
+      get ``topk_argsort_indices`` and ``index_math="count"`` (a partial inverse cannot be argsorted back).
 
   Returns:
     A tuple containing:
@@ -76,19 +148,37 @@ def ring_ragged_sort(
       - 1D tensor ``topk_argsort_revert_indices`` for inverse routing.
       - 1D tensor ``topk_argsort_indices``, the inverse of ``topk_argsort_revert_indices``.
   """
+  _check_index_math(index_math)
   # The routing index maps depend only on the (integer) expert ids, so they are computed outside the custom_vjp.
   # Only the gather below is differentiated.
   num_tokens_local = hidden_states_local.shape[0]
 
   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
+  if index_math == "count":
+    # The stable order by expert id as ONE single-operand sort of (expert << pos_bits | slot): the packed values
+    # are distinct, so the unstable sort has one result, the stable argsort's permutation. jnp.argsort sorts the
+    # (expert, iota) pair stably, which XLA lowers in the forward pass with a duplicated iota operand.
+    _, topk_argsort_indices = packed_stable_sort(topk_indices_flat.astype(jnp.int32), num_experts)
+    # Flat slot j belongs to token j // topk, so the sorted token indices need no gather of a repeated iota.
+    token_indices_sorted = (topk_argsort_indices // topk).astype(jnp.int32)  # num_tokens_local x topk
+  else:
+    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
+  if index_math == "count":
+    # Every reader of the inverse permutation masks it to [shard_output_start, shard_output_end), so it is
+    # computed on this shard's slots only, by counting instead of a second argsort, and is -1 elsewhere.
+    shard_idx = jax.lax.axis_index(ep_name)
+    local_pos = _local_slot_positions(topk_indices_local, num_experts, shard_idx, ep_size)
+    shard_start = jnp.sum(jnp.where(jnp.arange(num_experts) < shard_idx * (num_experts // ep_size), group_sizes_local, 0))
+    topk_argsort_revert_indices = jnp.where(local_pos >= 0, shard_start + local_pos, -1).astype(jnp.int32)
+  else:
+    topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices)  # num_tokens_local x topk
+  packed_preprocess = index_math == "count"
 
   if tag_routing_fn is not None:
     topk_argsort_indices, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices = (
@@ -205,6 +295,7 @@ def _ring_ragged_sort_bwd(res, g_out):
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
+          packed_preprocess=packed_preprocess,
       )
     else:
       # Buffering: g_x has size `local_buffer_size` (packed).
@@ -231,6 +322,7 @@ def _ring_ragged_sort_bwd(res, g_out):
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
+          packed_preprocess=packed_preprocess,
       )
     return grad_hidden_states, None, None, None
 
@@ -256,6 +348,7 @@ def ring_ragged_unsort(
     gather_reduce_bytes_accessed_override=-1,
     use_single_sparsecore=False,
     topk_argsort_indices=None,
+    index_math="sort",
 ):
   """Dual of :func:`ring_ragged_sort`.
 
@@ -285,11 +378,18 @@ def ring_ragged_unsort(
     topk_argsort_indices: optional inverse of ``topk_argsort_revert_indices`` (as returned by
       :func:`ring_ragged_sort`). The backward needs this permutation; when it is
       None, the backward re-derives it with an argsort. Both give the same permutation.
+      Required with ``index_math="count"``, whose inverse covers only this shard's slots.
+    index_math: must match the :func:`ring_ragged_sort` call that produced the routing. With "count" the
+      gather-reduce packs its validity sort, and in the backward the routing-weight gradient goes back to slot
+      order by a masked gather through the inverse permutation instead of a scatter (each valid slot reads its
+      own row; other shards' and dropped slots read 0). Same values either way.
 
   Returns:
     A 2D ``[num_tokens_local, hidden]`` tensor with expert outputs scattered back
     to their original global sequence locations, with unpopulated positions zeroed out.
   """
+  _check_index_math(index_math)
+  packed_preprocess = index_math == "count"
 
   @jax.custom_vjp
   def _ring_ragged_unsort(
@@ -352,6 +452,7 @@ def _ring_ragged_unsort_fwd(
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
+          packed_preprocess=packed_preprocess,
       )
     else:
       # Shift indices so they map to the packed local buffer [0, local_num_tokens).
@@ -371,6 +472,7 @@ def _ring_ragged_unsort_fwd(
           flops_override=gather_reduce_flops_override,
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
+          packed_preprocess=packed_preprocess,
       )
 
     res = (
@@ -418,6 +520,9 @@ def _ring_ragged_unsort_bwd(res, g_out):
     # topk_argsort_indices; reuse it when the caller has it rather than sorting again.
     if topk_argsort_indices is not None:
       idx_inv = topk_argsort_indices
+    elif index_math == "count":
+      # The "count" inverse holds only this shard's slots (-1 elsewhere), so it cannot be argsorted back.
+      raise ValueError("index_math='count' needs topk_argsort_indices (the forward permutation of ring_ragged_sort).")
     else:
       idx_inv = jnp.argsort(topk_argsort_revert_indices)
 
@@ -444,7 +549,14 @@ def _ring_ragged_unsort_bwd(res, g_out):
       grad_sorted_tokens = (gathered * weight_for_sorted[:, None]).astype(gathered.dtype)
       # Row-wise dot in sorted order, then permuted back to flat slot order.
       dot_sorted = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local[:n].astype(jnp.float32), axis=-1)
-      grad_topk_weights = dot_sorted[topk_argsort_revert_indices]
+      if index_math == "count":
+        # Other shards' slots hold -1, not a row outside the window, so mask them instead of reading row -1.
+        slot_valid = (topk_argsort_revert_indices >= shard_output_start) & (
+            topk_argsort_revert_indices < shard_output_end
+        )
+        grad_topk_weights = jnp.where(slot_valid, dot_sorted[jnp.where(slot_valid, topk_argsort_revert_indices, 0)], 0.0)
+      else:
+        grad_topk_weights = dot_sorted[topk_argsort_revert_indices]
     else:
       # Slice the inverse permutation to match the packed local buffer.
       padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size))
@@ -468,10 +580,21 @@ def _ring_ragged_unsort_bwd(res, g_out):
       mask = jnp.arange(buffer_size) < limit
       gathered = jnp.where(mask[:, None], gathered, 0.0)
       grad_sorted_tokens = (gathered * sliced_weights[:, None]).astype(gathered.dtype)
-      # 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)
-      slots = jnp.where(mask, sliced_idx_inv, n)
-      grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop")
+      if index_math == "count":
+        # Route the per-row dot back to flat slot order with a gather through the inverse permutation instead
+        # of a TensorCore scatter: each valid slot reads its own buffer row; other shards' slots (-1) and
+        # dropped slots read 0, as the scatter leaves them. The table is padded to n entries because XLA only
+        # offloads a 1-D gather to the SparseCore for a large table (v7x AOT: 225,280 entries -> TensorCore
+        # gather at ~7 ns/element; 655,360 -> SparseCore).
+        shifted = topk_argsort_revert_indices - shard_output_start
+        slot_valid = (shifted >= 0) & (shifted < limit)
+        dot_table = jnp.pad(dot_local, (0, max(0, n - buffer_size)))
+        grad_topk_weights = jnp.where(slot_valid, dot_table[jnp.where(slot_valid, shifted, 0)], 0.0)
+      else:
+        # Scatter the per-slot dot back to flat slot order; dropped slots stay zero.
+        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, None
 
   _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd)
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index fa08ed2702..1858c91abe 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -125,7 +125,7 @@ class RouteOutput:
   # moe_log_max_load_ratio=True on the ring-of-experts ragged path.
   max_load_ratio: Optional[jax.Array] = None
   # Inverse of sorted_selected_experts on the ring-of-experts ragged path, so the unsort backward does not re-sort to
-  # get it. Only set when moe_routing_maps is not "remat".
+  # get it. Set when moe_routing_maps is not "remat", and always with moe_routing_index_math="count".
   topk_argsort_indices: Optional[jax.Array] = None
   # moe_tc_ragged_sort: TcRouting shared by the TensorCore ragged sort and unsort of this chunk.
   tc_routing: Optional[Any] = None
@@ -464,6 +464,25 @@ def _top_2_in_group_sum(scores_grouped: jax.Array) -> jax.Array:
   return max_1.astype(jnp.float32) + max_2.astype(jnp.float32)
 
 
+def topk_indices_by_argmax(scores: jax.Array, k: int) -> jax.Array:
+  """`jax.lax.top_k(scores, k)[1]` without a sort (moe_routing_index_math="count").
+
+  k passes of argmax over the last axis, each masking its winner to -inf. argmax returns the first maximal index,
+  so ties go to the lower index, in the order top_k returns them. With k small against the expert count (10 of 512
+  for Qwen3.5-397B) the k reductions are expected to cost less than the sort top_k lowers to; this was timed only as
+  part of moe_routing_index_math="count", not on its own. Needs at least k entries above -inf per row, true of the
+  router probabilities it is used on.
+  """
+  lane = jax.lax.broadcasted_iota(jnp.int32, scores.shape, scores.ndim - 1)
+  masked = scores
+  picks = []
+  for _ in range(k):
+    pick = jnp.argmax(masked, axis=-1).astype(jnp.int32)
+    picks.append(pick)
+    masked = jnp.where(lane == pick[..., None], jnp.array(-jnp.inf, scores.dtype), masked)
+  return jnp.stack(picks, axis=-1)
+
+
 def calculate_load_balance_updates(top_k_indices, num_experts, rate, axis_names=None):
   """Computes a bias adjustment update based on expert load.
 
@@ -1455,7 +1474,12 @@ 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)
+      if self.config.moe_routing_index_math == "count":
+        # The router's own choice for tokens without (valid) forced ids: the same indices as top_k, from k argmax
+        # passes instead of a sort. Under moe_exclude_padding_from_routing every token takes this branch.
+        auto_indices = topk_indices_by_argmax(router_probs, self.num_experts_per_tok)
+      else:
+        _, 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
@@ -1834,9 +1858,12 @@ def permute(
           gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed,
           use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           tag_routing_fn=tag_routing_fn,
+          index_math=self.config.moe_routing_index_math,
       )
-      if tag_routing_fn is None:
+      if tag_routing_fn is None and self.config.moe_routing_index_math == "sort":
         # moe_routing_maps=remat: unpermute re-derives the permutation in its backward (existing behavior).
+        # moe_routing_index_math="count" always passes it: its inverse covers only this shard's slots, so it
+        # cannot be argsorted back, and under remat the replay recomputes it anyway (no memory is saved).
         topk_argsort_indices = None
     else:
       flatten_selected_experts = jnp.ravel(selected_experts)
@@ -2014,6 +2041,7 @@ def unpermute(
           gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed,
           use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           topk_argsort_indices=topk_argsort_indices,
+          index_math=self.config.moe_routing_index_math,
       )
     else:
       unsort_intermediate = _sort_activations(
diff --git a/tests/unit/moe_routing_index_math_test.py b/tests/unit/moe_routing_index_math_test.py
new file mode 100644
index 0000000000..1d79314ba4
--- /dev/null
+++ b/tests/unit/moe_routing_index_math_test.py
@@ -0,0 +1,502 @@
+# 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.
+"""moe_routing_index_math="count" against the default "sort", on CPU.
+
+"count" computes the ring-of-experts ragged sort's routing index maps without sorts where it can: the expert
+permutation by one packed single-operand sort (`packed_stable_sort`), the inverse permutation by counting on this
+shard's slots only (`ragged_sort._local_slot_positions`, -1 elsewhere), the router-weight gradient back to slot order
+by a masked gather, the gather-reduce's validity sort packed into one uint32 operand (`_preprocess(packed=True)`), and
+the router's top-k indices of non-forced tokens by k argmax passes (`moe.topk_indices_by_argmax`). The permutation is
+the same, so everything must be bit-identical to "sort".
+
+What a CPU run can and cannot reach:
+  - The gather-reduce's validity sort runs only on the SparseCore path. Off a TPU target `ragged_gather_reduce` takes
+    its jnp reference fallback, which does not sort, so `packed_preprocess` changes nothing in the layer runs below.
+    `PackedPreprocessTest` therefore calls `_preprocess` directly, the pure-JAX part of the SparseCore path.
+  - Likewise the ragged gather is the jnp fallback here. Everything else ("count"'s permutation, its counted inverse,
+    the masked weight-gradient gather, the argmax top-k) is plain JAX and runs in the layer below exactly as on TPU.
+
+In-process checks:
+  - `_local_slot_positions` equals argsort(argsort(ids)) - (this shard's window start) on this shard's slots and -1
+    elsewhere, for every shard, with token counts that are not a multiple of the block, the padding id, and tokens
+    that repeat an expert.
+  - `topk_indices_by_argmax` equals `jax.lax.top_k`'s indices, ties included.
+  - `_preprocess(packed=True)` equals `_preprocess(packed=False)`, every output.
+  - The config validator rejects "count" off the ring-of-experts ragged sort, one case per requirement.
+
+The layer check runs as a subprocess, so the forced device count takes effect before JAX initializes: the shrunk
+Qwen3.5 routed MoE of `moe_roe_local_routing_test.py` (16 experts, top 4, 8 CPU devices under cp-as-ep, ragged sort,
+local routing, two token chunks, megablox in interpret mode). Pre-registered: the forward output, the loss and every
+gradient bit-identical to "sort", in each case: the plain and the pipelined chunk loop, padding excluded from routing
+(every token takes the forced-routing branch, so the argmax top-k runs), forced ids with a buffer that drops tokens
+(asserted to overflow), a dropless buffer, saved routing maps, and a bfloat16 router. Two-sided:
+  - Structure: "count" must trace fewer `sort` equations than "sort" in the gradient and no argsort at all (a sort
+    of more than one operand: jnp.argsort sorts its keys with an iota), while "sort" traces argsorts; and with
+    padding excluded no `top_k`. So a switch that silently fell back, completely or for one map, cannot pass.
+  - Wiring of the SparseCore-only validity sort: CPU cannot run it, but the flag that selects it reaches
+    `ragged_gather_reduce` on every target. A spy records it: every call under "count" must pass
+    packed_preprocess=True and every call under "sort" False.
+  - The bench's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps saved on device,
+    ragged_buffer_factor 2.75, prefused weights with fp8 fake-quant, padding excluded) in one case.
+  - A deliberately wrong counted inverse (every local position shifted by one) must FAIL the comparison.
+  - "count" without the forward permutation must raise in the unsort backward instead of argsorting a partial inverse.
+"""
+
+import contextlib
+import os
+import re
+import subprocess
+import sys
+from unittest import mock
+
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from jax.sharding import PartitionSpec as P
+import numpy as np
+import pytest
+
+from maxtext.configs import pyconfig
+from maxtext.kernels.ragged import ragged_gather_reduce_v2
+from maxtext.kernels.ragged import ragged_sort
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+_NUM_DEVICES = 8
+_NUM_EXPERTS = 16
+_TOP_K = 4
+_EMB = 256
+_SEQ = 64
+_PASS_TOKEN = "MOE_ROUTING_INDEX_MATH_CHECKS_PASSED"
+_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false"
+
+
+def _expected_local_positions(ids, num_experts, shard, ep_size):
+  """Reference: each slot's rank in the stable sort by expert id, relative to this shard's window; -1 off-shard."""
+  flat = np.asarray(ids).reshape(-1)
+  rank = np.argsort(np.argsort(flat, kind="stable"), kind="stable")
+  local_num_experts = num_experts // ep_size
+  start = shard * local_num_experts
+  local = (flat >= start) & (flat < start + local_num_experts)
+  window_start = int(np.sum(flat < start))
+  return np.where(local, rank - window_start, -1).astype(np.int32)
+
+
+@pytest.mark.cpu_only
+class LocalSlotPositionsTest(parameterized.TestCase):
+  """`_local_slot_positions` is the inverse permutation restricted to this shard's window."""
+
+  @parameterized.named_parameters(
+      # (num_tokens, topk, num_experts, ep_size, block): one exact block, several blocks with a partial last one,
+      # fewer tokens than a block, and the production shape's ratios at a small scale.
+      ("one_block", 128, 4, 16, 8, 128),
+      ("partial_blocks", 300, 4, 16, 8, 64),
+      ("short", 5, 2, 8, 2, 128),
+      ("wide_topk", 200, 10, 64, 4, 32),
+  )
+  def test_matches_double_argsort(self, num_tokens, topk, num_experts, ep_size, block):
+    rng = np.random.default_rng(num_tokens)
+    # Ids in [0, num_experts]: num_experts is the padding id (never local). Distinct per token, as top-k gives.
+    ids = np.stack([rng.permutation(num_experts + 1)[:topk] for _ in range(num_tokens)]).astype(np.int32)
+    for shard in range(ep_size):
+      got = np.asarray(ragged_sort._local_slot_positions(jnp.asarray(ids), num_experts, shard, ep_size, block))  # pylint: disable=protected-access
+      np.testing.assert_array_equal(got, _expected_local_positions(ids, num_experts, shard, ep_size), err_msg=f"{shard=}")
+
+  def test_repeated_expert_within_a_token(self):
+    """A token that picks one expert twice: the stable order puts its earlier slot first."""
+    ids = np.array([[3, 3, 1], [1, 3, 3], [16, 16, 0], [0, 1, 3]], np.int32)
+    for shard in range(8):
+      got = np.asarray(ragged_sort._local_slot_positions(jnp.asarray(ids), 16, shard, 8, 2))  # pylint: disable=protected-access
+      np.testing.assert_array_equal(got, _expected_local_positions(ids, 16, shard, 8), err_msg=f"{shard=}")
+
+  def test_wrong_reference_is_detected(self):
+    """Two-sided: the comparison above fails on a deliberately wrong expectation (an unstable rank)."""
+    ids = np.array([[1, 0], [1, 0], [0, 1]], np.int32)
+    got = np.asarray(ragged_sort._local_slot_positions(jnp.asarray(ids), 2, 0, 1, 2))  # pylint: disable=protected-access
+    flat = ids.reshape(-1)
+    unstable_rank = np.argsort(np.argsort(-np.arange(flat.size) + flat * 100))  # ties broken by DEcreasing index
+    self.assertFalse(np.array_equal(got, unstable_rank.astype(np.int32)))
+
+
+@pytest.mark.cpu_only
+class TopkIndicesByArgmaxTest(parameterized.TestCase):
+  """`topk_indices_by_argmax` returns `jax.lax.top_k`'s indices."""
+
+  @parameterized.named_parameters(("distinct", False), ("ties", True))
+  def test_matches_top_k(self, ties):
+    scores = jax.random.uniform(jax.random.PRNGKey(0), (4, 33, 64), jnp.float32)
+    if ties:
+      # Few distinct values, so most rows tie at the top-k boundary; and exact zeros, like underflowed softmax.
+      scores = jnp.round(scores * 4) / 4
+    want = jax.lax.top_k(scores, 10)[1]
+    got = moe.topk_indices_by_argmax(scores, 10)
+    np.testing.assert_array_equal(np.asarray(got), np.asarray(want))
+
+  def test_on_softmax_probabilities(self):
+    logits = jax.random.normal(jax.random.PRNGKey(1), (2, 16, 512), jnp.float32) * 30.0  # many probs underflow to 0
+    probs = jax.nn.softmax(logits, axis=-1)
+    np.testing.assert_array_equal(
+        np.asarray(moe.topk_indices_by_argmax(probs, 10)), np.asarray(jax.lax.top_k(probs, 10)[1])
+    )
+
+
+@pytest.mark.cpu_only
+class PackedPreprocessTest(parameterized.TestCase):
+  """The SparseCore gather-reduce's validity sort, packed and unpacked.
+
+  Only the SparseCore path runs `_preprocess`; on a CPU target `ragged_gather_reduce` takes its reference fallback,
+  so this pure-JAX function is called directly.
+  """
+
+  @parameterized.named_parameters(
+      # (rows, reduce_group_size, num_row_partitions, num_simd_lanes, row_chunk_size)
+      ("one_partition", 96, 4, 1, 16, 16),
+      ("v7x_like", 1024, 8, 16, 16, 48),
+      ("padded_chunks", 640, 10, 8, 8, 24),
+  )
+  def test_packed_matches_argsort(self, rows, group, partitions, lanes, chunk):
+    for density in (0.0, 0.3, 0.9, 1.0):
+      mask = jax.random.uniform(jax.random.PRNGKey(rows), (rows,)) < density
+      want = ragged_gather_reduce_v2._preprocess(mask, group, partitions, lanes, chunk)  # pylint: disable=protected-access
+      got = ragged_gather_reduce_v2._preprocess(mask, group, partitions, lanes, chunk, packed=True)  # pylint: disable=protected-access
+      for name, a, b in zip(("sorted_by_validity", "rows_per_partition", "group_mask"), want, got):
+        self.assertEqual(a.dtype, b.dtype, name)
+        np.testing.assert_array_equal(np.asarray(b), np.asarray(a), err_msg=f"{name} {density=}")
+
+
+class RoutingIndexMathConfigTest(parameterized.TestCase):
+  """moe_routing_index_math="count" is rejected where no ring-of-experts ragged sort computes routing index maps."""
+
+  @parameterized.named_parameters(
+      # One case per requirement of the validator; each must be named in the error.
+      ("no_ragged_sort", {"use_ragged_sort": False}, "use_ragged_sort=True"),
+      (
+          "no_ring_of_experts",
+          {"use_ring_of_experts": False, "num_moe_token_chunks": 1, "ring_of_experts_local_routing": False},
+          "use_ring_of_experts=True",
+      ),
+      ("dense_matmul", {"sparse_matmul": False}, "sparse_matmul=True"),
+      ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"),
+  )
+  def test_rejected_off_the_ring_ragged_sort(self, overrides, requirement):
+    with self.assertRaisesRegex(ValueError, f"moe_routing_index_math='count' requires: .*{re.escape(requirement)}"):
+      _cfg(moe_routing_index_math="count", **overrides)
+
+  def test_accepted_on_the_ring_ragged_sort(self):
+    self.assertEqual(_cfg(moe_routing_index_math="count").moe_routing_index_math, "count")
+
+
+def _cfg(**overrides):
+  """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing)."""
+  kwargs = {
+      "run_name": "moe_routing_index_math",
+      "enable_checkpointing": False,
+      "log_config": False,
+      "skip_jax_distributed_system": True,
+      "model_name": "qwen3.5-397b-a17b",
+      "override_model_config": True,
+      "num_experts": _NUM_EXPERTS,
+      "num_experts_per_tok": _TOP_K,
+      "base_emb_dim": _EMB,
+      "base_mlp_dim": 256,
+      "base_moe_mlp_dim": 256,
+      "max_target_length": _SEQ,
+      "max_prefill_predict_length": _SEQ,
+      "per_device_batch_size": 0.25,
+      "dtype": "bfloat16",
+      "weight_dtype": "float32",
+      "float32_gate_logits": True,
+      "megablox": True,
+      "sparse_matmul": True,
+      # Tokamax gmm v2 is TPU-only; both settings share whichever GMM runs, so CPU uses megablox.
+      "use_tokamax_gmm": False,
+      "use_gmm_v2": False,
+      "use_ring_of_experts": True,
+      "use_ragged_sort": True,
+      "use_custom_sort_vjp": False,
+      "ragged_buffer_factor": 2.0,
+      "num_moe_token_chunks": 2,
+      "custom_mesh_and_rule": "cp-as-ep",
+      "ici_fsdp_parallelism": 1,
+      "ici_context_parallelism": 4,
+      "ici_expert_parallelism": 2,
+      "context_parallel_load_balance": False,
+      "ring_of_experts_local_routing": True,
+  }
+  kwargs.update(overrides)
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
+
+
+@pytest.mark.cpu_only
+def test_count_matches_sort_on_cpu_mesh():
+  """Runs `main` below in a child process with 8 CPU devices."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}"
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+  print(result.stdout)
+  assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+def _build(cfg, mesh):
+  return moe.RoutedMoE(
+      config=cfg,
+      num_experts=cfg.num_experts,
+      num_experts_per_tok=cfg.num_experts_per_tok,
+      mesh=mesh,
+      kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+      kernel_axes=("embed", "mlp"),
+      intermediate_dim=cfg.moe_mlp_dim,
+      weight_dtype=cfg.weight_dtype,
+      dtype=cfg.dtype,
+      rngs=nnx.Rngs(params=0),
+  )
+
+
+def _inputs(cfg, mode):
+  """The layer input, forced expert ids (padded and duplicate slots) and a token mask, per `mode`."""
+  batch = int(cfg.per_device_batch_size * _NUM_DEVICES)
+  x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32)
+  x = x if cfg.float32_gate_logits else x.astype(jnp.bfloat16)
+  forced = token_mask = None
+  if "forced" in mode:
+    keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ)
+    ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K)
+    pos = jnp.arange(_SEQ)[None, :, None]
+    ids = jnp.where((pos % 5 == 0) & (jnp.arange(_TOP_K) == _TOP_K - 1), -1, ids)
+    ids = jnp.where((pos % 7 == 3) & (jnp.arange(_TOP_K) == 1), ids[..., :1], ids)
+    forced = ids.astype(jnp.int32)
+  if "padmask" in mode:
+    # Each sequence ends in a run of padding of a different length.
+    lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)])
+    token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32)
+  return x, forced, token_mask
+
+
+def _count(jaxpr, name):
+  """Number of `name` equations in a jaxpr, including nested jaxprs."""
+  total = 0
+  for eqn in jaxpr.eqns:
+    total += eqn.primitive.name == name
+    for param in eqn.params.values():
+      for sub in param if isinstance(param, (list, tuple)) else (param,):
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        if hasattr(sub, "eqns"):
+          total += _count(sub, name)
+  return total
+
+
+def _count_argsorts(jaxpr):
+  """Number of `sort` equations with more than one operand: an argsort sorts its keys together with an iota.
+
+  "count" sorts one packed key; the forced-routing duplicate check's `jnp.sort` is single-operand too.
+  """
+  total = 0
+  for eqn in jaxpr.eqns:
+    total += eqn.primitive.name == "sort" and len(eqn.invars) > 1
+    for param in eqn.params.values():
+      for sub in param if isinstance(param, (list, tuple)) else (param,):
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        if hasattr(sub, "eqns"):
+          total += _count_argsorts(sub)
+  return total
+
+
+def _run(cfg, mesh, params, x, forced, token_mask):
+  """Loss, output, gradients, the overflow flag, the sort / top_k equation counts of the gradient, and the
+  packed_preprocess flag of every ragged_gather_reduce call traced."""
+  preprocess_flags = []
+  real_gather_reduce = ragged_sort.ragged_gather_reduce
+
+  def spy(*args, **kwargs):
+    preprocess_flags.append(kwargs.get("packed_preprocess", False))
+    return real_gather_reduce(*args, **kwargs)
+
+  with (
+      mock.patch.object(ragged_sort, "ragged_gather_reduce", spy),
+      jax.set_mesh(mesh),
+      nn_partitioning.axis_rules(cfg.logical_axis_rules),
+  ):
+    graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+
+    def loss_fn(params, x):
+      layer = nnx.merge(graphdef, params, rest)
+      out, lb_loss, _ = layer(x, forced_routed_experts=forced, token_mask=token_mask)
+      sown = nnx.state(layer, nnx.Intermediate).to_pure_dict()
+      flags = maxtext_utils.collect_intermediates_by_suffix(sown, "moe_has_overflow")
+      loss = jnp.mean(out.astype(jnp.float32) ** 2)
+      if lb_loss is not None:
+        loss = loss + lb_loss.astype(jnp.float32)
+      return loss, (out, jnp.any(jnp.concatenate([jnp.ravel(f) for f in flags])))
+
+    grad_fn = jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True)
+    jaxpr = jax.make_jaxpr(grad_fn)(params, x).jaxpr
+    (loss, (out, overflow)), grads = jax.jit(grad_fn)(params, x)
+  counts = {
+      "sort": _count(jaxpr, "sort"),
+      "argsort": _count_argsorts(jaxpr),
+      "top_k": _count(jaxpr, "top_k"),
+      "packed_preprocess": sorted(set(preprocess_flags)),
+  }
+  return jax.device_get((loss, out, grads)), bool(overflow), counts
+
+
+def _bit_identical(label, old, new):
+  """Prints the per-leaf report; True iff the output, the loss and every gradient are bit-identical."""
+  (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new
+  named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [
+      (jax.tree_util.keystr(path), a, b)
+      for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new))
+  ]
+  passed, lines = True, []
+  for name, a, b in named:
+    a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32)
+    same = bool(np.array_equal(a, b)) and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0
+    passed &= same
+    lines.append(f"  [{'ok' if same else 'FAIL'}] {name}: bit-identical={same} max|d|={np.max(np.abs(a - b)):.3e}")
+  print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+def _case(label, overrides, mode="real", control=None, expect_overflow=False):
+  """Runs "sort" and "count" on the same parameters; `control` is applied to the "count" run only."""
+  cfg_sort = _cfg(**overrides)
+  cfg_count = _cfg(**overrides, moe_routing_index_math="count")
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg_sort), cfg_sort.mesh_axes)
+  x, forced, token_mask = _inputs(cfg_sort, mode)
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_sort.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg_sort, mesh), nnx.Param, ...)
+  old, overflow_old, counts_old = _run(cfg_sort, mesh, params, x, forced, token_mask)
+  with control if control is not None else contextlib.nullcontext():
+    new, overflow_new, counts_new = _run(cfg_count, mesh, params, x, forced, token_mask)
+  structure = counts_new["sort"] < counts_old["sort"]
+  # "count" sorts only the packed key (one operand); "sort" argsorts with an iota operand.
+  structure &= counts_new["argsort"] == 0 < counts_old["argsort"]
+  # The SparseCore validity-sort flag, on every gather-reduce call (CPU runs the reference fallback, which ignores it).
+  structure &= counts_new["packed_preprocess"] == [True] and counts_old["packed_preprocess"] == [False]
+  if token_mask is not None:
+    structure &= counts_new["top_k"] == 0 < counts_old["top_k"]
+  overflow_ok = overflow_old == overflow_new == expect_overflow
+  print(
+      f"{label}: [{'ok' if structure else 'FAIL'}] sort/top_k equations sort={counts_old} count={counts_new};"
+      f" [{'ok' if overflow_ok else 'FAIL'}] overflow sort={overflow_old} count={overflow_new}"
+      f" (want {expect_overflow})",
+      flush=True,
+  )
+  return _bit_identical(label, old, new) and structure and overflow_ok
+
+
+def _shifted_local_positions():
+  """Deliberately wrong: every local slot's counted position is one too far."""
+  real = ragged_sort._local_slot_positions  # pylint: disable=protected-access
+
+  def wrong(*args, **kwargs):
+    pos = real(*args, **kwargs)
+    return jnp.where(pos >= 0, pos + 1, pos)
+
+  return mock.patch.object(ragged_sort, "_local_slot_positions", wrong)
+
+
+def _count_needs_forward_permutation():
+  """The unsort backward of "count" raises without the forward permutation (its inverse is partial)."""
+  mesh = Mesh(np.array(jax.devices()[:2]), ("expert",))
+  num_tokens, topk, num_experts = 8, 2, 4
+  ids = jnp.asarray(np.stack([np.random.default_rng(t).permutation(num_experts)[:topk] for t in range(num_tokens)]))
+
+  def unsort_loss(tokens):
+    def body(tokens, ids):
+      _, group_sizes, revert, _ = ragged_sort.ring_ragged_sort(
+          tokens, ids, num_experts, topk, "expert", 2, index_math="count"
+      )
+      out = ragged_sort.ring_ragged_unsort(
+          tokens.repeat(topk, axis=0),
+          group_sizes,
+          revert,
+          topk,
+          num_experts // 2,
+          "expert",
+          topk_weights=jnp.ones((num_tokens * topk,), jnp.float32),
+          index_math="count",
+      )
+      return jnp.sum(out.astype(jnp.float32))
+
+    return jax.shard_map(body, mesh=mesh, in_specs=(P(), P()), out_specs=P(), check_vma=False)(tokens, ids)
+
+  try:
+    jax.grad(unsort_loss)(jnp.ones((num_tokens, 128), jnp.bfloat16))
+  except ValueError as e:
+    raised = "needs topk_argsort_indices" in str(e)
+  else:
+    raised = False
+  print(f"count without the forward permutation: [{'ok' if raised else 'FAIL'}] the unsort backward raises", flush=True)
+  return raised
+
+
+# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+_BENCH_FLAGS = {
+    "moe_chunk_pipeline": True,
+    "moe_expert_weight_prefetch": "forward",
+    "ring_of_experts_row_major_reduce_scatter": True,
+    "moe_routing_maps": "device",
+    "ragged_buffer_factor": 2.75,
+    "prefuse_moe_weights": True,
+    "fp8_moe_fake_quant": True,
+    "moe_exclude_padding_from_routing": True,
+}
+
+
+def main():
+  assert jax.device_count() == _NUM_DEVICES, jax.devices()
+  pipelined = {"moe_chunk_pipeline": True}
+  padded = {**pipelined, "moe_exclude_padding_from_routing": True}
+  results = {
+      "plain_loop": _case("plain_loop", {}),
+      "pipelined": _case("pipelined", pipelined),
+      "padding_excluded": _case("padding_excluded", padded, mode="padmask"),
+      "forced_dropping": _case(
+          "forced_dropping", {**padded, "ragged_buffer_factor": 1.0}, mode="forced_padmask", expect_overflow=True
+      ),
+      "dropless": _case("dropless", {"ragged_buffer_factor": -1.0}, mode="forced"),
+      "saved_routing_maps": _case("saved_routing_maps", {**pipelined, "moe_routing_maps": "device"}),
+      "bf16_router": _case("bf16_router", {**padded, "float32_gate_logits": False}, mode="padmask"),
+      "bench_flags": _case("bench_flags", _BENCH_FLAGS, mode="padmask"),
+      "guard": _count_needs_forward_permutation(),
+  }
+  controls = {
+      "control_shifted_local_positions": _case(
+          "control_shifted_local_positions", pipelined, control=_shifted_local_positions()
+      ),
+  }
+  print("SUMMARY", {**results, **controls}, flush=True)
+  failed = [k for k, v in results.items() if not v]
+  undetected = [k for k, v in controls.items() if v]
+  if failed or undetected:
+    print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True)
+    sys.exit(1)
+  print(_PASS_TOKEN, flush=True)
+
+
+if __name__ == "__main__":
+  main()

From 048d617017479c3165961ad6cf806057dca9b03b Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 02:29:45 -0700
Subject: [PATCH 04/10] MoE: three exact glue rewrites
 (moe_unsplit_prefused_wi, moe_ragged_sort_unique_pair_argsort,
 moe_prefused_swiglu_vjp)

Three rewrites of MoE glue ops that leave every value unchanged, each
behind its own key (default off, program unchanged):

moe_unsplit_prefused_wi. With prefuse_moe_weights, the layer sliced the
fused wi = [w0 | w1] into its halves at entry, so the fp8 fake-quant ran
twice, the FSDP all-gather and the expert-weight prefetch barrier carried
two operands, and gmm_up concatenated the gathered halves back before the
up-projection. The key keeps wi whole up to the GEMM; w1_kernel=None marks
the layout in sparse_matmul. The fake-quant scales are per output channel,
so rounding the whole tensor gives the halves' values.

moe_ragged_sort_unique_pair_argsort. With moe_routing_index_math="sort",
the ring-of-experts ragged sort's two argsorts and the SparseCore
gather-reduce's validity sort run as unstable sorts of the unique (key,
index) pairs (argsort_unique_pairs). A stable jnp.argsort lowers in the
forward pass with a third, duplicated iota operand; the pair sort has two
operands and exactly one result, the stable permutation (on v7x, an
s32[655360] argsort takes 0.61-0.63 ms with three operands and 0.42 ms
with two).

moe_prefused_swiglu_vjp. With prefuse_moe_weights, the activation on the
whole [gate | up] up-projection output runs as one custom-VJP op whose
backward takes jax.vjp of the same activation and returns d[gate | up] as
one array. Split, the two [rows, mlp] cotangents leave their fusion
separately and XLA joins them (add of two pads) in an extra pass over the
expert buffer. Any gated activation works (SwiGLU, GPT-OSS's clipped
form). Same values; the gradient bits can differ only in the sign of an
exact zero (the add of zero pads gives -0.0 + 0.0 = +0.0, the concatenate
keeps -0.0; measured on XLA:CPU).

Which keys to use. With the rest of this series the recommended set is
moe_routing_index_math="count", moe_unsplit_prefused_wi,
moe_roe_packed_routing_gather, moe_router_grad_from_wo_input and
moe_wi_gmm_fused_swiglu. Of this commit's keys only the whole wi is in it:
"count" removes the sorts the pair argsort rewrites, and the fused SiLU
GMM never forms the [gate | up] output the fused activation reads, so the
validators reject the pair argsort with "count" and the fused activation
with the fused SiLU. The two stay as standalone alternatives for
configurations the recommended keys do not cover: the pair argsort for
runs that keep moe_routing_index_math="sort" (it keeps the maps as sorts,
the smaller change), and the fused activation where the fused SiLU GMM
does not apply (it needs tokamax gmm_v2 and plain silu + linear; the
fused activation works with any GMM backend and gated activation).

The validators, in the form moe_routing_index_math started, reject each
key for every unmet requirement: the pair argsort off the
ring-of-experts ragged sort or with "count"; the whole wi off prefused
sparse-matmul training, with weight quantization, non-default
expert-weight sharding, embedding chunks, chunk weight-gradient
accumulation, or an expert MLP dimension that tensor, tensor_sequence or
autoregressive parallelism shards (a device would then hold a slice of
[w0 | w1] instead of matching slices of w0 and w1); the fused activation
off prefused sparse matmul, with MLP biases or embedding chunks, or under
a remat policy that saves the routed experts' gate or up half
(_remat_keeps_moe_tensor_requirement, shared with the later keys of this
series: the minimal* presets, or custom with mlpwi_0, moe_mlpwi_0 or
moe_mlpwi_1 saved or offloaded). The fused activation never forms those
halves, so such a policy would silently stop saving the MoE tensor (the
shared expert's dense MLP, tagged mlpwi_0 too, would still be saved). An
MLP dimension sharded under custom logical axis rules, weight
quantization objects and sparsity, which only the built layer knows,
raise at trace time instead of falling back.

Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), all
three together, switched by an environment variable in a build that a
TPU lowering check shows compiling to the same program as this commit:
regular micro-batch 10.022 -> 9.852 s (-0.170 s; arm o275x-mf-gb-prof
against its in-pod control o275x-mf-ctl), 99.25 -> 97.70 s per step, 0
dropped tokens in both. The same control measured in different pods
spreads by +/-0.02-0.03 s per micro-batch. The parts were not timed
separately, and moe_routing_index_math was "sort" in that arm. The timing
arms run with the benchmark's log-prob gates on (loss and gradient norm
read 0), so they say nothing about numerics.

TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks for moe_unsplit_prefused_wi and
moe_prefused_swiglu_vjp stacked with moe_routing_index_math="count",
moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run
nxs6c: loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm
difference 7.0e-6 against a replicate spread of 1.9e-5), and for the whole
wi in the recommended set above (run nxs6silu: gradient-norm difference
2.5e-6). Caveats: stack level, not per key; the environment-switch build,
not this branch's binary; and the pair argsort is in neither run (it is
inert next to "count"; a permutation, exact by construction, and
bit-identical on CPU).

Tests (moe_glue_rewrites_test.py, new), on CPU:
- argsort_unique_pairs equals the stable argsort (ties, booleans, every
  axis, index dtype); _preprocess gives the same partition with it (the
  SparseCore gather-reduce is TPU-only, so its preprocessing is tested
  directly); _prefused_ffn_activation gives the split value bit for bit
  and the split VJP value for value, with bit differences only at zeros
  (SwiGLU and clipped GPT-OSS, bf16 and f32, a -0.0 cotangent row).
- On the shrunk Qwen3.5 routed MoE on 8 CPU devices with
  prefuse_moe_weights (cp-as-ep, local routing, two chunks), each key
  alone and all three are bit-identical to the plain program in the
  output, the loss and every gradient: plain and pipelined chunk loop,
  fp8 fake-quant with padding excluded, forced ids with drops and
  dropless, with "count", and with the benchmark's MoE settings
  (expert-weight prefetch, row-major reduce-scatter, routing maps on
  device, ragged_buffer_factor 2.75).
- Structure: fewer concatenates (and fake-quant barriers) with the whole
  wi; exactly two two-key sorts per chunk with the pair argsort, and
  unique_pair_argsort=True on every gather-reduce call (a spy, since only
  the TPU kernel reads it); one more custom_vjp per chunk with the fused
  activation. A fused backward that swaps d(gate) and d(up), and a pair
  argsort with reversed ties under a dropping buffer, are caught.
- On an expert 4 x tensor 2 mesh the layer runs without the key and
  raises at trace time when the key reaches it past the validator.
- Each validator rejects its key for each unmet requirement, naming it.
---
 src/maxtext/configs/base.yml                  |   9 +
 src/maxtext/configs/types.py                  | 160 +++++
 .../kernels/ragged/ragged_gather_reduce_v2.py |  31 +-
 src/maxtext/kernels/ragged/ragged_sort.py     |  20 +-
 src/maxtext/layers/moe.py                     | 121 +++-
 tests/unit/moe_glue_rewrites_test.py          | 579 ++++++++++++++++++
 6 files changed, 908 insertions(+), 12 deletions(-)
 create mode 100644 tests/unit/moe_glue_rewrites_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 98d0514258..de62223012 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -321,6 +321,12 @@ log_moe_drops: false # drop audit: sow moe_dropped_assignments / moe_max_load_ra
 moe_expert_input_dim: -1 # feature dimension of the tokens entering the MoE expert blocks.
 base_moe_mlp_dim: -1 # intermediate dimension at MoE layer. For a fully MoE model, base_mlp_dim must be equal to base_moe_mlp_dim.
 prefuse_moe_weights: false # whether to pre-fuse MoE weights (wi_0 and wi_1 into wi) during initialization
+# With prefuse_moe_weights on the sparse-matmul training path: keep wi = [w0 | w1] one tensor through the fp8
+# fake-quant, the FSDP all-gather and the prefetch barrier into the up-projection GEMM (no split + concatenate).
+moe_unsplit_prefused_wi: false
+# With prefuse_moe_weights on the sparse-matmul path: the activation on the whole [gate | up] up-projection output as
+# one custom-VJP op whose backward writes d[gate | up] in one pass. Same values.
+moe_prefused_swiglu_vjp: false
 fp8_moe: false # whether MoE weights are quantized to float8_e4m3fn in inference/rollout
 # Keep routed-expert weights in weight_dtype but round them in the forward pass to the per-channel FP8 grid
 # fp8_moe stores them on, with straight-through gradients: the trainer computes with the rollout's weights.
@@ -360,6 +366,9 @@ use_ragged_sort: false # whether to use the Pallas ragged-sort kernels in the Mo
 # router replay / padding exclusion). Same permutation, same values.
 # Needs sparse_matmul, use_ring_of_experts and use_ragged_sort (and te_moe_block=False).
 moe_routing_index_math: 'sort'
+# With moe_routing_index_math='sort': the ragged sort's argsorts (and the SparseCore gather-reduce validity sort) as
+# unstable two-operand sorts of unique (key, index) pairs instead of stable argsorts (three operands). Same values.
+moe_ragged_sort_unique_pair_argsort: false
 moe_tc_ragged_sort: false # ring-of-experts ragged sort/unsort on TensorCore (Pallas DMA kernels) instead of the SparseCore
                           # kernels; requires use_ring_of_experts, use_ragged_sort and ragged_buffer_factor > 0.
 moe_tc_ragged_gather_block_size: 1024 # moe_tc_ragged_sort: rows per grid step of the TC ragged gather.
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index e8aed40890..389c6f0271 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1113,6 +1113,22 @@ def _reject_unless(setting: str, requirements: list[tuple[bool, str]]) -> None:
     raise ValueError(f"{setting} requires: {'; '.join(unmet)}.")
 
 
+# The MLP checkpoint names that the remat_policy presets save or offload (Decoder.get_remat_policy, nnx_decoders.py).
+# remat_policy=custom names them per tensor instead (mlpwi_0, mlpwo, moe_mlpwi_0, ... = device / offload).
+_REMAT_PRESET_MLP_NAMES = {
+    "minimal": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
+    "minimal_with_context": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
+    "minimal_flash": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
+    "minimal_with_quantization": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
+    "minimal_with_context_and_quantization": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
+    "minimal_offloaded": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
+    "save_dot_except_mlpwi": ("mlpwo",),
+}
+# The checkpoint names of the routed experts' up-projection output in the sparse-matmul path: the gate half carries
+# mlpwi_0 and moe_mlpwi_0, the up half moe_mlpwi_1 (moe.py gmm_up).
+_MOE_UP_PROJECTION_NAMES = ("mlpwi_0", "moe_mlpwi_0", "moe_mlpwi_1")
+
+
 class MoEGeneral(BaseModel):
   """General configuration for Mixture of Experts (MoE) layers."""
 
@@ -1449,6 +1465,18 @@ class MoEGeneral(BaseModel):
           " use_ring_of_experts=True and use_ragged_sort=True."
       ),
   )
+  moe_ragged_sort_unique_pair_argsort: bool = Field(
+      False,
+      description=(
+          "With moe_routing_index_math='sort': run the ring-of-experts ragged sort's argsorts (the expert"
+          " permutation, its inverse, and the SparseCore gather-reduce's validity sort) as unstable sorts of the"
+          " unique (key, index) pairs instead of stable argsorts. XLA lowers a stable jnp.argsort in the forward"
+          " pass with a third, duplicated iota operand; the pair sort has two operands and exactly one result, the"
+          " same permutation. Exact (same values). An alternative to moe_routing_index_math='count', which removes"
+          " these sorts instead. Requires sparse_matmul=True, te_moe_block=False, use_ring_of_experts=True,"
+          " use_ragged_sort=True and moe_routing_index_math='sort'."
+      ),
+  )
   moe_quantize_combine_bwd_method: str | None = Field(
       "",
       description=(
@@ -1655,6 +1683,35 @@ class MoEGeneral(BaseModel):
       description="Whether to pre-fuse MoE weights (w0 and w1) during initialization. "
       "This enables a single FFN1 grouped GEMM in sparse MoE paths and passes fused weights directly in vllm_rpa mode.",
   )
+  moe_unsplit_prefused_wi: bool = Field(
+      False,
+      description=(
+          "With prefuse_moe_weights on the sparse-matmul training path: keep the prefused wi = [w0 | w1] one tensor"
+          " from the parameter to the up-projection GEMM, instead of slicing it into w0 and w1 at the layer entry and"
+          " concatenating the gathered halves back before the GEMM. One fp8 fake-quant (per output channel, so the"
+          " same values as rounding the halves), one FSDP all-gather, one operand through the expert-weight"
+          " prefetch barrier, and no post-gather concatenate. Exact (same values). Requires sparse_matmul=True,"
+          " te_moe_block=False, no weight quantization, an unsharded expert MLP dimension (tensor, tensor_sequence"
+          " and autoregressive parallelism 1), the default expert-weight sharding (not shard_exp_on_fsdp,"
+          " use_2d_fsdp_sharding, shard_embed_moe_on_fsdp, moe_flat_fsdp_weights or"
+          " moe_fsdp_use_two_stage_all_gather), num_moe_emb_chunks=0 and moe_accumulate_chunk_wgrad=False."
+      ),
+  )
+  moe_prefused_swiglu_vjp: bool = Field(
+      False,
+      description=(
+          "With prefuse_moe_weights on the sparse-matmul path: apply the MoE's gated activation (SwiGLU, or any"
+          " other mlp_activations pair, e.g. GPT-OSS's clipped form) to the whole [gate | up] up-projection output as"
+          " one custom-VJP op whose backward returns d[gate | up] as one array, computed by the same ops autodiff"
+          " emits. Split, the two [rows, mlp] cotangents leave their fusion separately and XLA concatenates them in"
+          " an extra full pass over the expert buffer. Exact (same values; the gradient bits can differ only in the"
+          " sign of an exact zero). An alternative to moe_wi_gmm_fused_swiglu for configurations that one does not"
+          " support (other GMM backends, other activations). Requires sparse_matmul=True, te_moe_block=False,"
+          " mlp_bias=False, num_moe_emb_chunks=0, and a remat policy that does not save the routed experts' gate or"
+          " up half (mlpwi_0, moe_mlpwi_0, moe_mlpwi_1; the minimal* policies save mlpwi_0): the fused activation"
+          " never forms them, so the policy would silently stop saving the MoE tensor."
+      ),
+  )
   fp8_moe: bool = Field(
       False,
       description="Whether MoE weights are quantized to float8_e4m3fn in inference/rollout.",
@@ -1742,6 +1799,109 @@ def validate_moe_routing_index_math(self) -> "MoEGeneral":
       )
     return self
 
+  def _prefused_sparse_matmul_requirements(self) -> list[tuple[bool, str]]:
+    """What the MoE layer needs to run the prefused [w0 | w1] sparse-matmul up projection in training."""
+    return [
+        (self.prefuse_moe_weights, "prefuse_moe_weights=True"),
+        (self.sparse_matmul, "sparse_matmul=True"),
+        (not self.te_moe_block, "te_moe_block=False"),
+        (
+            self.attention not in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"),
+            "an attention other than vllm_rpa* (the fused-MoE inference path)",
+        ),
+    ]
+
+  def _no_weight_quantization_requirement(self) -> tuple[bool, str]:
+    return (not self.quantization and not self.use_qwix_quantization, "no weight quantization")
+
+  def _remat_keeps_moe_tensor_requirement(self, tensor: str, names: tuple[str, ...]) -> tuple[bool, str]:
+    """Requirement: the remat policy saves none of `names`, the checkpoint names of the routed experts' `tensor`.
+
+    The keys that ask for this never form that tensor (it is fused into the op that reads it), so a policy that
+    saves or offloads it would silently stop saving the MoE tensor; the shared expert's dense MLP, which carries the
+    generic names (mlpwi_0, mlpwo) too, would still be saved.
+    """
+    if self.remat_policy == "custom":
+      saved = [name for name in names if getattr(self, name) != RematLocation.REMAT]
+    else:
+      saved = [name for name in names if name in _REMAT_PRESET_MLP_NAMES.get(self.remat_policy, ())]
+    return (
+        not saved,
+        f"a remat policy that does not save the routed experts' {tensor}, which this key never forms"
+        f" (remat_policy={self.remat_policy} saves it as {', '.join(saved)})",
+    )
+
+  def _expert_mlp_dim_shards(self) -> int:
+    """Devices that split the expert weights' MLP dimension inside the expert shard_map, by the default logical axis
+    rules ('mlp_no_fsdp' -> tensor, tensor_sequence, autoregressive)."""
+    degrees = (
+        self.ici_tensor_parallelism,
+        self.dcn_tensor_parallelism,
+        self.ici_tensor_sequence_parallelism,
+        self.dcn_tensor_sequence_parallelism,
+        self.ici_autoregressive_parallelism,
+        self.dcn_autoregressive_parallelism,
+    )
+    return math.prod(degrees)
+
+  @model_validator(mode="after")
+  def validate_moe_ragged_sort_unique_pair_argsort(self) -> "MoEGeneral":
+    """The pair argsort rewrites the ring-of-experts ragged sort's argsorts, which "count" removes."""
+    if self.moe_ragged_sort_unique_pair_argsort:
+      _reject_unless(
+          "moe_ragged_sort_unique_pair_argsort=True",
+          self._ring_of_experts_ragged_sort_requirements()
+          + [
+              (
+                  self.moe_routing_index_math == "sort",
+                  "moe_routing_index_math='sort' (with 'count' there are no argsorts left to rewrite; set one of them)",
+              )
+          ],
+      )
+    return self
+
+  @model_validator(mode="after")
+  def validate_moe_unsplit_prefused_wi(self) -> "MoEGeneral":
+    """The whole wi needs the prefused sparse-matmul path and readers of w0/w1 that all take the [w0 | w1] layout."""
+    if self.moe_unsplit_prefused_wi:
+      _reject_unless(
+          "moe_unsplit_prefused_wi=True",
+          self._prefused_sparse_matmul_requirements()
+          + [
+              self._no_weight_quantization_requirement(),
+              # The default expert-weight sharding: these split w0 and w1 differently from the whole [w0 | w1].
+              (not self.moe_fsdp_use_two_stage_all_gather, "moe_fsdp_use_two_stage_all_gather=False"),
+              (not self.moe_flat_fsdp_weights, "moe_flat_fsdp_weights=False"),
+              (not self.shard_exp_on_fsdp, "shard_exp_on_fsdp=False"),
+              (not self.use_2d_fsdp_sharding, "use_2d_fsdp_sharding=False"),
+              (not self.shard_embed_moe_on_fsdp, "shard_embed_moe_on_fsdp=False"),
+              (not self.moe_accumulate_chunk_wgrad, "moe_accumulate_chunk_wgrad=False"),
+              (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0"),
+              # Sharded, a device would hold a slice of [w0 | w1] instead of matching slices of w0 and w1. The layer
+              # checks the actual partition spec again when it is traced (custom logical axis rules).
+              (
+                  self._expert_mlp_dim_shards() == 1,
+                  "an unsharded expert MLP dimension (tensor, tensor_sequence and autoregressive parallelism 1)",
+              ),
+          ],
+      )
+    return self
+
+  @model_validator(mode="after")
+  def validate_moe_prefused_swiglu_vjp(self) -> "MoEGeneral":
+    """The fused activation needs the whole prefused [gate | up] output, unbiased and unchunked, never saved per half."""
+    if self.moe_prefused_swiglu_vjp:
+      _reject_unless(
+          "moe_prefused_swiglu_vjp=True",
+          self._prefused_sparse_matmul_requirements()
+          + [
+              (not self.mlp_bias, "mlp_bias=False (biases are added per half)"),
+              (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0 (partial sums per half)"),
+              self._remat_keeps_moe_tensor_requirement("gate and up halves", _MOE_UP_PROJECTION_NAMES),
+          ],
+      )
+    return self
+
 
 class MoEKernels(BaseModel):
   """Configuration for MoE-specific kernels like Megablox."""
diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
index 3e94b5b08b..0b8b3fa101 100644
--- a/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
+++ b/src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py
@@ -263,6 +263,25 @@ def _calculate_col_chunk_size(col_size: int, num_simd_lanes: int) -> int:
   return 128
 
 
+def argsort_unique_pairs(keys: jax.Array, axis: int = -1) -> jax.Array:
+  """``jnp.argsort(keys, axis=axis, stable=True)``, computed as an unstable sort of the pairs ``(key, index)``.
+
+  The pairs ``(keys[i], i)`` along ``axis`` are all distinct, so sorting them lexicographically has exactly one
+  result, the stable argsort's permutation: equal keys come out in increasing index order. Marked unstable, the
+  sort keeps two operands (keys, indices) and a two-key comparator. ``jnp.argsort`` emits a stable
+  ``sort_key_val(keys, iota)``, which XLA lowers in the forward pass with a THIRD operand (a second copy of the iota
+  as the tie-breaker). Measured on TPU v7x for an s32[655360] argsort: 0.61-0.63 ms for the three-operand sort,
+  0.42 ms for the two-operand one (moe_ragged_sort_unique_pair_argsort). The index dtype and sharding are the ones
+  ``jnp.argsort`` uses.
+  """
+  dim = axis % keys.ndim
+  # jnp.argsort's index dtype: the default integer type (int32 unless x64 is enabled).
+  idx_dtype = jax.dtypes.canonicalize_dtype(jnp.int64)
+  iota = jax.lax.broadcasted_iota(idx_dtype, keys.shape, dim, out_sharding=jax.typeof(keys).sharding)
+  _, indices = jax.lax.sort((keys, iota), dimension=dim, is_stable=False, num_keys=2)
+  return indices
+
+
 def _preprocess(
     valid_rows_mask: jax.Array,
     reduce_group_size: int,
@@ -270,11 +289,14 @@ def _preprocess(
     num_simd_lanes: int,
     row_chunk_size: int,
     packed: bool = False,
+    unique_pair_argsort: bool = False,
 ) -> tuple[jax.Array, jax.Array, jax.Array]:
   """Sorts valid source rows to the front of each row partition.
 
   With ``packed`` (``moe_routing_index_math="count"``) the stable partition is one single-operand sort of
-  packed uint32 keys instead of a stable argsort; the permutation is the same.
+  packed uint32 keys instead of a stable argsort; the permutation is the same. With ``unique_pair_argsort``
+  (``moe_ragged_sort_unique_pair_argsort``) it is ``argsort_unique_pairs``, the same permutation again.
+  ``packed`` takes precedence.
 
   Returns:
     sorted_by_validity: original row index of each slot after the stable
@@ -296,6 +318,8 @@ def _preprocess(
     key = key | jnp.arange(row_partition_size, dtype=jnp.uint32)[None, :]
     key = jax.lax.sort(key, dimension=1, is_stable=False)
     sorted_by_validity = (key & jnp.uint32((1 << idx_bits) - 1)).astype(jnp.int32)
+  elif unique_pair_argsort:
+    sorted_by_validity = argsort_unique_pairs(~valid_rows_mask_2d, axis=-1)
   else:
     # Stable sort of a boolean key is a stable partition: valid rows keep their
     # relative order and move ahead of the invalid ones.
@@ -652,6 +676,7 @@ def col_loop(col_compute_offset):
         "bytes_accessed_override",
         "use_single_sparsecore",
         "packed_preprocess",
+        "unique_pair_argsort",
     ),
 )
 def ragged_gather_reduce(
@@ -665,6 +690,7 @@ def ragged_gather_reduce(
     bytes_accessed_override: int = -1,
     use_single_sparsecore: bool = False,
     packed_preprocess: bool = False,
+    unique_pair_argsort: bool = False,
 ) -> jax.Array:
   """Gathers ``x`` by ``indices``, weights and masks, then reduces by group.
 
@@ -677,6 +703,8 @@ def ragged_gather_reduce(
     packed_preprocess: sort the valid rows to the front with one packed uint32 sort (``_preprocess(packed=True)``,
       ``moe_routing_index_math="count"``). Same permutation, so the same output. Only the SparseCore path sorts;
       the reference fallback (any non-TPU target) ignores it.
+    unique_pair_argsort: the same validity sort as a two-operand sort of unique (key, index) pairs
+      (``argsort_unique_pairs``, ``moe_ragged_sort_unique_pair_argsort``). Same permutation; SparseCore path only.
 
   Returns:
     Reduced output, ``(input_size // reduce_group_size, hidden_size)``.
@@ -741,6 +769,7 @@ def ragged_gather_reduce(
       num_simd_lanes,
       row_chunk_size,
       packed=packed_preprocess,
+      unique_pair_argsort=unique_pair_argsort,
   )
 
   # Step 4: Launch the SparseCore kernel.
diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index 9c57255ba0..b060d98fc6 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -17,6 +17,7 @@
 import jax
 import jax.numpy as jnp
 from maxtext.kernels.ragged.ragged_gather import ragged_gather
+from maxtext.kernels.ragged.ragged_gather_reduce_v2 import argsort_unique_pairs
 from maxtext.kernels.ragged.ragged_gather_reduce_v2 import ragged_gather_reduce
 from maxtext.kernels.ragged.ragged_gather_tc import packed_stable_sort
 
@@ -101,6 +102,7 @@ def ring_ragged_sort(
     use_single_sparsecore=False,
     tag_routing_fn=None,
     index_math="sort",
+    unique_pair_argsort=False,
 ):
   """Ragged-gather variant for AG-RS Expert Parallelism token routing.
 
@@ -138,6 +140,10 @@ def ring_ragged_sort(
       gather-reduce of the backward packs its validity sort too. Every reader of the inverse masks it to this
       shard's window, so the forward output and the gradients are the same. The unsort of these rows must then
       get ``topk_argsort_indices`` and ``index_math="count"`` (a partial inverse cannot be argsorted back).
+    unique_pair_argsort: with ``index_math="sort"``, run both argsorts, and the validity sort of the backward's
+      gather-reduce, as unstable two-operand sorts of the unique (key, index) pairs (``argsort_unique_pairs``,
+      ``moe_ragged_sort_unique_pair_argsort``) instead of stable argsorts, which XLA lowers with a duplicated iota
+      operand. Same permutations, same values.
 
   Returns:
     A tuple containing:
@@ -154,6 +160,7 @@ def ring_ragged_sort(
   num_tokens_local = hidden_states_local.shape[0]
 
   topk_indices_flat = topk_indices_local.flatten()  # num_tokens_local x topk
+  argsort = argsort_unique_pairs if unique_pair_argsort else jnp.argsort  # the "sort" maps' argsort
   if index_math == "count":
     # The stable order by expert id as ONE single-operand sort of (expert << pos_bits | slot): the packed values
     # are distinct, so the unstable sort has one result, the stable argsort's permutation. jnp.argsort sorts the
@@ -162,7 +169,7 @@ def ring_ragged_sort(
     # Flat slot j belongs to token j // topk, so the sorted token indices need no gather of a repeated iota.
     token_indices_sorted = (topk_argsort_indices // topk).astype(jnp.int32)  # num_tokens_local x topk
   else:
-    topk_argsort_indices = jnp.argsort(topk_indices_flat)  # num_tokens_local x topk
+    topk_argsort_indices = 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
@@ -177,8 +184,9 @@ def ring_ragged_sort(
     shard_start = jnp.sum(jnp.where(jnp.arange(num_experts) < shard_idx * (num_experts // ep_size), group_sizes_local, 0))
     topk_argsort_revert_indices = jnp.where(local_pos >= 0, shard_start + local_pos, -1).astype(jnp.int32)
   else:
-    topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices)  # num_tokens_local x topk
+    topk_argsort_revert_indices = argsort(topk_argsort_indices)  # num_tokens_local x topk
   packed_preprocess = index_math == "count"
+  unique_pair_preprocess = unique_pair_argsort and not packed_preprocess
 
   if tag_routing_fn is not None:
     topk_argsort_indices, token_indices_sorted, group_sizes_local, topk_argsort_revert_indices = (
@@ -296,6 +304,7 @@ def _ring_ragged_sort_bwd(res, g_out):
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
           packed_preprocess=packed_preprocess,
+          unique_pair_argsort=unique_pair_preprocess,
       )
     else:
       # Buffering: g_x has size `local_buffer_size` (packed).
@@ -323,6 +332,7 @@ def _ring_ragged_sort_bwd(res, g_out):
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
           packed_preprocess=packed_preprocess,
+          unique_pair_argsort=unique_pair_preprocess,
       )
     return grad_hidden_states, None, None, None
 
@@ -349,6 +359,7 @@ def ring_ragged_unsort(
     use_single_sparsecore=False,
     topk_argsort_indices=None,
     index_math="sort",
+    unique_pair_argsort=False,
 ):
   """Dual of :func:`ring_ragged_sort`.
 
@@ -383,6 +394,8 @@ def ring_ragged_unsort(
       gather-reduce packs its validity sort, and in the backward the routing-weight gradient goes back to slot
       order by a masked gather through the inverse permutation instead of a scatter (each valid slot reads its
       own row; other shards' and dropped slots read 0). Same values either way.
+    unique_pair_argsort: with ``index_math="sort"``, the forward gather-reduce's validity sort as a two-operand
+      sort of unique pairs (as in :func:`ring_ragged_sort`). Same values.
 
   Returns:
     A 2D ``[num_tokens_local, hidden]`` tensor with expert outputs scattered back
@@ -390,6 +403,7 @@ def ring_ragged_unsort(
   """
   _check_index_math(index_math)
   packed_preprocess = index_math == "count"
+  unique_pair_preprocess = unique_pair_argsort and not packed_preprocess
 
   @jax.custom_vjp
   def _ring_ragged_unsort(
@@ -453,6 +467,7 @@ def _ring_ragged_unsort_fwd(
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
           packed_preprocess=packed_preprocess,
+          unique_pair_argsort=unique_pair_preprocess,
       )
     else:
       # Shift indices so they map to the packed local buffer [0, local_num_tokens).
@@ -473,6 +488,7 @@ def _ring_ragged_unsort_fwd(
           bytes_accessed_override=gather_reduce_bytes_accessed_override,
           use_single_sparsecore=use_single_sparsecore,
           packed_preprocess=packed_preprocess,
+          unique_pair_argsort=unique_pair_preprocess,
       )
 
     res = (
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 1858c91abe..6ae4826001 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -281,6 +281,39 @@ def barrier_tie(x, dep, backward: bool = False):
   return barrier((x, dep))
 
 
+def _prefused_ffn_activation(act_fn, h01):
+  """`act_fn(h01[..., :n], h01[..., n:])` on the prefused up-projection output `h01 = [gate | up]`.
+
+  moe_prefused_swiglu_vjp. The forward is exactly the slice + `act_fn` the split path runs. The backward calls
+  `jax.vjp` of the same `act_fn` at the same primal point, so d(gate) and d(up) come from the ops autodiff emits for
+  it, and returns them as ONE `[rows, 2n]` array, so the concatenate fuses into the elementwise backward. Split, the
+  two `[rows, n]` cotangents leave the SwiGLU-backward fusion separately and XLA joins them in an extra pass over
+  the whole expert buffer (the transpose of the two slices, `add_any(pad, pad)`). The values are the same; the bits
+  can differ only in the sign of an exact zero, since the add of zero pads gives -0.0 + 0.0 = +0.0 where the
+  concatenate keeps -0.0 (measured on XLA:CPU). The residual is `h01` itself, from which the backward recomputes the
+  slices, so no separate gate slice is kept for the backward.
+  """
+  n = h01.shape[-1] // 2
+
+  def split(h):
+    return h[..., :n], h[..., n:]
+
+  @jax.custom_vjp
+  def fused_act(h):
+    return act_fn(*split(h))
+
+  def fused_act_fwd(h):
+    return act_fn(*split(h)), h
+
+  def fused_act_bwd(h, cotangent):
+    _, act_vjp = jax.vjp(act_fn, *split(h))
+    d_gate, d_up = act_vjp(cotangent)
+    return (jnp.concatenate([d_gate, d_up], axis=-1),)
+
+  fused_act.defvjp(fused_act_fwd, fused_act_bwd)
+  return fused_act(h01)
+
+
 @functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
 def _forward_layout_constraint(x, major_to_minor):
   """`with_layout_constraint` on the forward value and the identity on its cotangent."""
@@ -1685,6 +1718,52 @@ def apply_ffn_activation(self, layer_w0, layer_w1):
         intermediate_layer = jnp.multiply(layer_act, layer_w1)
       return intermediate_layer.astype(self.dtype)
 
+  def apply_ffn_activation_maybe_fused(self, layer_w0, layer_w1):
+    """`apply_ffn_activation`; `layer_w1=None` means `layer_w0` is the whole [gate | up] up-projection output.
+
+    gmm_up returns that layout under moe_prefused_swiglu_vjp, and the activation then runs as one custom-VJP op
+    (`_prefused_ffn_activation`) whose backward writes d[gate | up] in one pass.
+    """
+    if layer_w1 is None:
+      return _prefused_ffn_activation(self.apply_ffn_activation, layer_w0)
+    return self.apply_ffn_activation(layer_w0, layer_w1)
+
+  def _keep_prefused_wi_whole(self):
+    """moe_unsplit_prefused_wi: whether the prefused `wi = [w0 | w1]` stays one tensor up to the up-projection GEMM.
+
+    Then `wi` is fake-quantized, FSDP-all-gathered and carried through the expert-weight prefetch barrier as one
+    tensor, and the GEMM reads it directly (`w1_kernel=None` marks the layout), instead of being sliced into w0/w1
+    here and concatenated back after the gather. The config validator admits the key only where every reader of
+    w0/w1 handles that layout, including an expert MLP dimension that the default logical axis rules leave
+    unsharded; this checks what only the built layer knows (quantization objects, sparsity, the partition spec under
+    the actual axis rules), and raises rather than fall back.
+    """
+    if not self.config.moe_unsplit_prefused_wi:
+      return False
+    blockers = []
+    if self.quant is not None:
+      blockers.append("weight quantization")
+    if self.wi_0_sparsity_module is not None:
+      blockers.append("weight sparsity")
+    # The whole [w0 | w1] last dim must be on every device of the expert shard_map: sharded, a device would hold a
+    # slice of [w0 | w1] instead of matching slices of w0 and w1. Same pspec as sparse_matmul's default branch.
+    wi_pspec = remove_mesh_axes_from_partition_spec(self._logical_to_mesh_axes(("exp", None, "mlp_no_fsdp")), ("fsdp",))
+    f_shards = math.prod(self.mesh.shape[a] for a in mesh_axes_for_dim(wi_pspec[2]) if a in self.mesh.shape)
+    if f_shards != 1:
+      blockers.append(f"the MLP dimension sharded {f_shards}-way in the expert shard_map")
+    if blockers:
+      raise ValueError(f"moe_unsplit_prefused_wi=True cannot apply: {', '.join(blockers)}.")
+    return True
+
+  def _use_prefused_swiglu_vjp(self):
+    """moe_prefused_swiglu_vjp: whether gmm_up hands [gate | up] to the activation as one array.
+
+    The per-half outputs are never formed (the fused activation keeps [gate | up] as its residual), so the config
+    validator admits the key only where no remat policy saves them (mlpwi_0 / moe_mlpwi_0 / moe_mlpwi_1), without
+    MLP biases (added per half) and without embedding chunks (partial sums per half).
+    """
+    return self.config.moe_prefused_swiglu_vjp
+
   def permute(
       self,
       inputs,
@@ -1859,6 +1938,7 @@ def permute(
           use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           tag_routing_fn=tag_routing_fn,
           index_math=self.config.moe_routing_index_math,
+          unique_pair_argsort=self.config.moe_ragged_sort_unique_pair_argsort,
       )
       if tag_routing_fn is None and self.config.moe_routing_index_math == "sort":
         # moe_routing_maps=remat: unpermute re-derives the permutation in its backward (existing behavior).
@@ -2042,6 +2122,7 @@ def unpermute(
           use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
           topk_argsort_indices=topk_argsort_indices,
           index_math=self.config.moe_routing_index_math,
+          unique_pair_argsort=self.config.moe_ragged_sort_unique_pair_argsort,
       )
     else:
       unsort_intermediate = _sort_activations(
@@ -2471,7 +2552,12 @@ def sparse_matmul(
       input_ids=None,
       forced_routed_experts=None,
   ):
-    """Perform sparse matrix multiplication of inputs and Experts."""
+    """Perform sparse matrix multiplication of inputs and Experts.
+
+    `w1_kernel=None` (moe_unsplit_prefused_wi) means `w0_kernel` is the whole prefused [w0 | w1] weight.
+    """
+    if w1_kernel is None and not self.config.prefuse_moe_weights:
+      raise ValueError("sparse_matmul got w1_kernel=None (the whole [w0 | w1] layout), which needs prefuse_moe_weights.")
     # log_moe_drops only: the shard_map gets a 5th (moe_audit) output only when requested, so the
     # compiled program is unchanged with the flag off.
     log_moe_drops = getattr(self.config, "log_moe_drops", False)
@@ -3375,9 +3461,19 @@ def gmm_up(
       if not gather_weights:
         wi_gather_axes = []
       if self.config.prefuse_moe_weights:
-        # 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)
+        if w1 is None:
+          # moe_unsplit_prefused_wi: w0 is already the gathered (and fake-quantized) whole [w0 | w1] weight.
+          w_fused = w0
+        else:
+          # 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 partial_accum0 is None and partial_accum1 is None and self._use_prefused_swiglu_vjp():
+          # moe_prefused_swiglu_vjp: the whole [gate | up] output goes to the activation as one array
+          # (apply_ffn_activation_maybe_fused, which layer_w1=None selects), whose backward then writes d[gate | up]
+          # in one pass. There are no per-half outputs to tag; the config validator rejects a remat policy that
+          # would save them, since it would silently stop saving the MoE tensor.
+          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:
@@ -3691,7 +3787,7 @@ def _down_projection(
         w_rows = tc_buffer_row_weights(routing.tc_routing, jnp.ravel(routing.weights), output0.shape[0])
         if self.config.moe_tc_routing_checkpoint != "remat":
           w_rows = adc.checkpoint_name(w_rows, "moe_tc_routing_checkpoint")
-      intermediate_layer = self.apply_ffn_activation(output0, output1)
+      intermediate_layer = self.apply_ffn_activation_maybe_fused(output0, output1)
       if tc_prescaled:
         intermediate_layer = (intermediate_layer * w_rows[:, None]).astype(intermediate_layer.dtype)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
@@ -3875,7 +3971,7 @@ def _cn(t):
         w_rows = tc_buffer_row_weights(routing.tc_routing, jnp.ravel(routing.weights), output0.shape[0])
         if self.config.moe_tc_routing_checkpoint != "remat":
           w_rows = adc.checkpoint_name(w_rows, "moe_tc_routing_checkpoint")
-      intermediate_layer = self.apply_ffn_activation(output0, output1)
+      intermediate_layer = self.apply_ffn_activation_maybe_fused(output0, output1)
       if tc_prescaled:
         intermediate_layer = (intermediate_layer * w_rows[:, None]).astype(intermediate_layer.dtype)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
@@ -5640,9 +5736,15 @@ def _maybe_native_fused_weight(kernel, scale):
     elif cfg.prefuse_moe_weights:
       wi_scale = self.wi_scale[...] if self.wi_scale is not None else None
       wi = quantizations.dequantize_weight(self.wi[...], wi_scale, self.dtype)
-      n = wi.shape[-1] // 2
-      w0_kernel = wi[..., :n]
-      w1_kernel = wi[..., n:]
+      if self._keep_prefused_wi_whole():
+        # moe_unsplit_prefused_wi: w0_kernel carries the whole [w0 | w1] and w1_kernel=None marks the layout, so
+        # the fake-quant, the FSDP all-gather and the prefetch barrier see one tensor and the up-projection reads it
+        # without the concatenate that would otherwise put the gathered halves back together.
+        w0_kernel, w1_kernel = wi, None
+      else:
+        n = wi.shape[-1] // 2
+        w0_kernel = wi[..., :n]
+        w1_kernel = wi[..., n:]
     else:
       wi_0_scale = self.wi_0_scale[...] if self.wi_0_scale is not None else None
       wi_1_scale = self.wi_1_scale[...] if self.wi_1_scale is not None else None
@@ -5659,7 +5761,8 @@ def _maybe_native_fused_weight(kernel, scale):
       # Axis 2 is the output dim of wi_0/wi_1 (mlp) and wo (embed). An fp8_moe rollout quantizes the
       # fused wi = [wi_0 | wi_1] per output channel, so rounding the halves separately is identical.
       w0_kernel = quantizations.fake_quantize_weight_per_channel(w0_kernel)
-      w1_kernel = quantizations.fake_quantize_weight_per_channel(w1_kernel)
+      if w1_kernel is not None:  # None: moe_unsplit_prefused_wi, and the line above rounded both halves at once
+        w1_kernel = quantizations.fake_quantize_weight_per_channel(w1_kernel)
       wo_kernel = quantizations.fake_quantize_weight_per_channel(wo_kernel)
 
     # For fused MoE path (inference only), if we have not fused expert
diff --git a/tests/unit/moe_glue_rewrites_test.py b/tests/unit/moe_glue_rewrites_test.py
new file mode 100644
index 0000000000..ae98a6c562
--- /dev/null
+++ b/tests/unit/moe_glue_rewrites_test.py
@@ -0,0 +1,579 @@
+# 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.
+"""Three exact rewrites of MoE glue ops, each against the plain program, on CPU.
+
+  - moe_unsplit_prefused_wi: the prefused wi = [w0 | w1] stays one tensor through the fp8 fake-quant, the FSDP
+    all-gather and the prefetch barrier into the up-projection GEMM (no split at the layer entry, no concatenate).
+  - moe_ragged_sort_unique_pair_argsort: the ring-of-experts ragged sort's argsorts as unstable two-operand sorts of
+    unique (key, index) pairs (`argsort_unique_pairs`) instead of stable argsorts.
+  - moe_prefused_swiglu_vjp: the activation on the whole [gate | up] output as one custom-VJP op
+    (`moe._prefused_ffn_activation`) whose backward writes d[gate | up] as one array.
+
+None changes a value, so everything must be bit-identical.
+
+What a CPU run can and cannot reach: the SparseCore gather-reduce (and with it the validity sort that
+moe_ragged_sort_unique_pair_argsort also rewrites) runs only on a TPU target; on CPU `ragged_gather_reduce` takes its
+jnp reference fallback, which does not sort. `PreprocessTest` therefore calls `_preprocess` directly, the pure-JAX
+part of that path. The ragged sort's own argsorts, the weight layout and the activation are plain JAX and run in the
+layer below as on TPU (the GMM is megablox in interpret mode; tokamax gmm v2 is TPU-only).
+
+In-process checks: `argsort_unique_pairs` equals the stable argsort (ties, booleans, every axis, the index dtype);
+`_preprocess` with and without it agree; `_prefused_ffn_activation` gives the split activation's value bit for bit and
+its VJP value for value (SwiGLU and the clipped GPT-OSS form, bfloat16 and float32). A -0.0 cotangent is where the
+two can differ in bits: XLA:CPU turns it into +0.0 on the split path (an add of zero pads) and the fused concatenate
+keeps it, so the gradient bits may differ only in the sign of exact zeros, which the test pins down. The validators
+reject each key for each requirement it does not meet, naming that requirement; that includes moe_unsplit_prefused_wi
+with tensor parallelism, which shards the expert MLP dimension.
+
+The layer check runs as a subprocess with 8 CPU devices: the shrunk Qwen3.5 routed MoE of
+`moe_roe_local_routing_test.py` with prefuse_moe_weights (16 experts, top 4, cp-as-ep, ragged sort, local routing,
+two token chunks). Pre-registered: each key alone and all three together give an output, a loss and every gradient
+bit-identical to all three off, in the plain and the pipelined chunk loop, with fp8 fake-quant and padding excluded
+from routing, with forced ids and a dropless buffer, and with moe_routing_index_math="count". Two-sided:
+  - Structure, so a switch that silently does nothing, or does only part of its job, cannot pass:
+    moe_unsplit_prefused_wi traces fewer concatenates (and with fake-quant fewer barriers);
+    moe_ragged_sort_unique_pair_argsort traces exactly two two-key sorts per token chunk (the permutation and its
+    inverse) and the plain program none, and every ragged_gather_reduce call gets unique_pair_argsort=True (a spy:
+    CPU runs the reference fallback, which does not sort); moe_prefused_swiglu_vjp traces one more custom_vjp call
+    per token chunk.
+  - The trace-time guard behind the validator: on a mesh with tensor parallelism 2 (default axis rules, expert 4 x
+    tensor 2), a config that claims moe_unsplit_prefused_wi past the validator must make the layer raise rather
+    than compute with a sharded [w0 | w1] (it would give the wrong output), while the plain layer runs.
+  - The bench's MoE settings (expert-weight prefetch, row-major reduce-scatter, routing maps saved on device,
+    ragged_buffer_factor 2.75, fp8 fake-quant, padding excluded) with all three keys in one case.
+  - Deliberately wrong versions must FAIL: a fused activation whose backward swaps d(gate) and d(up) (the output
+    is unchanged, the gradients are not), and a pair argsort that breaks ties by decreasing index, under a buffer
+    that drops tokens (the tie order decides which tokens a full buffer drops, so the output changes).
+"""
+
+import contextlib
+import os
+import re
+import subprocess
+import sys
+from unittest import mock
+
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+import numpy as np
+import pytest
+
+from maxtext.configs import pyconfig
+from maxtext.kernels.ragged import ragged_gather_reduce_v2
+from maxtext.kernels.ragged import ragged_sort
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+_NUM_DEVICES = 8
+_NUM_EXPERTS = 16
+_TOP_K = 4
+_EMB = 256
+_SEQ = 64
+_PASS_TOKEN = "MOE_GLUE_REWRITES_CHECKS_PASSED"
+_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false"
+_KEYS = ("moe_unsplit_prefused_wi", "moe_ragged_sort_unique_pair_argsort", "moe_prefused_swiglu_vjp")
+
+
+@pytest.mark.cpu_only
+class ArgsortUniquePairsTest(parameterized.TestCase):
+  """`argsort_unique_pairs` is the stable argsort."""
+
+  @parameterized.named_parameters(
+      ("int_ties_1d", (257,), jnp.int32, -1),
+      ("int_ties_2d_last", (5, 96), jnp.int32, -1),
+      ("int_ties_2d_first", (96, 5), jnp.int32, 0),
+      ("bool_2d", (8, 100), jnp.bool_, 1),
+  )
+  def test_matches_stable_argsort(self, shape, dtype, axis):
+    keys = jax.random.randint(jax.random.PRNGKey(0), shape, 0, 7).astype(dtype)  # many ties
+    want = jnp.argsort(keys, axis=axis, stable=True)
+    got = ragged_gather_reduce_v2.argsort_unique_pairs(keys, axis=axis)
+    self.assertEqual(got.dtype, want.dtype)
+    np.testing.assert_array_equal(np.asarray(got), np.asarray(want))
+
+  def test_wrong_tie_order_is_detected(self):
+    """Two-sided: the same comparison fails for a pair sort that breaks ties by decreasing index."""
+    keys = jnp.array([1, 0, 1, 0, 1], jnp.int32)
+    wrong = _argsort_ties_reversed(keys)
+    self.assertFalse(np.array_equal(np.asarray(wrong), np.asarray(jnp.argsort(keys, stable=True))))
+
+
+@pytest.mark.cpu_only
+class PreprocessTest(parameterized.TestCase):
+  """The SparseCore gather-reduce's validity sort with and without the pair argsort.
+
+  Only the SparseCore path runs `_preprocess`; on a CPU target `ragged_gather_reduce` takes its reference fallback,
+  so this pure-JAX function is called directly.
+  """
+
+  @parameterized.named_parameters(("one_partition", 96, 4, 1, 16, 16), ("v7x_like", 1024, 8, 16, 16, 48))
+  def test_unique_pair_argsort_matches_argsort(self, rows, group, partitions, lanes, chunk):
+    for density in (0.0, 0.4, 1.0):
+      mask = jax.random.uniform(jax.random.PRNGKey(rows), (rows,)) < density
+      want = ragged_gather_reduce_v2._preprocess(mask, group, partitions, lanes, chunk)  # pylint: disable=protected-access
+      got = ragged_gather_reduce_v2._preprocess(  # pylint: disable=protected-access
+          mask, group, partitions, lanes, chunk, unique_pair_argsort=True
+      )
+      for a, b in zip(want, got):
+        self.assertEqual(a.dtype, b.dtype)
+        np.testing.assert_array_equal(np.asarray(b), np.asarray(a))
+
+
+def _swiglu(gate, up):
+  return (jax.nn.silu(gate) * up).astype(gate.dtype)
+
+
+def _gpt_oss(gate, up, limit=7.0):
+  gate = jnp.clip(gate, max=limit)
+  up = jnp.clip(up, -limit, limit)
+  return (gate * jax.nn.sigmoid(gate * 1.702) * (up + 1)).astype(gate.dtype)
+
+
+@pytest.mark.cpu_only
+class PrefusedFfnActivationTest(parameterized.TestCase):
+  """`_prefused_ffn_activation` gives the split activation's value bit for bit and its VJP value for value."""
+
+  @parameterized.product(act=("swiglu", "gpt_oss"), dtype=("bfloat16", "float32"))
+  def test_value_and_vjp_match_split(self, act, dtype):
+    act_fn = {"swiglu": _swiglu, "gpt_oss": _gpt_oss}[act]
+    h01 = (jax.random.normal(jax.random.PRNGKey(1), (64, 2 * 48), jnp.float32) * 4).astype(dtype)
+    cot = jax.random.normal(jax.random.PRNGKey(2), (64, 48), jnp.float32).astype(dtype)
+    cot = cot.at[0].set(-0.0)  # -0.0 cotangents must come back with the sign the split path gives them
+
+    def split(h):
+      n = h.shape[-1] // 2
+      return act_fn(h[:, :n], h[:, n:])
+
+    @jax.jit
+    def both(h, c):
+      y_split, vjp_split = jax.vjp(split, h)
+      y_fused, vjp_fused = jax.vjp(lambda h_: moe._prefused_ffn_activation(act_fn, h_), h)  # pylint: disable=protected-access
+      return y_split, vjp_split(c)[0], y_fused, vjp_fused(c)[0]
+
+    y_split, d_split, y_fused, d_fused = both(h01, cot)
+    np.testing.assert_array_equal(np.asarray(y_fused).view(np.uint8), np.asarray(y_split).view(np.uint8))
+    d_split, d_fused = np.asarray(d_split), np.asarray(d_fused)
+    # Equal values everywhere. Bits may differ only in the sign of an exact zero: the split path's gradient is the
+    # transpose of two slices, add(pad(d_gate), pad(d_up)), and XLA:CPU computes -0.0 + 0.0 = +0.0 there, while the
+    # fused backward's concatenate keeps -0.0.
+    np.testing.assert_array_equal(d_fused, d_split)
+    bits_differ = d_fused.view(np.uint16 if d_fused.itemsize == 2 else np.uint32) != d_split.view(
+        np.uint16 if d_split.itemsize == 2 else np.uint32
+    )
+    # Zeros come from the -0.0 cotangent row and, in the clipped GPT-OSS form, from clipped entries (derivative 0).
+    self.assertTrue(np.all(d_fused[bits_differ] == 0) and np.all(d_split[bits_differ] == 0))
+
+
+# A mesh with tensor parallelism: the default logical axis rules (cp-as-ep has no tensor axis), expert 4 x tensor 2.
+_TP_MESH = {
+    "custom_mesh_and_rule": "",
+    "ici_context_parallelism": 1,
+    "ici_expert_parallelism": 4,
+    "ici_tensor_parallelism": 2,
+    "per_device_batch_size": 0.5,
+}
+_OFF_PREFUSED_PATH = (
+    ("prefuse", {"prefuse_moe_weights": False}, "prefuse_moe_weights=True"),
+    ("dense", {"sparse_matmul": False}, "sparse_matmul=True"),
+    ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"),
+    ("vllm_rpa", {"attention": "vllm_rpa"}, "an attention other than vllm_rpa*"),
+)
+_OFF_RING_RAGGED_SORT = (
+    ("dense", {"sparse_matmul": False}, "sparse_matmul=True"),
+    ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"),
+    (
+        "no_ring",
+        {"use_ring_of_experts": False, "num_moe_token_chunks": 1, "ring_of_experts_local_routing": False},
+        "use_ring_of_experts=True",
+    ),
+    ("no_ragged_sort", {"use_ragged_sort": False}, "use_ragged_sort=True"),
+)
+_REJECTED = (
+    [(f"pair_argsort_{name}", "moe_ragged_sort_unique_pair_argsort", o, m) for name, o, m in _OFF_RING_RAGGED_SORT]
+    + [
+        (
+            "pair_argsort_with_count",
+            "moe_ragged_sort_unique_pair_argsort",
+            {"moe_routing_index_math": "count"},
+            "moe_routing_index_math='sort'",
+        )
+    ]
+    + [(f"unsplit_wi_{name}", "moe_unsplit_prefused_wi", o, m) for name, o, m in _OFF_PREFUSED_PATH]
+    + [
+        (f"unsplit_wi_{name}", "moe_unsplit_prefused_wi", o, m)
+        for name, o, m in (
+            ("quantization", {"quantization": "int8"}, "no weight quantization"),
+            ("qwix", {"use_qwix_quantization": True}, "no weight quantization"),
+            (
+                "two_stage_ag",
+                {"moe_fsdp_use_two_stage_all_gather": True},
+                "moe_fsdp_use_two_stage_all_gather=False",
+            ),
+            ("flat_fsdp", {"moe_flat_fsdp_weights": True}, "moe_flat_fsdp_weights=False"),
+            ("shard_exp_on_fsdp", {"shard_exp_on_fsdp": True}, "shard_exp_on_fsdp=False"),
+            ("2d_fsdp", {"use_2d_fsdp_sharding": True}, "use_2d_fsdp_sharding=False"),
+            ("shard_embed", {"shard_embed_moe_on_fsdp": True}, "shard_embed_moe_on_fsdp=False"),
+            ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"),
+            ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"),
+            ("tensor_parallel", _TP_MESH, "an unsharded expert MLP dimension"),
+        )
+    ]
+    + [(f"swiglu_vjp_{name}", "moe_prefused_swiglu_vjp", o, m) for name, o, m in _OFF_PREFUSED_PATH]
+    + [
+        (f"swiglu_vjp_{name}", "moe_prefused_swiglu_vjp", o, m)
+        for name, o, m in (
+            ("mlp_bias", {"mlp_bias": True}, "mlp_bias=False"),
+            ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"),
+            ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwi_0"),
+            ("offloaded_remat", {"remat_policy": "minimal_offloaded"}, "remat_policy=minimal_offloaded saves it as"),
+            ("saved_gate", {"remat_policy": "custom", "mlpwi_0": "device"}, "remat_policy=custom saves it as mlpwi_0"),
+            ("saved_moe_gate", {"remat_policy": "custom", "moe_mlpwi_0": "device"}, "saves it as moe_mlpwi_0"),
+            ("offloaded_moe_up", {"remat_policy": "custom", "moe_mlpwi_1": "offload"}, "saves it as moe_mlpwi_1"),
+        )
+    ]
+)
+
+
+class GlueRewritesConfigTest(parameterized.TestCase):
+  """Each key is rejected for each requirement it does not meet, and accepted on the path it rewrites."""
+
+  @parameterized.named_parameters(*_REJECTED)
+  def test_rejected(self, key, overrides, requirement):
+    overrides = {"prefuse_moe_weights": True, **overrides}
+    with self.assertRaisesRegex(ValueError, f"{key}=True requires: .*{re.escape(requirement)}"):
+      _cfg(**{key: True}, **overrides)
+
+  def test_accepted(self):
+    cfg = _cfg(prefuse_moe_weights=True, remat_policy="custom", **{k: True for k in _KEYS})
+    self.assertTrue(all(getattr(cfg, k) for k in _KEYS))
+
+  def test_accepted_where_the_policy_saves_only_other_tensors(self):
+    """Policies that save no routed-expert gate/up tensor: full, a dense-only name, the shared expert's mlpwi_1."""
+    for overrides in ({"remat_policy": "full"}, {"remat_policy": "custom", "mlpwi_1": "device", "mlpwo": "device"}):
+      self.assertTrue(_cfg(prefuse_moe_weights=True, moe_prefused_swiglu_vjp=True, **overrides).moe_prefused_swiglu_vjp)
+
+
+def _cfg(**overrides):
+  """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing)."""
+  kwargs = {
+      "run_name": "moe_glue_rewrites",
+      "enable_checkpointing": False,
+      "log_config": False,
+      "skip_jax_distributed_system": True,
+      "model_name": "qwen3.5-397b-a17b",
+      "override_model_config": True,
+      "num_experts": _NUM_EXPERTS,
+      "num_experts_per_tok": _TOP_K,
+      "base_emb_dim": _EMB,
+      "base_mlp_dim": 256,
+      "base_moe_mlp_dim": 256,
+      "max_target_length": _SEQ,
+      "max_prefill_predict_length": _SEQ,
+      "per_device_batch_size": 0.25,
+      "dtype": "bfloat16",
+      "weight_dtype": "float32",
+      "float32_gate_logits": True,
+      "megablox": True,
+      "sparse_matmul": True,
+      # Tokamax gmm v2 is TPU-only; both programs share whichever GMM runs, so CPU uses megablox.
+      "use_tokamax_gmm": False,
+      "use_gmm_v2": False,
+      "use_ring_of_experts": True,
+      "use_ragged_sort": True,
+      "use_custom_sort_vjp": False,
+      "ragged_buffer_factor": 2.0,
+      "num_moe_token_chunks": 2,
+      "custom_mesh_and_rule": "cp-as-ep",
+      "ici_fsdp_parallelism": 1,
+      "ici_context_parallelism": 4,
+      "ici_expert_parallelism": 2,
+      "context_parallel_load_balance": False,
+      "ring_of_experts_local_routing": True,
+  }
+  kwargs.update(overrides)
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
+
+
+def _argsort_ties_reversed(keys, axis=-1):
+  """Deliberately wrong: a two-operand pair sort that breaks ties by DEcreasing index."""
+  dim = axis % keys.ndim
+  iota = jax.lax.broadcasted_iota(jnp.int32, keys.shape, dim)
+  _, rev = jax.lax.sort((keys, -iota), dimension=dim, is_stable=False, num_keys=2)
+  return -rev
+
+
+@pytest.mark.cpu_only
+def test_glue_rewrites_match_plain_program_on_cpu_mesh():
+  """Runs `main` below in a child process with 8 CPU devices."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}"
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+  print(result.stdout)
+  assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+def _build(cfg, mesh):
+  return moe.RoutedMoE(
+      config=cfg,
+      num_experts=cfg.num_experts,
+      num_experts_per_tok=cfg.num_experts_per_tok,
+      mesh=mesh,
+      kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+      kernel_axes=("embed", "mlp"),
+      intermediate_dim=cfg.moe_mlp_dim,
+      weight_dtype=cfg.weight_dtype,
+      dtype=cfg.dtype,
+      rngs=nnx.Rngs(params=0),
+  )
+
+
+def _inputs(cfg, mode):
+  """The layer input, forced expert ids (padded and duplicate slots) and a token mask, per `mode`."""
+  batch = int(cfg.per_device_batch_size * _NUM_DEVICES)
+  x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32)
+  forced = token_mask = None
+  if "forced" in mode:
+    keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ)
+    ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K)
+    pos = jnp.arange(_SEQ)[None, :, None]
+    ids = jnp.where((pos % 5 == 0) & (jnp.arange(_TOP_K) == _TOP_K - 1), -1, ids)
+    ids = jnp.where((pos % 7 == 3) & (jnp.arange(_TOP_K) == 1), ids[..., :1], ids)
+    forced = ids.astype(jnp.int32)
+  if "padmask" in mode:
+    lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)])
+    token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32)
+  return x, forced, token_mask
+
+
+def _eqns(jaxpr):
+  """Every equation of a jaxpr, including those of nested jaxprs."""
+  for eqn in jaxpr.eqns:
+    yield eqn
+    for param in eqn.params.values():
+      for sub in param if isinstance(param, (list, tuple)) else (param,):
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        if hasattr(sub, "eqns"):
+          yield from _eqns(sub)
+
+
+def _structure(jaxpr):
+  eqns = list(_eqns(jaxpr))
+  return {
+      "concatenate": sum(e.primitive.name == "concatenate" for e in eqns),
+      "barrier": sum(e.primitive.name == "optimization_barrier" for e in eqns),
+      "pair_sort": sum(e.primitive.name == "sort" and e.params.get("num_keys") == 2 for e in eqns),
+      "custom_vjp": sum(e.primitive.name.startswith("custom_vjp_call") for e in eqns),
+  }
+
+
+def _run(cfg, mesh, params, x, forced, token_mask):
+  """Loss, output, gradients, the structure of the traced forward, and the unique_pair_argsort flag of every
+  ragged_gather_reduce call traced."""
+  pair_flags = []
+  real_gather_reduce = ragged_sort.ragged_gather_reduce
+
+  def spy(*args, **kwargs):
+    pair_flags.append(kwargs.get("unique_pair_argsort", False))
+    return real_gather_reduce(*args, **kwargs)
+
+  with (
+      mock.patch.object(ragged_sort, "ragged_gather_reduce", spy),
+      jax.set_mesh(mesh),
+      nn_partitioning.axis_rules(cfg.logical_axis_rules),
+  ):
+    graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+
+    def loss_fn(params, x):
+      out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask)
+      loss = jnp.mean(out.astype(jnp.float32) ** 2)
+      if lb_loss is not None:
+        loss = loss + lb_loss.astype(jnp.float32)
+      return loss, out
+
+    structure = _structure(jax.make_jaxpr(loss_fn)(params, x).jaxpr)
+    (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x)
+  structure["gather_reduce_pair_flags"] = sorted(set(pair_flags))
+  return jax.device_get((loss, out, grads)), structure
+
+
+def _bit_identical(label, old, new):
+  """Prints the per-leaf report; True iff the output, the loss and every gradient are bit-identical."""
+  (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new
+  named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [
+      (jax.tree_util.keystr(path), a, b)
+      for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new))
+  ]
+  passed, lines = True, []
+  for name, a, b in named:
+    a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32)
+    same = bool(np.array_equal(a, b)) and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0
+    passed &= same
+    lines.append(f"  [{'ok' if same else 'FAIL'}] {name}: bit-identical={same} max|d|={np.max(np.abs(a - b)):.3e}")
+  print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+def _case(label, keys, overrides, mode="real", control=None):
+  """Runs the plain program and the one with `keys` on, on the same parameters; `control` wraps the second run."""
+  overrides = {"prefuse_moe_weights": True, **overrides}
+  cfg_old = _cfg(**overrides)
+  cfg_new = _cfg(**overrides, **{k: True for k in keys})
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes)
+  x, forced, token_mask = _inputs(cfg_old, mode)
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...)
+  old, s_old = _run(cfg_old, mesh, params, x, forced, token_mask)
+  with control if control is not None else contextlib.nullcontext():
+    new, s_new = _run(cfg_new, mesh, params, x, forced, token_mask)
+  n_chunks = cfg_new.num_moe_token_chunks
+  want = {}
+  if "moe_unsplit_prefused_wi" in keys:
+    want["concatenate"] = s_new["concatenate"] < s_old["concatenate"]
+    if cfg_new.fp8_moe_fake_quant:
+      want["barrier"] = s_new["barrier"] < s_old["barrier"]
+  if "moe_ragged_sort_unique_pair_argsort" in keys:
+    # Both argsorts of every chunk's ragged sort (the forward traced here), and the backward's gather-reduce flag.
+    want["pair_sort"] = s_new["pair_sort"] == 2 * n_chunks and s_old["pair_sort"] == 0
+    flags = (s_new["gather_reduce_pair_flags"], s_old["gather_reduce_pair_flags"])
+    want["gather_reduce_pair_flags"] = flags == ([True], [False])
+  if "moe_prefused_swiglu_vjp" in keys:
+    want["custom_vjp"] = s_new["custom_vjp"] == s_old["custom_vjp"] + n_chunks
+  structure = all(want.values())
+  print(f"{label}: [{'ok' if structure else 'FAIL'}] structure plain={s_old} new={s_new} checks={want}", flush=True)
+  return _bit_identical(label, old, new) and structure
+
+
+def _swapped_activation_cotangents():
+  """Deliberately wrong: the fused activation's backward returns [d(up) | d(gate)]."""
+  real = moe._prefused_ffn_activation  # pylint: disable=protected-access
+
+  def wrong(act_fn, h01):
+    n = h01.shape[-1] // 2
+    swapped = jnp.concatenate([h01[..., n:], h01[..., :n]], axis=-1)
+    # Same forward value (act_fn reads the halves in the right order), gradient of the halves exchanged.
+    return real(lambda up, gate: act_fn(gate, up), jax.lax.stop_gradient(swapped) + (h01 - jax.lax.stop_gradient(h01)))
+
+  return mock.patch.object(moe, "_prefused_ffn_activation", wrong)
+
+
+def _ties_reversed_pair_argsort():
+  return mock.patch.object(ragged_sort, "argsort_unique_pairs", _argsort_ties_reversed)
+
+
+class _ConfigOverride:
+  """`config` with some attributes replaced: reaches the layer's trace-time checks past the config validator."""
+
+  def __init__(self, config, **overrides):
+    self._config = config
+    self._overrides = overrides
+
+  def __getattr__(self, name):  # only called for names that are not set on the instance
+    if name.startswith("__") or "_config" not in self.__dict__:
+      raise AttributeError(name)
+    return self._overrides[name] if name in self._overrides else getattr(self._config, name)
+
+
+def _unsplit_wi_tensor_parallel_guard():
+  """On a tensor-parallel mesh the layer refuses moe_unsplit_prefused_wi at trace time; the plain layer runs."""
+  cfg = _cfg(prefuse_moe_weights=True, **_TP_MESH)
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+  x, forced, token_mask = _inputs(cfg, "real")
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+  (loss, _, _), _ = _run(cfg, mesh, params, x, forced, token_mask)
+  plain_runs = bool(np.isfinite(loss))
+  try:
+    _run(_ConfigOverride(cfg, moe_unsplit_prefused_wi=True), mesh, params, x, forced, token_mask)
+  except ValueError as e:
+    raised = "the MLP dimension sharded 2-way in the expert shard_map" in str(e)
+  else:
+    raised = False
+  ok = plain_runs and raised
+  print(
+      f"unsplit_wi_tensor_parallel_guard: [{'ok' if ok else 'FAIL'}] mesh {dict(mesh.shape)}: plain layer runs:"
+      f" {plain_runs}; moe_unsplit_prefused_wi past the validator raises at trace time: {raised}",
+      flush=True,
+  )
+  return ok
+
+
+# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+_BENCH_FLAGS = {
+    "moe_chunk_pipeline": True,
+    "moe_expert_weight_prefetch": "forward",
+    "ring_of_experts_row_major_reduce_scatter": True,
+    "moe_routing_maps": "device",
+    "ragged_buffer_factor": 2.75,
+    "fp8_moe_fake_quant": True,
+    "moe_exclude_padding_from_routing": True,
+}
+
+
+def main():
+  assert jax.device_count() == _NUM_DEVICES, jax.devices()
+  wi, pairs, swiglu = _KEYS
+  pipelined = {"moe_chunk_pipeline": True}
+  fake_quant = {**pipelined, "fp8_moe_fake_quant": True, "moe_exclude_padding_from_routing": True}
+  results = {
+      "unsplit_wi_plain_loop": _case("unsplit_wi_plain_loop", (wi,), {}),
+      "unsplit_wi_fake_quant": _case("unsplit_wi_fake_quant", (wi,), fake_quant, mode="padmask"),
+      "pair_argsort_pipelined": _case("pair_argsort_pipelined", (pairs,), pipelined),
+      "pair_argsort_dropping": _case(
+          "pair_argsort_dropping", (pairs,), {**pipelined, "ragged_buffer_factor": 1.0}, mode="forced"
+      ),
+      "swiglu_vjp_pipelined": _case("swiglu_vjp_pipelined", (swiglu,), pipelined),
+      "all_three_plain_loop": _case("all_three_plain_loop", _KEYS, {}),
+      "all_three_fake_quant": _case("all_three_fake_quant", _KEYS, fake_quant, mode="padmask"),
+      "all_three_dropless": _case("all_three_dropless", _KEYS, {"ragged_buffer_factor": -1.0}, mode="forced"),
+      "wi_and_swiglu_with_count": _case(
+          "wi_and_swiglu_with_count", (wi, swiglu), {**pipelined, "moe_routing_index_math": "count"}
+      ),
+      "all_three_bench_flags": _case("all_three_bench_flags", _KEYS, _BENCH_FLAGS, mode="padmask"),
+      "unsplit_wi_tensor_parallel_guard": _unsplit_wi_tensor_parallel_guard(),
+  }
+  controls = {
+      "control_swapped_activation_cotangents": _case(
+          "control_swapped_activation_cotangents", (swiglu,), pipelined, control=_swapped_activation_cotangents()
+      ),
+      "control_ties_reversed_under_drops": _case(
+          "control_ties_reversed_under_drops",
+          (pairs,),
+          {**pipelined, "ragged_buffer_factor": 1.0},
+          mode="forced",
+          control=_ties_reversed_pair_argsort(),
+      ),
+  }
+  print("SUMMARY", {**results, **controls}, flush=True)
+  failed = [k for k, v in results.items() if not v]
+  undetected = [k for k, v in controls.items() if v]
+  if failed or undetected:
+    print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True)
+    sys.exit(1)
+  print(_PASS_TOKEN, flush=True)
+
+
+if __name__ == "__main__":
+  main()

From c88674bbca16e96896934f47e948940aac4f573f Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 02:42:25 -0700
Subject: [PATCH 05/10] MoE: one packed routing all-gather per layer in the
 pipelined ring of experts (moe_roe_packed_routing_gather)

With ring_of_experts_local_routing and moe_chunk_pipeline, every token
chunk routes its own tokens and all-gathers its top-k weights and expert
ids over the expert-parallel axis: two all-gathers, two SparseCore
offloads, per chunk, and in the backward one weight-cotangent
reduce-scatter per chunk.

moe_roe_packed_routing_gather routes all of the layer's local tokens once
and all-gathers weights and ids as ONE float32 [batch, seq, 2 * top_k]
array (pack_routing_for_all_gather: ids are < num_experts or -1, exact in
float32), which each chunk slices: one gather per layer and one
reduce-scatter in the backward. Routing is per token, so each slice holds
what the chunk would have computed: bit-identical with float32 routing
weights (float32_weight_sum, the default). With bfloat16 weights the
values round-trip exactly and their cotangent is reduce-scattered in
float32 instead of bfloat16. The all-gather is the ring of experts' own
(_ep_all_gather, now defined once for roe_ag_and_route and this gather),
so a SparseCore pin applies to both alike.

The per-chunk load-balance loss is computed as without the key, from this
shard's own ids of the chunk, so the loss and the output are
bit-identical; the router gradient then agrees to float32 reduction order
(the top-k and the loss terms reach the logits in another order;
1.2e-7 relative on CPU). The gathered array is a routing map: tagged
moe_routing_maps, so moe_routing_maps=device saves it and the
rematerialized forward skips the gather; untagged under remat.

The validator (the shared form) requires sparse_matmul, te_moe_block=False,
ring_of_experts_local_routing and moe_chunk_pipeline (the packed gather
replaces the pipelined loop's per-chunk gathers only), and rejects
use_random_routing (random routing draws per call, so one draw for the
layer differs from one per chunk). If XLA strips optimization barriers
the pipeline is skipped with a warning, and the packed gather too.

Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step,
load_balance_loss_weight 0, moe_routing_maps=device), switched by an
environment variable in a build that a TPU lowering check shows compiling
to the same program as this commit: regular micro-batch 9.938 s (arm
o275x-mf-pra-prof), 98.68 s per step, 0 dropped tokens. That pod ran no
flags-off control; against the flags-off control of another pod
(o275x-gdn-ctl, 10.033 s, 99.50 s per step) that is -0.095 s per
micro-batch, with a spread of the same control across pods of
+/-0.02-0.03 s per micro-batch. The environment-switch build dropped the
load-balance loss (it returned None for it); this change computes it,
which the measured run, at weight 0, does not exercise. The timing arms
run with the benchmark's log-prob gates on (loss and gradient norm read
0), so they say nothing about numerics.

TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks with this key in both measured
stacks: with moe_routing_index_math="count", moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp and moe_router_grad_from_wo_input (run nxs6c:
loss relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference
7.0e-6 against a replicate spread of 1.9e-5), and with the fused SiLU GMM
in place of the fused activation (run nxs6silu: gradient-norm difference
2.5e-6). Caveats: stack level, not this key alone; the environment-switch
build, not this branch's binary.

Tests (moe_roe_packed_routing_gather_test.py, new):
- pack_routing_for_all_gather / unpack_routing round-trip bit for bit
  with ids up to 4095 and -1 (bfloat16 holds every integer only up to
  256), float32 and bfloat16 weights, and the packed array is float32.
- On the shrunk Qwen3.5 routed MoE on 8 CPU devices (cp-as-ep, local
  routing, pipelined loop, two chunks), bit-identical output, loss and
  every gradient against the per-chunk gathers: real routing, padding
  excluded from routing, forced ids with a dropping buffer, saved routing
  maps, with moe_routing_index_math=count, and with the benchmark's MoE
  settings (expert-weight prefetch, row-major reduce-scatter,
  ragged_buffer_factor 2.75, prefused weights with fp8 fake-quant). With
  the load-balance loss on: output and loss bit-identical, gradients
  within 1e-5 of max (measured 1.2e-7). bf16 routing weights: within
  relative L2 1e-2 (bit-identical on CPU).
- Structure: the plain loop all-gathers two [..., top_k] arrays per chunk
  and the packed one exactly one [..., 2 * top_k] array, tagged
  moe_routing_maps only under moe_routing_maps=device; nothing is tagged
  moe_route.
- A packed gather shifted along the sequence, and the packed path with
  its load-balance loss dropped, are caught.
- The validator rejects the key for each unmet requirement, naming it,
  and accepts the load-balance loss. The SparseCore-pinned all-gather is
  a TPU-only placement and is not reached on CPU.
---
 src/maxtext/configs/base.yml                  |   3 +
 src/maxtext/configs/types.py                  |  31 ++
 src/maxtext/layers/moe.py                     | 105 ++++-
 .../moe_roe_packed_routing_gather_test.py     | 370 ++++++++++++++++++
 4 files changed, 493 insertions(+), 16 deletions(-)
 create mode 100644 tests/unit/moe_roe_packed_routing_gather_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index de62223012..33f49d1644 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -340,6 +340,9 @@ use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backw
 router_topk_matmul_vjp: false # deepseek routing: dense one-hot VJP for the top-k weight gather (no scatter-add in the backward)
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
 ring_of_experts_local_routing: false # ring of experts: route each shard's own tokens, then all-gather only the top-k expert ids and weights (not the router logits)
+# With ring_of_experts_local_routing + moe_chunk_pipeline: route every token chunk once and all-gather ids + weights as
+# ONE float32 array per layer (tagged moe_routing_maps) instead of two gathers per chunk. Same routing.
+moe_roe_packed_routing_gather: false
 ring_of_experts_row_major_reduce_scatter: false # ring of experts: lay out the expert-parallel reduce-scatter input row-major in the forward pass, so XLA does not copy the combine output into another layout and the result back
 moe_topk_before_ep_all_gather: false # ring of experts: run top-k on local tokens and all-gather top-k weights/ids instead of router logits
 te_moe_block: false # whether to use TransformerEngine's MoEBlock for routed MoE layers
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 389c6f0271..262edf331c 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1400,6 +1400,20 @@ class MoEGeneral(BaseModel):
           " Requires use_ring_of_experts=True."
       ),
   )
+  moe_roe_packed_routing_gather: bool = Field(
+      False,
+      description=(
+          "With ring_of_experts_local_routing and moe_chunk_pipeline: route the layer's local tokens once (every"
+          " token chunk) and all-gather the top-k weights and expert ids as ONE float32 [batch, seq, 2*top_k] array"
+          " (ids are exact in float32), instead of two all-gathers (ids, weights) per token chunk: one SparseCore"
+          " offload per layer instead of 2 x num_moe_token_chunks, and one weight-gradient reduce-scatter instead of"
+          " one per chunk. Same routing; bit-identical with float32 routing weights (float32_weight_sum=True), while"
+          " bfloat16 weights get their cotangent reduce-scattered in float32. The per-chunk load-balance loss is"
+          " computed as without it. Tagged moe_routing_maps. Requires sparse_matmul=True, te_moe_block=False,"
+          " ring_of_experts_local_routing=True, moe_chunk_pipeline=True and use_random_routing=False (random routing"
+          " draws per call, so one draw for the layer differs from one per chunk)."
+      ),
+  )
   ring_of_experts_row_major_reduce_scatter: bool = Field(
       False,
       description=(
@@ -1790,6 +1804,23 @@ def _ring_of_experts_ragged_sort_requirements(self) -> list[tuple[bool, str]]:
         (self.use_ragged_sort, "use_ragged_sort=True"),
     ]
 
+  @model_validator(mode="after")
+  def validate_moe_roe_packed_routing_gather(self) -> "MoEGeneral":
+    """The packed gather replaces the per-chunk routing gathers of the pipelined local-routing chunk loop."""
+    if self.moe_roe_packed_routing_gather:
+      _reject_unless(
+          "moe_roe_packed_routing_gather=True",
+          [
+              (self.sparse_matmul, "sparse_matmul=True"),
+              (not self.te_moe_block, "te_moe_block=False"),
+              (self.ring_of_experts_local_routing, "ring_of_experts_local_routing=True"),
+              (self.moe_chunk_pipeline, "moe_chunk_pipeline=True (only the pipelined chunk loop gathers per chunk)"),
+              # Random routing draws per get_topk call: one draw for the layer is not one per chunk.
+              (not self.use_random_routing, "use_random_routing=False"),
+          ],
+      )
+    return self
+
   @model_validator(mode="after")
   def validate_moe_routing_index_math(self) -> "MoEGeneral":
     """moe_routing_index_math='count' needs the ring-of-experts ragged sort, the only user of these index maps."""
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 6ae4826001..41e385579f 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -612,6 +612,21 @@ def _fix(path, leaf):
   return jax.tree_util.tree_map_with_path(_fix, intermediate_outputs)
 
 
+def pack_routing_for_all_gather(weights: jax.Array, ids: jax.Array) -> jax.Array:
+  """`[..., k]` routing weights and expert ids as ONE float32 `[..., 2k]` array (moe_roe_packed_routing_gather).
+
+  Exact: expert ids are integers in [-1, num_experts), and float32 holds every integer up to 2**24; float32 weights
+  are kept as they are and narrower ones convert exactly. `unpack_routing` inverts it.
+  """
+  return jnp.concatenate([weights.astype(jnp.float32), ids.astype(jnp.float32)], axis=-1)
+
+
+def unpack_routing(packed: jax.Array, weights_dtype, ids_dtype) -> tuple[jax.Array, jax.Array]:
+  """The routing weights and expert ids of a `pack_routing_for_all_gather` array, in their own dtypes."""
+  k = packed.shape[-1] // 2
+  return packed[..., :k].astype(weights_dtype), packed[..., k:].astype(ids_dtype)
+
+
 def routing_tag_fn(config):
   """Returns the checkpoint_name tagger for the routing arrays, or None under moe_routing_maps=remat.
 
@@ -3038,6 +3053,25 @@ def quantize_and_all_gather_tokens(
       x_gathered_qval = all_gather_fn(x_input.qvalue)
       return dataclasses.replace(x_input, qvalue=x_gathered_qval)
 
+    # The ring of experts' all-gather over the expert-parallel axis (tokens and routing), on a SparseCore with
+    # moe_pin_sparse_core_all_gathers / moe_pin_sparse_core_ep_all_gathers. roe_ag_and_route and the packed routing
+    # gather of the pipelined chunk loop (moe_roe_packed_routing_gather) both use it.
+    if self.config.moe_pin_sparse_core_all_gathers or getattr(self.config, "moe_pin_sparse_core_ep_all_gathers", False):
+
+      @functools.partial(
+          compute_on,
+          compute_type="tpu_sparsecore",
+          out_memory_spaces=jax.memory.Space.Device,
+          compiler_options={"sparse_core_config": {"core_ids": [self.config.moe_ep_all_gather_sparse_core_id]}},
+      )
+      def _ep_all_gather(z):
+        return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True)
+
+    else:
+
+      def _ep_all_gather(z):
+        return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True)
+
     def roe_ag_and_route(
         x,
         logits,
@@ -3048,23 +3082,14 @@ def roe_ag_and_route(
         input_ids=None,
         forced_routed_experts=None,
         force_dropless=False,
+        gathered_topk=None,
     ):
-      if self.config.moe_pin_sparse_core_all_gathers or getattr(self.config, "moe_pin_sparse_core_ep_all_gathers", False):
-
-        @functools.partial(
-            compute_on,
-            compute_type="tpu_sparsecore",
-            out_memory_spaces=jax.memory.Space.Device,
-            compiler_options={"sparse_core_config": {"core_ids": [self.config.moe_ep_all_gather_sparse_core_id]}},
-        )
-        def _ep_all_gather(z):
-          return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True)
-
-      else:
-
-        def _ep_all_gather(z):
-          return jax.lax.all_gather(z, axis_name=self._expert_parallelism_name, tiled=True)
+      """Ring of experts: all-gathers this chunk's tokens over the expert-parallel axis and routes them.
 
+      `gathered_topk` (moe_roe_packed_routing_gather): `(weights, expert ids)` of these tokens, already routed and
+      all-gathered over the expert-parallel axis, plus this shard's own (local, not gathered) expert ids for the
+      load-balance loss. The local top-k and the two routing all-gathers of this chunk are then skipped.
+      """
       if use_tc_3d_dispatch():
         # moe_tc_ragged_3d_dispatch: all-gather the tokens in the TC kernels' (emb // 128, 128) layout; the
         # reshape is on the local tokens, so the gathered tensor needs no relayout before the TC sort.
@@ -3097,7 +3122,21 @@ def token_all_gather(z):
           and not self.is_hash_routing
           and not self.config.use_random_routing
       )
-      if self.config.ring_of_experts_local_routing:
+      if gathered_topk is not None:
+        # moe_roe_packed_routing_gather: the chunk pipeline routed this shard's tokens of every chunk at once and
+        # all-gathered weights and ids as one packed array; this chunk's slice arrives here. Same values as the
+        # local-routing branch below (routing is per token), including its per-chunk load-balance loss, which is
+        # computed from this shard's own ids of these tokens exactly as there.
+        weights, selected_experts, local_selected_experts = gathered_topk
+        lb_loss = None
+        if self.config.load_balance_loss_weight > 0.0 and not self.is_hash_routing and not use_megatron_seq_aux_loss:
+          lb_loss = jax.lax.pmean(
+              self.router_load_balance_loss(local_selected_experts, logits, pre_bias_logits),
+              self._expert_parallelism_name,
+          )
+        topk = (weights, selected_experts, lb_loss)
+        logits, pre_bias_logits = None, None
+      elif self.config.ring_of_experts_local_routing:
         # (68792f5 pin) Routing is per token, so each shard routes only its own tokens and all-gathers the top-k
         # ids and weights the way the tokens were gathered. The [tokens, num_experts] logits never
         # cross shards, so neither does their gradient, and the gradient of the top-k gather is a
@@ -3334,6 +3373,7 @@ def route(
         input_ids=None,
         forced_routed_experts=None,
         force_dropless=False,
+        gathered_topk=None,
     ):
       """Performs both across device and within device token routing/sorting"""
       num_ep = self.get_expert_parallelism_size()
@@ -3350,6 +3390,7 @@ def route(
             input_ids=input_ids,
             forced_routed_experts=forced_routed_experts,
             force_dropless=force_dropless,
+            gathered_topk=gathered_topk,
         )
       else:
         return ra2a_and_route(
@@ -4198,6 +4239,10 @@ def sparse_matmul_route_and_compute(
             "moe_chunk_pipeline=True is skipped: --xla_tpu_aggressive_opt_barrier_removal=true (or ENABLED)"
             " removes optimization barriers before scheduling. Set it to false to use the pipeline."
         )
+        if self.config.moe_roe_packed_routing_gather:
+          max_logging.warning(
+              "moe_roe_packed_routing_gather=True is skipped too: it only runs in the pipelined chunk loop."
+          )
         pipeline_enabled = False
 
       def _pipelined_route_and_compute(force_dropless):
@@ -4235,6 +4280,33 @@ def _pipelined_route_and_compute(force_dropless):
         seq_len = x.shape[1]
         chunk = seq_len // n_chunks
 
+        packed_topk = None
+        if self.config.moe_roe_packed_routing_gather:
+          # moe_roe_packed_routing_gather: route ALL of this shard's tokens once (every chunk) and all-gather the
+          # top-k weights and expert ids as ONE float32 [batch, seq, 2k] array: one all-gather (one SparseCore
+          # offload) per layer instead of two per chunk, and in the backward one reduce-scatter of the weight
+          # cotangent instead of one per chunk. Routing is per token, so each chunk's slice holds the values the
+          # per-chunk local routing would compute (pack_routing_for_all_gather: exact).
+          num_expert_shards = self.get_expert_parallelism_size()
+          expert_shard = jax.lax.axis_index(self._expert_parallelism_name) if num_expert_shards > 1 else 0
+          local_weights, local_ids = self.get_topk(
+              logits,
+              pre_bias_logits,
+              rngs,
+              sharded_input_ids,
+              forced_routed_experts,
+              num_token_shards=num_expert_shards,
+              token_shard_index=expert_shard,
+          )
+          packed = _ep_all_gather(pack_routing_for_all_gather(local_weights, local_ids))
+          # A routing map: with moe_routing_maps=device (or offload) the backward reuses the gathered array and the
+          # rematerialized forward skips the gather. Untagged under moe_routing_maps=remat.
+          tag_fn = routing_tag_fn(self.config)
+          if tag_fn is not None:
+            packed = tag_fn(packed)
+          gathered_weights, gathered_ids = unpack_routing(packed, local_weights.dtype, local_ids.dtype)
+          packed_topk = (gathered_weights, gathered_ids, local_ids)
+
         def _dispatch(c, x_c):
           sl = slice(c * chunk, (c + 1) * chunk)
           return route(
@@ -4245,6 +4317,7 @@ def _dispatch(c, x_c):
               input_ids=None if sharded_input_ids is None else sharded_input_ids[:, sl],
               forced_routed_experts=None if forced_routed_experts is None else forced_routed_experts[:, sl, :],
               force_dropless=force_dropless,
+              gathered_topk=None if packed_topk is None else tuple(t[:, sl, :] for t in packed_topk),
           )
 
         def _with_routed(dispatched_c, routed):
diff --git a/tests/unit/moe_roe_packed_routing_gather_test.py b/tests/unit/moe_roe_packed_routing_gather_test.py
new file mode 100644
index 0000000000..b5ece5ac61
--- /dev/null
+++ b/tests/unit/moe_roe_packed_routing_gather_test.py
@@ -0,0 +1,370 @@
+# 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.
+"""moe_roe_packed_routing_gather against the per-chunk routing gathers of the pipelined chunk loop, on a CPU mesh.
+
+With ring_of_experts_local_routing and moe_chunk_pipeline, every token chunk routes its own tokens and all-gathers
+its top-k weights and expert ids: two all-gathers per chunk. The key routes all of the layer's local tokens once and
+all-gathers weights and ids as ONE float32 [batch, seq, 2 * top_k] array, which each chunk slices. Routing is per
+token, so each slice holds what the chunk would have computed.
+
+In-process checks: `pack_routing_for_all_gather` / `unpack_routing` round-trip routing weights and expert ids bit
+for bit, ids up to 4095 and the -1 padding id included (the layer below has 16 experts, and bfloat16 holds every
+integer up to 256, so only ids beyond 256 can tell an exact float32 packing from a lossy one), and the packed array
+is float32. The validator rejects the key for each requirement it does not meet, naming it.
+
+Runs as a subprocess so the forced device count takes effect before JAX initializes: the shrunk Qwen3.5 routed MoE
+of `moe_roe_local_routing_test.py` (16 experts, top 4, 8 CPU devices under cp-as-ep, ragged sort, local routing,
+two token chunks, megablox in interpret mode), with the pipelined loop in both runs. Not reachable on CPU: the
+all-gather pinned to a SparseCore (moe_pin_sparse_core_all_gathers / moe_pin_sparse_core_ep_all_gathers), a TPU-only
+compute_on placement; the packed gather shares it with the per-chunk gathers, and the Qwen3.5-397B RL trainer
+benchmark does not pin (it measured +0.363 s per micro-batch slower with the pin).
+
+Pre-registered comparisons, per case:
+  - Bit-identical output, loss and every gradient with float32 routing weights (float32_weight_sum, the default) and
+    no load-balance loss: real routing, padding excluded from routing, forced ids with a buffer that drops tokens,
+    saved routing maps, with moe_routing_index_math="count", and with the benchmark's MoE settings (expert-weight
+    prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, prefused weights with
+    fp8 fake-quant, padding excluded).
+  - With the load-balance loss on, the loss and the output stay bit-identical (the per-chunk loss is computed from
+    this shard's own ids of the chunk, as without the key) and the gradients agree to float32 reduction order,
+    max |new - old| <= 1e-5 * max |old|: the router gradient sums the top-k and the loss terms in another order.
+  - bfloat16 routing weights (float32_weight_sum=False): the output is bit-identical; their cotangent is
+    reduce-scattered in float32 instead of bfloat16, so the gradients agree to relative L2 < 1e-2.
+Two-sided:
+  - Structure: the plain loop must all-gather two [..., top_k] arrays per chunk and the packed one none of those
+    but exactly one [..., 2 * top_k] float32 array; under moe_routing_maps=device that array is tagged
+    moe_routing_maps, under remat it is untagged, and nothing is ever tagged moe_route.
+  - Deliberately wrong versions must FAIL: a packed gather whose result is shifted along the sequence (each chunk
+    gets another chunk's routing), and the packed path with its load-balance loss dropped (what an earlier version
+    of this change did).
+"""
+
+import contextlib
+import os
+import re
+import subprocess
+import sys
+from unittest import mock
+
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+import numpy as np
+import pytest
+
+from maxtext.configs import pyconfig
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+_NUM_DEVICES = 8
+_NUM_EXPERTS = 16
+_TOP_K = 4
+_EMB = 256
+_SEQ = 64
+_PASS_TOKEN = "MOE_ROE_PACKED_ROUTING_GATHER_CHECKS_PASSED"
+_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false"
+_REL_TOL = 1e-5
+_BF16_REL_NORM_TOL = 1e-2
+
+
+def _cfg(**overrides):
+  """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh, local routing, pipelined chunk loop."""
+  kwargs = {
+      "run_name": "moe_roe_packed_routing_gather",
+      "enable_checkpointing": False,
+      "log_config": False,
+      "skip_jax_distributed_system": True,
+      "model_name": "qwen3.5-397b-a17b",
+      "override_model_config": True,
+      "num_experts": _NUM_EXPERTS,
+      "num_experts_per_tok": _TOP_K,
+      "base_emb_dim": _EMB,
+      "base_mlp_dim": 256,
+      "base_moe_mlp_dim": 256,
+      "max_target_length": _SEQ,
+      "max_prefill_predict_length": _SEQ,
+      "per_device_batch_size": 0.25,
+      "dtype": "bfloat16",
+      "weight_dtype": "float32",
+      "float32_gate_logits": True,
+      "megablox": True,
+      "sparse_matmul": True,
+      # Tokamax gmm v2 is TPU-only; both runs share whichever GMM runs, so CPU uses megablox.
+      "use_tokamax_gmm": False,
+      "use_gmm_v2": False,
+      "use_ring_of_experts": True,
+      "use_ragged_sort": True,
+      "use_custom_sort_vjp": False,
+      "ragged_buffer_factor": 2.0,
+      "num_moe_token_chunks": 2,
+      "custom_mesh_and_rule": "cp-as-ep",
+      "ici_fsdp_parallelism": 1,
+      "ici_context_parallelism": 4,
+      "ici_expert_parallelism": 2,
+      "context_parallel_load_balance": False,
+      "ring_of_experts_local_routing": True,
+      "moe_chunk_pipeline": True,
+  }
+  kwargs.update(overrides)
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
+
+
+class PackRoutingTest(parameterized.TestCase):
+  """The packed routing array carries weights and expert ids exactly, ids beyond bfloat16's integers included."""
+
+  @parameterized.named_parameters(("float32_weights", jnp.float32), ("bfloat16_weights", jnp.bfloat16))
+  def test_round_trip_is_exact(self, weights_dtype):
+    k0, k1 = jax.random.split(jax.random.PRNGKey(0))
+    weights = jax.random.uniform(k0, (2, 64, 10), jnp.float32).astype(weights_dtype)
+    # Every id from -1 (padding) to 4095, so the odd ids above 256 that bfloat16 cannot hold are all there.
+    ids = (jax.random.permutation(k1, 2 * 64 * 10) % 4097 - 1).reshape(2, 64, 10).astype(jnp.int32)
+    packed = moe.pack_routing_for_all_gather(weights, ids)
+    self.assertEqual(packed.dtype, jnp.float32)
+    self.assertEqual(packed.shape, (2, 64, 20))
+    got_weights, got_ids = moe.unpack_routing(packed, weights.dtype, ids.dtype)
+    self.assertEqual((got_weights.dtype, got_ids.dtype), (weights.dtype, ids.dtype))
+    np.testing.assert_array_equal(np.asarray(got_ids), np.asarray(ids))
+    np.testing.assert_array_equal(np.asarray(got_weights).view(np.uint8), np.asarray(weights).view(np.uint8))
+
+
+class PackedRoutingGatherConfigTest(parameterized.TestCase):
+  """The key is rejected for each requirement it does not meet; the load-balance loss is allowed."""
+
+  @parameterized.named_parameters(
+      ("dense", {"sparse_matmul": False}, "sparse_matmul=True"),
+      ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"),
+      ("no_local_routing", {"ring_of_experts_local_routing": False}, "ring_of_experts_local_routing=True"),
+      ("no_pipeline", {"moe_chunk_pipeline": False}, "moe_chunk_pipeline=True"),
+      ("random_routing", {"use_random_routing": True}, "use_random_routing=False"),
+  )
+  def test_rejected(self, overrides, requirement):
+    with self.assertRaisesRegex(ValueError, f"moe_roe_packed_routing_gather=True requires: .*{re.escape(requirement)}"):
+      _cfg(moe_roe_packed_routing_gather=True, **overrides)
+
+  def test_accepted_with_load_balance_loss(self):
+    cfg = _cfg(moe_roe_packed_routing_gather=True, load_balance_loss_weight=0.01)
+    self.assertTrue(cfg.moe_roe_packed_routing_gather)
+
+
+@pytest.mark.cpu_only
+def test_packed_routing_gather_matches_per_chunk_gathers_on_cpu_mesh():
+  """Runs `main` below in a child process with 8 CPU devices."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}"
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+  print(result.stdout)
+  assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+def _build(cfg, mesh):
+  return moe.RoutedMoE(
+      config=cfg,
+      num_experts=cfg.num_experts,
+      num_experts_per_tok=cfg.num_experts_per_tok,
+      mesh=mesh,
+      kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+      kernel_axes=("embed", "mlp"),
+      intermediate_dim=cfg.moe_mlp_dim,
+      weight_dtype=cfg.weight_dtype,
+      dtype=cfg.dtype,
+      rngs=nnx.Rngs(params=0),
+  )
+
+
+def _inputs(cfg, mode):
+  """The layer input, forced expert ids (padded and duplicate slots) and a token mask, per `mode`."""
+  batch = int(cfg.per_device_batch_size * _NUM_DEVICES)
+  x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32)
+  forced = token_mask = None
+  if "forced" in mode:
+    keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ)
+    ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K)
+    pos = jnp.arange(_SEQ)[None, :, None]
+    ids = jnp.where((pos % 5 == 0) & (jnp.arange(_TOP_K) == _TOP_K - 1), -1, ids)
+    ids = jnp.where((pos % 7 == 3) & (jnp.arange(_TOP_K) == 1), ids[..., :1], ids)
+    forced = ids.astype(jnp.int32)
+  if "padmask" in mode:
+    lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)])
+    token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32)
+  return x, forced, token_mask
+
+
+def _eqns(jaxpr):
+  """Every equation of a jaxpr, including those of nested jaxprs."""
+  for eqn in jaxpr.eqns:
+    yield eqn
+    for param in eqn.params.values():
+      for sub in param if isinstance(param, (list, tuple)) else (param,):
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        if hasattr(sub, "eqns"):
+          yield from _eqns(sub)
+
+
+def _structure(jaxpr):
+  """Routing all-gathers by last dim, and the checkpoint names on [..., 2 * top_k] arrays and anywhere."""
+  eqns = list(_eqns(jaxpr))
+  gathers = [tuple(v.aval.shape) for e in eqns if e.primitive.name == "all_gather" for v in e.outvars]
+  names = [(e.params.get("name"), tuple(e.invars[0].aval.shape)) for e in eqns if e.primitive.name == "name"]
+  return {
+      "topk_gathers": sum(s[-1:] == (_TOP_K,) for s in gathers),
+      "packed_gathers": sum(s[-1:] == (2 * _TOP_K,) for s in gathers),
+      "packed_tags": sorted({n for n, s in names if s[-1:] == (2 * _TOP_K,)}),
+      "moe_route_tags": sum(n == "moe_route" for n, _ in names),
+  }
+
+
+def _run(cfg, mesh, params, x, forced, token_mask):
+  """Loss, output, gradients, and the structure of the traced forward."""
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules):
+    graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+
+    def loss_fn(params, x):
+      out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask)
+      loss = jnp.mean(out.astype(jnp.float32) ** 2)
+      if lb_loss is not None:
+        loss = loss + lb_loss.astype(jnp.float32)
+      return loss, out
+
+    structure = _structure(jax.make_jaxpr(loss_fn)(params, x).jaxpr)
+    (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x)
+  return jax.device_get((loss, out, grads)), structure
+
+
+def _compare(label, old, new, mode):
+  """Prints the per-leaf report. mode: "bitwise", "reduction_order" (gradients) or "bf16" (gradients)."""
+  (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new
+  named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [
+      (jax.tree_util.keystr(path), a, b)
+      for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new))
+  ]
+  passed, lines = True, []
+  for name, a, b in named:
+    a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32)
+    diff, scale = float(np.max(np.abs(b - a))), float(np.max(np.abs(a)))
+    rel_l2 = float(np.linalg.norm(b - a)) / max(float(np.linalg.norm(a)), 1e-30)
+    same = bool(np.array_equal(a, b))
+    if mode == "bitwise" or name in ("output", "loss"):
+      ok = same
+    elif mode == "reduction_order":
+      ok = diff <= _REL_TOL * scale
+    else:
+      ok = rel_l2 < _BF16_REL_NORM_TOL
+    ok = ok and bool(np.isfinite(b).all()) and scale > 0
+    passed &= ok
+    lines.append(
+        f"  [{'ok' if ok else 'FAIL'}] {name}: bit-identical={same} max|d|={diff:.3e} max|old|={scale:.3e}"
+        f" rel_l2={rel_l2:.3e}"
+    )
+  print(f"{label}: {'PASS' if passed else 'FAIL'} ({mode})", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+def _case(label, overrides, mode="real", compare="bitwise", control=None):
+  """The per-chunk routing gathers against the packed one, on the same parameters; `control` wraps the packed run."""
+  cfg_old = _cfg(**overrides)
+  cfg_new = _cfg(**overrides, moe_roe_packed_routing_gather=True)
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes)
+  x, forced, token_mask = _inputs(cfg_old, mode)
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...)
+  old, s_old = _run(cfg_old, mesh, params, x, forced, token_mask)
+  with control if control is not None else contextlib.nullcontext():
+    new, s_new = _run(cfg_new, mesh, params, x, forced, token_mask)
+  n_chunks = cfg_new.num_moe_token_chunks
+  want_tags = ["moe_routing_maps"] if cfg_new.moe_routing_maps != "remat" else []
+  structure = (
+      s_old["topk_gathers"] == 2 * n_chunks
+      and s_old["packed_gathers"] == 0
+      and s_new["topk_gathers"] == 0
+      and s_new["packed_gathers"] == 1
+      and s_new["packed_tags"] == want_tags
+      and s_old["moe_route_tags"] == s_new["moe_route_tags"] == 0
+  )
+  print(f"{label}: [{'ok' if structure else 'FAIL'}] structure plain={s_old} packed={s_new}", flush=True)
+  return _compare(label, old, new, compare) and structure
+
+
+def _shifted_packed_gather():
+  """Deliberately wrong: the packed routing array comes back rolled along the sequence by one chunk."""
+  real = jax.lax.all_gather
+
+  def wrong(x, axis_name, **kwargs):
+    out = real(x, axis_name, **kwargs)
+    if out.ndim == 3 and out.shape[-1] == 2 * _TOP_K:
+      # Half the LOCAL sequence: inside the shard_map the sequence is sharded over the context axis.
+      out = jnp.roll(out, out.shape[1] // 2, axis=1)
+    return out
+
+  return mock.patch.object(jax.lax, "all_gather", wrong)
+
+
+def _dropped_load_balance_loss():
+  """Deliberately wrong: the packed path without its load-balance loss (pmean of the loss returns zeros)."""
+  return mock.patch.object(jax.lax, "pmean", lambda x, axis_name, **kwargs: jnp.zeros_like(x))
+
+
+# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+_BENCH_FLAGS = {
+    "moe_expert_weight_prefetch": "forward",
+    "ring_of_experts_row_major_reduce_scatter": True,
+    "moe_routing_maps": "device",
+    "ragged_buffer_factor": 2.75,
+    "prefuse_moe_weights": True,
+    "fp8_moe_fake_quant": True,
+    "moe_exclude_padding_from_routing": True,
+}
+
+
+def main():
+  assert jax.device_count() == _NUM_DEVICES, jax.devices()
+  padded = {"moe_exclude_padding_from_routing": True}
+  lb = {"load_balance_loss_weight": 0.01}
+  results = {
+      "real": _case("real", {}),
+      "padding_excluded": _case("padding_excluded", padded, mode="padmask"),
+      "forced_dropping": _case("forced_dropping", {"ragged_buffer_factor": 1.0}, mode="forced"),
+      "saved_routing_maps": _case("saved_routing_maps", {"moe_routing_maps": "device"}),
+      "with_count": _case("with_count", {**padded, "moe_routing_index_math": "count"}, mode="padmask"),
+      "load_balance_loss": _case("load_balance_loss", lb, compare="reduction_order"),
+      "bf16_routing_weights": _case("bf16_routing_weights", {"float32_weight_sum": False}, compare="bf16"),
+      "bench_flags": _case("bench_flags", _BENCH_FLAGS, mode="padmask"),
+  }
+  controls = {
+      "control_shifted_packed_gather": _case("control_shifted_packed_gather", {}, control=_shifted_packed_gather()),
+      "control_dropped_load_balance_loss": _case(
+          "control_dropped_load_balance_loss", lb, compare="reduction_order", control=_dropped_load_balance_loss()
+      ),
+  }
+  print("SUMMARY", {**results, **controls}, flush=True)
+  failed = [k for k, v in results.items() if not v]
+  undetected = [k for k, v in controls.items() if v]
+  if failed or undetected:
+    print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True)
+    sys.exit(1)
+  print(_PASS_TOKEN, flush=True)
+
+
+if __name__ == "__main__":
+  main()

From 4ab9f123d74fb9b516d3665fca4605f43b5afb9c Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 02:51:50 -0700
Subject: [PATCH 06/10] MoE: SwiGLU in the wi GMM epilogue
 (moe_wi_gmm_fused_swiglu)

The up projection wrote the [buffer, 2 * mlp] GMM output and a separate
elementwise pass read it back to apply the SwiGLU, over every buffer row,
the ragged buffer's padding included.

moe_wi_gmm_fused_swiglu calls megablox ops.gmm_silu instead: the tokamax
gmm_v2 kernel with fuse_act="silu" applies the activation to the float32
accumulator of the rows it computes and writes the [buffer, mlp]
activation. When differentiated it also writes the gate and up halves
([buffer, 3 * mlp] in all), rounded to the output dtype exactly as the
unfused GMM rounds its output (the new return_pre_act option of gmm_v2:
both halves go through the store path of the 2D output, with its
zero-init DMAs and partial-row carry), and the custom VJP is the unfused
backward: the SwiGLU VJP on those bf16 halves, then gmm's own DLHS / TGMM
on [d gate | d up]. ops._gmm_silu_bwd calls gmm's backward with keyword
arguments, as gmm's custom VJP would for a plain bf16 gmm_v2 GEMM. Works
with split weights and with the whole [w0 | w1] of moe_unsplit_prefused_wi.
The layer hands the fused GMM the routing operands explicitly
(get_gmm_operands_for_local_experts, which get_gmm_for_local_experts now
wraps; the key-off program is unchanged).

Numerics change: silu reads the f32 accumulator instead of its bf16
rounding (the more precise value), so the activation, and through it
everything downstream, differs by rounding.

The validator (the shared form) requires prefused sparse-matmul training,
use_tokamax_gmm and use_gmm_v2, silu + linear activations without
clipping (and not GPT-OSS), no weight quantization, MLP bias, embedding
chunks, flat FSDP weights, 3D TC GMM, token all-gather quantization or
chunk weight-gradient accumulation, and a remat policy that saves neither
half of the routed experts' up projection (the rule
moe_prefused_swiglu_vjp uses: the fused GMM never forms the
[buffer, 2 * mlp] output, so such a policy would silently stop saving the
MoE tensor; remat_policy=full and the shared expert's own names are
fine). moe_prefused_swiglu_vjp is rejected next to it, as it would do
nothing. Weight quantization objects and an explicit weight all-gather,
which only the built layer knows, raise at trace time.

Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), switched
by an environment variable in a build whose kernel files differ from these
only in docstrings, comments and line breaks (AST-identical) and whose
program a TPU lowering check matches to this commit. Same pod, stacked on
moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather and
moe_router_grad_from_wo_input (the fused activation has nothing to act on
next to this key, and the validator rejects the two together): regular
micro-batch 9.494 -> 9.418 s (-0.0755 s), 86.970 -> 86.253 s per step
(-0.716 s), and the step's second micro-batch, a proxy for the step
boundary, 9.450 -> 9.375 s, so the boundary pays nothing extra (arms
o275x-s6silu-l3-prof against o275x-s6c-l3-prof, one pod, on a build with
the staged step boundary, which is why its steps are shorter than the
other arms of this series). The same control measured in different pods
spreads by +/-0.02-0.03 s per micro-batch. The timing arms run with the
benchmark's log-prob gates on (loss and gradient norm read 0), so they say
nothing about numerics.

TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks for this key stacked with
moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_roe_packed_routing_gather and moe_router_grad_from_wo_input (run
nxs6silu: N1 1.4e-6, gradient-norm difference 2.5e-6, within the
replicate spread of 1.9e-5 measured with the same reference). Caveats:
stack level, not this key alone; the environment-switch build, not this
branch's binary. megablox_gmm_silu_test.py, the kernel's own TPU test,
has not been run yet, so that stack run is the kernel's only TPU
evidence.

Tests:
- megablox_gmm_silu_test.py (new, TPU only, skipped on CPU): gmm_v2's
  gate and up equal the unfused output bit for bit, rows outside the
  groups are zero, act is silu of the f32 accumulator, the act-only call
  matches, with an EP shard (group_offset); gmm_silu's VJP equals the
  unfused path's value for value.
- moe_wi_gmm_fused_swiglu_test.py (new, CPU):
  - ops.gmm_silu with stand-ins for the Pallas kernel and gmm's TPU
    backward only: the forward tile is tiling[2] halved, act alone is
    written undifferentiated and act + gate + up differentiated, the
    backward hands gmm's backward [d gate | d up] of the SwiGLU on the
    bf16 halves bit for bit, with this op's operands as residuals and its
    DLHS options on the right keywords (each option set alone, both
    group offsets); an odd forward tile is rejected.
  - The shrunk Qwen3.5 MoE on 8 CPU devices, with jnp stand-ins for gmm
    and gmm_silu (gmm_silu as the unfused composition of the same
    stand-in): the layer with the key is bit-identical to the layer
    without it (output, loss, every gradient): plain and pipelined loop,
    dropless (non-zero group offsets), padding excluded, the whole
    prefused wi, with the other MoE keys of this series, and with those
    and the benchmark's MoE settings. One gmm_silu call per chunk, with
    the plain path's wi tiling and the config's DLHS options
    (moe_gmm_v2_dlhs_transpose_rhs True in the benchmark case). Stand-ins
    that swap the gate and up halves, or ignore the group offset, are
    caught.
  - Key off: the dropless ring-of-experts ragged sort matches the ring of
    experts without it (use_ragged_sort=False), and fails to once its GMM
    ignores the group offset, so the routing operands that both paths now
    take from get_gmm_operands_for_local_experts are checked against a
    reference this change does not touch.
  - The validator rejects the key for each unmet requirement, naming it.
---
 src/maxtext/configs/base.yml                  |   5 +
 src/maxtext/configs/types.py                  |  49 ++
 src/maxtext/kernels/megablox/__init__.py      |   2 +-
 src/maxtext/kernels/megablox/ops.py           | 137 +++++
 .../pallas_mosaic_tpu_v2_gmm_kernel.py        | 154 ++++-
 src/maxtext/layers/moe.py                     | 123 +++-
 tests/unit/megablox_gmm_silu_test.py          | 149 +++++
 tests/unit/moe_wi_gmm_fused_swiglu_test.py    | 549 ++++++++++++++++++
 8 files changed, 1142 insertions(+), 26 deletions(-)
 create mode 100644 tests/unit/megablox_gmm_silu_test.py
 create mode 100644 tests/unit/moe_wi_gmm_fused_swiglu_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 33f49d1644..60cbcfc1bc 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -327,6 +327,11 @@ moe_unsplit_prefused_wi: false
 # With prefuse_moe_weights on the sparse-matmul path: the activation on the whole [gate | up] up-projection output as
 # one custom-VJP op whose backward writes d[gate | up] in one pass. Same values.
 moe_prefused_swiglu_vjp: false
+# The MoE SwiGLU in the wi GMM epilogue (megablox gmm_silu, tokamax gmm_v2): the forward kernel writes the
+# [buffer, mlp] activation instead of the [buffer, 2 * mlp] GMM output (differentiated: act, gate and up, i.e.
+# [buffer, 3 * mlp]). Numerics change: silu reads the float32 accumulator.
+# Needs prefuse_moe_weights, use_tokamax_gmm and use_gmm_v2, silu+linear without clipping, no bias or quantization.
+moe_wi_gmm_fused_swiglu: false
 fp8_moe: false # whether MoE weights are quantized to float8_e4m3fn in inference/rollout
 # Keep routed-expert weights in weight_dtype but round them in the forward pass to the per-channel FP8 grid
 # fp8_moe stores them on, with straight-through gradients: the trainer computes with the rollout's weights.
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 262edf331c..0be85e3df3 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1726,6 +1726,24 @@ class MoEGeneral(BaseModel):
           " never forms them, so the policy would silently stop saving the MoE tensor."
       ),
   )
+  moe_wi_gmm_fused_swiglu: bool = Field(
+      False,
+      description=(
+          "Apply the MoE SwiGLU in the epilogue of the wi (up-projection) GMM (megablox ops.gmm_silu, the tokamax"
+          " gmm_v2 kernel with fuse_act): the forward kernel writes the [buffer, mlp] activation of the rows it"
+          " computes instead of the [buffer, 2 * mlp] GMM output that a separate elementwise pass reads back over every"
+          " buffer row, padding included. When differentiated it writes the activation and the gate and up halves"
+          " ([buffer, 3 * mlp] in all, gate and up rounded as the unfused GMM rounds its output), and the backward is"
+          " the unfused one. NUMERICS CHANGE: silu reads the float32 accumulator instead of its bfloat16 rounding (the"
+          " more precise value), so the activation differs by rounding. Requires prefuse_moe_weights=True,"
+          " sparse_matmul=True, te_moe_block=False, use_tokamax_gmm=True, use_gmm_v2=True, no weight quantization,"
+          " mlp_activations=['silu', 'linear'] without clipping, mlp_bias=False, num_moe_emb_chunks=0,"
+          " moe_flat_fsdp_weights, moe_tc_ragged_3d_gmm, moe_quantize_token_all_gather and moe_accumulate_chunk_wgrad"
+          " off, and a remat policy that does not save the routed experts' gate or up half (mlpwi_0, moe_mlpwi_0,"
+          " moe_mlpwi_1; the minimal* policies save mlpwi_0): the fused GMM never forms them, so the policy would"
+          " silently stop saving the MoE tensor."
+      ),
+  )
   fp8_moe: bool = Field(
       False,
       description="Whether MoE weights are quantized to float8_e4m3fn in inference/rollout.",
@@ -1929,6 +1947,37 @@ def validate_moe_prefused_swiglu_vjp(self) -> "MoEGeneral":
               (not self.mlp_bias, "mlp_bias=False (biases are added per half)"),
               (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0 (partial sums per half)"),
               self._remat_keeps_moe_tensor_requirement("gate and up halves", _MOE_UP_PROJECTION_NAMES),
+              (
+                  not self.moe_wi_gmm_fused_swiglu,
+                  "moe_wi_gmm_fused_swiglu=False (the wi GMM already applies the activation)",
+              ),
+          ],
+      )
+    return self
+
+  @model_validator(mode="after")
+  def validate_moe_wi_gmm_fused_swiglu(self) -> "MoEGeneral":
+    """The fused gmm_v2 SwiGLU epilogue replaces the up projection only on the plain bf16 prefused gmm_v2 path."""
+    if self.moe_wi_gmm_fused_swiglu:
+      _reject_unless(
+          "moe_wi_gmm_fused_swiglu=True",
+          self._prefused_sparse_matmul_requirements()
+          + [
+              # gmm_silu is the bf16 tokamax gmm_v2 kernel only.
+              (self.use_tokamax_gmm and self.use_gmm_v2, "use_tokamax_gmm=True and use_gmm_v2=True"),
+              self._no_weight_quantization_requirement(),
+              # The epilogue computes silu(gate) * up; GPT-OSS and clipped activations differ.
+              (list(self.mlp_activations) == ["silu", "linear"], "mlp_activations=['silu', 'linear']"),
+              (self.decoder_block != DecoderBlockType.GPT_OSS, "a decoder_block other than gpt_oss"),
+              (self.mlp_activations_limit <= 0.0, "mlp_activations_limit <= 0 (no clipping)"),
+              (not self.mlp_bias, "mlp_bias=False"),
+              (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0"),
+              (not self.moe_flat_fsdp_weights, "moe_flat_fsdp_weights=False"),
+              # A 2D bf16 token buffer.
+              (not self.moe_tc_ragged_3d_gmm, "moe_tc_ragged_3d_gmm=False"),
+              (not self.moe_quantize_token_all_gather, "moe_quantize_token_all_gather=False"),
+              (not self.moe_accumulate_chunk_wgrad, "moe_accumulate_chunk_wgrad=False"),
+              self._remat_keeps_moe_tensor_requirement("gate and up halves", _MOE_UP_PROJECTION_NAMES),
           ],
       )
     return self
diff --git a/src/maxtext/kernels/megablox/__init__.py b/src/maxtext/kernels/megablox/__init__.py
index b912dc316b..ff2a18776e 100644
--- a/src/maxtext/kernels/megablox/__init__.py
+++ b/src/maxtext/kernels/megablox/__init__.py
@@ -13,4 +13,4 @@
 # limitations under the License.
 """Megablox kernel"""
 
-from maxtext.kernels.megablox.ops import TcUnsortCfg, gmm
+from maxtext.kernels.megablox.ops import TcUnsortCfg, gmm, gmm_silu
diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py
index cd32ac3f92..2ca87711fe 100644
--- a/src/maxtext/kernels/megablox/ops.py
+++ b/src/maxtext/kernels/megablox/ops.py
@@ -1395,3 +1395,140 @@ def _drhs_run_megablox(
       interpret=interpret,
       varying_axes=rhs_vma_axes,
   )
+
+
+# ==============================================================================
+# GMM + SwiGLU: the activation in the gmm_v2 epilogue
+# ==============================================================================
+
+
+def gmm_silu(
+    lhs: jnp.ndarray,
+    rhs: jnp.ndarray,
+    group_sizes: jnp.ndarray,
+    preferred_element_type: jnp.dtype,
+    tiling: tuple[int, int, int, int, int, int, int, int, int],
+    group_offset: jnp.ndarray | int | None = None,
+    use_gmm_v2_heuristic_tiling: bool = False,
+    use_dlhs_transpose_rhs: bool = False,
+    fuse_dlhs_scale: bool = False,
+    fuse_3d_dlhs_ps: bool = False,
+) -> jnp.ndarray:
+  """silu(lhs @ gate) * (lhs @ up) per group, with the activation in the gmm_v2 epilogue.
+
+  `rhs` is the prefused [g, k, 2n] weight, gate = rhs[..., :n] and up = rhs[..., n:]. The unfused path writes
+  the [m, 2n] GMM output and runs the activation over every row in a separate elementwise pass; here the kernel
+  applies it to the f32 accumulator of the rows it computes and writes [m, n].
+
+  Numerics: the activation reads the f32 accumulator instead of its bf16 rounding, so `act` differs from the
+  unfused value by rounding (it is the more precise of the two). The backward is the unfused one: it takes
+  gate and up rounded exactly as the unfused GMM rounds its output (the kernel writes them next to act when
+  differentiated), the SwiGLU VJP of `jax.nn.silu(gate) * up`, and the unchanged tokamax gmm_v2 DLHS / TGMM
+  DRHS of `gmm`.
+
+  Only the bf16 tokamax gmm_v2 path: no quantization, no partial_sum, no return_lhs / return_rhs, 2D lhs and
+  output, rhs not transposed. tiling is gmm's 9-tuple; the fused kernel's tile_n counts the columns of each
+  half, so tiling[2] is halved and every grid step does the unfused kernel's MXU work in the same VMEM.
+  """
+  if isinstance(lhs, qpl.QArray) or isinstance(rhs, qpl.QArray) or lhs.ndim != 2 or rhs.ndim != 3:
+    raise NotImplementedError("gmm_silu: plain 2D lhs and [g, k, 2n] rhs only.")
+  return _gmm_silu(
+      lhs,
+      rhs,
+      group_sizes,
+      group_offset,
+      preferred_element_type,
+      tuple(tiling),
+      use_gmm_v2_heuristic_tiling,
+      use_dlhs_transpose_rhs,
+      fuse_dlhs_scale,
+      fuse_3d_dlhs_ps,
+  )
+
+
+def _gmm_silu_kernel(lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, return_pre_act):
+  """The fused forward: act, or (act, gate, up) with return_pre_act."""
+  if heuristic:
+    tile_info = gmm_v2.calculate_tiling
+  else:
+    half_n = rhs.shape[-1] // 2
+    tile_n = min(tiling[2], rhs.shape[-1]) // 2
+    if tiling[2] % 2 or (tile_n % 128 and tile_n != half_n):
+      raise ValueError(f"gmm_silu: fwd tile_n {tiling[2]} must halve to a multiple of 128 (or to n={half_n}).")
+    tile_info = gmm_v2.TileSizes(tile_m=tiling[0], tile_k=tiling[1], tile_n=tile_n)
+  return gmm_v2.gmm_v2(
+      lhs=lhs,
+      rhs=rhs,
+      group_sizes=group_sizes,
+      tile_info=tile_info,
+      preferred_element_type=preferred_element_type,
+      group_offset=group_offset,
+      fuse_act="silu",
+      return_pre_act=return_pre_act,
+  )
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(4, 5, 6, 7, 8, 9))
+def _gmm_silu(
+    lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, dlhs_transpose, fuse_scale, fuse_3d
+):
+  # Not differentiated: the activation alone, nothing else is written.
+  del dlhs_transpose, fuse_scale, fuse_3d
+  return _gmm_silu_kernel(lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, False)
+
+
+def _gmm_silu_fwd(
+    lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, dlhs_transpose, fuse_scale, fuse_3d
+):
+  """Differentiated: the kernel also writes gate and up, the residuals of the unfused backward."""
+  del dlhs_transpose, fuse_scale, fuse_3d
+  act, gate, up = _gmm_silu_kernel(lhs, rhs, group_sizes, group_offset, preferred_element_type, tiling, heuristic, True)
+  return act, (lhs, rhs, group_sizes, group_offset, gate, up)
+
+
+def _swiglu(gate, up):
+  """MaxText's MoE activation on the GMM output (`apply_ffn_activation`, mlp_activations silu + linear)."""
+  return (jax.nn.silu(gate) * up).astype(gate.dtype)
+
+
+def _gmm_silu_bwd(preferred_element_type, tiling, heuristic, dlhs_transpose, fuse_scale, fuse_3d, residual, dact):
+  """The unfused backward from the saved bf16 gate and up: the SwiGLU VJP, then gmm's DLHS / TGMM."""
+  lhs, rhs, group_sizes, group_offset, gate, up = residual
+  # The unfused backward: the SwiGLU VJP on the bf16 GMM output, then gmm's own backward on [d gate | d up].
+  _, swiglu_vjp = jax.vjp(_swiglu, gate, up)
+  dgate, dup = swiglu_vjp(dact.astype(gate.dtype))
+  grad = jnp.concatenate([dgate, dup], axis=-1)
+  # gmm's own backward, as gmm's custom VJP calls it for a plain bf16 tokamax gmm_v2 GMM (no quantization, no
+  # partial sum, no returned operands), with this op's DLHS options.
+  dlhs, drhs, *_ = _gmm_bwd(
+      lhs_dtype=lhs.dtype,
+      rhs_dtype=rhs.dtype,
+      preferred_element_type=preferred_element_type,
+      tiling=tiling,
+      transpose_rhs=False,
+      interpret=False,  # the megablox backend only
+      quantization_rule=None,
+      use_tokamax_backend=True,
+      weight_gather_axes=None,
+      use_manual_quantization=False,
+      lhs_vma_axes=(),
+      rhs_vma_axes=(),
+      use_gmm_v2=True,
+      use_gmm_v2_heuristic_tiling=heuristic,
+      return_rhs=False,
+      return_lhs=False,
+      fuse_dlhs_scale=fuse_scale,
+      use_dlhs_transpose_rhs=dlhs_transpose,
+      flat_fsdp_shard=None,
+      out_is_3d=False,
+      fuse_3d_dlhs_ps=fuse_3d,
+      tc_unsort_cfg=None,
+      # (lhs, rhs, group_sizes, group_offset, partial_sum_fwd, lhs_is_qarray, rhs_is_qarray, tc_routing)
+      residual=(lhs, rhs, group_sizes, group_offset, None, False, False, None),
+      # (d out, d returned lhs, d returned rhs)
+      grad=(grad, None, None),
+  )
+  return dlhs, drhs, None, None
+
+
+_gmm_silu.defvjp(_gmm_silu_fwd, _gmm_silu_bwd)
diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
index 634270419c..f02f6452d6 100644
--- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
+++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
@@ -254,6 +254,10 @@ class GmmConfigs:
   # Output is stored as [size_m, out_size_n // 128, 128] (one contiguous slab
   # per row); tiles are relayouted from [tile_m, tile_n] inside VMEM.
   out_is_3d: bool = False
+  # With fuse_act: also write the two halves of the accumulator before the
+  # activation (gate and up), each [size_m, out_size_n] in out_dtype, next to
+  # the activation output. The backward of the activation needs them.
+  return_pre_act: bool = False
 
   @property
   def lhs_is_3d(self) -> bool:
@@ -538,8 +542,17 @@ def inner_kernel(
     acc_ref: Reference to the accumulator.
     metadata_ref: Reference to the metadata.
     cfgs: GmmConfigs.
+
+  With cfgs.return_pre_act, tiled_out_ref and partial_out_ref are (act, gate,
+  up) triples. The activation takes the regular store path; gate and up are
+  stored from the same accumulator before the activation is applied.
   """
 
+  pre_act_refs = ()
+  if cfgs.return_pre_act:
+    (tiled_out_ref, *tiled_pre_act_refs), (partial_out_ref, *partial_pre_act_refs) = tiled_out_ref, partial_out_ref
+    pre_act_refs = tuple(zip(tiled_pre_act_refs, partial_pre_act_refs))
+
   def _matmul(is_first_k_step: bool, is_last_k_step: bool):
     tpu_info = pltpu.get_tpu_info()
     mxu_size = tpu_info.mxu_column_size
@@ -773,6 +786,12 @@ def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int):
           ps_tile = tiled_ps_ref[...].reshape(acc.shape)  # pyrefly: ignore[unsupported-operation]
         acc += ps_tile.astype(acc.dtype)
 
+      if pre_act_refs:
+        # acc is [tile_m, 2 * tile_n] = gate | up. Store both halves before the activation, rounded to the
+        # output dtype exactly as the unfused kernel rounds its output; the activation below stays in f32.
+        for (tiled_pre_ref, partial_pre_ref), acc_half in zip(pre_act_refs, jnp.split(acc, 2, -1)):
+          _store_group_rows(tiled_pre_ref, partial_pre_ref, acc_half, metadata_ref, cfgs)
+
       acc = apply_act_fn(acc, cfgs.fuse_act)
 
       gm_id = pl.program_id(cfgs.gm_axis)
@@ -881,6 +900,42 @@ def matmul_last():
   )
 
 
+def _store_group_rows(
+    tiled_ref: jax.Array,  # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n]
+    partial_ref: jax.Array,  # [size_lhs_sublane, tile_n]
+    acc: jax.Array,  # [tile_m, tile_n]
+    metadata_ref: MetadataRef,
+    cfgs: GmmConfigs,
+):
+  """Stores the current group's rows of `acc` into a 2D output tile, with the partial-row carry.
+
+  The same steps as the 2D store at the end of inner_kernel's last k step: zero
+  the rows of other groups, round to the output dtype, add the sublane rows the
+  previous gm tile left in `partial_ref`, and leave this tile's unfinished last
+  sublane rows there for the next one. Used for the extra (return_pre_act)
+  outputs written from the same accumulator.
+  """
+  gm_id = pl.program_id(cfgs.gm_axis)
+  m_start = metadata_ref.gm_id_to_m_offset[gm_id]
+  m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1]
+  m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane
+  m_start_local = m_start - m_offset
+  m_end_local = m_end - m_offset
+
+  iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0)
+  mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local)
+  tiled_ref[...] = jnp.where(mask, acc, 0).reshape(tiled_ref.shape).astype(tiled_ref.dtype)
+
+  partial_zeros = jnp.zeros_like(partial_ref)
+  tiled_ref[0] += jnp.where(gm_id == 0, partial_zeros, partial_ref[...])
+  last_row = m_end_local // cfgs.dims.size_lhs_sublane
+  partial_ref[...] = jnp.where(
+      m_end_local % cfgs.dims.size_lhs_sublane == 0,
+      partial_zeros,
+      tiled_ref[last_row],
+  )
+
+
 def fill_metadata(
     lhs_group_sizes_ref: jax.Array,  # int32[size_lhs_group]
     group_offset_ref: jax.Array,  # int32[1]
@@ -1077,6 +1132,8 @@ def kernel_main(
     semaphore_ref: jax.Array | None,  # [1]
     *,
     cfgs: GmmConfigs,
+    pre_act_out_refs: tuple[jax.Array, ...] = (),  # return_pre_act: (gate, up) outputs, each shaped like out_ref
+    pre_act_partial_refs: tuple[jax.Array, ...] = (),  # return_pre_act: their partial-row scratch
 ):
   """Entry point for GMM kernel.
 
@@ -1132,8 +1189,22 @@ def kernel_main(
         num_gm,
         dims=cfgs.dims,
     )
+    # The pre-activation outputs have the act output's shape, so they zero the same rows with DMAs of the
+    # same size on the same semaphore; zero_out_end below waits for them one output at a time.
+    for pre_act_out_ref in pre_act_out_refs:
+      zero_out_start(
+          pre_act_out_ref,
+          zero_ref,  # pyrefly: ignore[bad-argument-type]
+          semaphore_ref,  # pyrefly: ignore[bad-argument-type]
+          metadata_ref,
+          num_gm,
+          dims=cfgs.dims,
+      )
 
   (lhs_spec, rhs_spec, ps_spec), out_spec = generate_block_specs(metadata_ref, cfgs)
+  if pre_act_out_refs:
+    # (act, gate, up) tiles share one index map; inner_kernel receives them as one tuple.
+    out_spec = ((out_spec,) * (1 + len(pre_act_out_refs)),)
 
   if cfgs.fuse_act is not None:
     if cfgs.transpose_rhs:
@@ -1174,10 +1245,54 @@ def kernel_main(
   # Works for both [size_m, size_n] and the 3D [size_m, D0, 128] output.
   out_in = out_ref.reshape(-1, cfgs.dims.size_lhs_sublane, *out_ref.shape[1:])
   scratches = [partial_out_ref, acc_ref, metadata_ref]
+  if pre_act_out_refs:
+    out_in = (out_in, *(r.reshape(-1, cfgs.dims.size_lhs_sublane, *r.shape[1:]) for r in pre_act_out_refs))
+    scratches[0] = (partial_out_ref, *pre_act_partial_refs)
   pipeline_fn(lhs_in, rhs_ref, ps_in, out_in, scratches=scratches)
 
   if cfgs.zero_init:
     zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims)  # pyrefly: ignore[bad-argument-type, unbound-name]
+    for pre_act_out_ref in pre_act_out_refs:
+      # pyrefly: ignore[bad-argument-type, unbound-name]
+      zero_out_end(pre_act_out_ref, semaphore_ref, zero_size, dims=cfgs.dims)
+
+
+def kernel_main_pre_act(
+    lhs_group_sizes_ref: jax.Array,
+    group_offset_ref: jax.Array,
+    lhs_ref: LhsRef,
+    rhs_ref: WeightsRef,
+    partial_sum_ref: jax.Array,
+    out_ref: jax.Array,  # act: [size_m, size_n // 2]
+    gate_out_ref: jax.Array,  # gate before the activation: [size_m, size_n // 2]
+    up_out_ref: jax.Array,  # up before the activation: [size_m, size_n // 2]
+    partial_out_ref: jax.Array,
+    acc_ref: jax.Array,
+    metadata_ref: MetadataRef,
+    zero_ref: jax.Array | None,
+    semaphore_ref: jax.Array | None,
+    partial_gate_ref: jax.Array,  # [size_lhs_sublane, tile_n]
+    partial_up_ref: jax.Array,  # [size_lhs_sublane, tile_n]
+    *,
+    cfgs: GmmConfigs,
+):
+  """kernel_main for return_pre_act: pallas_call passes the gate and up outputs right after act."""
+  kernel_main(
+      lhs_group_sizes_ref,
+      group_offset_ref,
+      lhs_ref,
+      rhs_ref,
+      partial_sum_ref,
+      out_ref,
+      partial_out_ref,
+      acc_ref,
+      metadata_ref,
+      zero_ref,
+      semaphore_ref,
+      cfgs=cfgs,
+      pre_act_out_refs=(gate_out_ref, up_out_ref),
+      pre_act_partial_refs=(partial_gate_ref, partial_up_ref),
+  )
 
 
 def calculate_tiling(
@@ -1395,6 +1510,8 @@ def get_cost_estimate(cfgs: GmmConfigs):
     rhs_bytes += dims.size_group * dims.size_n * fp32_bytes
 
   out_bytes = dims.size_m * cfgs.out_size_n * cfgs.out_dtype.itemsize
+  if cfgs.return_pre_act:
+    out_bytes *= 3  # act, gate and up
 
   total_bytes = lhs_bytes + rhs_bytes + out_bytes
 
@@ -1406,6 +1523,7 @@ def get_cost_estimate(cfgs: GmmConfigs):
 
 
 def get_scope_name(cfgs: GmmConfigs) -> str:
+  """The kernel's named scope: its dims, tiles and the flags that change the kernel."""
   dims = cfgs.dims
   tiles = cfgs.tiles
   scope_name = (
@@ -1414,6 +1532,8 @@ def get_scope_name(cfgs: GmmConfigs) -> str:
   )
   if cfgs.transpose_rhs:
     scope_name += "-transpose_rhs"
+  if cfgs.return_pre_act:
+    scope_name += "-pre_act"
   return scope_name
 
 
@@ -1436,6 +1556,7 @@ def make_gmm_configs(
     lhs_scale: jax.Array | None = None,
     transpose_rhs: bool = False,
     out_is_3d: bool = False,
+    return_pre_act: bool = False,
 ):
   """Fills the GMM config for the GMM kernel."""
 
@@ -1562,6 +1683,8 @@ def make_gmm_configs(
       raise ValueError(f"3D out requires an even tile_m; got {tiles=}.")
   elif partial_sum is not None and partial_sum.ndim == 3:
     raise ValueError("A 3D partial_sum requires out_is_3d=True.")
+  if return_pre_act and (fuse_act is None or out_is_3d or partial_sum is not None):
+    raise NotImplementedError("return_pre_act needs fuse_act and a 2D output, and no partial_sum.")
 
   return GmmConfigs(
       dims=dims,
@@ -1575,11 +1698,15 @@ def make_gmm_configs(
       fuse_act=fuse_act,
       transpose_rhs=transpose_rhs,
       out_is_3d=out_is_3d,
+      return_pre_act=return_pre_act,
   )
 
 
 def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]:
+  """The kernel config flattened to the pallas_call metadata dict."""
   cfgs_dict = dataclasses.asdict(cfgs)
+  if not cfgs.return_pre_act:
+    del cfgs_dict["return_pre_act"]  # keeps the metadata of every other gmm_v2 call unchanged
   ret = {}
   for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict):
     key = jax.tree_util.keystr(path, simple=True, separator=".")
@@ -1601,6 +1728,7 @@ def get_metadata(cfgs: GmmConfigs) -> dict[str, str | int | float]:
         "fuse_act",
         "transpose_rhs",
         "out_is_3d",
+        "return_pre_act",
     ]
 )
 def gmm_v2(
@@ -1623,7 +1751,8 @@ def gmm_v2(
     fuse_act: str | None = None,
     transpose_rhs: bool = False,
     out_is_3d: bool = False,
-) -> jax.Array:
+    return_pre_act: bool = False,
+) -> jax.Array | tuple[jax.Array, jax.Array, jax.Array]:
   """GMM kernel implemented with emit_pipeline.
 
   Dynamically calculate offset lhs/out tiles to reduce redundant computations.
@@ -1661,9 +1790,16 @@ def gmm_v2(
       Both require tile_k / tile_n to be a multiple of 128 that divides the
       dim. Tiles whose D0 (tile // 128) is not a multiple of 8 use full-D0
       blocks and relayout the tile window in VMEM.
+    return_pre_act: With fuse_act, also return the gate and up halves of the
+      matmul before the activation (what the activation's backward needs), each
+      rounded to the output dtype as the unfused kernel rounds them. The
+      activation itself is computed from the f32 accumulator. 2D output only,
+      no partial_sum.
 
   Returns:
-    Output of shape [size_m, size_n] (or [size_m, size_n // 128, 128]).
+    Output of shape [size_m, size_n] (or [size_m, size_n // 128, 128]); with
+    fuse_act the output width is size_n // 2. With return_pre_act, the tuple
+    (act, gate, up), each of shape [size_m, size_n // 2].
   """
 
   del precision
@@ -1695,6 +1831,7 @@ def gmm_v2(
       lhs_scale=lhs_scale,
       transpose_rhs=transpose_rhs,
       out_is_3d=out_is_3d,
+      return_pre_act=return_pre_act,
   )
   dims = cfgs.dims
   tiles = cfgs.tiles
@@ -1775,6 +1912,9 @@ def gmm_v2(
     ]
   else:
     scratch_shapes += [None, None]
+  if cfgs.return_pre_act:
+    # Partial-row carries of the gate and up outputs (kernel_main_pre_act takes them last).
+    scratch_shapes += [pltpu.VMEM(partial_out_shape, cfgs.out_dtype) for _ in range(2)]
 
   aligned_n = align_to(cfgs.out_size_n, num_lanes)
   if cfgs.out_is_3d:
@@ -1794,9 +1934,11 @@ def gmm_v2(
     partial_sum_idx = sum(1 for x in leaves if x is not None)
     input_output_aliases = {partial_sum_idx: 0}
 
+  out_spec = pl.BlockSpec(memory_space=pltpu.HBM)
   out = pl.pallas_call(
-      functools.partial(kernel_main, cfgs=cfgs),
-      out_shape=out_init,
+      functools.partial(kernel_main_pre_act if cfgs.return_pre_act else kernel_main, cfgs=cfgs),
+      # return_pre_act: (act, gate, up), all of act's shape.
+      out_shape=(out_init,) * 3 if cfgs.return_pre_act else out_init,
       grid_spec=pltpu.PrefetchScalarGridSpec(
           num_scalar_prefetch=2,
           in_specs=[
@@ -1811,7 +1953,7 @@ def gmm_v2(
               ),
               partial_sum_spec,
           ],
-          out_specs=pl.BlockSpec(memory_space=pltpu.HBM),
+          out_specs=(out_spec,) * 3 if cfgs.return_pre_act else out_spec,
           scratch_shapes=scratch_shapes,  # pyrefly: ignore[bad-argument-type]
       ),
       compiler_params=pltpu.CompilerParams(
@@ -1823,6 +1965,8 @@ def gmm_v2(
       metadata=get_metadata(cfgs),  # pyrefly: ignore[bad-argument-type]
       input_output_aliases=input_output_aliases,
   )(group_sizes, group_offset, lhs_in, rhs_weights, partial_sum)
+  if cfgs.return_pre_act:
+    return tuple(o[:, : cfgs.out_size_n] for o in out)
   if cfgs.out_is_3d:
     return out
   return out[:, : cfgs.out_size_n]
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 41e385579f..1e8c412b0b 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -3583,28 +3583,105 @@ def valid_token_count(x, routing, route_metadata):
         return jnp.sum(jax.lax.dynamic_slice_in_dim(routing.group_sizes, experts_start, num_experts_per_shard))
       return jnp.sum(routing.group_sizes)
 
-    def get_gmm_for_local_experts(x, routing, route_metadata):
-      """Return a partial GMM function with preconfigured routing params."""
+    # moe_wi_gmm_fused_swiglu: apply the SwiGLU in the wi GMM's epilogue (megablox ops.gmm_silu). The unfused path
+    # writes the [buffer, 2 * mlp] GMM output and reads it back in an elementwise pass over every buffer row, padding
+    # included; the fused GMM writes the [buffer, mlp] activation of the rows it computes. When differentiated it also
+    # writes gate and up, and the backward is the unfused one. Works with split (w0, w1) and with the whole prefused
+    # wi of moe_unsplit_prefused_wi (w0 = [w0 | w1], w1 = None). Off, gmm_up_or_silu is gmm_up and nothing below runs.
+    fuse_silu_in_wi_gmm = self.config.moe_wi_gmm_fused_swiglu
+    if fuse_silu_in_wi_gmm:
+      # The config validator checks the config-level requirements; these two only the built layer knows.
+      unmet = [
+          name
+          for name, hit in (
+              ("weight quantization", self.quant is not None),
+              ("an explicit expert-weight all-gather", weight_gather),
+          )
+          if hit
+      ]
+      if unmet:
+        raise ValueError(f"moe_wi_gmm_fused_swiglu=True does not support {' or '.join(unmet)}.")
+
+    def gmm_up_or_silu(
+        x,
+        w0,
+        w1,
+        w0_bias,
+        w1_bias,
+        gmm_fn,
+        gmm_operands,
+        _weight_gather,
+        mask=None,
+        gather_weights=True,
+        return_weights=False,
+    ):
+      """gmm_up; with moe_wi_gmm_fused_swiglu, (silu(x @ w0) * (x @ w1), None) from one fused GMM.
+
+      `gmm_operands` are the routing operands that `gmm_fn` is bound to (get_gmm_operands_for_local_experts); the
+      fused GMM takes them explicitly. An output1 of None tells the caller the activation is already applied.
+      Biases, the mask and gather_weights only matter in configurations the key does not support (mlp_bias, an
+      explicit weight all-gather).
+      """
+      if not fuse_silu_in_wi_gmm:
+        return gmm_up(
+            x,
+            w0,
+            w1,
+            w0_bias,
+            w1_bias,
+            gmm_fn,
+            _weight_gather,
+            mask=mask,
+            gather_weights=gather_weights,
+            return_weights=return_weights,
+        )
+      if return_weights:
+        raise NotImplementedError("moe_wi_gmm_fused_swiglu does not support moe_accumulate_chunk_wgrad.")
+      del w0_bias, w1_bias, mask, gather_weights
+      _, wi_tile_size = get_wi_gmm_params()
+      # As gmm(): pad the rows to the forward tile and cast to the compute dtype.
+      inputs, padding_amount = max_utils.maybe_pad(x, fwd_tile("wi_tile_fwd_batch_seq"))
+      # Weights are stored as (G,K,2N): w0/w1 are adjacent slices (the same concatenate as gmm_up), or w0 is already
+      # the whole [w0 | w1] (moe_unsplit_prefused_wi).
+      w_fused = w0 if w1 is None else jnp.concatenate([w0, w1], axis=-1)
+      act = mblx.gmm_silu(
+          inputs.astype(self.dtype),
+          w_fused.astype(self.dtype),
+          gmm_operands["group_sizes"],
+          preferred_element_type=self.dtype,
+          tiling=wi_tile_size,
+          group_offset=gmm_operands["group_offset"],
+          use_gmm_v2_heuristic_tiling=self.config.use_gmm_v2_heuristic_tiling,
+          use_dlhs_transpose_rhs=self.config.moe_gmm_v2_dlhs_transpose_rhs,
+          fuse_dlhs_scale=self.config.moe_accumulate_wi_dlhs,
+          fuse_3d_dlhs_ps=self.config.moe_3d_dlhs_partial_sum,
+      )
+      return (act[: x.shape[0]] if padding_amount else act), None
+
+    def get_gmm_operands_for_local_experts(x, routing, route_metadata):
+      """The routing operands of the local experts' GMMs: group sizes, expert assignments, group offset."""
       num_ep = self.get_expert_parallelism_size()
       num_experts_per_shard = self.config.num_experts // num_ep
       if self.config.use_ring_of_experts and x.shape[0] < routing.sorted_selected_experts.shape[0]:
         local_group_sizes = routing.local_group_sizes
-        return functools.partial(
-            gmm,
-            group_sizes=local_group_sizes,
-            expert_assignments=routing.selected_experts,
-            group_offset=0,
-        )
+        return {
+            "group_sizes": local_group_sizes,
+            "expert_assignments": routing.selected_experts,
+            "group_offset": 0,
+        }
       if self.config.use_ragged_sort and self.config.use_ring_of_experts:
         experts_start = route_metadata.expert_shard_id * num_experts_per_shard
       else:
         experts_start = 0
-      return functools.partial(
-          gmm,
-          group_sizes=routing.group_sizes,
-          expert_assignments=routing.selected_experts,
-          group_offset=experts_start,
-      )
+      return {
+          "group_sizes": routing.group_sizes,
+          "expert_assignments": routing.selected_experts,
+          "group_offset": experts_start,
+      }
+
+    def get_gmm_for_local_experts(x, routing, route_metadata):
+      """Return a partial GMM function with preconfigured routing params."""
+      return functools.partial(gmm, **get_gmm_operands_for_local_experts(x, routing, route_metadata))
 
     def unsort_output_and_ra2a(
         intermediate_output,
@@ -3793,14 +3870,16 @@ def _cn(t):
       if self.config.mlp_bias:
         w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias)
 
-      gmm_fn = get_gmm_for_local_experts(x, routing, route_metadata)
-      up_res = gmm_up(
+      gmm_operands = get_gmm_operands_for_local_experts(x, routing, route_metadata)
+      gmm_fn = functools.partial(gmm, **gmm_operands)  # get_gmm_for_local_experts
+      up_res = gmm_up_or_silu(
           x,
           w0,
           w1,
           w0_bias,
           w1_bias,
           gmm_fn,
+          gmm_operands,
           weight_gather,
           mask=mask,
           gather_weights=gather_weights,
@@ -3828,7 +3907,8 @@ def _down_projection(
         w_rows = tc_buffer_row_weights(routing.tc_routing, jnp.ravel(routing.weights), output0.shape[0])
         if self.config.moe_tc_routing_checkpoint != "remat":
           w_rows = adc.checkpoint_name(w_rows, "moe_tc_routing_checkpoint")
-      intermediate_layer = self.apply_ffn_activation_maybe_fused(output0, output1)
+      # moe_wi_gmm_fused_swiglu: the wi GMM already applied the activation (output1 is None).
+      intermediate_layer = output0 if fuse_silu_in_wi_gmm else self.apply_ffn_activation_maybe_fused(output0, output1)
       if tc_prescaled:
         intermediate_layer = (intermediate_layer * w_rows[:, None]).astype(intermediate_layer.dtype)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
@@ -3983,14 +4063,16 @@ def _cn(t):
         if self.config.mlp_bias:
           w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias)
 
-        gmm_fn = get_gmm_for_local_experts(x, routing, route_metadata)
-        up_res = gmm_up(
+        gmm_operands = get_gmm_operands_for_local_experts(x, routing, route_metadata)
+        gmm_fn = functools.partial(gmm, **gmm_operands)  # get_gmm_for_local_experts
+        up_res = gmm_up_or_silu(
             x,
             w0,
             w1,
             w0_bias,
             w1_bias,
             gmm_fn,
+            gmm_operands,
             weight_gather,
             mask=mask,
             gather_weights=gather_weights,
@@ -4012,7 +4094,8 @@ def _cn(t):
         w_rows = tc_buffer_row_weights(routing.tc_routing, jnp.ravel(routing.weights), output0.shape[0])
         if self.config.moe_tc_routing_checkpoint != "remat":
           w_rows = adc.checkpoint_name(w_rows, "moe_tc_routing_checkpoint")
-      intermediate_layer = self.apply_ffn_activation_maybe_fused(output0, output1)
+      # moe_wi_gmm_fused_swiglu: the wi GMM already applied the activation (output1 is None).
+      intermediate_layer = output0 if fuse_silu_in_wi_gmm else self.apply_ffn_activation_maybe_fused(output0, output1)
       if tc_prescaled:
         intermediate_layer = (intermediate_layer * w_rows[:, None]).astype(intermediate_layer.dtype)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
diff --git a/tests/unit/megablox_gmm_silu_test.py b/tests/unit/megablox_gmm_silu_test.py
new file mode 100644
index 0000000000..df23bc29db
--- /dev/null
+++ b/tests/unit/megablox_gmm_silu_test.py
@@ -0,0 +1,149 @@
+# 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.
+
+"""TPU tests: the SwiGLU in the gmm_v2 epilogue (megablox ops.gmm_silu, moe_wi_gmm_fused_swiglu) against the unfused GMM.
+
+gmm_v2(fuse_act="silu", return_pre_act=True) returns (act, gate, up) from one accumulator: gate and up must be the
+unfused gmm_v2 output's halves bit for bit (they are rounded the same way), every row outside the computed groups must
+be zero in all three, and act must be silu(gate) * up computed from the float32 accumulator (close to, not equal to,
+the same activation of the bf16-rounded halves: that is the numerics change). The act-only call must give the same
+act. ops.gmm_silu's VJP is the unfused one: its gradients must equal, value for value, the VJP of the unfused path
+(ops.gmm, then silu(gate) * up on its bf16 output). Groups are skewed, with an empty group, padding rows past the
+last group, boundaries inside a sublane, and an expert-parallel shard (group_offset).
+
+Pallas TPU kernels: these run only on a TPU. On CPU the moe.py wiring of the key is covered by
+moe_wi_gmm_fused_swiglu_test.py with a jnp stand-in for the kernel.
+"""
+
+import functools
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.kernels.megablox import ops
+from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2
+
+pytestmark = pytest.mark.tpu_only  # Pallas TPU kernels.
+
+M, K, N = 1024, 512, 512  # rhs is [groups, K, 2 * N] = [gate | up]
+GROUP_SIZES = np.array([300, 0, 517, 133], np.int32)  # sum 950 < M: padding rows; group ends inside sublanes
+TILING = (256, 512, 1024, 256, 512, 512, 256, 512, 512)  # gmm's 9-tuple; gmm_silu halves the forward tile_n
+
+
+def _rand(key, shape, scale=1.0):
+  return (jax.random.normal(key, shape, jnp.float32) * scale).astype(jnp.bfloat16)
+
+
+def _reference_accumulator(lhs, rhs, group_sizes, group_offset):
+  """f32 lhs @ rhs[expert of row] for the rows of the local groups [group_offset, group_offset + groups), else 0."""
+  ends = jnp.cumsum(group_sizes)
+  local = jnp.sum(jnp.arange(lhs.shape[0])[:, None] >= ends[None, :], axis=1) - group_offset
+  out = jnp.zeros((lhs.shape[0], rhs.shape[-1]), jnp.float32)
+  for e in range(rhs.shape[0]):
+    out = out + jnp.where((local == e)[:, None], lhs.astype(jnp.float32) @ rhs[e].astype(jnp.float32), 0.0)
+  return out
+
+
+def _f32(x):
+  return np.asarray(x.astype(jnp.float32))
+
+
+class GmmV2PreActTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.assertEqual(jax.devices()[0].platform, "tpu")
+
+  @parameterized.product(group_offset=[0, 2], tile_k=[256, 512])
+  def test_pre_act_halves_are_the_unfused_output(self, group_offset, tile_k):
+    k0, k1 = jax.random.split(jax.random.key(group_offset * 7 + tile_k))
+    lhs = _rand(k0, (M, K))
+    rhs = _rand(k1, (len(GROUP_SIZES) - group_offset, K, 2 * N), 0.05)
+    gs, go = jnp.asarray(GROUP_SIZES), jnp.array([group_offset], jnp.int32)
+    unfused = jax.jit(
+        functools.partial(
+            gmm_v2.gmm_v2, tile_info=gmm_v2.TileSizes(256, tile_k, 512), preferred_element_type=jnp.bfloat16
+        )
+    )(lhs, rhs, gs, group_offset=go)
+    fused = functools.partial(
+        gmm_v2.gmm_v2, tile_info=gmm_v2.TileSizes(256, tile_k, 256), preferred_element_type=jnp.bfloat16, fuse_act="silu"
+    )
+    act, gate, up = jax.jit(functools.partial(fused, return_pre_act=True))(lhs, rhs, gs, group_offset=go)
+    act_only = jax.jit(fused)(lhs, rhs, gs, group_offset=go)
+
+    np.testing.assert_array_equal(_f32(gate), _f32(unfused[:, :N]))
+    np.testing.assert_array_equal(_f32(up), _f32(unfused[:, N:]))
+    np.testing.assert_array_equal(_f32(act_only), _f32(act))
+    # Rows outside the local groups: zero in every output.
+    computed = np.zeros(M, bool)
+    start = int(np.sum(GROUP_SIZES[:group_offset]))
+    computed[start : start + int(np.sum(GROUP_SIZES[group_offset:]))] = True
+    for out in (act, gate, up):
+      self.assertTrue(np.all(_f32(out)[~computed] == 0))
+    # act is silu on the f32 accumulator, rounded once to bf16.
+    acc = _reference_accumulator(lhs, rhs, gs, group_offset)
+    want = (jax.nn.silu(acc[:, :N]) * acc[:, N:]).astype(jnp.bfloat16)
+    np.testing.assert_allclose(_f32(act), _f32(want), rtol=1e-2, atol=1e-3)
+
+
+class GmmSiluTest(parameterized.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    self.assertEqual(jax.devices()[0].platform, "tpu")
+
+  @parameterized.parameters(0, 2)
+  def test_forward_and_vjp_match_the_unfused_path(self, group_offset):
+    k0, k1, k2 = jax.random.split(jax.random.key(group_offset), 3)
+    lhs = _rand(k0, (M, K))
+    rhs = _rand(k1, (len(GROUP_SIZES) - group_offset, K, 2 * N), 0.05)
+    cot = _rand(k2, (M, N))
+    gs, go = jnp.asarray(GROUP_SIZES), jnp.array([group_offset], jnp.int32)
+
+    def unfused(lhs, rhs):
+      out = ops.gmm(
+          lhs,
+          rhs,
+          gs,
+          preferred_element_type=jnp.bfloat16,
+          tiling=TILING,
+          group_offset=go,
+          use_tokamax_backend=True,
+          use_gmm_v2=True,
+      )
+      return (jax.nn.silu(out[:, :N]) * out[:, N:]).astype(jnp.bfloat16)
+
+    def fused(lhs, rhs):
+      return ops.gmm_silu(lhs, rhs, gs, preferred_element_type=jnp.bfloat16, tiling=TILING, group_offset=go)
+
+    @jax.jit
+    def run(lhs, rhs, cot):
+      y_u, vjp_u = jax.vjp(unfused, lhs, rhs)
+      y_f, vjp_f = jax.vjp(fused, lhs, rhs)
+      return y_u, vjp_u(cot), y_f, vjp_f(cot)
+
+    y_u, (dl_u, dr_u), y_f, (dl_f, dr_f) = run(lhs, rhs, cot)
+    # Forward: the f32-accumulator activation, close to the unfused one.
+    np.testing.assert_allclose(_f32(y_f), _f32(y_u), rtol=2e-2, atol=2e-3)
+    # Backward: the unfused backward on the same bf16 gate and up, so the same values.
+    np.testing.assert_array_equal(_f32(dl_f), _f32(dl_u))
+    np.testing.assert_array_equal(_f32(dr_f), _f32(dr_u))
+
+
+if __name__ == "__main__":
+  absltest.main()
diff --git a/tests/unit/moe_wi_gmm_fused_swiglu_test.py b/tests/unit/moe_wi_gmm_fused_swiglu_test.py
new file mode 100644
index 0000000000..9d75356fb6
--- /dev/null
+++ b/tests/unit/moe_wi_gmm_fused_swiglu_test.py
@@ -0,0 +1,549 @@
+# 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.
+"""moe_wi_gmm_fused_swiglu on CPU: megablox ops.gmm_silu's own code, its wiring in the MoE layer, and the validator.
+
+With the key, the up projection calls megablox `gmm_silu` (the tokamax gmm_v2 kernel with the SwiGLU in its epilogue)
+instead of `gmm` followed by the activation. The kernels are Pallas TPU kernels and do not run on CPU; their numerics
+are tested on TPU in megablox_gmm_silu_test.py. Two CPU checks cover everything around them:
+
+ops.gmm_silu (in-process, `GmmSiluVjpTest`). Only the Pallas entry `gmm_v2.gmm_v2` and gmm's TPU backward
+`ops._gmm_bwd` are replaced: the first by a jnp grouped matmul that honours fuse_act and return_pre_act and records the
+tile it is given, the second by a stand-in that records how it is called and passes its cotangent through as the input
+gradient. The rest is the real ops.py: gmm_silu, its custom VJP, the forward tile (tiling[2] halved: the fused kernel's
+tile_n counts the columns of each half), the residuals, and the backward, which must hand gmm's backward [d gate |
+d up] of silu(gate) * up on the bf16 halves, bit for bit, with this op's DLHS options bound to the right keyword
+arguments.
+
+The layer (a subprocess, so the forced device count applies): the shrunk Qwen3.5 routed MoE of
+moe_roe_local_routing_test.py with prefuse_moe_weights on 8 CPU devices (cp-as-ep, ragged sort, local routing, two
+token chunks). `gmm` is replaced by a grouped matmul over the same group sizes and group offset, and `gmm_silu` by the
+unfused composition of that same stand-in, i.e. silu(gate) * up of its bf16 output. The layer with the key must then
+reproduce the layer without it bit for bit (the output, the loss and every gradient): that checks what the key
+changes in moe.py (the rows and the routing operands handed to the fused GMM, the padding to the forward tile and its
+removal, the whole [w0 | w1] weight, and the skipped activation), not the kernel. The fused kernel's real numerics
+differ by rounding (silu reads the f32 accumulator). Cases: the plain and the pipelined chunk loop, a dropless buffer
+(where each shard's group offset is not 0), padding excluded from routing, the whole prefused wi
+(moe_unsplit_prefused_wi), with the other MoE keys of this series, and with those and the benchmark's MoE settings
+(expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, fp8 fake-quant,
+moe_gmm_v2_dlhs_transpose_rhs). Two-sided: the fused path must trace exactly one gmm_silu call per chunk and the plain
+one none; gmm_silu must get the routing operands, the wi tiling the plain path gives gmm, and the DLHS options of the
+config (moe_gmm_v2_dlhs_transpose_rhs is True in the benchmark case and False elsewhere); stand-ins that swap the gate
+and up halves, or that ignore the group offset, must FAIL.
+
+Key off, not compared with anything this change touches: the full-buffer ring-of-experts ragged sort (dropless), whose
+GMM routing operands get_gmm_operands_for_local_experts now builds for both paths, must match the ring of experts
+without the ragged sort (use_ragged_sort=False), and a GMM that ignores the group offset must fail that comparison.
+
+The validator rejects the key for each requirement it does not meet, naming it.
+"""
+
+import os
+import re
+import subprocess
+import sys
+from unittest import mock
+
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+import numpy as np
+import pytest
+
+from maxtext.configs import pyconfig
+from maxtext.kernels import megablox
+from maxtext.kernels.megablox import ops as megablox_ops
+from maxtext.kernels.megablox import pallas_mosaic_tpu_v2_gmm_kernel as gmm_v2
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+_NUM_DEVICES = 8
+_NUM_EXPERTS = 16
+_TOP_K = 4
+_EMB = 256
+_SEQ = 64
+_PASS_TOKEN = "MOE_WI_GMM_FUSED_SWIGLU_CHECKS_PASSED"
+_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false"
+# The ragged-sort and plain ring paths combine the same expert outputs, possibly in another order: bf16 rounding.
+_UNSORTED_REL_L2 = 1e-2
+
+
+def _cfg(**overrides):
+  """The shrunk Qwen3.5 routed-MoE config with prefused weights and the tokamax gmm_v2 path the key requires."""
+  kwargs = {
+      "run_name": "moe_wi_gmm_fused_swiglu",
+      "enable_checkpointing": False,
+      "log_config": False,
+      "skip_jax_distributed_system": True,
+      "model_name": "qwen3.5-397b-a17b",
+      "override_model_config": True,
+      "num_experts": _NUM_EXPERTS,
+      "num_experts_per_tok": _TOP_K,
+      "base_emb_dim": _EMB,
+      "base_mlp_dim": 256,
+      "base_moe_mlp_dim": 256,
+      "max_target_length": _SEQ,
+      "max_prefill_predict_length": _SEQ,
+      "per_device_batch_size": 0.25,
+      "dtype": "bfloat16",
+      "weight_dtype": "float32",
+      "float32_gate_logits": True,
+      "megablox": True,
+      "sparse_matmul": True,
+      "use_tokamax_gmm": True,
+      "use_gmm_v2": True,
+      "prefuse_moe_weights": True,
+      "use_ring_of_experts": True,
+      "use_ragged_sort": True,
+      "use_custom_sort_vjp": False,
+      "ragged_buffer_factor": 2.0,
+      "num_moe_token_chunks": 2,
+      "custom_mesh_and_rule": "cp-as-ep",
+      "ici_fsdp_parallelism": 1,
+      "ici_context_parallelism": 4,
+      "ici_expert_parallelism": 2,
+      "context_parallel_load_balance": False,
+      "ring_of_experts_local_routing": True,
+  }
+  kwargs.update(overrides)
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
+
+
+class FusedSwigluConfigTest(parameterized.TestCase):
+  """The key is rejected for each requirement of the fused gmm_v2 epilogue it does not meet."""
+
+  @parameterized.named_parameters(
+      ("not_prefused", {"prefuse_moe_weights": False}, "prefuse_moe_weights=True"),
+      ("dense", {"sparse_matmul": False}, "sparse_matmul=True"),
+      ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"),
+      ("vllm_rpa", {"attention": "vllm_rpa"}, "an attention other than vllm_rpa*"),
+      ("no_gmm_v2", {"use_gmm_v2": False}, "use_tokamax_gmm=True and use_gmm_v2=True"),
+      ("no_tokamax", {"use_tokamax_gmm": False}, "use_tokamax_gmm=True and use_gmm_v2=True"),
+      ("quantization", {"quantization": "int8"}, "no weight quantization"),
+      ("gelu", {"mlp_activations": ["gelu", "linear"]}, "mlp_activations=['silu', 'linear']"),
+      ("gpt_oss", {"decoder_block": "gpt_oss"}, "a decoder_block other than gpt_oss"),
+      ("clipped", {"mlp_activations_limit": 7.0}, "mlp_activations_limit <= 0 (no clipping)"),
+      ("mlp_bias", {"mlp_bias": True}, "mlp_bias=False"),
+      ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"),
+      ("flat_fsdp", {"moe_flat_fsdp_weights": True}, "moe_flat_fsdp_weights=False"),
+      ("tc_3d_gmm", {"moe_tc_ragged_3d_gmm": True}, "moe_tc_ragged_3d_gmm=False"),
+      ("quantized_token_ag", {"moe_quantize_token_all_gather": True}, "moe_quantize_token_all_gather=False"),
+      ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"),
+      ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwi_0"),
+      ("saved_gate", {"remat_policy": "custom", "mlpwi_0": "device"}, "remat_policy=custom saves it as mlpwi_0"),
+      ("offloaded_moe_up", {"remat_policy": "custom", "moe_mlpwi_1": "offload"}, "saves it as moe_mlpwi_1"),
+  )
+  def test_rejected(self, overrides, requirement):
+    with self.assertRaisesRegex(ValueError, f"moe_wi_gmm_fused_swiglu=True requires: .*{re.escape(requirement)}"):
+      _cfg(moe_wi_gmm_fused_swiglu=True, **overrides)
+
+  def test_fused_activation_rejected_next_to_it(self):
+    with self.assertRaisesRegex(ValueError, "moe_prefused_swiglu_vjp=True requires: .*moe_wi_gmm_fused_swiglu=False"):
+      _cfg(moe_wi_gmm_fused_swiglu=True, moe_prefused_swiglu_vjp=True)
+
+  def test_accepted(self):
+    self.assertTrue(_cfg(moe_wi_gmm_fused_swiglu=True, moe_unsplit_prefused_wi=True).moe_wi_gmm_fused_swiglu)
+    # remat_policy=full saves nothing, and the shared expert's mlpwi_1 is not a routed-expert tensor.
+    for overrides in ({"remat_policy": "full"}, {"remat_policy": "custom", "mlpwi_1": "device"}):
+      self.assertTrue(_cfg(moe_wi_gmm_fused_swiglu=True, **overrides).moe_wi_gmm_fused_swiglu)
+
+
+_OP_M, _OP_K, _OP_N = 64, 256, 128  # the op case: lhs [M, K], rhs [groups, K, 2N]; K == 2N, so a pass-through dlhs fits
+_OP_GROUP_SIZES = np.array([20, 0, 30, 10], np.int32)  # sum 60 < M: padding rows past the last group
+_OP_TILING = (64, 256, 256, 64, 256, 256, 64, 256, 256)  # gmm's 9-tuple; the fused forward tile is (64, 256, 128)
+
+
+def _grouped_f32(lhs, rhs, group_sizes, group_offset):
+  """f32 lhs @ rhs[local expert of row]: rows of the groups [group_offset, group_offset + groups), others 0."""
+  ends = jnp.cumsum(group_sizes)
+  local = jnp.sum(jnp.arange(lhs.shape[0])[:, None] >= ends[None, :], axis=1) - jnp.reshape(group_offset, ())
+  out = jnp.zeros((lhs.shape[0], rhs.shape[-1]), jnp.float32)
+  for e in range(rhs.shape[0]):
+    out = out + jnp.where((local == e)[:, None], lhs.astype(jnp.float32) @ rhs[e].astype(jnp.float32), 0.0)
+  return out
+
+
+@pytest.mark.cpu_only
+class GmmSiluVjpTest(parameterized.TestCase):
+  """ops.gmm_silu around stand-ins for the TPU kernel and gmm's TPU backward (see the module docstring)."""
+
+  @parameterized.product(
+      group_offset=(0, 1),
+      # (use_dlhs_transpose_rhs, fuse_dlhs_scale, fuse_3d_dlhs_ps): each alone, so a swapped binding shows.
+      dlhs_options=((False, False, False), (True, False, False), (False, True, False), (False, False, True)),
+  )
+  def test_custom_vjp(self, group_offset, dlhs_options):
+    dlhs_transpose, fuse_scale, fuse_3d = dlhs_options
+    seen = {"tile_n": [], "return_pre_act": [], "bwd": []}
+
+    def kernel_stand_in(
+        lhs,
+        rhs,
+        group_sizes,
+        tile_info,
+        preferred_element_type,
+        group_offset=None,
+        fuse_act=None,
+        return_pre_act=False,
+        **kwargs,
+    ):
+      del kwargs
+      self.assertEqual(fuse_act, "silu")
+      seen["tile_n"].append(tile_info.tile_n)
+      seen["return_pre_act"].append(return_pre_act)
+      acc = _grouped_f32(lhs, rhs, group_sizes, 0 if group_offset is None else group_offset)
+      n = acc.shape[-1] // 2
+      act = (jax.nn.silu(acc[:, :n]) * acc[:, n:]).astype(preferred_element_type)
+      if not return_pre_act:
+        return act
+      return act, acc[:, :n].astype(preferred_element_type), acc[:, n:].astype(preferred_element_type)
+
+    def gmm_bwd_stand_in(**kwargs):  # keyword arguments only: ops._gmm_silu_bwd must name every one
+      seen["bwd"].append(kwargs)
+      lhs, rhs = kwargs["residual"][:2]
+      return kwargs["grad"][0].astype(lhs.dtype), jnp.zeros_like(rhs), None, None, None, None, None
+
+    k0, k1, k2 = jax.random.split(jax.random.key(group_offset), 3)
+    lhs = jax.random.normal(k0, (_OP_M, _OP_K), jnp.float32).astype(jnp.bfloat16)
+    rhs = (jax.random.normal(k1, (len(_OP_GROUP_SIZES) - group_offset, _OP_K, 2 * _OP_N)) * 0.05).astype(jnp.bfloat16)
+    cotangent = jax.random.normal(k2, (_OP_M, _OP_N), jnp.float32).astype(jnp.bfloat16)
+    group_sizes, offset = jnp.asarray(_OP_GROUP_SIZES), jnp.array([group_offset], jnp.int32)
+
+    def gmm_silu(l, r):
+      return megablox_ops.gmm_silu(
+          l,
+          r,
+          group_sizes,
+          preferred_element_type=jnp.bfloat16,
+          tiling=_OP_TILING,
+          group_offset=offset,
+          use_dlhs_transpose_rhs=dlhs_transpose,
+          fuse_dlhs_scale=fuse_scale,
+          fuse_3d_dlhs_ps=fuse_3d,
+      )
+
+    with (
+        mock.patch.object(gmm_v2, "gmm_v2", kernel_stand_in),
+        mock.patch.object(megablox_ops, "_gmm_bwd", gmm_bwd_stand_in),
+    ):
+      gmm_silu(lhs, rhs)  # not differentiated: the activation alone
+      _, vjp = jax.vjp(gmm_silu, lhs, rhs)
+      dlhs, _ = vjp(cotangent)
+
+    # The forward tile: tiling[2] halved, every call; act alone undifferentiated, act + gate + up differentiated.
+    self.assertEqual(seen["tile_n"], [_OP_TILING[2] // 2] * 2)
+    self.assertEqual(seen["return_pre_act"], [False, True])
+    # The backward: gmm's backward gets [d gate | d up] of silu(gate) * up on the bf16 halves, bit for bit.
+    acc = _grouped_f32(lhs, rhs, group_sizes, offset)
+    gate, up = acc[:, :_OP_N].astype(jnp.bfloat16), acc[:, _OP_N:].astype(jnp.bfloat16)
+    _, swiglu_vjp = jax.vjp(lambda g, u: (jax.nn.silu(g) * u).astype(g.dtype), gate, up)
+    want = jnp.concatenate(swiglu_vjp(cotangent), axis=-1)
+    np.testing.assert_array_equal(np.asarray(dlhs, np.float32), np.asarray(want, np.float32))
+    self.assertLen(seen["bwd"], 1)
+    bwd = seen["bwd"][0]
+    np.testing.assert_array_equal(np.asarray(bwd["grad"][0], np.float32), np.asarray(want, np.float32))
+    self.assertEqual(bwd["grad"][1:], (None, None))
+    # Its residuals are this op's operands, and its options are this op's (the plain bf16 tokamax gmm_v2 GEMM).
+    res_lhs, res_rhs, res_group_sizes, res_offset, *res_rest = bwd["residual"]
+    for got, want_operand in ((res_lhs, lhs), (res_rhs, rhs), (res_group_sizes, group_sizes), (res_offset, offset)):
+      np.testing.assert_array_equal(np.asarray(got), np.asarray(want_operand))
+    self.assertEqual(tuple(res_rest), (None, False, False, None))
+    expected_options = {
+        "lhs_dtype": jnp.bfloat16,
+        "rhs_dtype": jnp.bfloat16,
+        "preferred_element_type": jnp.bfloat16,
+        "tiling": _OP_TILING,
+        "transpose_rhs": False,
+        "quantization_rule": None,
+        "use_tokamax_backend": True,
+        "use_gmm_v2": True,
+        "use_gmm_v2_heuristic_tiling": False,
+        "return_rhs": False,
+        "return_lhs": False,
+        "fuse_dlhs_scale": fuse_scale,
+        "use_dlhs_transpose_rhs": dlhs_transpose,
+        "out_is_3d": False,
+        "fuse_3d_dlhs_ps": fuse_3d,
+        "tc_unsort_cfg": None,
+    }
+    self.assertEqual({name: bwd[name] for name in expected_options}, expected_options)
+
+  def test_odd_forward_tile_is_rejected(self):
+    lhs = jnp.zeros((_OP_M, _OP_K), jnp.bfloat16)
+    rhs = jnp.zeros((len(_OP_GROUP_SIZES), _OP_K, 2 * _OP_N), jnp.bfloat16)
+    with self.assertRaisesRegex(ValueError, "must halve to a multiple of 128"):
+      megablox_ops.gmm_silu(lhs, rhs, jnp.asarray(_OP_GROUP_SIZES), jnp.bfloat16, (64, 256, 192) + _OP_TILING[3:])
+
+
+@pytest.mark.cpu_only
+def test_fused_swiglu_wiring_on_cpu_mesh():
+  """Runs `main` below in a child process with 8 CPU devices."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}"
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+  print(result.stdout)
+  assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+def _grouped_matmul(lhs, rhs, group_sizes, group_offset):
+  """f32 lhs @ rhs[local expert of row]: rows of the groups [group_offset, group_offset + groups), others 0."""
+  ends = jnp.cumsum(group_sizes)
+  local = jnp.sum(jnp.arange(lhs.shape[0])[:, None] >= ends[None, :], axis=1) - group_offset
+  out = jnp.zeros((lhs.shape[0], rhs.shape[-1]), jnp.float32)
+  for e in range(rhs.shape[0]):
+    out = out + jnp.where((local == e)[:, None], lhs.astype(jnp.float32) @ rhs[e].astype(jnp.float32), 0.0)
+  return out
+
+
+def _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type=jnp.float32, tiling=None, group_offset=None, **kwargs):
+  """jnp stand-in for megablox gmm (the TPU kernel), for the calls the layer makes here."""
+  del tiling
+  for name in ("partial_sum", "return_lhs", "return_rhs"):
+    assert not kwargs.get(name), name
+  offset = 0 if group_offset is None else group_offset
+  return _grouped_matmul(lhs, rhs, group_sizes, offset).astype(preferred_element_type)
+
+
+def _gmm_silu_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset=None, **kwargs):
+  """megablox gmm_silu as the unfused composition of the gmm stand-in: silu(gate) * up of its rounded output."""
+  del kwargs
+  out = _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset)
+  n = out.shape[-1] // 2
+  return (jax.nn.silu(out[:, :n]) * out[:, n:]).astype(preferred_element_type)
+
+
+def _halves_swapped_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset=None, **kwargs):
+  """Deliberately wrong: treats the up half as the gate and the gate half as up."""
+  del kwargs
+  out = _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset)
+  n = out.shape[-1] // 2
+  return (jax.nn.silu(out[:, n:]) * out[:, :n]).astype(preferred_element_type)
+
+
+def _offset_ignored_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, group_offset=None, **kwargs):
+  """Deliberately wrong: ignores the group offset, i.e. computes the first shard's experts on every shard."""
+  del group_offset
+  return _gmm_silu_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, None, **kwargs)
+
+
+def _build(cfg, mesh):
+  return moe.RoutedMoE(
+      config=cfg,
+      num_experts=cfg.num_experts,
+      num_experts_per_tok=cfg.num_experts_per_tok,
+      mesh=mesh,
+      kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+      kernel_axes=("embed", "mlp"),
+      intermediate_dim=cfg.moe_mlp_dim,
+      weight_dtype=cfg.weight_dtype,
+      dtype=cfg.dtype,
+      rngs=nnx.Rngs(params=0),
+  )
+
+
+def _inputs(cfg, mode):
+  batch = int(cfg.per_device_batch_size * _NUM_DEVICES)
+  x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32)
+  token_mask = None
+  if mode == "padmask":
+    lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)])
+    token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32)
+  return x, token_mask
+
+
+def _run(cfg, mesh, params, x, token_mask, gmm_silu, gmm=None):
+  """Loss, output and gradients with the stand-ins; the keyword arguments of every gmm_silu call traced, and the
+  tiling of the plain path's up-projection gmm calls ([K, 2N] weight)."""
+  silu_calls, wi_tilings = [], []
+  gmm = gmm or _gmm_stand_in
+
+  def recorded_gmm_silu(*args, **kwargs):
+    silu_calls.append(kwargs)
+    return gmm_silu(*args, **kwargs)
+
+  def recorded_gmm(lhs, rhs, *args, **kwargs):
+    if rhs.shape[-1] == 2 * cfg.moe_mlp_dim:
+      wi_tilings.append(tuple(kwargs["tiling"]))
+    return gmm(lhs, rhs, *args, **kwargs)
+
+  with (
+      mock.patch.object(megablox, "gmm", recorded_gmm),
+      mock.patch.object(megablox, "gmm_silu", recorded_gmm_silu),
+      jax.set_mesh(mesh),
+      nn_partitioning.axis_rules(cfg.logical_axis_rules),
+  ):
+    graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+
+    def loss_fn(params, x):
+      out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, token_mask=token_mask)
+      loss = jnp.mean(out.astype(jnp.float32) ** 2)
+      if lb_loss is not None:
+        loss = loss + lb_loss.astype(jnp.float32)
+      return loss, out
+
+    jax.make_jaxpr(loss_fn)(params, x)
+    traced = list(silu_calls)
+    (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x)
+  return jax.device_get((loss, out, grads)), traced, sorted(set(wi_tilings))
+
+
+def _bit_identical(label, old, new):
+  """Prints the per-leaf report; True iff the output, the loss and every gradient are bit-identical."""
+  (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new
+  named = [("output", out_old, out_new), ("loss", loss_old, loss_new)] + [
+      (jax.tree_util.keystr(path), a, b)
+      for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new))
+  ]
+  passed, lines = True, []
+  for name, a, b in named:
+    a, b = np.asarray(a).astype(np.float32), np.asarray(b).astype(np.float32)
+    same = bool(np.array_equal(a, b)) and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0
+    passed &= same
+    lines.append(f"  [{'ok' if same else 'FAIL'}] {name}: bit-identical={same} max|d|={np.max(np.abs(a - b)):.3e}")
+  print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+def _case(label, overrides, mode="real", gmm_silu=_gmm_silu_stand_in):
+  """The layer without and with the key on the same parameters; `gmm_silu` is the fused GMM's stand-in."""
+  cfg_old = _cfg(**overrides)
+  cfg_new = _cfg(**overrides, moe_wi_gmm_fused_swiglu=True)
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes)
+  x, token_mask = _inputs(cfg_old, mode)
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...)
+  old, calls_old, wi_tilings = _run(cfg_old, mesh, params, x, token_mask, gmm_silu)
+  new, calls_new, _ = _run(cfg_new, mesh, params, x, token_mask, gmm_silu)
+  structure = not calls_old and len(calls_new) == cfg_new.num_moe_token_chunks
+  # gmm_silu gets the tiling the plain path gives the up-projection gmm, and the config's DLHS options.
+  want = {
+      "tiling": wi_tilings[0] if len(wi_tilings) == 1 else wi_tilings,
+      "use_gmm_v2_heuristic_tiling": cfg_new.use_gmm_v2_heuristic_tiling,
+      "use_dlhs_transpose_rhs": cfg_new.moe_gmm_v2_dlhs_transpose_rhs,
+      "fuse_dlhs_scale": cfg_new.moe_accumulate_wi_dlhs,
+      "fuse_3d_dlhs_ps": cfg_new.moe_3d_dlhs_partial_sum,
+      "preferred_element_type": jnp.bfloat16,
+  }
+  got = [{name: tuple(c[name]) if name == "tiling" else c[name] for name in want} for c in calls_new]
+  bindings = bool(got) and all(g == want for g in got)
+  print(
+      f"{label}: [{'ok' if structure else 'FAIL'}] gmm_silu calls traced plain={len(calls_old)}"
+      f" fused={len(calls_new)}; [{'ok' if bindings else 'FAIL'}] gmm_silu arguments {got[:1]} (want {want})",
+      flush=True,
+  )
+  return _bit_identical(label, old, new) and structure and bindings
+
+
+def _gmm_offset_ignored_stand_in(
+    lhs, rhs, group_sizes, preferred_element_type=jnp.float32, tiling=None, group_offset=None, **kwargs
+):
+  """Deliberately wrong: a gmm that ignores the group offset, i.e. computes the first shard's experts everywhere."""
+  del group_offset
+  return _gmm_stand_in(lhs, rhs, group_sizes, preferred_element_type, tiling, None, **kwargs)
+
+
+def _ragged_dropless_matches_unsorted(label, ragged_gmm=None):
+  """Key off: the full-buffer (dropless) ring-of-experts ragged sort against the ring of experts without it.
+
+  Both take their GMM routing operands from get_gmm_operands_for_local_experts, but only the ragged path uses a
+  non-zero group offset on a full buffer. Not a comparison of the key: a reference the key does not touch.
+  `ragged_gmm` replaces the gmm stand-in in the ragged run only.
+  """
+  cfg_ragged = _cfg(ragged_buffer_factor=-1.0)
+  cfg_unsorted = _cfg(ragged_buffer_factor=-1.0, use_ragged_sort=False)
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg_ragged), cfg_ragged.mesh_axes)
+  x, token_mask = _inputs(cfg_ragged, "real")
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_ragged.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg_ragged, mesh), nnx.Param, ...)
+  reference, _, _ = _run(cfg_unsorted, mesh, params, x, token_mask, _gmm_silu_stand_in)
+  ragged, _, _ = _run(cfg_ragged, mesh, params, x, token_mask, _gmm_silu_stand_in, gmm=ragged_gmm)
+  (loss_ref, out_ref, grads_ref), (loss_new, out_new, grads_new) = reference, ragged
+  named = [("output", out_ref, out_new), ("loss", loss_ref, loss_new)] + [
+      (jax.tree_util.keystr(path), a, b)
+      for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_ref)[0], jax.tree_util.tree_leaves(grads_new))
+  ]
+  passed, lines = True, []
+  for name, a, b in named:
+    a, b = np.asarray(a, np.float32), np.asarray(b, np.float32)
+    rel = float(np.linalg.norm(b - a) / max(np.linalg.norm(a), 1e-30))
+    ok = rel < _UNSORTED_REL_L2 and bool(np.isfinite(b).all()) and float(np.max(np.abs(a))) > 0
+    passed &= ok
+    lines.append(f"  [{'ok' if ok else 'FAIL'}] {name}: rel_l2 vs use_ragged_sort=False {rel:.2e}")
+  print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+_BENCH_FLAGS = {
+    "moe_chunk_pipeline": True,
+    "moe_expert_weight_prefetch": "forward",
+    "ring_of_experts_row_major_reduce_scatter": True,
+    "moe_routing_maps": "device",
+    "ragged_buffer_factor": 2.75,
+    "fp8_moe_fake_quant": True,
+    "moe_exclude_padding_from_routing": True,
+    "moe_gmm_v2_dlhs_transpose_rhs": True,
+}
+
+
+def main():
+  assert jax.device_count() == _NUM_DEVICES, jax.devices()
+  pipelined = {"moe_chunk_pipeline": True}
+  other_keys = {
+      "moe_unsplit_prefused_wi": True,
+      "moe_routing_index_math": "count",
+      "moe_roe_packed_routing_gather": True,
+  }
+  results = {
+      "plain_loop": _case("plain_loop", {}),
+      "pipelined": _case("pipelined", pipelined),
+      "dropless": _case("dropless", {"ragged_buffer_factor": -1.0}),
+      "padding_excluded": _case("padding_excluded", {**pipelined, "moe_exclude_padding_from_routing": True}, "padmask"),
+      "unsplit_wi": _case("unsplit_wi", {**pipelined, "moe_unsplit_prefused_wi": True}),
+      "with_other_keys": _case("with_other_keys", {**pipelined, **other_keys}),
+      "with_other_keys_bench_flags": _case("with_other_keys_bench_flags", {**_BENCH_FLAGS, **other_keys}, "padmask"),
+      "key_off_ragged_dropless_matches_unsorted": _ragged_dropless_matches_unsorted(
+          "key_off_ragged_dropless_matches_unsorted"
+      ),
+  }
+  controls = {
+      "control_halves_swapped": _case("control_halves_swapped", pipelined, gmm_silu=_halves_swapped_stand_in),
+      "control_offset_ignored": _case(
+          "control_offset_ignored", {"ragged_buffer_factor": -1.0}, gmm_silu=_offset_ignored_stand_in
+      ),
+      "control_key_off_ragged_offset_ignored": _ragged_dropless_matches_unsorted(
+          "control_key_off_ragged_offset_ignored", ragged_gmm=_gmm_offset_ignored_stand_in
+      ),
+  }
+  print("SUMMARY", {**results, **controls}, flush=True)
+  failed = [k for k, v in results.items() if not v]
+  undetected = [k for k, v in controls.items() if v]
+  if failed or undetected:
+    print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True)
+    sys.exit(1)
+  print(_PASS_TOKEN, flush=True)
+
+
+if __name__ == "__main__":
+  main()

From 63cba2632c123a94d10486b84872270fec7a86e1 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 03:10:35 -0700
Subject: [PATCH 07/10] MoE: fuse the ring-of-experts down projection and
 unsort, with a backward that reads the GEMM input
 (moe_router_grad_from_wo_input)

On the ring-of-experts ragged path the down projection y = h @ wo (per
expert, over the sorted rows) feeds ring_ragged_unsort, whose backward
needs y for the routing-weight gradient <g, y>. Under remat the layer's
backward therefore recomputes the down-projection GEMM only to form y,
and the combine backward makes a pass over two [rows, hidden] tensors that
also writes dy = w * g.

moe_router_grad_from_wo_input runs the activation, the down projection and
the unsort as one custom-VJP op (ragged_sort.ring_ragged_unsort_down_projection).
Its forward is the GEMM followed by ring_ragged_unsort's forward, so the
output is bit-identical. Its backward never reads y: row j of the GEMM is
linear in h_j and the weight is a scalar per row, so
<g, h wo> = <g wo^T, h>, (w g) wo^T = w (g wo^T) and h^T (w g) = (w h)^T g.
It takes the GEMM's own VJP at the input w * h with the unweighted
cotangent g (u = g wo^T, dwo = (w h)^T g), then dh = w u and the routing
gradient <u, h>. The GEMMs keep their shapes, the recomputed forward GEMM
has no reader and XLA drops it, and the weighting and the dot run over
[rows, mlp] tensors. In the pipelined chunk loop, tie 2 moves into the op
as a forward-only barrier on the GEMM output. The cotangent rows outside
this shard's groups are left unmasked: the GEMM backward never mixes them
into group rows and every output is masked.

The op shares its code with ring_ragged_unsort: the unsort's forward and
the fan-out half of its backward are now helpers (_ring_unsort_window,
_ring_unsort_combine, _ring_unsort_fan_out, with _unsort_kwargs for their
kernel options) that both call, so the fused forward is the unsort's by
construction. The helpers keep the unsort's ops in their order: lowered
for TPU v7x (forward plus backward of a shrunk Qwen3.5 MoE layer, the
benchmark's MoE settings, also the dropless retry program), the key-off
program is the same as the parent commit's and every key-on program the
same as before this refactor.

Numerics change: each gradient keeps one rounding to the activation dtype,
at a different point (w * h instead of w * g; dh after the product with w
instead of before the GEMM; the routing gradient from g wo^T instead of
y), so the gradients agree with the two ops apart to bf16 rounding, not bit
for bit. On the CPU layer below they differ by 1e-3 to 4.2e-3 relative L2.
On the op alone, against the same function in float32, the fused
gradients are as far as the two ops' are: the ratio of the two errors is
0.99-1.01 for dh and dwo and 0.94-1.07 for the routing gradient, over 6
seeds with a full and a dropping buffer.

The validator (the shared form) requires the ring-of-experts ragged sort
(sparse_matmul, te_moe_block=False, use_ring_of_experts and
use_ragged_sort) and a down projection that is the bare GEMM of the
sorted rows: no tensor parallelism (also checked on the mesh at trace
time), MLP bias, weight quantization, TC ragged sort, chunk weight-gradient
accumulation, pre-quantized unsort backward, token all-gather quantization
or embedding chunks; and a remat policy that saves no name of the routed
experts' down-projection output (mlpwo, moe_mlpwo: the minimal* and
save_dot_except_mlpwi presets, or custom with either saved or offloaded),
the rule of the series: the op never forms that output, so such a policy
would silently stop saving the MoE tensor. It composes with
moe_routing_index_math, the prefused-weight keys and
moe_wi_gmm_fused_swiglu. The layer raises if it is asked to return the
weights for chunk weight-gradient accumulation; _moe_body's three exits
now share one construction of its outputs.

Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per step), switched
by an environment variable in a build that a TPU lowering check shows
compiling to the same program as this commit: regular micro-batch 10.022
-> 9.765 s (-0.257 s; arm o275x-mf-a14-prof against its in-pod control
o275x-mf-ctl), 99.25 -> 97.05 s per step, 0 dropped tokens in both. All
five MoE keys of this series together (moe_routing_index_math=count,
moe_unsplit_prefused_wi, moe_roe_packed_routing_gather, this key and
moe_wi_gmm_fused_swiglu), in one pod on a build with the staged step
boundary: 10.0445 -> 9.510 s per micro-batch (-0.535 s), 91.87 -> 87.04 s
per step (o275x-s5moe5-l3-prof against o275x-s5c-l3-prof). The same
control measured in different pods spreads by +/-0.02-0.03 s per
micro-batch. The timing arms run with the benchmark's log-prob gates on
(loss and gradient norm read 0), so they say nothing about numerics.

TPU numerics: a 3-step A/B with the gates off, against a reference run and
a replicate of it, passes all six checks with this key in both measured
stacks: with moe_routing_index_math=count, moe_unsplit_prefused_wi,
moe_prefused_swiglu_vjp and moe_roe_packed_routing_gather (run nxs6c: loss
relative difference 4.6e-8, N1 2.9e-6, gradient-norm difference 7.0e-6
against a replicate spread of 1.9e-5), and with the fused SiLU GMM in
place of the fused activation (run nxs6silu: gradient-norm difference
2.5e-6). Caveats: stack level, not this key alone; the environment-switch
build, not this branch's binary.

Tests (moe_router_grad_from_wo_input_test.py, new), on 8 CPU devices:
- The op on a 4-device EP mesh with a jnp grouped matmul as the down
  projection, the full buffer and a packed buffer that drops tokens,
  "sort" and "count": output bit-identical to the two ops apart, every
  gradient within relative L2 1e-2 of theirs; a forward tie comes back
  unchanged with its cotangent passed through and is tied to the GEMM
  output by a barrier; "count" without the forward permutation raises.
- Precision of the op alone (its input h given, so only the backward
  rounds), against the same function written densely in float32 (the
  ragged gather-reduce returns bf16, so the two ops cannot be the
  reference): each fused gradient at most 1.10x as far as the two ops',
  on 6 seeds x {full buffer, dropping buffer}, with routing weights where
  bf16 rounds them worst. A backward that rounds the routing weight to
  bf16 measures 1.31-1.34 and one that shrinks the routing gradient by
  0.5% 3.0-3.3; both are caught.
- The shrunk Qwen3.5 routed MoE: output and loss bit-identical to the key
  off, gradients within relative L2 1e-2: plain and pipelined loop,
  dropless, "count" with padding excluded, saved routing maps, prefused
  weights with the whole wi and the fused SwiGLU VJP, and the benchmark's
  MoE settings with every compatible key of the series on against all
  off. The compiled rematerialized backward has exactly
  num_moe_token_chunks GEMM loops fewer (19 -> 17 for 2 chunks), the
  pipelined forward keeps its four barriers, and the next chunk's routed
  tokens are read from the op's tie output.
- On an expert 4 x tensor 2 mesh the layer runs without the key and
  raises at trace time when the key reaches it past the validator.
- A fused backward that drops the routing weight is caught at the op and
  in the layer; the validator rejects the key for each unmet requirement,
  naming it.
---
 src/maxtext/configs/base.yml                  |   3 +
 src/maxtext/configs/types.py                  |  44 +
 src/maxtext/kernels/ragged/ragged_sort.py     | 538 ++++++++----
 src/maxtext/layers/moe.py                     | 177 +++-
 .../moe_router_grad_from_wo_input_test.py     | 803 ++++++++++++++++++
 5 files changed, 1377 insertions(+), 188 deletions(-)
 create mode 100644 tests/unit/moe_router_grad_from_wo_input_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 60cbcfc1bc..a43f4198cf 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -348,6 +348,9 @@ ring_of_experts_local_routing: false # ring of experts: route each shard's own t
 # With ring_of_experts_local_routing + moe_chunk_pipeline: route every token chunk once and all-gather ids + weights as
 # ONE float32 array per layer (tagged moe_routing_maps) instead of two gathers per chunk. Same routing.
 moe_roe_packed_routing_gather: false
+# Ring of experts with ragged sort: activation + down projection + unsort as one op whose backward reads the down
+# projection's input, not its output (no recomputed wo GEMM under remat). Numerics change: gradients to bf16 rounding.
+moe_router_grad_from_wo_input: false
 ring_of_experts_row_major_reduce_scatter: false # ring of experts: lay out the expert-parallel reduce-scatter input row-major in the forward pass, so XLA does not copy the combine output into another layout and the result back
 moe_topk_before_ep_all_gather: false # ring of experts: run top-k on local tokens and all-gather top-k weights/ids instead of router logits
 te_moe_block: false # whether to use TransformerEngine's MoEBlock for routed MoE layers
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 0be85e3df3..405cb657d8 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1127,6 +1127,8 @@ def _reject_unless(setting: str, requirements: list[tuple[bool, str]]) -> None:
 # The checkpoint names of the routed experts' up-projection output in the sparse-matmul path: the gate half carries
 # mlpwi_0 and moe_mlpwi_0, the up half moe_mlpwi_1 (moe.py gmm_up).
 _MOE_UP_PROJECTION_NAMES = ("mlpwi_0", "moe_mlpwi_0", "moe_mlpwi_1")
+# ... and of its down-projection output: mlpwo and moe_mlpwo.
+_MOE_DOWN_PROJECTION_NAMES = ("mlpwo", "moe_mlpwo")
 
 
 class MoEGeneral(BaseModel):
@@ -1414,6 +1416,24 @@ class MoEGeneral(BaseModel):
           " draws per call, so one draw for the layer differs from one per chunk)."
       ),
   )
+  moe_router_grad_from_wo_input: bool = Field(
+      False,
+      description=(
+          "Ring of Experts with the ragged sort: run the activation, the down projection (wo GEMM) and the unsort"
+          " as one op whose backward does not read the down projection's output. The routing-weight gradient"
+          " <g, h wo> is computed as <g wo^T, h> from the down projection's input and the routing weight is applied"
+          " on the [rows, mlp] side, so a rematerialized layer no longer recomputes the down-projection GEMM in the"
+          " backward, and the combine backward's pass over [rows, hidden] tensors moves to [rows, mlp] ones. The"
+          " forward is bit-identical. NUMERICS CHANGE: each gradient keeps one rounding to the activation dtype, at"
+          " a different point, so the gradients agree with the unfused ops to bf16 rounding, not bit for bit."
+          " Requires sparse_matmul=True, te_moe_block=False, use_ring_of_experts=True and use_ragged_sort=True, and"
+          " the down projection to be the bare GEMM of the sorted rows: no tensor parallelism, mlp_bias=False, no"
+          " weight quantization, moe_tc_ragged_sort, moe_accumulate_chunk_wgrad, moe_bwd_prequant_before_unsort and"
+          " moe_quantize_token_all_gather off, num_moe_emb_chunks=0; and a remat policy that does not save the routed"
+          " experts' down-projection output (mlpwo, moe_mlpwo; the minimal* and save_dot_except_mlpwi policies save"
+          " mlpwo): the fused op never forms it, so the policy would silently stop saving the MoE tensor."
+      ),
+  )
   ring_of_experts_row_major_reduce_scatter: bool = Field(
       False,
       description=(
@@ -1839,6 +1859,30 @@ def validate_moe_roe_packed_routing_gather(self) -> "MoEGeneral":
       )
     return self
 
+  @model_validator(mode="after")
+  def validate_moe_router_grad_from_wo_input(self) -> "MoEGeneral":
+    """The fused op needs the ring-of-experts ragged sort and a down projection that is the bare GEMM of its rows."""
+    if self.moe_router_grad_from_wo_input:
+      _reject_unless(
+          "moe_router_grad_from_wo_input=True",
+          self._ring_of_experts_ragged_sort_requirements()
+          + [
+              (
+                  self.ici_tensor_parallelism * self.dcn_tensor_parallelism == 1,
+                  "no tensor parallelism (the layer checks the mesh again when it is traced)",
+              ),
+              (not self.mlp_bias, "mlp_bias=False"),
+              self._no_weight_quantization_requirement(),
+              (not self.moe_tc_ragged_sort, "moe_tc_ragged_sort=False"),
+              (not self.moe_accumulate_chunk_wgrad, "moe_accumulate_chunk_wgrad=False"),
+              (not self.moe_bwd_prequant_before_unsort, "moe_bwd_prequant_before_unsort=False"),
+              (not self.moe_quantize_token_all_gather, "moe_quantize_token_all_gather=False"),
+              (self.num_moe_emb_chunks <= 0, "num_moe_emb_chunks=0"),
+              self._remat_keeps_moe_tensor_requirement("down-projection output", _MOE_DOWN_PROJECTION_NAMES),
+          ],
+      )
+    return self
+
   @model_validator(mode="after")
   def validate_moe_routing_index_math(self) -> "MoEGeneral":
     """moe_routing_index_math='count' needs the ring-of-experts ragged sort, the only user of these index maps."""
diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index b060d98fc6..5c7469e7d2 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -342,6 +342,201 @@ def _ring_ragged_sort_bwd(res, g_out):
   return x, group_sizes_local, topk_argsort_revert_indices, topk_argsort_indices
 
 
+def _unsort_kwargs(
+    enforce_gather_fallback,
+    enforce_gather_reduce_fallback,
+    gather_flops_override,
+    gather_reduce_flops_override,
+    gather_bytes_accessed_override,
+    gather_reduce_bytes_accessed_override,
+    use_single_sparsecore,
+    index_math,
+    unique_pair_argsort,
+):
+  """The keyword arguments of the unsort's ragged gather (backward) and ragged gather-reduce (forward)."""
+  _check_index_math(index_math)
+  gather_kwargs = {
+      "enforce_fallback": enforce_gather_fallback,
+      "flops_override": gather_flops_override,
+      "bytes_accessed_override": gather_bytes_accessed_override,
+      "use_single_sparsecore": use_single_sparsecore,
+  }
+  gather_reduce_kwargs = {
+      "enforce_fallback": enforce_gather_reduce_fallback,
+      "flops_override": gather_reduce_flops_override,
+      "bytes_accessed_override": gather_reduce_bytes_accessed_override,
+      "use_single_sparsecore": use_single_sparsecore,
+      "packed_preprocess": index_math == "count",
+      "unique_pair_argsort": unique_pair_argsort and index_math != "count",
+  }
+  return gather_kwargs, gather_reduce_kwargs
+
+
+def _ring_unsort_window(group_sizes_local, local_num_experts, ep_name):
+  """This shard's ``[start, end)`` window in the globally expert-sorted slot order."""
+  group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True)
+
+  shard_idx = jax.lax.axis_index(ep_name)
+  experts_start = shard_idx * local_num_experts
+  experts_end = experts_start + local_num_experts
+
+  return group_offsets[experts_start], group_offsets[experts_end]
+
+
+def _ring_unsort_combine(
+    sorted_tokens_local,
+    topk_argsort_revert_indices,
+    topk_weights_flat,
+    shard_output_start,
+    shard_output_end,
+    topk,
+    gather_reduce_kwargs,
+):
+  """The unsort's forward: each token's routing-weighted sum over its slots on this shard's experts."""
+  buffer_size = sorted_tokens_local.shape[0]
+  num_tokens = topk_argsort_revert_indices.shape[0]
+
+  # We support two buffering modes:
+  # 1. buffer_size >= num_tokens: sorted_tokens_local has size >= num_tokens,
+  #    and local tokens are at their global positions [start, end).
+  # 2. buffer_size < num_tokens: sorted_tokens_local has size < num_tokens,
+  #    and local tokens are packed at [0, local_num_tokens).
+  if buffer_size >= num_tokens:
+    # Express the scatter as a gather-reduce: each output row pulls
+    # from sorted_tokens_local at position `topk_argsort_revert_indices[i]` if
+    # that position is within this shard's [start, end) range, else zero.
+    # The routing weights are applied per-row before the topk reduction.
+    valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & (
+        topk_argsort_revert_indices < shard_output_end
+    )
+    return ragged_gather_reduce(
+        sorted_tokens_local,
+        topk_argsort_revert_indices,
+        topk_weights=topk_weights_flat,
+        valid_rows_mask=valid_rows_mask,
+        reduce_group_size=topk,
+        **gather_reduce_kwargs,
+    )
+  # Shift indices so they map to the packed local buffer [0, local_num_tokens).
+  shifted_indices = topk_argsort_revert_indices - shard_output_start
+  local_num_tokens = shard_output_end - shard_output_start
+  limit = jnp.minimum(local_num_tokens, buffer_size)
+  valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit)
+  safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0)
+  return ragged_gather_reduce(
+      sorted_tokens_local,
+      safe_indices,
+      topk_weights=topk_weights_flat,
+      valid_rows_mask=valid_rows_mask,
+      reduce_group_size=topk,
+      **gather_reduce_kwargs,
+  )
+
+
+def _ring_unsort_fan_out(
+    g_hidden_states_local,
+    topk_argsort_revert_indices,
+    topk_weights_flat,
+    shard_output_start,
+    shard_output_end,
+    buffer_size,
+    topk,
+    topk_argsort_indices,
+    gather_kwargs,
+    mask_rows,
+    index_math="sort",
+):
+  """The first half of the unsort's backward: the output cotangent sent to the rows of the sorted buffer.
+
+  Returns ``(gathered, weight_for_sorted, mask, to_slot_order)``: row ``j`` of ``gathered`` is the unweighted
+  cotangent of the token that sorted row ``j`` came from; ``weight_for_sorted[j]`` is that slot's routing
+  weight; ``mask`` marks the rows of this shard's experts; ``to_slot_order`` maps a per-row value that is
+  zero outside ``mask`` back to flat slot order, with zero for the slots of other shards and dropped slots.
+
+  With ``mask_rows`` the rows of ``gathered`` outside ``mask`` are zeroed (:func:`ring_ragged_unsort`, whose
+  cotangent output is ``gathered`` scaled). Without it they are left as the ragged gather wrote them (unwritten by
+  the SparseCore kernel, other slots' rows on the reference fallback), which saves a pass over ``[rows, hidden]``
+  for a caller whose consumers never read those rows (:func:`ring_ragged_unsort_down_projection`).
+  """
+  n = topk_argsort_revert_indices.shape[0]
+  # Build the inverse permutation idx_inv such that idx_inv[j] = i
+  # where revert[i] = j. revert is itself argsort(topk_argsort_indices), so its inverse is
+  # topk_argsort_indices; reuse it when the caller has it rather than sorting again.
+  if topk_argsort_indices is not None:
+    idx_inv = topk_argsort_indices
+  elif index_math == "count":
+    # The "count" inverse holds only this shard's slots (-1 elsewhere), so it cannot be argsorted back.
+    raise ValueError("index_math='count' needs topk_argsort_indices (the forward permutation of ring_ragged_sort).")
+  else:
+    idx_inv = jnp.argsort(topk_argsort_revert_indices)
+
+  # Handle the same two buffering modes for backward pass.
+  # ragged_gather does the fan-out, by indexing into the un-expanded
+  # g_hidden_states_local via idx_inv // topk. It gathers unweighted so the same rows
+  # feed both gradients: the activation one after scaling, the weight one after a dot.
+  if buffer_size >= n:
+    weight_for_sorted = topk_weights_flat[idx_inv]
+    gathered = ragged_gather(
+        g_hidden_states_local,
+        idx_inv // topk,
+        shard_output_start[None],
+        shard_output_end[None],
+        **gather_kwargs,
+    )
+    # The gradients of the elements outside the valid shard output range are zero.
+    mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end)
+    if mask_rows:
+      gathered = jnp.where(mask[:, None], gathered, 0.0)
+
+    def from_buffer_rows(per_row):
+      # Row j of the sorted buffer is slot idx_inv[j], so slot i reads row revert[i].
+      if index_math == "count":
+        # Other shards' slots hold -1, not a row outside the window, so mask them instead of reading row -1.
+        slot_valid = (topk_argsort_revert_indices >= shard_output_start) & (
+            topk_argsort_revert_indices < shard_output_end
+        )
+        return jnp.where(slot_valid, per_row[jnp.where(slot_valid, topk_argsort_revert_indices, 0)], 0.0)
+      return per_row[topk_argsort_revert_indices]
+
+    return gathered, weight_for_sorted, mask, from_buffer_rows
+
+  # Slice the inverse permutation and the per-slot routing weights to match the packed local buffer.
+  padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size))
+  sliced_idx_inv = jax.lax.dynamic_slice_in_dim(padded_idx_inv, shard_output_start, buffer_size, axis=0)
+  gather_end = jnp.minimum(shard_output_end - shard_output_start, buffer_size)
+  padded_weights = jnp.pad(topk_weights_flat[idx_inv], (0, buffer_size))
+  sliced_weights = jax.lax.dynamic_slice_in_dim(padded_weights, shard_output_start, buffer_size, axis=0)
+  gathered = ragged_gather(
+      g_hidden_states_local,
+      sliced_idx_inv // topk,
+      jnp.int32(0)[None],
+      gather_end[None],
+      **gather_kwargs,
+  )
+  # The gradients of the elements beyond the valid limit of the local buffer are zero.
+  limit = jnp.minimum(shard_output_end - shard_output_start, buffer_size)
+  mask = jnp.arange(buffer_size) < limit
+  if mask_rows:
+    gathered = jnp.where(mask[:, None], gathered, 0.0)
+
+  def from_packed_rows(per_row):
+    if index_math == "count":
+      # Route the per-row value back to flat slot order with a gather through the inverse permutation instead
+      # of a TensorCore scatter: each valid slot reads its own buffer row; other shards' slots (-1) and
+      # dropped slots read 0, as the scatter leaves them. The table is padded to n entries because XLA only
+      # offloads a 1-D gather to the SparseCore for a large table (v7x AOT: 225,280 entries -> TensorCore
+      # gather at ~7 ns/element; 655,360 -> SparseCore).
+      shifted = topk_argsort_revert_indices - shard_output_start
+      slot_valid = (shifted >= 0) & (shifted < limit)
+      dot_table = jnp.pad(per_row, (0, max(0, n - buffer_size)))
+      return jnp.where(slot_valid, dot_table[jnp.where(slot_valid, shifted, 0)], 0.0)
+    # Scatter the per-row value back to flat slot order; dropped slots stay zero.
+    slots = jnp.where(mask, sliced_idx_inv, n)
+    return jnp.zeros((n,), jnp.float32).at[slots].set(per_row, mode="drop")
+
+  return gathered, sliced_weights, mask, from_packed_rows
+
+
 def ring_ragged_unsort(
     sorted_tokens_local,
     group_sizes_local,
@@ -401,9 +596,17 @@ def ring_ragged_unsort(
     A 2D ``[num_tokens_local, hidden]`` tensor with expert outputs scattered back
     to their original global sequence locations, with unpopulated positions zeroed out.
   """
-  _check_index_math(index_math)
-  packed_preprocess = index_math == "count"
-  unique_pair_preprocess = unique_pair_argsort and not packed_preprocess
+  gather_kwargs, gather_reduce_kwargs = _unsort_kwargs(
+      enforce_gather_fallback,
+      enforce_gather_reduce_fallback,
+      gather_flops_override,
+      gather_reduce_flops_override,
+      gather_bytes_accessed_override,
+      gather_reduce_bytes_accessed_override,
+      use_single_sparsecore,
+      index_math,
+      unique_pair_argsort,
+  )
 
   @jax.custom_vjp
   def _ring_ragged_unsort(
@@ -431,76 +634,25 @@ def _ring_ragged_unsort_fwd(
       topk_argsort_indices,
   ):
     """Executes unsorting sending tokens back."""
-    group_offsets = jnp.cumulative_sum(group_sizes_local, include_initial=True)
-
-    shard_idx = jax.lax.axis_index(ep_name)
-    experts_start = shard_idx * local_num_experts
-    experts_end = experts_start + local_num_experts
-
-    shard_output_start = group_offsets[experts_start]
-    shard_output_end = group_offsets[experts_end]
-
-    buffer_size = sorted_tokens_local.shape[0]
-    num_tokens = topk_argsort_revert_indices.shape[0]
-
-    # We support two buffering modes:
-    # 1. buffer_size >= num_tokens: sorted_tokens_local has size >= num_tokens,
-    #    and local tokens are at their global positions [start, end).
-    # 2. buffer_size < num_tokens: sorted_tokens_local has size < num_tokens,
-    #    and local tokens are packed at [0, local_num_tokens).
-    if buffer_size >= num_tokens:
-      # Express the scatter as a gather-reduce: each output row pulls
-      # from sorted_tokens_local at position `topk_argsort_revert_indices[i]` if
-      # that position is within this shard's [start, end) range, else zero.
-      # The routing weights are applied per-row before the topk reduction.
-      valid_rows_mask = (topk_argsort_revert_indices >= shard_output_start) & (
-          topk_argsort_revert_indices < shard_output_end
-      )
-      out = ragged_gather_reduce(
-          sorted_tokens_local,
-          topk_argsort_revert_indices,
-          topk_weights=topk_weights_flat,
-          valid_rows_mask=valid_rows_mask,
-          reduce_group_size=topk,
-          enforce_fallback=enforce_gather_reduce_fallback,
-          flops_override=gather_reduce_flops_override,
-          bytes_accessed_override=gather_reduce_bytes_accessed_override,
-          use_single_sparsecore=use_single_sparsecore,
-          packed_preprocess=packed_preprocess,
-          unique_pair_argsort=unique_pair_preprocess,
-      )
-    else:
-      # Shift indices so they map to the packed local buffer [0, local_num_tokens).
-      shifted_indices = topk_argsort_revert_indices - shard_output_start
-      local_num_tokens = shard_output_end - shard_output_start
-      limit = jnp.minimum(local_num_tokens, buffer_size)
-      valid_rows_mask = (shifted_indices >= 0) & (shifted_indices < limit)
-      safe_indices = jnp.where(valid_rows_mask, shifted_indices, 0)
-
-      out = ragged_gather_reduce(
-          sorted_tokens_local,
-          safe_indices,
-          topk_weights=topk_weights_flat,
-          valid_rows_mask=valid_rows_mask,
-          reduce_group_size=topk,
-          enforce_fallback=enforce_gather_reduce_fallback,
-          flops_override=gather_reduce_flops_override,
-          bytes_accessed_override=gather_reduce_bytes_accessed_override,
-          use_single_sparsecore=use_single_sparsecore,
-          packed_preprocess=packed_preprocess,
-          unique_pair_argsort=unique_pair_preprocess,
-      )
-
+    shard_output_start, shard_output_end = _ring_unsort_window(group_sizes_local, local_num_experts, ep_name)
+    out = _ring_unsort_combine(
+        sorted_tokens_local,
+        topk_argsort_revert_indices,
+        topk_weights_flat,
+        shard_output_start,
+        shard_output_end,
+        topk,
+        gather_reduce_kwargs,
+    )
     res = (
         sorted_tokens_local,
         topk_argsort_revert_indices,
         topk_weights_flat,
         shard_output_start,
         shard_output_end,
-        buffer_size,
+        sorted_tokens_local.shape[0],
         topk_argsort_indices,
     )
-
     return out, res
 
   @jax.named_scope("ragged-unsort-bwd")
@@ -528,89 +680,25 @@ def _ring_ragged_unsort_bwd(res, g_out):
         buffer_size,
         topk_argsort_indices,
     ) = res
-    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
-    # where revert[i] = j. revert is itself argsort(topk_argsort_indices), so its inverse is
-    # topk_argsort_indices; reuse it when the caller has it rather than sorting again.
-    if topk_argsort_indices is not None:
-      idx_inv = topk_argsort_indices
-    elif index_math == "count":
-      # The "count" inverse holds only this shard's slots (-1 elsewhere), so it cannot be argsorted back.
-      raise ValueError("index_math='count' needs topk_argsort_indices (the forward permutation of ring_ragged_sort).")
-    else:
-      idx_inv = jnp.argsort(topk_argsort_revert_indices)
-
-    # Handle the same two buffering modes for backward pass.
-    # ragged_gather does the fan-out, by indexing into the un-expanded
-    # g_hidden_states_local via idx_inv // topk. It gathers unweighted so the same rows
-    # feed both gradients: the activation one after scaling, the weight one after a dot.
-    if buffer_size >= n:
-      weight_for_sorted = topk_weights_flat[idx_inv]
-      gathered = ragged_gather(
-          g_hidden_states_local,
-          idx_inv // topk,
-          shard_output_start[None],
-          shard_output_end[None],
-          enforce_fallback=enforce_gather_fallback,
-          flops_override=gather_flops_override,
-          bytes_accessed_override=gather_bytes_accessed_override,
-          use_single_sparsecore=use_single_sparsecore,
-      )
-      # Mask out gradients that correspond to elements outside the valid shard
-      # output range.
-      mask = (jnp.arange(n) >= shard_output_start) & (jnp.arange(n) < shard_output_end)
-      gathered = jnp.where(mask[:, None], gathered, 0.0)
-      grad_sorted_tokens = (gathered * weight_for_sorted[:, None]).astype(gathered.dtype)
-      # Row-wise dot in sorted order, then permuted back to flat slot order.
-      dot_sorted = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local[:n].astype(jnp.float32), axis=-1)
-      if index_math == "count":
-        # Other shards' slots hold -1, not a row outside the window, so mask them instead of reading row -1.
-        slot_valid = (topk_argsort_revert_indices >= shard_output_start) & (
-            topk_argsort_revert_indices < shard_output_end
-        )
-        grad_topk_weights = jnp.where(slot_valid, dot_sorted[jnp.where(slot_valid, topk_argsort_revert_indices, 0)], 0.0)
-      else:
-        grad_topk_weights = dot_sorted[topk_argsort_revert_indices]
-    else:
-      # Slice the inverse permutation to match the packed local buffer.
-      padded_idx_inv = jnp.pad(idx_inv, (0, buffer_size))
-      sliced_idx_inv = jax.lax.dynamic_slice_in_dim(padded_idx_inv, shard_output_start, buffer_size, axis=0)
-      gather_end = jnp.minimum(shard_output_end - shard_output_start, buffer_size)
-      # Slice the per-slot routing weights to match the packed local buffer.
-      padded_weights = jnp.pad(topk_weights_flat[idx_inv], (0, buffer_size))
-      sliced_weights = jax.lax.dynamic_slice_in_dim(padded_weights, shard_output_start, buffer_size, axis=0)
-      gathered = ragged_gather(
-          g_hidden_states_local,
-          sliced_idx_inv // topk,
-          jnp.int32(0)[None],
-          gather_end[None],
-          enforce_fallback=enforce_gather_fallback,
-          flops_override=gather_flops_override,
-          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.
-      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)
-      dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1)
-      if index_math == "count":
-        # Route the per-row dot back to flat slot order with a gather through the inverse permutation instead
-        # of a TensorCore scatter: each valid slot reads its own buffer row; other shards' slots (-1) and
-        # dropped slots read 0, as the scatter leaves them. The table is padded to n entries because XLA only
-        # offloads a 1-D gather to the SparseCore for a large table (v7x AOT: 225,280 entries -> TensorCore
-        # gather at ~7 ns/element; 655,360 -> SparseCore).
-        shifted = topk_argsort_revert_indices - shard_output_start
-        slot_valid = (shifted >= 0) & (shifted < limit)
-        dot_table = jnp.pad(dot_local, (0, max(0, n - buffer_size)))
-        grad_topk_weights = jnp.where(slot_valid, dot_table[jnp.where(slot_valid, shifted, 0)], 0.0)
-      else:
-        # Scatter the per-slot dot back to flat slot order; dropped slots stay zero.
-        slots = jnp.where(mask, sliced_idx_inv, n)
-        grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop")
+    gathered, weight_for_sorted, _, to_slot_order = _ring_unsort_fan_out(
+        g_out,
+        topk_argsort_revert_indices,
+        topk_weights_flat,
+        shard_output_start,
+        shard_output_end,
+        buffer_size,
+        topk,
+        topk_argsort_indices,
+        gather_kwargs,
+        mask_rows=True,
+        index_math=index_math,
+    )
+    grad_sorted_tokens = (gathered * weight_for_sorted[:, None]).astype(gathered.dtype)
+    # Row-wise dot in sorted order, then back to flat slot order (a full buffer has its rows at [0, n)).
+    gathered_f32 = gathered.astype(jnp.float32)
+    sorted_rows = sorted_tokens_local[:n] if buffer_size >= n else sorted_tokens_local
+    grad_topk_weights = to_slot_order(jnp.sum(gathered_f32 * sorted_rows.astype(jnp.float32), axis=-1))
     return grad_sorted_tokens, None, None, grad_topk_weights, None
 
   _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd)
@@ -627,6 +715,162 @@ def _ring_ragged_unsort_bwd(res, g_out):
   )
 
 
+def ring_ragged_unsort_down_projection(
+    intermediate_local,
+    wo_local,
+    down_projection,
+    down_projection_operands,
+    group_sizes_local,
+    topk_argsort_revert_indices,
+    topk,
+    local_num_experts,
+    ep_name,
+    topk_weights,
+    topk_argsort_indices=None,
+    forward_tie=None,
+    enforce_gather_fallback=False,
+    enforce_gather_reduce_fallback=False,
+    gather_flops_override=-1,
+    gather_reduce_flops_override=-1,
+    gather_bytes_accessed_override=-1,
+    gather_reduce_bytes_accessed_override=-1,
+    use_single_sparsecore=False,
+    index_math="sort",
+    unique_pair_argsort=False,
+):
+  """:func:`ring_ragged_unsort` of the down projection, with a backward that does not read its output.
+
+  Forward: ``y = down_projection(intermediate_local, wo_local, *down_projection_operands)``, the
+  per-expert GEMM ``y_j = h_j @ wo[e(j)]`` over the sorted rows ``h_j``, then the forward of
+  ``ring_ragged_unsort(y, ...)``: the same helpers (``_ring_unsort_window``, ``_ring_unsort_combine``), so the
+  output is bit-identical. The backward's fan-out is ``ring_ragged_unsort``'s too (``_ring_unsort_fan_out``).
+
+  Backward. Let ``g_j`` be the unweighted output cotangent of the token that sorted row ``j`` came
+  from and ``w_j`` that slot's routing weight. The two ops apart compute ``dy_j = w_j g_j``,
+  ``dh = dy @ wo^T`` (dlhs), ``dwo = h^T @ dy`` (drhs) and the routing-weight gradient
+  ``g_w_j = <g_j, y_j>``. That reads ``y``, so a rematerialized layer recomputes the down-projection
+  GEMM in the backward only for this dot, and the dot takes a pass over two ``[rows, hidden]``
+  tensors that also writes ``dy``. Row ``j`` of the GEMM is linear in ``h_j`` and the weight is a
+  scalar per row, so::
+
+      <g_j, h_j wo> = <g_j wo^T, h_j>,   (w_j g_j) wo^T = w_j (g_j wo^T),   h^T (w * g) = (w * h)^T g.
+
+  This backward therefore takes the down projection's own VJP at the input ``w * h`` with the
+  cotangent ``g``. Its two GEMMs give ``u = g wo^T`` and ``dwo = (w * h)^T g``, and then::
+
+      dh_j = w_j u_j,     g_w_j = <u_j, h_j>.
+
+  The GEMMs keep their shapes and cost; ``y`` is no longer read, so XLA drops its recomputation; and
+  the weighting and the dot run over ``[rows, mlp]`` tensors instead of ``[rows, hidden]`` ones.
+  Each gradient still has one rounding to the activation dtype of a weighted operand or GEMM output,
+  but at a different point (``w * h`` instead of ``w * g``; ``dh`` after the product with ``w``
+  instead of before the GEMM; ``g_w`` from ``g wo^T`` instead of ``y``). So the gradients agree with
+  the two ops apart to that rounding, not bit for bit.
+
+  Rows of the sorted buffer outside this shard's experts' groups: the cotangent is left unmasked
+  there, which saves a pass over ``[rows, hidden]``. The down projection's backward GEMMs do not use
+  those rows for any row inside a group (the dlhs is row-wise and the drhs masks both operands of a
+  partial group tile), and every output of this backward is masked to the group rows.
+
+  Args:
+    intermediate_local: ``[buffer_size, mlp]`` down-projection input, rows grouped by local expert.
+    wo_local: the down-projection kernel of the local experts, as ``down_projection`` takes it.
+    down_projection: ``(h, wo, *down_projection_operands) -> [buffer_size, hidden]``, a GEMM with its
+      own VJP. It must not close over traced arrays; pass them in ``down_projection_operands``.
+    down_projection_operands: tuple of non-differentiated operands of ``down_projection`` (group
+      sizes and the like).
+    group_sizes_local, topk_argsort_revert_indices, topk, local_num_experts, ep_name, topk_weights,
+      topk_argsort_indices: as for :func:`ring_ragged_unsort`.
+    forward_tie: optional array tied to ``y`` by an optimization barrier in the forward pass only,
+      so everything that reads either comes after both. It is returned as a second output and its
+      cotangent passes through unchanged.
+    index_math, unique_pair_argsort: as for :func:`ring_ragged_unsort` (they select how the combine's
+      gather-reduce sorts its valid rows, and with "count" how the weight gradient returns to slot order).
+
+  Returns:
+    The unsorted ``[num_tokens_local, hidden]`` output, and ``forward_tie`` if one was given.
+  """
+  num_slots = topk_argsort_revert_indices.shape[0]
+  if intermediate_local.shape[0] > num_slots:
+    # ring_ragged_sort's buffer has exactly num_slots rows (full mode) or fewer (packed mode).
+    raise ValueError(
+        f"ring_ragged_unsort_down_projection: {intermediate_local.shape[0]} buffer rows > {num_slots} slots."
+    )
+  gather_kwargs, gather_reduce_kwargs = _unsort_kwargs(
+      enforce_gather_fallback,
+      enforce_gather_reduce_fallback,
+      gather_flops_override,
+      gather_reduce_flops_override,
+      gather_bytes_accessed_override,
+      gather_reduce_bytes_accessed_override,
+      use_single_sparsecore,
+      index_math,
+      unique_pair_argsort,
+  )
+
+  @jax.custom_vjp
+  def _unsort_down_projection(h, wo, operands, group_sizes, revert, weights, argsort_indices, tie):
+    """Down projection, then unsort and combine."""
+    return _unsort_down_projection_fwd(h, wo, operands, group_sizes, revert, weights, argsort_indices, tie)[0]
+
+  @jax.named_scope("ragged-unsort-down-projection-fwd")
+  def _unsort_down_projection_fwd(h, wo, operands, group_sizes, revert, weights, argsort_indices, tie):
+    y = down_projection(h, wo, *operands)
+    if tie is not None:
+      y, tie = jax.lax.optimization_barrier((y, tie))
+    start, end = _ring_unsort_window(group_sizes, local_num_experts, ep_name)
+    out = _ring_unsort_combine(y, revert, weights, start, end, topk, gather_reduce_kwargs)
+    # y is not a residual: the backward below never reads it.
+    res = (h, wo, operands, revert, weights, start, end, argsort_indices)
+    return (out if tie is None else (out, tie)), res
+
+  @jax.named_scope("ragged-unsort-down-projection-bwd")
+  def _unsort_down_projection_bwd(res, g_out):
+    h, wo, operands, revert, weights, start, end, argsort_indices = res
+    g_unsorted, g_tie = (g_out, None) if forward_tie is None else g_out
+    gathered, weight_for_sorted, mask, to_slot_order = _ring_unsort_fan_out(
+        g_unsorted,
+        revert,
+        weights,
+        start,
+        end,
+        h.shape[0],
+        topk,
+        argsort_indices,
+        gather_kwargs,
+        mask_rows=False,
+        index_math=index_math,
+    )
+    weight = weight_for_sorted[:, None]
+    h32 = h.astype(jnp.float32)
+    # The routing weight goes on the narrow [rows, mlp] side: dwo = (w * h)^T g.
+    weighted_h = jnp.where(mask[:, None], h32 * weight, 0.0).astype(h.dtype)
+    # The down projection's own VJP at the input w * h: its input gradient g wo^T does not depend on the
+    # input, and its weight gradient is (w * h)^T g. Its forward GEMM has no reader, so XLA removes it.
+    _, project_vjp = jax.vjp(lambda h_, wo_: down_projection(h_, wo_, *operands), weighted_h, wo)
+    u, grad_wo = project_vjp(gathered)
+    u = u.astype(jnp.float32)
+    grad_h = jnp.where(mask[:, None], u * weight, 0.0).astype(h.dtype)
+    # The routing-weight gradient <g, h wo> = <g wo^T, h>, from the down projection's input.
+    dot = jnp.where(mask, jnp.sum(u * h32, axis=-1), 0.0)
+    grad_topk_weights = to_slot_order(dot)
+    grad_operands = jax.tree.map(lambda _: None, operands)
+    return grad_h, grad_wo, grad_operands, None, None, grad_topk_weights, None, g_tie
+
+  _unsort_down_projection.defvjp(_unsort_down_projection_fwd, _unsort_down_projection_bwd)
+
+  return _unsort_down_projection(
+      intermediate_local,
+      wo_local,
+      tuple(down_projection_operands),
+      group_sizes_local,
+      topk_argsort_revert_indices,
+      topk_weights.astype(jnp.float32),
+      topk_argsort_indices,
+      forward_tie,
+  )
+
+
 def a2a_ragged_sort(
     inputs,
     sort_indices,
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index 1e8c412b0b..cee0570dc3 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -47,6 +47,7 @@
 from maxtext.kernels.ragged.ragged_sort import a2a_ragged_unsort
 from maxtext.kernels.ragged.ragged_sort import ring_ragged_sort
 from maxtext.kernels.ragged.ragged_sort import ring_ragged_unsort
+from maxtext.kernels.ragged.ragged_sort import ring_ragged_unsort_down_projection
 from maxtext.kernels.ragged.ragged_sort_tc import ring_ragged_sort_tc
 from maxtext.kernels.ragged.ragged_sort_tc import ring_ragged_unsort_tc
 from maxtext.kernels.ragged.ragged_sort_tc import tc_buffer_row_weights
@@ -2576,6 +2577,14 @@ def sparse_matmul(
     # log_moe_drops only: the shard_map gets a 5th (moe_audit) output only when requested, so the
     # compiled program is unchanged with the flag off.
     log_moe_drops = getattr(self.config, "log_moe_drops", False)
+    # moe_router_grad_from_wo_input: the ring-of-experts down projection and unsort run as one op
+    # (_ring_down_project_and_unpermute). The config validator checks the rest; the mesh decides this.
+    router_grad_from_wo_input = self.config.moe_router_grad_from_wo_input
+    if router_grad_from_wo_input and self.get_tensor_parallelism_size() > 1:
+      raise ValueError(
+          "moe_router_grad_from_wo_input=True does not support tensor parallelism in the MoE layer: the down"
+          " projection must be the bare GEMM of the sorted rows, without its reduce-scatter."
+      )
 
     def fwd_tile(name):
       """Reads a forward GMM tile size, preferring `eval_*` under eval axis rules."""
@@ -3972,6 +3981,64 @@ def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length, t
           ),
       )
 
+    def _ring_down_project_and_unpermute(
+        output0, output1, wo, routing, route_metadata, batch_size, sequence_length, forward_tie=None
+    ):
+      """moe_router_grad_from_wo_input: `_ring_unpermute(_down_projection(...))` as one op.
+
+      The fused op's backward does not read the down projection's output: the routing-weight gradient
+      <g, h wo> is computed as <g wo^T, h> from the down projection's input, and the routing weight is applied
+      on the [rows, mlp] side, so a rematerialized layer no longer recomputes the down-projection GEMM in the
+      backward (ring_ragged_unsort_down_projection). The config validator admits the key only where the down
+      projection is the bare GEMM of the sorted rows (no bias, TC sort, tensor parallelism or quantization).
+      `forward_tie` is the pipelined loop's tie 2, applied inside the op to the down-projection output.
+      Returns the output as [expert shards x batch, seq, embed], as `_ring_unpermute`, and the tie's output.
+      """
+      # moe_wi_gmm_fused_swiglu: the wi GMM already applied the activation (output1 is None).
+      intermediate_layer = output0 if fuse_silu_in_wi_gmm else self.apply_ffn_activation_maybe_fused(output0, output1)
+      wo_gather_axes, wo_tile_size = get_wo_gmm_params()
+
+      def down_projection(h, kernel, group_sizes, expert_assignments, group_offset):
+        return gmm(
+            h,
+            kernel,
+            tiling=wo_tile_size,
+            group_sizes=group_sizes,
+            expert_assignments=expert_assignments,
+            weight_gather_axes=wo_gather_axes,
+            group_offset=group_offset,
+        )
+
+      operands = get_gmm_operands_for_local_experts(intermediate_layer, routing, route_metadata)
+      output = ring_ragged_unsort_down_projection(
+          intermediate_layer,
+          wo,
+          down_projection,
+          (operands["group_sizes"], operands["expert_assignments"], operands["group_offset"]),
+          routing.group_sizes,
+          routing.sorted_selected_experts,
+          self.num_experts_per_tok,
+          self.config.num_experts // self.get_expert_parallelism_size(),
+          self._expert_parallelism_name,
+          topk_weights=jnp.ravel(routing.weights).astype(jnp.float32),
+          topk_argsort_indices=routing.topk_argsort_indices,
+          forward_tie=forward_tie,
+          enforce_gather_fallback=self.config.ragged_gather_fallback,
+          enforce_gather_reduce_fallback=self.config.ragged_gather_reduce_fallback,
+          gather_flops_override=self.config.ragged_gather_cost_estimate_flops,
+          gather_reduce_flops_override=self.config.ragged_gather_reduce_cost_estimate_flops,
+          gather_bytes_accessed_override=self.config.ragged_gather_cost_estimate_bytes_accessed,
+          gather_reduce_bytes_accessed_override=self.config.ragged_gather_reduce_cost_estimate_bytes_accessed,
+          use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
+          index_math=self.config.moe_routing_index_math,
+          unique_pair_argsort=self.config.moe_ragged_sort_unique_pair_argsort,
+      )
+      if forward_tie is not None:
+        output, forward_tie = output
+      # The reshapes of unpermute and _ring_unpermute (the embed dim is whole: no tensor parallelism here).
+      output = output.reshape(batch_size, sequence_length, -1).astype(self.dtype)
+      return jnp.reshape(output, (-1, sequence_length, self.moe_expert_input_dim)), forward_tie
+
     def _ring_reduce_scatter(output):
       """Ring of experts: sum the partial outputs across the expert shards."""
       if self.config.ring_of_experts_row_major_reduce_scatter:
@@ -4000,6 +4067,20 @@ def _ring_reduce_scatter(output):
         output = jnp.reshape(output, (*output.shape[:2], -1))
       return output
 
+    def _moe_body_outputs(output, routing, weights):
+      """What _moe_body returns: the layer output, the routing's side outputs, and the weights `(w0, w1, wo)`
+      (as returned by the GMMs under moe_accumulate_chunk_wgrad)."""
+      return (
+          output,
+          routing.lb_loss,
+          routing.bias_updates,
+          routing.has_overflow,
+          routing.required_rbf,
+          routing.max_load_ratio,
+          routing.moe_audit,
+          weights,
+      )
+
     def _moe_body(
         x,
         logits,
@@ -4083,6 +4164,17 @@ def _cn(t):
         else:
           output0, output1 = up_res  # pylint: disable=unbalanced-tuple-unpacking
 
+      if router_grad_from_wo_input:
+        # moe_router_grad_from_wo_input: the activation, the down projection and the unsort as one op whose
+        # backward does not read the down projection's output; then the same reduce-scatter as below. The op does
+        # not return the weights, so it cannot serve moe_accumulate_chunk_wgrad (the config validator rejects it).
+        if return_weights or not gather_weights:
+          raise NotImplementedError("moe_router_grad_from_wo_input does not support moe_accumulate_chunk_wgrad.")
+        output, _ = _ring_down_project_and_unpermute(
+            output0, output1, wo, routing, route_metadata, batch_size, sequence_length
+        )
+        return _moe_body_outputs(_ring_reduce_scatter(output), routing, (w0, w1, wo))
+
       tc_prescaled = (
           self.config.use_ring_of_experts
           and getattr(routing, "tc_routing", None) is not None
@@ -4196,16 +4288,7 @@ def _cn(t):
             )
           # Same reduce-scatter as the pipelined chunk loop (adds ring_of_experts_row_major_reduce_scatter).
           output = _ring_reduce_scatter(output)
-        return (
-            output,
-            routing.lb_loss,
-            routing.bias_updates,
-            routing.has_overflow,
-            routing.required_rbf,
-            routing.max_load_ratio,
-            routing.moe_audit,
-            (w0, w1, wo),
-        )
+        return _moe_body_outputs(output, routing, (w0, w1, wo))
 
       if self.get_expert_parallelism_size() > 1:
         original_inputs_first_dim = batch_size * sequence_length * self.config.num_experts_per_tok
@@ -4237,16 +4320,7 @@ def _cn(t):
           group_sizes=routing.group_sizes,
       )
 
-      return (
-          output,
-          routing.lb_loss,
-          routing.bias_updates,
-          routing.has_overflow,
-          routing.required_rbf,
-          routing.max_load_ratio,
-          routing.moe_audit,
-          (w0, w1, wo),
-      )
+      return _moe_body_outputs(output, routing, (w0, w1, wo))
 
     in_specs = (
         input_partition_pspec,
@@ -4439,24 +4513,28 @@ def _with_routed(dispatched_c, routed):
             #    so the TensorCore runs that GEMM while the gather-reduce is on the SparseCore.
             unpermuted, output0 = barrier_tie(unpermuted, output0)
             outs.append(_ring_reduce_scatter(unpermuted))
-          intermediate_output, wo_c, tc_prescaled, x_is_3d = _down_projection(
-              x_c,
-              output0,
-              output1,
-              gmm_fn,
-              cur_wo,
-              wo_bias_c,
-              routing,
-              route_metadata,
-              gather_weights=gather_c,
-              return_weights=return_c,
-          )
-          if return_c:
-            cur_w0, cur_w1, cur_wo = w0_c, w1_c, wo_c
-          if c + 1 < n_chunks:
-            # 2. This chunk's gather-reduce waits for the next chunk's routed tokens.
-            intermediate_output, routed = barrier_tie(intermediate_output, dispatched[c + 1][0])
-            dispatched[c + 1] = _with_routed(dispatched[c + 1], routed)
+          if router_grad_from_wo_input:
+            # moe_router_grad_from_wo_input: the fused op below runs the down projection, after tie 3.
+            intermediate_output, tc_prescaled, x_is_3d = None, False, False
+          else:
+            intermediate_output, wo_c, tc_prescaled, x_is_3d = _down_projection(
+                x_c,
+                output0,
+                output1,
+                gmm_fn,
+                cur_wo,
+                wo_bias_c,
+                routing,
+                route_metadata,
+                gather_weights=gather_c,
+                return_weights=return_c,
+            )
+            if return_c:
+              cur_w0, cur_w1, cur_wo = w0_c, w1_c, wo_c
+            if c + 1 < n_chunks:
+              # 2. This chunk's gather-reduce waits for the next chunk's routed tokens.
+              intermediate_output, routed = barrier_tie(intermediate_output, dispatched[c + 1][0])
+              dispatched[c + 1] = _with_routed(dispatched[c + 1], routed)
           if outs:
             # 3. This chunk's gather-reduce waits for the previous chunk's reduce-scatter. The tie holds the
             #    gather-reduce's routing weights, a small input: holding the down-projection output instead
@@ -4467,9 +4545,26 @@ def _with_routed(dispatched_c, routed):
             #    run before the previous chunk's reduce-scatter.
             weights, outs[-1] = barrier_tie(routing.weights, outs[-1])
             routing = routing.replace(weights=weights)
-          unpermuted = _ring_unpermute(
-              intermediate_output, routing, x.shape[0], chunk, tc_prescaled=tc_prescaled, x_is_3d=x_is_3d
-          )
+          if router_grad_from_wo_input:
+            # moe_router_grad_from_wo_input: the down projection and the unsort as one op. Tie 2 moves inside it
+            # (forward only, as barrier_tie's backward=False): this chunk's down-projection output and the next
+            # chunk's routed tokens. Tie 3 above already holds the routing weights the op's gather-reduce reads.
+            unpermuted, routed = _ring_down_project_and_unpermute(
+                output0,
+                output1,
+                cur_wo,
+                routing,
+                route_metadata,
+                x.shape[0],
+                chunk,
+                forward_tie=dispatched[c + 1][0] if c + 1 < n_chunks else None,
+            )
+            if c + 1 < n_chunks:
+              dispatched[c + 1] = _with_routed(dispatched[c + 1], routed)
+          else:
+            unpermuted = _ring_unpermute(
+                intermediate_output, routing, x.shape[0], chunk, tc_prescaled=tc_prescaled, x_is_3d=x_is_3d
+            )
           lb_losses.append(routing.lb_loss)
           bias_updates_list.append(routing.bias_updates)
           has_overflows.append(routing.has_overflow)
diff --git a/tests/unit/moe_router_grad_from_wo_input_test.py b/tests/unit/moe_router_grad_from_wo_input_test.py
new file mode 100644
index 0000000000..329e5816b4
--- /dev/null
+++ b/tests/unit/moe_router_grad_from_wo_input_test.py
@@ -0,0 +1,803 @@
+# 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.
+"""moe_router_grad_from_wo_input: the down projection and the unsort as one op, against the two apart, on CPU.
+
+`ragged_sort.ring_ragged_unsort_down_projection` runs the down-projection GEMM and then the forward of
+`ring_ragged_unsort` (the same helpers), so the output must be bit-identical. Its backward does not read the GEMM's
+output: it takes the GEMM's own VJP at the input w * h with the unweighted cotangent g, giving u = g wo^T and
+dwo = (w * h)^T g, then dh = w * u and the routing-weight gradient <u, h> (= <g, h wo>). Each gradient keeps one
+rounding to the activation dtype, at a different point, so the gradients agree with the two ops apart to bf16
+rounding, not bit for bit.
+
+What a CPU run can and cannot reach: the ragged gather and gather-reduce take their jnp reference fallbacks off a TPU
+target (the SparseCore kernels, and the validity sort of the gather-reduce, are TPU-only), and the layer's GMM is
+megablox in interpret mode. The fused op's own math (the custom VJP, the weighting, the dot, the masks, the slot-order
+return under both index maps and both buffering modes) is plain JAX and runs here as on TPU.
+
+Pre-registered, all in a subprocess with 8 CPU devices:
+  - The op on a 4-device expert-parallel mesh, with a jnp grouped matmul as the down projection, the full buffer and
+    a packed buffer that drops tokens, moe_routing_index_math "sort" and "count": the output bit-identical to the
+    two ops apart, every gradient (dx, dw, dwi, dwo) within relative L2 1e-2 of theirs. With a forward tie: the tie
+    comes back unchanged, its cotangent passes through, and the traced forward ties it to the GEMM output with an
+    optimization barrier. "count" without the forward permutation must raise.
+  - Precision, the op alone (h an input, so only the backward rounds): against the same function written densely in
+    float32 (slot by slot, plain autodiff, no ragged kernels; the ragged gather-reduce returns bf16, so the two-ops
+    path cannot serve as a float32 reference) at the same bf16-representable inputs, each fused gradient (dh, dw,
+    dwo) must be no further than _PRECISION_RATIO = 1.10 times as far as the two-ops gradient, on 6 seeds x {full
+    buffer "sort", dropping packed buffer "count"}. The routing weights sit at the bottom of a binade
+    ([0.25, 0.3)), where bf16 rounds them worst, so a backward that rounds them is visible. Measured over the 12
+    runs: dh 0.996-1.005, dwo 0.992-1.003, dw 0.941-1.065 (dw, one value per slot, is the noisiest), so 1.10 sits
+    above the spread. A backward that rounds the routing weight to bf16 measures 1.31-1.34 on dh and dwo on every
+    seed, and one that shrinks the routing-weight gradient by 0.5% measures 3.0-3.3 on dw; both must fail.
+  - The shrunk Qwen3.5 routed MoE of moe_roe_local_routing_test.py (16 experts, top 4, cp-as-ep, ragged sort, local
+    routing, two token chunks): output and loss bit-identical to the key off, every gradient within relative L2 1e-2
+    (measured <= 4.2e-3): plain and pipelined chunk loop, dropless, padding excluded with "count", saved routing
+    maps, the prefused-weight keys, and the benchmark's MoE settings with every compatible key of this series on
+    against all off (moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp,
+    moe_roe_packed_routing_gather; expert-weight prefetch, row-major reduce-scatter, routing maps on device,
+    ragged_buffer_factor 2.75, fp8 fake-quant, padding excluded). A 0.5% error in the routing-weight gradient sits
+    inside the layer's 1e-2 tolerance; the precision check above is what catches it. Structure: the compiled
+    jax.grad(jax.checkpoint(loss)) has exactly num_moe_token_chunks GEMM loops fewer (the down-projection GEMM is no
+    longer recomputed in the backward), the pipelined forward keeps the same number of optimization barriers (tie 2
+    moves into the op), and every chunk but the last reads the next chunk's routed tokens from the op's tie output.
+  - On an expert 4 x tensor 2 mesh the layer runs without the key and raises at trace time when the key reaches it
+    past the validator.
+  - Deliberately wrong versions must FAIL: the fused backward with the routing weight dropped from the input and
+    weight gradients (at the op and in the layer).
+  - The validator rejects the key for each requirement it does not meet, naming it.
+"""
+
+import contextlib
+import os
+import re
+import subprocess
+import sys
+from unittest import mock
+
+from absl.testing import parameterized
+from flax import nnx
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+from jax.sharding import PartitionSpec as P
+import numpy as np
+import pytest
+
+from maxtext.configs import pyconfig
+from maxtext.kernels.ragged import ragged_sort
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils.test_helpers import get_test_config_path
+
+_NUM_DEVICES = 8
+_NUM_EXPERTS = 16
+_TOP_K = 4
+_EMB = 256
+_SEQ = 64
+_PASS_TOKEN = "MOE_ROUTER_GRAD_FROM_WO_INPUT_CHECKS_PASSED"
+_XLA_FLAGS = f"--xla_force_host_platform_device_count={_NUM_DEVICES} --xla_allow_excess_precision=false"
+_GRAD_REL_L2 = 1e-2
+# The fused backward's error against a float32 run, over the two-ops backward's error (see the module docstring).
+_PRECISION_RATIO = 1.10
+_PRECISION_SEEDS = range(6)
+
+
+def _cfg(**overrides):
+  """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing)."""
+  kwargs = {
+      "run_name": "moe_router_grad_from_wo_input",
+      "enable_checkpointing": False,
+      "log_config": False,
+      "skip_jax_distributed_system": True,
+      "model_name": "qwen3.5-397b-a17b",
+      "override_model_config": True,
+      "num_experts": _NUM_EXPERTS,
+      "num_experts_per_tok": _TOP_K,
+      "base_emb_dim": _EMB,
+      "base_mlp_dim": 256,
+      "base_moe_mlp_dim": 256,
+      "max_target_length": _SEQ,
+      "max_prefill_predict_length": _SEQ,
+      "per_device_batch_size": 0.25,
+      "dtype": "bfloat16",
+      "weight_dtype": "float32",
+      "float32_gate_logits": True,
+      "megablox": True,
+      "sparse_matmul": True,
+      # Tokamax gmm v2 is TPU-only; both runs share whichever GMM runs, so CPU uses megablox.
+      "use_tokamax_gmm": False,
+      "use_gmm_v2": False,
+      "use_ring_of_experts": True,
+      "use_ragged_sort": True,
+      "use_custom_sort_vjp": False,
+      "ragged_buffer_factor": 2.0,
+      "num_moe_token_chunks": 2,
+      "custom_mesh_and_rule": "cp-as-ep",
+      "ici_fsdp_parallelism": 1,
+      "ici_context_parallelism": 4,
+      "ici_expert_parallelism": 2,
+      "context_parallel_load_balance": False,
+      "ring_of_experts_local_routing": True,
+  }
+  kwargs.update(overrides)
+  return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs)
+
+
+# A mesh with tensor parallelism: the default logical axis rules (cp-as-ep has no tensor axis), expert 4 x tensor 2.
+_TP_MESH = {
+    "custom_mesh_and_rule": "",
+    "ici_context_parallelism": 1,
+    "ici_expert_parallelism": 4,
+    "ici_tensor_parallelism": 2,
+    "per_device_batch_size": 0.5,
+}
+
+
+class RouterGradFromWoInputConfigTest(parameterized.TestCase):
+  """The key is rejected for each requirement it does not meet."""
+
+  @parameterized.named_parameters(
+      ("dense", {"sparse_matmul": False}, "sparse_matmul=True"),
+      ("te_moe_block", {"te_moe_block": True}, "te_moe_block=False"),
+      (
+          "no_ring",
+          {"use_ring_of_experts": False, "num_moe_token_chunks": 1, "ring_of_experts_local_routing": False},
+          "use_ring_of_experts=True",
+      ),
+      ("no_ragged_sort", {"use_ragged_sort": False}, "use_ragged_sort=True"),
+      ("tensor_parallel", _TP_MESH, "no tensor parallelism"),
+      ("mlp_bias", {"mlp_bias": True}, "mlp_bias=False"),
+      ("quantization", {"quantization": "int8"}, "no weight quantization"),
+      ("tc_ragged_sort", {"moe_tc_ragged_sort": True}, "moe_tc_ragged_sort=False"),
+      ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"),
+      ("prequant_unsort", {"moe_bwd_prequant_before_unsort": True}, "moe_bwd_prequant_before_unsort=False"),
+      ("quantized_token_ag", {"moe_quantize_token_all_gather": True}, "moe_quantize_token_all_gather=False"),
+      ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"),
+      ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwo"),
+      ("save_dot_remat", {"remat_policy": "save_dot_except_mlpwi"}, "remat_policy=save_dot_except_mlpwi saves it"),
+      ("saved_mlpwo", {"remat_policy": "custom", "mlpwo": "device"}, "remat_policy=custom saves it as mlpwo"),
+      ("saved_moe_mlpwo", {"remat_policy": "custom", "moe_mlpwo": "offload"}, "saves it as moe_mlpwo"),
+  )
+  def test_rejected(self, overrides, requirement):
+    with self.assertRaisesRegex(ValueError, f"moe_router_grad_from_wo_input=True requires: .*{re.escape(requirement)}"):
+      _cfg(moe_router_grad_from_wo_input=True, **overrides)
+
+  def test_accepted(self):
+    self.assertTrue(_cfg(moe_router_grad_from_wo_input=True, moe_chunk_pipeline=True).moe_router_grad_from_wo_input)
+    # A policy that saves only other tensors (here the dense MLP input of the shared expert) is fine.
+    cfg = _cfg(moe_router_grad_from_wo_input=True, remat_policy="custom", mlpwi_1="device")
+    self.assertTrue(cfg.moe_router_grad_from_wo_input)
+
+
+@pytest.mark.cpu_only
+def test_router_grad_from_wo_input_on_cpu_mesh():
+  """Runs `main` below in a child process with 8 CPU devices."""
+  env = os.environ.copy()
+  env["XLA_FLAGS"] = env.get("XLA_FLAGS", "") + f" {_XLA_FLAGS}"
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+  print(result.stdout)
+  assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+# --------------------------------------------------------------------------------------------------- the op alone
+_OP_EP, _OP_EXPERTS, _OP_TOPK, _OP_TOKENS, _OP_D, _OP_F = 4, 8, 2, 64, 128, 64
+# A packed buffer smaller than the busiest shard's share of the 128 slots (32 on average), so tokens are dropped.
+_OP_DROPPING_BUFFER = 24
+# The precision check's op: more slots (1024, 256 per shard on average), so each error estimate is tight.
+_PRECISION_TOKENS = 512
+_PRECISION_DROPPING_BUFFER = 192
+
+
+def _grouped_matmul(h, w, group_sizes, expert_assignments, group_offset):
+  """A down projection with its own (autodiff) VJP: row r uses expert (global group of r) - group_offset."""
+  del expert_assignments
+  ends = jnp.cumsum(group_sizes)
+  local = jnp.sum(jnp.arange(h.shape[0])[:, None] >= ends[None, :], axis=1) - group_offset
+  out = jnp.zeros((h.shape[0], w.shape[-1]), jnp.float32)
+  for e in range(w.shape[0]):
+    out = out + jnp.where((local == e)[:, None], h.astype(jnp.float32) @ w[e].astype(jnp.float32), 0.0)
+  return out.astype(h.dtype)
+
+
+def _op_inputs(seed=0, num_tokens=_OP_TOKENS):
+  """Tokens, top-k expert ids, routing weights, wi, wo, the output cotangent and per-shard down-projection inputs."""
+  k = jax.random.split(jax.random.PRNGKey(seed), 7)
+  x = jax.random.normal(k[0], (num_tokens, _OP_D), jnp.float32)
+  ids = jnp.stack([jax.random.permutation(kk, _OP_EXPERTS)[:_OP_TOPK] for kk in jax.random.split(k[1], num_tokens)])
+  w = jax.nn.softmax(jax.random.normal(k[2], (num_tokens, _OP_TOPK)), axis=-1).reshape(-1)
+  wi = jax.random.normal(k[3], (_OP_EXPERTS, _OP_D, _OP_F)) * 0.1
+  wo = jax.random.normal(k[4], (_OP_EXPERTS, _OP_F, _OP_D)) * 0.1
+  cot = jax.random.normal(k[5], (num_tokens, _OP_D), jnp.float32)
+  # For the precision check: the down-projection input of every slot, on every shard.
+  h = jax.random.normal(k[6], (_OP_EP, num_tokens * _OP_TOPK, _OP_F), jnp.float32)
+  return x, ids.astype(jnp.int32), w, wi, wo, cot, h
+
+
+def _expert_operands(gs, buffer_rows, num_slots, local):
+  """The grouped matmul's operands for this shard's buffer: (group sizes, None, group offset)."""
+  shard = jax.lax.axis_index("ep")
+  if buffer_rows < num_slots:
+    # Packed buffer: this shard's groups from row 0, clipped to the buffer.
+    ends = jnp.minimum(jnp.cumsum(jax.lax.dynamic_slice_in_dim(gs, shard * local, local)), buffer_rows)
+    return (ends - jnp.concatenate([jnp.zeros((1,), ends.dtype), ends[:-1]]), None, 0)
+  return (gs, None, shard * local)
+
+
+def _op_run(buffer_size, index_math, fused, tie=False):
+  """bf16 run on a 4-device EP mesh; returns (output, dx, dw, dwi, dwo) summed / gathered over the shards, and the
+  forward-tie checks (tie unchanged, cotangent passed through)."""
+  x, ids, w, wi, wo, cot, _ = _op_inputs()
+  ep, local = _OP_EP, _OP_EXPERTS // _OP_EP
+  mesh = Mesh(np.array(jax.devices()[:ep]), ("ep",))
+  wi_s = wi.reshape(ep, local, _OP_D, _OP_F).astype(jnp.bfloat16)
+  wo_s = wo.reshape(ep, local, _OP_F, _OP_D).astype(jnp.bfloat16)
+  tie_value = jnp.arange(8.0, dtype=jnp.float32)
+
+  def body(x, w, wi, wo, cot):
+    wi, wo = wi[0], wo[0]
+
+    def f(x, w, wi, wo, t):
+      sorted_x, gs, revert, perm = ragged_sort.ring_ragged_sort(
+          x, ids, _OP_EXPERTS, _OP_TOPK, "ep", ep, buffer_size=buffer_size, index_math=index_math
+      )
+      operands = _expert_operands(gs, sorted_x.shape[0], _OP_TOKENS * _OP_TOPK, local)
+      h = jax.nn.silu(_grouped_matmul(sorted_x, wi, *operands).astype(jnp.float32)).astype(sorted_x.dtype)
+      perm_arg = perm if index_math == "count" else None
+      if fused:
+        return ragged_sort.ring_ragged_unsort_down_projection(
+            h,
+            wo,
+            _grouped_matmul,
+            operands,
+            gs,
+            revert,
+            _OP_TOPK,
+            local,
+            "ep",
+            topk_weights=w,
+            topk_argsort_indices=perm_arg,
+            forward_tie=t if tie else None,
+            index_math=index_math,
+        )
+      y = _grouped_matmul(h, wo, *operands)
+      out = ragged_sort.ring_ragged_unsort(
+          y, gs, revert, _OP_TOPK, local, "ep", topk_weights=w, topk_argsort_indices=perm_arg, index_math=index_math
+      )
+      return (out, t) if tie else out
+
+    def g(x, w, wi, wo, t):
+      res = f(x, w, wi, wo, t)
+      return res if tie else (res, t)
+
+    (out, t_out), vjp = jax.vjp(g, x, w, wi, wo, tie_value)
+    gx, gw, gwi, gwo, gt = vjp((cot.astype(out.dtype), 2.0 * t_out))
+    return out[None], gx[None], gw[None], gwi[None], gwo[None], t_out[None], gt[None]
+
+  sm = jax.shard_map(
+      body, mesh=mesh, in_specs=(P(), P(), P("ep"), P("ep"), P()), out_specs=(P("ep"),) * 7, check_vma=False
+  )
+  args = (x.astype(jnp.bfloat16), w, wi_s, wo_s, cot)
+  outs = [np.asarray(jnp.asarray(a).astype(jnp.float32)) for a in jax.jit(sm)(*args)]
+  out, gx, gw, gwi, gwo, t_out, gt = outs
+  arrays = [
+      out.sum(0),
+      gx.sum(0),
+      gw.sum(0),
+      gwi.reshape(_OP_EXPERTS, _OP_D, _OP_F),
+      gwo.reshape(_OP_EXPERTS, _OP_F, _OP_D),
+  ]
+  tie_ok = bool(np.array_equal(t_out[0], np.asarray(tie_value)))
+  tie_ok &= bool(np.array_equal(gt[0], 2.0 * np.asarray(tie_value)))
+  barrier = None
+  if fused and tie:
+    jaxpr = jax.make_jaxpr(sm)(*args).jaxpr
+    barrier = any(_barrier_ties_shapes(eqn) for eqn in _eqns(jaxpr))
+  return arrays, tie_ok, barrier
+
+
+def _barrier_ties_shapes(eqn):
+  """An optimization barrier whose operands include the tie (shape (8,)) next to a [rows, D] GEMM output."""
+  if eqn.primitive.name != "optimization_barrier":
+    return False
+  shapes = [tuple(v.aval.shape) for v in eqn.invars]
+  return (8,) in shapes and any(len(s) == 2 and s[-1] == _OP_D for s in shapes)
+
+
+def _rel(value, reference):
+  """Relative L2 distance of `value` from `reference`."""
+  return float(np.linalg.norm(value - reference) / max(np.linalg.norm(reference), 1e-30))
+
+
+def _op_case(label, buffer_size, index_math, tie=False, control=None):
+  """The fused op against the two ops apart; `control` wraps the fused run."""
+  base, _, _ = _op_run(buffer_size, index_math, False, tie)
+  with control if control is not None else contextlib.nullcontext():
+    new, tie_ok, barrier = _op_run(buffer_size, index_math, True, tie)
+  passed = bool(np.array_equal(new[0], base[0]))
+  lines = [f"  [{'ok' if passed else 'FAIL'}] output bit-identical={passed}"]
+  for name, a, b in zip(("dx", "dw", "dwi", "dwo"), base[1:], new[1:]):
+    rel = _rel(b, a)
+    ok = rel < _GRAD_REL_L2 and bool(np.isfinite(b).all())
+    passed &= ok
+    lines.append(f"  [{'ok' if ok else 'FAIL'}] {name}: rel_l2 vs two ops={rel:.2e}")
+  if tie:
+    passed &= tie_ok and bool(barrier)
+    lines.append(
+        f"  [{'ok' if tie_ok and barrier else 'FAIL'}] forward tie unchanged, cotangent passed: {tie_ok};"
+        f" barrier ties it to the GEMM output: {barrier}"
+    )
+  print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+def _dense_unsort_down_projection(h, w, wo, ids, gs, revert, shard, local):
+  """The fused op's function written densely, slot by slot (no ragged kernels, no rounding in float32).
+
+  Slot i (token i // top_k) reads buffer row `revert[i]` (packed buffer: minus this shard's window start, and only
+  below the buffer's row count), multiplies it by its expert's down projection, and adds it, weighted, to its token.
+  """
+  num_slots, buffer_rows = revert.shape[0], h.shape[0]
+  local_expert = ids.reshape(-1) - shard * local
+  is_local = (local_expert >= 0) & (local_expert < local)
+  start = jnp.sum(jnp.where(jnp.arange(gs.shape[0]) < shard * local, gs, 0))
+  row = revert if buffer_rows >= num_slots else revert - start
+  limit = buffer_rows if buffer_rows >= num_slots else jnp.minimum(jnp.sum(jnp.where(is_local, 1, 0)), buffer_rows)
+  valid = is_local & (row >= 0) & (row < limit)
+  h_slot = jnp.where(valid[:, None], h[jnp.where(valid, row, 0)], 0.0)
+  y_slot = jnp.einsum("nf,nfd->nd", h_slot, wo[jnp.clip(local_expert, 0, local - 1)])
+  weighted = jnp.where(valid[:, None], (w[:, None] * y_slot), 0.0)
+  return jnp.sum(weighted.reshape(num_slots // _OP_TOPK, _OP_TOPK, -1), axis=1)
+
+
+def _op_backward_run(seed, buffer_size, index_math, path):
+  """The op (path "fused"), the two ops apart ("two_ops", both bf16) or the dense float32 function ("reference"),
+  differentiated at (h, w, wo) with the cotangent g on a 4-device EP mesh. h is an input, so the only roundings are
+  the backward's own (and, in the two-ops path, the forward GEMM output's). Returns (dh, dw, dwo), summed or gathered
+  over the shards. The routing weights are in [0.25, 0.3), the bottom of a binade, where bf16 rounds them worst."""
+  x, ids, _, _, wo, cot, h = _op_inputs(seed, _PRECISION_TOKENS)
+  w = jax.random.uniform(jax.random.PRNGKey(1000 + seed), (_PRECISION_TOKENS * _OP_TOPK,), jnp.float32, 0.25, 0.3)
+  ep, local = _OP_EP, _OP_EXPERTS // _OP_EP
+  mesh = Mesh(np.array(jax.devices()[:ep]), ("ep",))
+  # The same bf16-representable inputs on every path; the reference only computes in float32.
+  dtype = jnp.float32 if path == "reference" else jnp.bfloat16
+  wo_s = wo.reshape(ep, local, _OP_F, _OP_D).astype(jnp.bfloat16).astype(dtype)
+  h_s = h.astype(jnp.bfloat16).astype(dtype)
+  cot = cot.astype(jnp.bfloat16).astype(dtype)
+
+  def body(h, w, wo, cot):
+    h, wo = h[0], wo[0]
+    sorted_x, gs, revert, perm = ragged_sort.ring_ragged_sort(
+        x.astype(jnp.bfloat16), ids, _OP_EXPERTS, _OP_TOPK, "ep", ep, buffer_size=buffer_size, index_math=index_math
+    )
+    rows = sorted_x.shape[0]
+    operands = _expert_operands(gs, rows, _PRECISION_TOKENS * _OP_TOPK, local)
+    perm_arg = perm if index_math == "count" else None
+
+    def f(h, w, wo):
+      if path == "reference":
+        return _dense_unsort_down_projection(h, w, wo, ids, gs, revert, jax.lax.axis_index("ep"), local)
+      if path == "fused":
+        return ragged_sort.ring_ragged_unsort_down_projection(
+            h,
+            wo,
+            _grouped_matmul,
+            operands,
+            gs,
+            revert,
+            _OP_TOPK,
+            local,
+            "ep",
+            topk_weights=w,
+            topk_argsort_indices=perm_arg,
+            index_math=index_math,
+        )
+      y = _grouped_matmul(h, wo, *operands)
+      return ragged_sort.ring_ragged_unsort(
+          y, gs, revert, _OP_TOPK, local, "ep", topk_weights=w, topk_argsort_indices=perm_arg, index_math=index_math
+      )
+
+    out, vjp = jax.vjp(f, h[:rows], w, wo)
+    dh, dw, dwo = vjp(cot.astype(out.dtype))
+    return dh[None], dw[None], dwo[None]
+
+  sm = jax.shard_map(body, mesh=mesh, in_specs=(P("ep"), P(), P("ep"), P()), out_specs=(P("ep"),) * 3, check_vma=False)
+  dh, dw, dwo = [np.asarray(jnp.asarray(a).astype(jnp.float32)) for a in jax.jit(sm)(h_s, w, wo_s, cot)]
+  return dh, dw.sum(0), dwo.reshape(_OP_EXPERTS, _OP_F, _OP_D)
+
+
+def _op_precision_case(label, buffer_size, index_math, control=None):
+  """Fused vs two-ops error against a float32 run, over _PRECISION_SEEDS; `control` wraps the fused runs."""
+  passed, lines, ratios = True, [], []
+  for seed in _PRECISION_SEEDS:
+    reference = _op_backward_run(seed, buffer_size, index_math, "reference")
+    two_ops = _op_backward_run(seed, buffer_size, index_math, "two_ops")
+    with control if control is not None else contextlib.nullcontext():
+      fused = _op_backward_run(seed, buffer_size, index_math, "fused")
+    parts = []
+    for name, r, a, b in zip(("dh", "dw", "dwo"), reference, two_ops, fused):
+      err_two_ops, err_fused = _rel(a, r), _rel(b, r)
+      ratio = err_fused / err_two_ops
+      ok = ratio <= _PRECISION_RATIO and bool(np.isfinite(b).all()) and err_two_ops > 0
+      passed &= ok
+      ratios.append(ratio)
+      parts.append(f"{name} {err_two_ops:.2e}/{err_fused:.2e}={ratio:.3f}{'' if ok else ' FAIL'}")
+    lines.append(f"  seed {seed}: error vs float32, two ops / fused = ratio: {'; '.join(parts)}")
+  print(
+      f"{label}: {'PASS' if passed else 'FAIL'} (ratio <= {_PRECISION_RATIO}; over the seeds min"
+      f" {min(ratios):.3f} max {max(ratios):.3f})",
+      flush=True,
+  )
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+def _count_needs_forward_permutation():
+  try:
+    _op_run_count_without_permutation()
+  except ValueError as e:
+    raised = "needs topk_argsort_indices" in str(e)
+  else:
+    raised = False
+  print(f"count without the forward permutation: [{'ok' if raised else 'FAIL'}] the fused backward raises", flush=True)
+  return raised
+
+
+def _op_run_count_without_permutation():
+  """Differentiates the fused op under "count" without passing the forward permutation (must raise)."""
+  x, ids, w, _, wo, cot, _ = _op_inputs()
+  mesh = Mesh(np.array(jax.devices()[:2]), ("ep",))
+
+  def body(x, wo, cot):
+    def f(x, wo):
+      sorted_x, gs, revert, _ = ragged_sort.ring_ragged_sort(x, ids, _OP_EXPERTS, _OP_TOPK, "ep", 2, index_math="count")
+      return ragged_sort.ring_ragged_unsort_down_projection(
+          sorted_x[:, :_OP_F],
+          wo[:4],
+          _grouped_matmul,
+          (gs, None, 0),
+          gs,
+          revert,
+          _OP_TOPK,
+          4,
+          "ep",
+          topk_weights=w,
+          index_math="count",
+      )
+
+    out, vjp = jax.vjp(f, x, wo)
+    return vjp(cot.astype(out.dtype))[1]
+
+  jax.shard_map(body, mesh=mesh, in_specs=(P(), P(), P()), out_specs=P(), check_vma=False)(
+      x.astype(jnp.bfloat16), wo.astype(jnp.bfloat16), cot
+  )
+
+
+def _fan_out_changed(weight=None, slot_value=None):
+  """Deliberately wrong: the fused op's backward fan-out with its routing weights or its per-slot result changed.
+
+  ring_ragged_unsort uses the same helper, so only wrap the fused run with this.
+  """
+  real = ragged_sort._ring_unsort_fan_out  # pylint: disable=protected-access
+
+  def wrong(*args, **kwargs):
+    gathered, weight_for_sorted, mask, to_slot_order = real(*args, **kwargs)
+    if weight is not None:
+      weight_for_sorted = weight(weight_for_sorted)
+    if slot_value is not None:
+      return gathered, weight_for_sorted, mask, lambda per_row: to_slot_order(slot_value(per_row))
+    return gathered, weight_for_sorted, mask, to_slot_order
+
+  return mock.patch.object(ragged_sort, "_ring_unsort_fan_out", wrong)
+
+
+def _weight_dropped_from_backward():
+  """Deliberately wrong: the fused backward's fan-out reports every routing weight as 1."""
+  return _fan_out_changed(weight=jnp.ones_like)
+
+
+def _weight_rounded_to_bf16():
+  """Deliberately wrong: the fused backward uses the routing weight rounded to bf16."""
+  return _fan_out_changed(weight=lambda w: w.astype(jnp.bfloat16).astype(jnp.float32))
+
+
+def _router_grad_shrunk():
+  """Deliberately wrong: the fused backward's routing-weight gradient 0.5% too small."""
+  return _fan_out_changed(slot_value=lambda dot: dot * 0.995)
+
+
+# --------------------------------------------------------------------------------------------------- the layer
+def _build(cfg, mesh):
+  return moe.RoutedMoE(
+      config=cfg,
+      num_experts=cfg.num_experts,
+      num_experts_per_tok=cfg.num_experts_per_tok,
+      mesh=mesh,
+      kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+      kernel_axes=("embed", "mlp"),
+      intermediate_dim=cfg.moe_mlp_dim,
+      weight_dtype=cfg.weight_dtype,
+      dtype=cfg.dtype,
+      rngs=nnx.Rngs(params=0),
+  )
+
+
+def _inputs(cfg, mode):
+  """The layer input, forced expert ids and a token mask, per `mode`."""
+  batch = int(cfg.per_device_batch_size * _NUM_DEVICES)
+  x = jax.random.normal(jax.random.PRNGKey(7), (batch, _SEQ, _EMB), jnp.float32)
+  forced = token_mask = None
+  if "forced" in mode:
+    keys = jax.random.split(jax.random.PRNGKey(11), batch * _SEQ)
+    ids = jax.vmap(lambda k: jax.random.permutation(k, _NUM_EXPERTS)[:_TOP_K])(keys).reshape(batch, _SEQ, _TOP_K)
+    forced = ids.astype(jnp.int32)
+  if "padmask" in mode:
+    lens = jnp.array([_SEQ - 5 * (b + 1) for b in range(batch)])
+    token_mask = (jnp.arange(_SEQ)[None, :] < lens[:, None]).astype(jnp.int32)
+  return x, forced, token_mask
+
+
+def _eqns(jaxpr):
+  for eqn in jaxpr.eqns:
+    yield eqn
+    for param in eqn.params.values():
+      for sub in param if isinstance(param, (list, tuple)) else (param,):
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+        if hasattr(sub, "eqns"):
+          yield from _eqns(sub)
+
+
+def _compiled_gemm_loops(hlo_text):
+  """Number of while instructions in a compiled HLO module.
+
+  The GMMs are megablox kernels in interpret mode, each a loop over its tiles; identical GEMMs share one body
+  computation, so the loop instructions (one per call), not the dots in the bodies, count the GEMMs.
+  """
+  return sum(" while(" in line for line in hlo_text.splitlines())
+
+
+def _sub_jaxprs(eqn):
+  for param in eqn.params.values():
+    for sub in param if isinstance(param, (list, tuple)) else (param,):
+      sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+      sub = sub if hasattr(sub, "eqns") else getattr(sub, "jaxpr", None)
+      if hasattr(sub, "eqns"):
+        yield sub
+
+
+def _tie_reads(jaxpr):
+  """(fused ops with a forward tie, how many of them have their tie output read later), over all sub-jaxprs.
+
+  A fused op with a tie is a custom_vjp call with two outputs whose primal holds an optimization barrier and the
+  unsort's ragged_gather_reduce (barrier_tie's own custom_vjp holds only the barrier). Its second output is the next
+  chunk's routed tokens after the barrier: if nothing reads it, that chunk is no longer ordered after this chunk's down
+  projection.
+  """
+  ties = read = 0
+  for i, eqn in enumerate(jaxpr.eqns):
+    call = eqn.params.get("call_jaxpr")
+    call = getattr(call, "jaxpr", call)
+    if (
+        eqn.primitive.name.startswith("custom_vjp_call")
+        and len(eqn.outvars) == 2
+        and hasattr(call, "eqns")
+        and any(e.primitive.name == "optimization_barrier" for e in _eqns(call))
+        and any(e.params.get("name") == "ragged_gather_reduce" for e in _eqns(call))
+    ):
+      ties += 1
+      tie_out = eqn.outvars[-1]
+      read += any(tie_out in later.invars for later in jaxpr.eqns[i + 1 :]) or tie_out in jaxpr.outvars
+    for sub in _sub_jaxprs(eqn):
+      sub_ties, sub_read = _tie_reads(sub)
+      ties, read = ties + sub_ties, read + sub_read
+  return ties, read
+
+
+def _run(cfg, mesh, params, x, forced, token_mask):
+  """Loss, output, gradients; forward barriers and fused-op ties; GEMM loops of the compiled rematerialized
+  backward."""
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules):
+    graphdef, _, rest = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+
+    def loss_fn(params, x):
+      out, lb_loss, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask)
+      loss = jnp.mean(out.astype(jnp.float32) ** 2)
+      if lb_loss is not None:
+        loss = loss + lb_loss.astype(jnp.float32)
+      return loss, out
+
+    forward = jax.make_jaxpr(loss_fn)(params, x).jaxpr
+    barriers = sum(e.primitive.name == "optimization_barrier" for e in _eqns(forward))
+    ties = _tie_reads(forward)
+
+    def linear_loss(params, x):
+      # Linear in the layer output, as behind the decoder's residual add: the layer's own backward does not need its
+      # output, so only the backward's own reads of the down-projection output keep that GEMM in the replay.
+      out, _, _ = nnx.merge(graphdef, params, rest)(x, forced_routed_experts=forced, token_mask=token_mask)
+      return jnp.sum(out.astype(jnp.float32) * jax.random.normal(jax.random.PRNGKey(3), out.shape, jnp.float32))
+
+    remat = jax.grad(jax.checkpoint(linear_loss, policy=jax.checkpoint_policies.nothing_saveable), argnums=(0, 1))
+    loops = _compiled_gemm_loops(jax.jit(remat).lower(params, x).compile().as_text())
+    (loss, out), grads = jax.jit(jax.value_and_grad(loss_fn, argnums=(0, 1), has_aux=True))(params, x)
+  return jax.device_get((loss, out, grads)), barriers, ties, loops
+
+
+def _layer_case(label, overrides, mode="real", control=None, keys=None, same_barriers=True):
+  """The layer with `keys` off and on (default: this key alone), on the same parameters; `control` wraps the run
+  with them on. `same_barriers`: the forward keeps its optimization barriers (False where other keys change them)."""
+  keys = {"moe_router_grad_from_wo_input": True} if keys is None else keys
+  cfg_old = _cfg(**overrides)
+  cfg_new = _cfg(**overrides, **keys)
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg_old), cfg_old.mesh_axes)
+  x, forced, token_mask = _inputs(cfg_old, mode)
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg_old.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg_old, mesh), nnx.Param, ...)
+  old, barriers_old, ties_old, loops_old = _run(cfg_old, mesh, params, x, forced, token_mask)
+  with control if control is not None else contextlib.nullcontext():
+    new, barriers_new, ties_new, loops_new = _run(cfg_new, mesh, params, x, forced, token_mask)
+  (loss_old, out_old, grads_old), (loss_new, out_new, grads_new) = old, new
+  n_chunks = cfg_new.num_moe_token_chunks
+  # One GEMM loop fewer per chunk in the rematerialized backward: the down projection is not replayed.
+  loops_ok = loops_old - loops_new == n_chunks
+  barriers_ok = barriers_new == barriers_old or not same_barriers
+  # Pipelined: every chunk but the last ties its down projection to the next chunk's routed tokens, which are read.
+  want_ties = (n_chunks - 1, n_chunks - 1) if cfg_new.moe_chunk_pipeline else (0, 0)
+  ties_ok = ties_old == (0, 0) and ties_new == want_ties
+  passed = loops_ok and barriers_ok and ties_ok
+  lines = [
+      f"  [{'ok' if loops_ok else 'FAIL'}] compiled remat-backward GEMM loops off/on={loops_old}/{loops_new}"
+      f" (want {n_chunks} fewer)",
+      f"  [{'ok' if barriers_ok else 'FAIL'}] forward barriers off/on={barriers_old}/{barriers_new}"
+      f"{'' if same_barriers else ' (not compared)'}",
+      f"  [{'ok' if ties_ok else 'FAIL'}] fused-op ties (count, read) off/on={ties_old}/{ties_new} (want {want_ties})",
+  ]
+  for name, a, b in (("output", out_old, out_new), ("loss", loss_old, loss_new)):
+    same = bool(np.array_equal(np.asarray(a, np.float32), np.asarray(b, np.float32)))
+    passed &= same
+    lines.append(f"  [{'ok' if same else 'FAIL'}] {name}: bit-identical={same}")
+  for (path, a), b in zip(jax.tree_util.tree_flatten_with_path(grads_old)[0], jax.tree_util.tree_leaves(grads_new)):
+    a, b = np.asarray(a, np.float32), np.asarray(b, np.float32)
+    rel = _rel(b, a)
+    ok = rel < _GRAD_REL_L2 and bool(np.isfinite(b).all())
+    passed &= ok
+    lines.append(f"  [{'ok' if ok else 'FAIL'}] {jax.tree_util.keystr(path)}: rel_l2={rel:.2e}")
+  print(f"{label}: {'PASS' if passed else 'FAIL'}", flush=True)
+  print("\n".join(lines), flush=True)
+  return passed
+
+
+class _ConfigOverride:
+  """`config` with some attributes replaced: reaches the layer's trace-time checks past the config validator."""
+
+  def __init__(self, config, **overrides):
+    self._config = config
+    self._overrides = overrides
+
+  def __getattr__(self, name):  # only called for names that are not set on the instance
+    if name.startswith("__") or "_config" not in self.__dict__:
+      raise AttributeError(name)
+    return self._overrides[name] if name in self._overrides else getattr(self._config, name)
+
+
+def _tensor_parallel_guard():
+  """On a tensor-parallel mesh the layer refuses the key at trace time; without it the layer runs."""
+  cfg = _cfg(**_TP_MESH)
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+  x, forced, token_mask = _inputs(cfg, "real")
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules):
+    _, params, _ = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+  (loss, _, _), _, _, _ = _run(cfg, mesh, params, x, forced, token_mask)
+  plain_runs = bool(np.isfinite(loss))
+  try:
+    _run(_ConfigOverride(cfg, moe_router_grad_from_wo_input=True), mesh, params, x, forced, token_mask)
+  except ValueError as e:
+    raised = "does not support tensor parallelism in the MoE layer" in str(e)
+  else:
+    raised = False
+  ok = plain_runs and raised
+  print(
+      f"tensor_parallel_guard: [{'ok' if ok else 'FAIL'}] mesh {dict(mesh.shape)}: plain layer runs: {plain_runs};"
+      f" the key past the validator raises at trace time: {raised}",
+      flush=True,
+  )
+  return ok
+
+
+# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+_BENCH_FLAGS = {
+    "moe_chunk_pipeline": True,
+    "moe_expert_weight_prefetch": "forward",
+    "ring_of_experts_row_major_reduce_scatter": True,
+    "moe_routing_maps": "device",
+    "ragged_buffer_factor": 2.75,
+    "prefuse_moe_weights": True,
+    "fp8_moe_fake_quant": True,
+    "moe_exclude_padding_from_routing": True,
+}
+# Every key of this series that composes with the others on the megablox path the CPU runs (the fused SiLU GMM is
+# tokamax-gmm_v2-only; moe_wi_gmm_fused_swiglu_test.py stacks it with the others using kernel stand-ins).
+_ALL_KEYS = {
+    "moe_routing_index_math": "count",
+    "moe_unsplit_prefused_wi": True,
+    "moe_prefused_swiglu_vjp": True,
+    "moe_roe_packed_routing_gather": True,
+    "moe_router_grad_from_wo_input": True,
+}
+
+
+def main():
+  assert jax.device_count() == _NUM_DEVICES, jax.devices()
+  pipelined = {"moe_chunk_pipeline": True}
+  dropping, precision_dropping = _OP_DROPPING_BUFFER, _PRECISION_DROPPING_BUFFER
+  results = {
+      "op_full_buffer_sort": _op_case("op_full_buffer_sort", None, "sort"),
+      "op_full_buffer_count": _op_case("op_full_buffer_count", None, "count"),
+      "op_dropping_buffer_sort": _op_case("op_dropping_buffer_sort", dropping, "sort"),
+      "op_dropping_buffer_count": _op_case("op_dropping_buffer_count", dropping, "count"),
+      "op_forward_tie": _op_case("op_forward_tie", dropping, "sort", tie=True),
+      "op_count_guard": _count_needs_forward_permutation(),
+      "op_precision_full_buffer_sort": _op_precision_case("op_precision_full_buffer_sort", None, "sort"),
+      "op_precision_dropping_buffer_count": _op_precision_case(
+          "op_precision_dropping_buffer_count", precision_dropping, "count"
+      ),
+      "layer_plain_loop": _layer_case("layer_plain_loop", {}),
+      "layer_pipelined": _layer_case("layer_pipelined", pipelined),
+      "layer_dropless": _layer_case("layer_dropless", {"ragged_buffer_factor": -1.0}, mode="forced"),
+      "layer_count_padding_excluded": _layer_case(
+          "layer_count_padding_excluded",
+          {**pipelined, "moe_routing_index_math": "count", "moe_exclude_padding_from_routing": True},
+          mode="padmask",
+      ),
+      "layer_saved_routing_maps": _layer_case("layer_saved_routing_maps", {**pipelined, "moe_routing_maps": "device"}),
+      "layer_prefused_keys": _layer_case(
+          "layer_prefused_keys",
+          {**pipelined, "prefuse_moe_weights": True, "moe_unsplit_prefused_wi": True, "moe_prefused_swiglu_vjp": True},
+      ),
+      "layer_bench_flags_all_keys": _layer_case(
+          "layer_bench_flags_all_keys", _BENCH_FLAGS, mode="padmask", keys=_ALL_KEYS, same_barriers=False
+      ),
+      "tensor_parallel_guard": _tensor_parallel_guard(),
+  }
+  controls = {
+      "control_op_weight_dropped": _op_case(
+          "control_op_weight_dropped", dropping, "sort", control=_weight_dropped_from_backward()
+      ),
+      "control_op_precision_weight_rounded_to_bf16": _op_precision_case(
+          "control_op_precision_weight_rounded_to_bf16", None, "sort", control=_weight_rounded_to_bf16()
+      ),
+      "control_op_precision_router_grad_shrunk": _op_precision_case(
+          "control_op_precision_router_grad_shrunk", precision_dropping, "count", control=_router_grad_shrunk()
+      ),
+      "control_layer_weight_dropped": _layer_case(
+          "control_layer_weight_dropped", pipelined, control=_weight_dropped_from_backward()
+      ),
+  }
+  print("SUMMARY", {**results, **controls}, flush=True)
+  failed = [k for k, v in results.items() if not v]
+  undetected = [k for k, v in controls.items() if v]
+  if failed or undetected:
+    print(f"cases that failed: {failed}; positive controls that were NOT detected: {undetected}", flush=True)
+    sys.exit(1)
+  print(_PASS_TOKEN, flush=True)
+
+
+if __name__ == "__main__":
+  main()

From cb709c93cfad39cdfcb9be59f12647990db67785 Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 04:14:18 -0700
Subject: [PATCH 08/10] GDN: keep the Conv1D halo input in qkv's dtype
 (gdn_conv1d_input_in_qkv_dtype)

The GDN kernel path's Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py)
builds its halo-extended input [conv_state | qkv] in float32 and slices the
gdn_conv_kernel_dim taps out of it. On the Qwen3.5-397B RL trainer that is an
f32[1, 16387, 12288] = 805.5 MB copy per GDN layer in the backward pass,
written once and read once per tap.

With gdn_conv1d_input_in_qkv_dtype=True the input stays in qkv's dtype and
each tap is cast to float32 where it is read, so XLA fuses the cast into the
consumer and the float32 copy is never written. No optimization barrier.

The option reaches the Conv1D through gdn_decoupled_conv1d's new static
argument conv_input_in_qkv_dtype, which scopes it around the op's primal,
forward and backward rules (the backward rule is traced after the model's
forward returns, so only the custom_vjp's static arguments reach it).

Numerics: the same float32 values in the same order, so op by op every
output is bit-identical. Compiled, the forward outputs and the qkv,
conv-state and bias gradients stay bit-identical; the conv weight gradient
(a float32 reduction over batch x sequence) can be summed in a different
order once XLA fuses the cast into it (CPU: 11 of 64 elements, <= 5.1e-6
relative, in 2 of 16 configurations).

Default off. Requires use_gdn_kernel=True (validated).
---
 src/maxtext/configs/base.yml                  |   5 +
 src/maxtext/configs/types.py                  |  21 +
 src/maxtext/kernels/gdn/gdn_bwd/api.py        | 848 +++++++++---------
 .../kernels/gdn/gdn_bwd/compute_conv1d_bwd.py | 179 +++-
 src/maxtext/kernels/gdn/model_runner.py       |   5 +
 .../gdn_conv1d_input_in_qkv_dtype_test.py     | 289 ++++++
 6 files changed, 931 insertions(+), 416 deletions(-)
 create mode 100644 tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index a43f4198cf..09f8b45af3 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1638,6 +1638,11 @@ gdn_cp_mode: "auto"
 # ('highest' = 6 bf16 passes on TPU, 'high' = 3 passes). Only used with use_gdn_kernel: true and
 # sequence-sharded GDN CP (gdn_cp_mode 'seq', or 'auto' when it does not pick head sharding); a no-op otherwise.
 gdn_cp_matmul_precision: "highest"
+# GDN kernel path: build the Conv1D + SiLU halo-extended input [conv_state | qkv] in qkv's dtype and cast each tap
+# to float32 where it is read, instead of materializing a float32 copy (805.5 MB per GDN layer in the backward of
+# the Qwen3.5-397B RL trainer). Same values: outputs and the qkv gradient bit-identical, the conv weight gradient up to
+# float32 summation order. Requires use_gdn_kernel: true.
+gdn_conv1d_input_in_qkv_dtype: false
 
 use_tokamax_splash: false
 # Setting this flag will use a non-pallas implementation.
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 405cb657d8..db5ed19e63 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2235,6 +2235,22 @@ class Qwen3Next(BaseModel):
           " HIGHEST) or with head-sharded context parallelism."
       ),
   )
+  gdn_conv1d_input_in_qkv_dtype: bool = Field(
+      False,
+      description=(
+          "GDN kernel path: build the Conv1D + SiLU input, the halo-extended [conv_state | qkv] (or qkv left-padded"
+          " by gdn_conv_kernel_dim - 1 zeros), in qkv's dtype and cast each of the gdn_conv_kernel_dim taps to"
+          " float32 where it is read, instead of materializing a float32 copy of the whole input that every tap"
+          " then slices (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py). The copy is f32[1, 16387, 12288] = 805.5 MB"
+          " per GDN layer in the backward pass of the Qwen3.5-397B RL trainer (bf16 activations, sequence-sharded"
+          " context parallelism). bfloat16 -> float32 is lossless and every product and sum is formed in the same"
+          " order, so the forward outputs and the qkv / conv-state / bias gradients are bit-identical; the conv weight"
+          " gradient, a float32 reduction over batch x sequence, may be summed in a different order once XLA fuses"
+          " the cast into it (CPU: a few ulp). Only which tensors live in HBM changes; no optimization barrier is"
+          " involved. Applies to the Conv1D forward and backward that gdn_decoupled_conv1d and its custom-VJP rules"
+          " trace. Requires use_gdn_kernel=True (the only path that runs that Conv1D)."
+      ),
+  )
 
 
 # ----------------------------------------------------------------------------
@@ -5292,6 +5308,11 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None:
         raise ValueError("GatedDeltaNet context parallelism requires context_parallel_load_balance=False.")
 
     if not cfg.use_gdn_kernel:
+      if getattr(cfg, "gdn_conv1d_input_in_qkv_dtype", False):
+        raise ValueError(
+            "gdn_conv1d_input_in_qkv_dtype=True requires use_gdn_kernel=True: only the GDN kernel path runs the"
+            " Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py) whose input it keeps in qkv's dtype."
+        )
       return
     if cfg.ici_tensor_parallelism > 1 or cfg.dcn_tensor_parallelism > 1:
       raise ValueError(
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index f703c1dda2..ca8fcb8550 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -26,6 +26,7 @@
 from .. import compute_gdn as local_compute_gdn
 from .. import wrapper as local_gdn_wrapper
 from . import cp_gdn
+from . import compute_conv1d_bwd
 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
@@ -473,7 +474,7 @@ def _run_cp_gdn_decoupled_fwd_impl(
   return (out, states), t_inv_2, chunk_states, conv_halo, s_in_r, m_local
 
 
-@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19))
+@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20))
 def gdn_decoupled_conv1d(
     qkv: jax.Array,
     b: jax.Array,
@@ -495,14 +496,48 @@ def gdn_decoupled_conv1d(
     cp_axis_name: str | tuple[str, ...] | None = None,
     segment_ids: Optional[jax.Array] = None,
     cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION,
+    conv_input_in_qkv_dtype: bool = False,
 ) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]:
   """Decoupled Conv1D + GDN with Pallas backward pass and optional sequence-sharded CP.
 
   `cp_matmul_precision` sets the precision of the f32 matmuls that compose the recurrent state across
   sequence-sharded CP ranks (`cp_gdn`); it is unused without CP.
+
+  `conv_input_in_qkv_dtype` (config gdn_conv1d_input_in_qkv_dtype) makes every Conv1D + SiLU forward and backward
+  traced by this op and its forward and backward rules keep the halo-extended conv input in qkv's dtype instead of
+  materializing it in float32 (compute_conv1d_bwd.py). Same values (the conv weight gradient up to float32
+  summation order); only HBM traffic changes.
   """
-  if _is_cp_active(cp_axis_name):
-    (out, states), _, _, _, _, _ = _run_cp_gdn_decoupled_fwd_impl(
+  with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
+    if _is_cp_active(cp_axis_name):
+      (out, states), _, _, _, _, _ = _run_cp_gdn_decoupled_fwd_impl(
+          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,
+          cp_axis_name=cp_axis_name,
+          segment_ids=segment_ids,
+          cp_matmul_precision=cp_matmul_precision,
+      )
+      return out, states
+
+    segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
+        segment_ids, (conv_state is not None) or (recurrent_state is not None), conv_kernel_size
+    )
+    (out, states), _, _ = _run_local_gdn_decoupled_fwd(
         qkv,
         b,
         a,
@@ -520,39 +555,12 @@ def gdn_decoupled_conv1d(
         chunk_size=chunk_size,
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
         compute_dtype=compute_dtype,
-        cp_axis_name=cp_axis_name,
         segment_ids=segment_ids,
-        cp_matmul_precision=cp_matmul_precision,
+        conv_halo_seg=conv_halo_seg,
+        init_seg=init_seg,
     )
     return out, states
 
-  segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
-      segment_ids, (conv_state is not None) or (recurrent_state is not None), conv_kernel_size
-  )
-  (out, states), _, _ = _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 out, states
-
 
 def _unwrap_primal(x):
   return jax.tree.map(lambda v: v.value if hasattr(v, "value") else v, x)
@@ -585,42 +593,88 @@ def _gdn_decoupled_conv1d_fwd(
     cp_axis_name: str | tuple[str, ...] | None = None,
     segment_ids: Optional[jax.Array] = None,
     cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION,
+    conv_input_in_qkv_dtype: bool = False,
 ):
   """Forward rule for custom_vjp registration of gdn_decoupled_conv1d."""
-  (
-      qkv,
-      b,
-      a,
-      conv_weight,
-      conv_bias,
-      a_log,
-      dt_bias,
-      conv_state,
-      recurrent_state,
-      segment_ids,
-  ) = _unwrap_primal((
-      qkv,
-      b,
-      a,
-      conv_weight,
-      conv_bias,
-      a_log,
-      dt_bias,
-      conv_state,
-      recurrent_state,
-      segment_ids,
-  ))
-  qkv = checkpoint_name(qkv, "gdn_fwd_conv")
-  qkv = checkpoint_name(qkv, "gdn_conv_out")
-  has_initial_state = (conv_state is not None) or (recurrent_state is not None)
-  if _is_cp_active(cp_axis_name):
-    if segment_ids is not None:
-      seg_metadata = cp_gdn.gather_cp_segment_metadata(
-          segment_ids, cp_axis_name, conv_kernel_size, has_initial_state=has_initial_state
+  with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
+    (
+        qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        segment_ids,
+    ) = _unwrap_primal((
+        qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        segment_ids,
+    ))
+    qkv = checkpoint_name(qkv, "gdn_fwd_conv")
+    qkv = checkpoint_name(qkv, "gdn_conv_out")
+    has_initial_state = (conv_state is not None) or (recurrent_state is not None)
+    if _is_cp_active(cp_axis_name):
+      if segment_ids is not None:
+        seg_metadata = cp_gdn.gather_cp_segment_metadata(
+            segment_ids, cp_axis_name, conv_kernel_size, has_initial_state=has_initial_state
+        )
+      else:
+        seg_metadata = None
+      (out, states), t_inv, chunk_states, conv_halo, s_in_r, m_local = _run_cp_gdn_decoupled_fwd_impl(
+          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,
+          cp_axis_name=cp_axis_name,
+          segment_ids=segment_ids,
+          seg_metadata=seg_metadata,
+          cp_matmul_precision=cp_matmul_precision,
       )
-    else:
-      seg_metadata = None
-    (out, states), t_inv, chunk_states, conv_halo, s_in_r, m_local = _run_cp_gdn_decoupled_fwd_impl(
+      out = checkpoint_name(out, "gdn_core_attn_out")
+      residuals = (
+          checkpoint_name(qkv, "gdn_qkv") if qkv is not None else None,
+          checkpoint_name(b, "gdn_b") if b is not None else None,
+          checkpoint_name(a, "gdn_a") if a is not None else None,
+          conv_weight,
+          conv_bias,
+          a_log,
+          dt_bias,
+          checkpoint_name(conv_halo, "gdn_conv_state"),
+          checkpoint_name(s_in_r, "gdn_recurrent_state"),
+          checkpoint_name(t_inv, "gdn_t_inv") if t_inv is not None else None,
+          checkpoint_name(chunk_states, "gdn_chunk_states") if chunk_states is not None else None,
+          checkpoint_name(m_local, "gdn_m_local") if m_local is not None else None,
+          conv_state is not None,
+          recurrent_state is not None,
+          seg_metadata if seg_metadata is not None else segment_ids,
+      )
+      return (out, states), residuals
+
+    segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(segment_ids, has_initial_state, conv_kernel_size)
+    (out, states), t_inv, chunk_states = _run_local_gdn_decoupled_fwd(
         qkv,
         b,
         a,
@@ -638,12 +692,12 @@ def _gdn_decoupled_conv1d_fwd(
         chunk_size=chunk_size,
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
         compute_dtype=compute_dtype,
-        cp_axis_name=cp_axis_name,
         segment_ids=segment_ids,
-        seg_metadata=seg_metadata,
-        cp_matmul_precision=cp_matmul_precision,
+        conv_halo_seg=conv_halo_seg,
+        init_seg=init_seg,
     )
     out = checkpoint_name(out, "gdn_core_attn_out")
+    seg_bundle = (segment_ids, conv_halo_seg, init_seg) if segment_ids is not None else None
     residuals = (
         checkpoint_name(qkv, "gdn_qkv") if qkv is not None else None,
         checkpoint_name(b, "gdn_b") if b is not None else None,
@@ -652,58 +706,14 @@ def _gdn_decoupled_conv1d_fwd(
         conv_bias,
         a_log,
         dt_bias,
-        checkpoint_name(conv_halo, "gdn_conv_state"),
-        checkpoint_name(s_in_r, "gdn_recurrent_state"),
+        checkpoint_name(conv_state, "gdn_conv_state") if conv_state is not None else None,
+        checkpoint_name(recurrent_state, "gdn_recurrent_state") if recurrent_state is not None else None,
         checkpoint_name(t_inv, "gdn_t_inv") if t_inv is not None else None,
         checkpoint_name(chunk_states, "gdn_chunk_states") if chunk_states is not None else None,
-        checkpoint_name(m_local, "gdn_m_local") if m_local is not None else None,
-        conv_state is not None,
-        recurrent_state is not None,
-        seg_metadata if seg_metadata is not None else segment_ids,
+        seg_bundle,
     )
     return (out, states), residuals
 
-  segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(segment_ids, has_initial_state, conv_kernel_size)
-  (out, states), t_inv, chunk_states = _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,
-  )
-  out = checkpoint_name(out, "gdn_core_attn_out")
-  seg_bundle = (segment_ids, conv_halo_seg, init_seg) if segment_ids is not None else None
-  residuals = (
-      checkpoint_name(qkv, "gdn_qkv") if qkv is not None else None,
-      checkpoint_name(b, "gdn_b") if b is not None else None,
-      checkpoint_name(a, "gdn_a") if a is not None else None,
-      conv_weight,
-      conv_bias,
-      a_log,
-      dt_bias,
-      checkpoint_name(conv_state, "gdn_conv_state") if conv_state is not None else None,
-      checkpoint_name(recurrent_state, "gdn_recurrent_state") if recurrent_state is not None else None,
-      checkpoint_name(t_inv, "gdn_t_inv") if t_inv is not None else None,
-      checkpoint_name(chunk_states, "gdn_chunk_states") if chunk_states is not None else None,
-      seg_bundle,
-  )
-  return (out, states), residuals
-
 
 def _gdn_decoupled_conv1d_bwd(
     num_k_heads: int,
@@ -716,222 +726,320 @@ def _gdn_decoupled_conv1d_bwd(
     compute_dtype: jnp.dtype,
     cp_axis_name: str | tuple[str, ...] | None = None,
     cp_matmul_precision: Any = cp_gdn.DEFAULT_PRECISION,
+    conv_input_in_qkv_dtype: bool = False,
     residuals: tuple[Any, ...] | None = None,
     cotangents: tuple[Any, ...] | None = None,
 ):
   """Backward rule for custom_vjp registration of gdn_decoupled_conv1d."""
   # Support positional calls that omit the trailing nondiff arguments: 10 args (no cp_axis_name,
-  # no cp_matmul_precision) or 11 args (no cp_matmul_precision).
+  # no cp_matmul_precision, no conv_input_in_qkv_dtype), 11 args (no cp_matmul_precision, no
+  # conv_input_in_qkv_dtype) or 12 args (no conv_input_in_qkv_dtype). jax.custom_vjp itself always
+  # passes all 11 nondiff arguments.
   if residuals is None and cotangents is None and isinstance(cp_axis_name, tuple) and len(cp_axis_name) >= 10:
     residuals, cotangents = cp_axis_name, cp_matmul_precision
     cp_axis_name, cp_matmul_precision = None, cp_gdn.DEFAULT_PRECISION
-  elif cotangents is None and isinstance(cp_matmul_precision, tuple) and len(cp_matmul_precision) >= 10:
-    residuals, cotangents = cp_matmul_precision, residuals
-    cp_matmul_precision = cp_gdn.DEFAULT_PRECISION
+  elif (
+      residuals is None
+      and cotangents is None
+      and isinstance(cp_matmul_precision, tuple)
+      and len(cp_matmul_precision) >= 10
+  ):
+    residuals, cotangents = cp_matmul_precision, conv_input_in_qkv_dtype
+    cp_matmul_precision, conv_input_in_qkv_dtype = cp_gdn.DEFAULT_PRECISION, False
+  elif cotangents is None and isinstance(conv_input_in_qkv_dtype, tuple) and len(conv_input_in_qkv_dtype) >= 10:
+    residuals, cotangents = conv_input_in_qkv_dtype, residuals
+    conv_input_in_qkv_dtype = False
 
-  m_local_fwd = None
-  seg_slot = None
-  include_segment_ids_grad = False
-  if len(residuals) == 15:
-    (
-        pre_conv_qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        t_inv_fwd,
-        chunk_states,
-        m_local_fwd,
-        has_user_conv_state,
-        has_user_recurrent_state,
-        seg_slot,
-    ) = residuals
-    include_segment_ids_grad = True
-  elif len(residuals) == 14:
-    (
-        pre_conv_qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        t_inv_fwd,
-        chunk_states,
-        m_local_fwd,
-        has_user_conv_state,
-        has_user_recurrent_state,
-    ) = residuals
-  elif len(residuals) == 12:
-    (
-        pre_conv_qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        t_inv_fwd,
-        chunk_states,
-        seg_slot,
-    ) = residuals
-    has_user_conv_state = conv_state is not None
-    has_user_recurrent_state = recurrent_state is not None
-    include_segment_ids_grad = True
-  elif len(residuals) == 11:
-    (
-        pre_conv_qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        t_inv_fwd,
-        chunk_states,
-    ) = residuals
-    has_user_conv_state = conv_state is not None
-    has_user_recurrent_state = recurrent_state is not None
-  else:
-    (
-        pre_conv_qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        t_inv_fwd,
-    ) = residuals
-    chunk_states = None
-    has_user_conv_state = conv_state is not None
-    has_user_recurrent_state = recurrent_state is not None
+  with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
+    m_local_fwd = None
+    seg_slot = None
+    include_segment_ids_grad = False
+    if len(residuals) == 15:
+      (
+          pre_conv_qkv,
+          b,
+          a,
+          conv_weight,
+          conv_bias,
+          a_log,
+          dt_bias,
+          conv_state,
+          recurrent_state,
+          t_inv_fwd,
+          chunk_states,
+          m_local_fwd,
+          has_user_conv_state,
+          has_user_recurrent_state,
+          seg_slot,
+      ) = residuals
+      include_segment_ids_grad = True
+    elif len(residuals) == 14:
+      (
+          pre_conv_qkv,
+          b,
+          a,
+          conv_weight,
+          conv_bias,
+          a_log,
+          dt_bias,
+          conv_state,
+          recurrent_state,
+          t_inv_fwd,
+          chunk_states,
+          m_local_fwd,
+          has_user_conv_state,
+          has_user_recurrent_state,
+      ) = residuals
+    elif len(residuals) == 12:
+      (
+          pre_conv_qkv,
+          b,
+          a,
+          conv_weight,
+          conv_bias,
+          a_log,
+          dt_bias,
+          conv_state,
+          recurrent_state,
+          t_inv_fwd,
+          chunk_states,
+          seg_slot,
+      ) = residuals
+      has_user_conv_state = conv_state is not None
+      has_user_recurrent_state = recurrent_state is not None
+      include_segment_ids_grad = True
+    elif len(residuals) == 11:
+      (
+          pre_conv_qkv,
+          b,
+          a,
+          conv_weight,
+          conv_bias,
+          a_log,
+          dt_bias,
+          conv_state,
+          recurrent_state,
+          t_inv_fwd,
+          chunk_states,
+      ) = residuals
+      has_user_conv_state = conv_state is not None
+      has_user_recurrent_state = recurrent_state is not None
+    else:
+      (
+          pre_conv_qkv,
+          b,
+          a,
+          conv_weight,
+          conv_bias,
+          a_log,
+          dt_bias,
+          conv_state,
+          recurrent_state,
+          t_inv_fwd,
+      ) = residuals
+      chunk_states = None
+      has_user_conv_state = conv_state is not None
+      has_user_recurrent_state = recurrent_state is not None
 
-  use_cp = _is_cp_active(cp_axis_name)
-  if isinstance(seg_slot, tuple):
-    segment_ids, conv_halo_seg, init_seg = seg_slot
-  elif use_cp:
-    segment_ids, conv_halo_seg, init_seg = seg_slot, None, None
-  else:
-    segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
-        seg_slot, bool(has_user_conv_state) or bool(has_user_recurrent_state), conv_kernel_size
-    )
+    use_cp = _is_cp_active(cp_axis_name)
+    if isinstance(seg_slot, tuple):
+      segment_ids, conv_halo_seg, init_seg = seg_slot
+    elif use_cp:
+      segment_ids, conv_halo_seg, init_seg = seg_slot, None, None
+    else:
+      segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
+          seg_slot, bool(has_user_conv_state) or bool(has_user_recurrent_state), conv_kernel_size
+      )
 
-  d_out_raw, d_states = cotangents
-  d_conv_state_raw, d_recurrent_state_raw = d_states
-  d_out = _unwrap_cotangent(d_out_raw, like_zero=True)
-  # next_conv_state / next_recurrent_state depend on the inputs even without
-  # caller initial states, so their cotangents are always propagated.
-  d_conv_state = _unwrap_cotangent(d_conv_state_raw, like_zero=False)
-  d_recurrent_state = _unwrap_cotangent(d_recurrent_state_raw, like_zero=False)
-  need_dh0 = bool(has_user_recurrent_state)
+    d_out_raw, d_states = cotangents
+    d_conv_state_raw, d_recurrent_state_raw = d_states
+    d_out = _unwrap_cotangent(d_out_raw, like_zero=True)
+    # next_conv_state / next_recurrent_state depend on the inputs even without
+    # caller initial states, so their cotangents are always propagated.
+    d_conv_state = _unwrap_cotangent(d_conv_state_raw, like_zero=False)
+    d_recurrent_state = _unwrap_cotangent(d_recurrent_state_raw, like_zero=False)
+    need_dh0 = bool(has_user_recurrent_state)
 
-  if use_cp and segment_ids is not None and conv_halo_seg is None and init_seg is None:
-    segment_ids, conv_halo_seg, init_seg = cp_gdn.gather_cp_segment_metadata(
-        segment_ids,
-        cp_axis_name,
-        conv_kernel_size,
-        has_initial_state=bool(has_user_conv_state) or bool(has_user_recurrent_state),
-    )
+    if use_cp and segment_ids is not None and conv_halo_seg is None and init_seg is None:
+      segment_ids, conv_halo_seg, init_seg = cp_gdn.gather_cp_segment_metadata(
+          segment_ids,
+          cp_axis_name,
+          conv_kernel_size,
+          has_initial_state=bool(has_user_conv_state) or bool(has_user_recurrent_state),
+      )
 
-  # When chunk_states is cached in residuals and t_inv was omitted (chunk_size > 64),
-  # the Pallas backward kernel rebuilds T^-1 per chunk in VMEM using xor_dnc instead
-  # of reading it from HBM or recomputing it with XLA.
-  recompute_t_inv_in_kernel = chunk_states is not None and t_inv_fwd is None
-  conv_out_cached = None
-  if not recompute_t_inv_in_kernel and (chunk_states is None or t_inv_fwd is None):
-    qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states(
-        qkv=pre_conv_qkv,
-        b=b,
-        a=a,
-        conv_weight=conv_weight,
-        conv_bias=conv_bias,
-        a_log=a_log,
-        dt_bias=dt_bias,
-        recurrent_state=recurrent_state,
-        conv_state=conv_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,
-        cached_t_inv=t_inv_fwd,
-        segment_ids=segment_ids,
-        conv_halo_seg=conv_halo_seg,
-        init_seg=init_seg,
-    )
-    if chunk_states is None:
-      chunk_states = chunk_states_recomputed
-    if t_inv_fwd is None:
-      t_inv_fwd = t_inv_recomputed
-  else:
-    conv_out_cached, qkv_conv = conv1d_silu_fwd(
-        qkv=pre_conv_qkv,
-        conv_weight=conv_weight,
-        conv_bias=conv_bias,
-        kernel_size=conv_kernel_size,
-        conv_state=conv_state,
-        segment_ids=segment_ids,
-        conv_halo_seg=conv_halo_seg,
-    )
-  t_inv = t_inv_fwd
+    # When chunk_states is cached in residuals and t_inv was omitted (chunk_size > 64),
+    # the Pallas backward kernel rebuilds T^-1 per chunk in VMEM using xor_dnc instead
+    # of reading it from HBM or recomputing it with XLA.
+    recompute_t_inv_in_kernel = chunk_states is not None and t_inv_fwd is None
+    conv_out_cached = None
+    if not recompute_t_inv_in_kernel and (chunk_states is None or t_inv_fwd is None):
+      qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states(
+          qkv=pre_conv_qkv,
+          b=b,
+          a=a,
+          conv_weight=conv_weight,
+          conv_bias=conv_bias,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          recurrent_state=recurrent_state,
+          conv_state=conv_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,
+          cached_t_inv=t_inv_fwd,
+          segment_ids=segment_ids,
+          conv_halo_seg=conv_halo_seg,
+          init_seg=init_seg,
+      )
+      if chunk_states is None:
+        chunk_states = chunk_states_recomputed
+      if t_inv_fwd is None:
+        t_inv_fwd = t_inv_recomputed
+    else:
+      conv_out_cached, qkv_conv = conv1d_silu_fwd(
+          qkv=pre_conv_qkv,
+          conv_weight=conv_weight,
+          conv_bias=conv_bias,
+          kernel_size=conv_kernel_size,
+          conv_state=conv_state,
+          segment_ids=segment_ids,
+          conv_halo_seg=conv_halo_seg,
+      )
+    t_inv = t_inv_fwd
+
+    if use_cp:
+      # Under jax.shard_map(..., check_vma=False), _shard_map_transpose (shard_map.py:1942)
+      # divides cotangents of unmapped out_specs (next_conv_state, next_recurrent_state) by D,
+      # and (shard_map.py:1959) emits a single lax.psum across cp_axis_name for cotangents of
+      # unmapped in_specs (conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state),
+      # avoiding redundant All-Reduce collectives inside the custom_vjp rule.
+      dht_final = (
+          jax.lax.psum(d_recurrent_state.astype(jnp.float32), cp_axis_name)
+          if d_recurrent_state is not None
+          else jnp.zeros(
+              (pre_conv_qkv.shape[0], num_v_heads, head_k_dim, head_v_dim),
+              dtype=jnp.float32,
+          )
+      )
+      d_cs_ext_unscaled = (
+          jax.lax.psum(d_conv_state.astype(jnp.float32), cp_axis_name).astype(pre_conv_qkv.dtype)
+          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,
+      )
+      dht_local, _ = cp_gdn.incoming_grad_state(
+          dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision
+      )
+      bwd_out = pallas_gdn_bwd_kernel(
+          qkv_conv=qkv_conv,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          do=d_out,
+          chunk_states=chunk_states,
+          t_inv=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,
+          d_recurrent_state=dht_local,
+          return_dh0=need_dh0,
+          recompute_t_inv=recompute_t_inv_in_kernel,
+      )
+      if need_dh0:
+        dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
+      else:
+        dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out  # pylint: disable=unbalanced-tuple-unpacking
+        dh0_local = None
 
-  if use_cp:
-    # Under jax.shard_map(..., check_vma=False), _shard_map_transpose (shard_map.py:1942)
-    # divides cotangents of unmapped out_specs (next_conv_state, next_recurrent_state) by D,
-    # and (shard_map.py:1959) emits a single lax.psum across cp_axis_name for cotangents of
-    # unmapped in_specs (conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state),
-    # avoiding redundant All-Reduce collectives inside the custom_vjp rule.
-    dht_final = (
-        jax.lax.psum(d_recurrent_state.astype(jnp.float32), cp_axis_name)
-        if d_recurrent_state is not None
-        else jnp.zeros(
-            (pre_conv_qkv.shape[0], num_v_heads, head_k_dim, head_v_dim),
-            dtype=jnp.float32,
+      d_pre_conv_qkv, d_conv_weight, d_conv_bias, d_cs_local = conv1d_silu_bwd(
+          qkv=pre_conv_qkv,
+          conv_weight=conv_weight,
+          conv_bias=conv_bias,
+          dy=dy_conv,
+          kernel_size=conv_kernel_size,
+          conv_out=conv_out_cached,
+          conv_state=conv_state,
+          return_d_conv_state=True,
+          segment_ids=segment_ids,
+          conv_halo_seg=conv_halo_seg,
+      )
+      d_cs_ext_for_halo = d_cs_ext_unscaled
+      if d_cs_ext_unscaled is not None and segment_ids is not None:
+        # With packing, next_conv_state is the window ending at the last valid
+        # token of the rank that owns it (see broadcast_end_conv_state). Scatter
+        # its cotangent into that rank's tokens / halo instead of the last rank's tail.
+        cs_has_valid = jnp.any(segment_ids > 0, axis=1)
+        is_owner = cp_gdn.select_end_conv_state_rank(cs_has_valid, cp_axis_name)
+        d_cs_owned = jnp.where(is_owner[:, None, None], d_cs_ext_unscaled, jnp.zeros_like(d_cs_ext_unscaled))
+        d_pre_conv_qkv, d_cs_local = local_compute_conv1d.extract_segment_conv_state_split_adjoint(
+            d_cs_owned,
+            d_pre_conv_qkv,
+            d_cs_local,
+            segment_ids,
+            conv_kernel_size,
+            conv_halo_seg,
         )
-    )
-    d_cs_ext_unscaled = (
-        jax.lax.psum(d_conv_state.astype(jnp.float32), cp_axis_name).astype(pre_conv_qkv.dtype)
-        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,
-    )
-    dht_local, _ = cp_gdn.incoming_grad_state(
-        dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision
-    )
+        d_cs_ext_for_halo = None
+      d_pre_conv_qkv, _ = cp_gdn.halo_exchange_for_conv_bwd(
+          dx=d_pre_conv_qkv,
+          d_conv_state=d_cs_local,
+          d_conv_state_ext=d_cs_ext_for_halo,
+          kernel_size=conv_kernel_size,
+          cp_axis=cp_axis_name,
+          segment_ids=segment_ids,
+      )
+
+      idx = jax.lax.axis_index(cp_axis_name)
+      d_init_cs_rank0 = jnp.where(idx == 0, d_cs_local, jnp.zeros_like(d_cs_local))
+      d_conv_state_out = d_init_cs_rank0 if has_user_conv_state else None
+      if need_dh0 and dh0_local is not None:
+        d_recurrent_state_out = jnp.where(idx == 0, dh0_local, jnp.zeros_like(dh0_local))
+      else:
+        d_recurrent_state_out = None
+      out_grads = (
+          d_pre_conv_qkv,
+          d_b,
+          d_a,
+          d_conv_weight,
+          d_conv_bias,
+          d_a_log,
+          d_dt_bias,
+          d_conv_state_out,
+          d_recurrent_state_out,
+      )
+      if include_segment_ids_grad:
+        return out_grads + (None,)
+      return out_grads
+
     bwd_out = pallas_gdn_bwd_kernel(
         qkv_conv=qkv_conv,
         b=b,
@@ -948,14 +1056,15 @@ def _gdn_decoupled_conv1d_bwd(
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
         segment_ids=segment_ids,
         init_seg=init_seg,
-        d_recurrent_state=dht_local,
+        d_recurrent_state=d_recurrent_state,
         return_dh0=need_dh0,
         recompute_t_inv=recompute_t_inv_in_kernel,
     )
+    # pylint: disable=unbalanced-tuple-unpacking
     if need_dh0:
       dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
     else:
-      dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out  # pylint: disable=unbalanced-tuple-unpacking
+      dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out
       dh0_local = None
 
     d_pre_conv_qkv, d_conv_weight, d_conv_bias, d_cs_local = conv1d_silu_bwd(
@@ -970,39 +1079,21 @@ def _gdn_decoupled_conv1d_bwd(
         segment_ids=segment_ids,
         conv_halo_seg=conv_halo_seg,
     )
-    d_cs_ext_for_halo = d_cs_ext_unscaled
-    if d_cs_ext_unscaled is not None and segment_ids is not None:
-      # With packing, next_conv_state is the window ending at the last valid
-      # token of the rank that owns it (see broadcast_end_conv_state). Scatter
-      # its cotangent into that rank's tokens / halo instead of the last rank's tail.
-      cs_has_valid = jnp.any(segment_ids > 0, axis=1)
-      is_owner = cp_gdn.select_end_conv_state_rank(cs_has_valid, cp_axis_name)
-      d_cs_owned = jnp.where(is_owner[:, None, None], d_cs_ext_unscaled, jnp.zeros_like(d_cs_ext_unscaled))
+    # pylint: enable=unbalanced-tuple-unpacking
+    if d_conv_state is not None and segment_ids is not None:
       d_pre_conv_qkv, d_cs_local = local_compute_conv1d.extract_segment_conv_state_split_adjoint(
-          d_cs_owned,
+          d_conv_state,
           d_pre_conv_qkv,
-          d_cs_local,
+          d_cs_local if has_user_conv_state else None,
           segment_ids,
           conv_kernel_size,
           conv_halo_seg,
       )
-      d_cs_ext_for_halo = None
-    d_pre_conv_qkv, _ = cp_gdn.halo_exchange_for_conv_bwd(
-        dx=d_pre_conv_qkv,
-        d_conv_state=d_cs_local,
-        d_conv_state_ext=d_cs_ext_for_halo,
-        kernel_size=conv_kernel_size,
-        cp_axis=cp_axis_name,
-        segment_ids=segment_ids,
-    )
+    elif d_conv_state is not None and pre_conv_qkv.shape[1] >= conv_kernel_size - 1:
+      d_pre_conv_qkv = d_pre_conv_qkv.at[:, -(conv_kernel_size - 1) :, :].add(d_conv_state.astype(d_pre_conv_qkv.dtype))
 
-    idx = jax.lax.axis_index(cp_axis_name)
-    d_init_cs_rank0 = jnp.where(idx == 0, d_cs_local, jnp.zeros_like(d_cs_local))
-    d_conv_state_out = d_init_cs_rank0 if has_user_conv_state else None
-    if need_dh0 and dh0_local is not None:
-      d_recurrent_state_out = jnp.where(idx == 0, dh0_local, jnp.zeros_like(dh0_local))
-    else:
-      d_recurrent_state_out = None
+    d_conv_state_out = d_cs_local if has_user_conv_state else None
+    d_recurrent_state_out = dh0_local if need_dh0 else None
     out_grads = (
         d_pre_conv_qkv,
         d_b,
@@ -1018,75 +1109,6 @@ def _gdn_decoupled_conv1d_bwd(
       return out_grads + (None,)
     return out_grads
 
-  bwd_out = pallas_gdn_bwd_kernel(
-      qkv_conv=qkv_conv,
-      b=b,
-      a=a,
-      a_log=a_log,
-      dt_bias=dt_bias,
-      do=d_out,
-      chunk_states=chunk_states,
-      t_inv=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,
-      d_recurrent_state=d_recurrent_state,
-      return_dh0=need_dh0,
-      recompute_t_inv=recompute_t_inv_in_kernel,
-  )
-  # pylint: disable=unbalanced-tuple-unpacking
-  if need_dh0:
-    dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
-  else:
-    dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out
-    dh0_local = None
-
-  d_pre_conv_qkv, d_conv_weight, d_conv_bias, d_cs_local = conv1d_silu_bwd(
-      qkv=pre_conv_qkv,
-      conv_weight=conv_weight,
-      conv_bias=conv_bias,
-      dy=dy_conv,
-      kernel_size=conv_kernel_size,
-      conv_out=conv_out_cached,
-      conv_state=conv_state,
-      return_d_conv_state=True,
-      segment_ids=segment_ids,
-      conv_halo_seg=conv_halo_seg,
-  )
-  # pylint: enable=unbalanced-tuple-unpacking
-  if d_conv_state is not None and segment_ids is not None:
-    d_pre_conv_qkv, d_cs_local = local_compute_conv1d.extract_segment_conv_state_split_adjoint(
-        d_conv_state,
-        d_pre_conv_qkv,
-        d_cs_local if has_user_conv_state else None,
-        segment_ids,
-        conv_kernel_size,
-        conv_halo_seg,
-    )
-  elif d_conv_state is not None and pre_conv_qkv.shape[1] >= conv_kernel_size - 1:
-    d_pre_conv_qkv = d_pre_conv_qkv.at[:, -(conv_kernel_size - 1) :, :].add(d_conv_state.astype(d_pre_conv_qkv.dtype))
-
-  d_conv_state_out = d_cs_local if has_user_conv_state else None
-  d_recurrent_state_out = dh0_local if need_dh0 else None
-  out_grads = (
-      d_pre_conv_qkv,
-      d_b,
-      d_a,
-      d_conv_weight,
-      d_conv_bias,
-      d_a_log,
-      d_dt_bias,
-      d_conv_state_out,
-      d_recurrent_state_out,
-  )
-  if include_segment_ids_grad:
-    return out_grads + (None,)
-  return out_grads
-
 
 gdn_decoupled_conv1d.defvjp(
     _gdn_decoupled_conv1d_fwd,
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py
index e30fc6bf28..7628934cbe 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py
@@ -14,11 +14,40 @@
 
 """Conv1D + SiLU forward and unrolled shift-multiply backward primitives."""
 
-from typing import Optional, Tuple
+import contextlib
+import contextvars
+from typing import Iterator, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
 
+# gdn_conv1d_input_in_qkv_dtype (see the comment above _conv1d_silu_fwd_narrow_input at the end of this file).
+# conv1d_silu_fwd / conv1d_silu_bwd take it as an argument; left at None they read this trace-time setting, which
+# gdn_decoupled_conv1d (gdn_bwd/api.py) scopes around its primal, forward and backward rules from its own static
+# argument. A scoped setting rather than an argument on every internal helper because the backward rule is traced
+# when the cotangent is pulled back, after the model's forward has returned, so only the custom_vjp's static
+# arguments reach it; and the rules reach these functions through several helpers (jax_compute_gdn_states, the CP
+# forward) that do not otherwise need it.
+_CONV_INPUT_IN_QKV_DTYPE: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "gdn_conv1d_input_in_qkv_dtype", default=False
+)
+
+
+@contextlib.contextmanager
+def conv_input_in_qkv_dtype_scope(enabled: bool) -> Iterator[None]:
+  """Scopes the default of conv1d_silu_fwd / conv1d_silu_bwd's `conv_input_in_qkv_dtype` while tracing."""
+  token = _CONV_INPUT_IN_QKV_DTYPE.set(bool(enabled))
+  try:
+    yield
+  finally:
+    _CONV_INPUT_IN_QKV_DTYPE.reset(token)
+
+
+def _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype: Optional[bool]) -> bool:
+  if conv_input_in_qkv_dtype is None:
+    return _CONV_INPUT_IN_QKV_DTYPE.get()
+  return bool(conv_input_in_qkv_dtype)
+
 
 def _build_conv_tap_masks(
     segment_ids: jax.Array,
@@ -53,8 +82,17 @@ def conv1d_silu_fwd(
     segment_ids: Optional[jax.Array] = None,
     conv_halo_seg: Optional[jax.Array] = None,
     output_dtype: Optional[jnp.dtype] = None,
+    conv_input_in_qkv_dtype: Optional[bool] = None,
 ) -> Tuple[jax.Array, jax.Array]:
-  """Forward Conv1D + SiLU returning (conv_out, qkv_conv)."""
+  """Forward Conv1D + SiLU returning (conv_out, qkv_conv).
+
+  `conv_input_in_qkv_dtype` (None: the conv_input_in_qkv_dtype_scope() setting, default False) keeps the halo-extended input
+  in qkv's dtype and casts each tap to float32 where it is read. Same values; see _conv1d_silu_fwd_narrow_input.
+  """
+  if _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype):
+    return _conv1d_silu_fwd_narrow_input(
+        qkv, conv_weight, conv_bias, kernel_size, conv_state, segment_ids, conv_halo_seg, output_dtype
+    )
   _, seq_len, _ = qkv.shape
   if conv_weight.ndim == 3:
     conv_weight_3d = conv_weight.astype(jnp.float32)
@@ -92,8 +130,18 @@ def conv1d_silu_bwd(
     return_d_conv_state: bool = False,
     segment_ids: Optional[jax.Array] = None,
     conv_halo_seg: Optional[jax.Array] = None,
+    conv_input_in_qkv_dtype: Optional[bool] = None,
 ):
-  """Dedicated Conv1D + SiLU backward pass using JAX primitives."""
+  """Dedicated Conv1D + SiLU backward pass using JAX primitives.
+
+  `conv_input_in_qkv_dtype` (None: the conv_input_in_qkv_dtype_scope() setting, default False) keeps the halo-extended input
+  in qkv's dtype and casts each tap to float32 where it is read. Same values; see _conv1d_silu_bwd_narrow_input.
+  """
+  if _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype):
+    return _conv1d_silu_bwd_narrow_input(
+        qkv, conv_weight, conv_bias, dy, kernel_size, conv_out, conv_state, return_d_conv_state, segment_ids,
+        conv_halo_seg
+    )
   _, seq_len, _ = qkv.shape
   if conv_weight.ndim == 3:
     conv_weight_3d = conv_weight.astype(jnp.float32)
@@ -179,3 +227,128 @@ def conv1d_silu_bwd(
     return dx, dw, db, d_conv_state
 
   return dx, dw, db
+
+
+# --------------------------------------------------------------------------------------------------------------------
+# gdn_conv1d_input_in_qkv_dtype: the halo-extended conv input in qkv's dtype.
+#
+# The default path above builds conv_input = [conv_state | qkv] (or qkv left-padded by kernel_size - 1 zeros) in
+# float32 and slices the kernel_size taps out of it. That materializes a float32 copy of the whole input: on the
+# Qwen3.5-397B RL trainer (bf16 qkv, sequence-sharded context parallelism) one f32[1, 16387, 12288] = 805.5 MB buffer
+# per GDN layer in the backward pass, written once and read kernel_size times.
+#
+# The functions below build conv_input in qkv's dtype (bf16 for bf16 qkv and conv_state; the promoted dtype if they
+# differ) and cast each tap slice to float32 where it is read, so XLA fuses the cast into the consumer and the
+# float32 copy is never written. Every product and sum is formed in the same order, on the same float32 values, as in
+# the default path (bf16 -> float32 is exact), so op by op every output is bit-identical; only which tensors are
+# materialized in HBM changes. Compiled, the forward outputs and dx / d_conv_state / db stay bit-identical, while the
+# weight gradient dw, a float32 reduction over batch x sequence, can come out in a different summation order when XLA
+# fuses the tap's cast into the reduction (CPU: a few ulp, only when the pre-activation is recomputed without packing
+# or a conv state). No jax.lax.optimization_barrier is used: the narrow buffer alone is what removes the copy.
+# Tests: tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py.
+
+
+def _conv_input_narrow(qkv, conv_state, kernel_size):
+  """Halo-extended conv input in the narrowest exact dtype (bf16 for bf16 qkv and conv_state)."""
+  if conv_state is not None:
+    cdt = jnp.promote_types(qkv.dtype, conv_state.dtype)
+    return jnp.concatenate([conv_state.astype(cdt), qkv.astype(cdt)], axis=1)
+  return jnp.pad(qkv, ((0, 0), (kernel_size - 1, 0), (0, 0)))
+
+
+def _conv_weight_3d_f32(conv_weight):
+  return conv_weight.astype(jnp.float32) if conv_weight.ndim == 3 else conv_weight[:, None, :].astype(jnp.float32)
+
+
+def _preact_from_narrow_input(conv_input, w3d, conv_bias, kernel_size, seq_len, tap_masks, valid_mask):
+  """f32 pre-activation z with the default path's op order: sum over taps k of (x_k * mask_k) * w_k, + b, * valid."""
+  if tap_masks is not None:
+    z = sum(
+        (conv_input[:, k : k + seq_len, :].astype(jnp.float32) * tap_masks[k]) * w3d[k, 0, :] for k in range(kernel_size)
+    )
+    if conv_bias is not None:
+      z = z + conv_bias.astype(jnp.float32)
+    return z * valid_mask
+  z = sum(conv_input[:, k : k + seq_len, :].astype(jnp.float32) * w3d[k, 0, :] for k in range(kernel_size))
+  if conv_bias is not None:
+    z = z + conv_bias.astype(jnp.float32)
+  return z
+
+
+def _conv1d_silu_fwd_narrow_input(
+    qkv, conv_weight, conv_bias, kernel_size, conv_state, segment_ids, conv_halo_seg, output_dtype
+):
+  """conv1d_silu_fwd reading the conv input in qkv's dtype. Returns (f32 conv_out, qkv_conv) like the default path."""
+  _, seq_len, _ = qkv.shape
+  w3d = _conv_weight_3d_f32(conv_weight)
+  conv_input = _conv_input_narrow(qkv, conv_state, kernel_size)
+  valid_mask = tap_masks = None
+  if segment_ids is not None:
+    valid_mask, tap_masks = _build_conv_tap_masks(segment_ids, kernel_size, conv_halo_seg)
+  conv_out = _preact_from_narrow_input(conv_input, w3d, conv_bias, kernel_size, seq_len, tap_masks, valid_mask)
+  qkv_conv = jax.nn.silu(conv_out)
+  return conv_out, qkv_conv.astype(output_dtype or qkv.dtype)
+
+
+def _conv1d_silu_bwd_narrow_input(
+    qkv, conv_weight, conv_bias, dy, kernel_size, conv_out, conv_state, return_d_conv_state, segment_ids, conv_halo_seg
+):
+  """conv1d_silu_bwd reading the conv input in qkv's dtype; the same steps 1-4 as the default path."""
+  _, seq_len, _ = qkv.shape
+  w3d = _conv_weight_3d_f32(conv_weight)
+  conv_input = _conv_input_narrow(qkv, conv_state, kernel_size)
+  valid_mask = tap_masks = None
+  if segment_ids is not None:
+    valid_mask, tap_masks = _build_conv_tap_masks(segment_ids, kernel_size, conv_halo_seg)
+  if conv_out is None:
+    conv_out = _preact_from_narrow_input(conv_input, w3d, conv_bias, kernel_size, seq_len, tap_masks, valid_mask)
+  z = conv_out.astype(jnp.float32)
+
+  # dz = dy * SiLU'(z), as in the default path.
+  sig_z = jax.nn.sigmoid(z)
+  silu_prime = sig_z * (1.0 + z * (1.0 - sig_z))
+  dz = dy.astype(jnp.float32) * silu_prime
+  if valid_mask is not None:
+    dz = dz * valid_mask
+
+  if conv_bias is not None:
+    db = jnp.sum(dz, axis=(0, 1)).astype(conv_bias.dtype)
+    if conv_bias.ndim != 1:
+      db = db.reshape(conv_bias.shape)
+  else:
+    db = None
+
+  # dw[k] = sum_{b,t} dz[b,t] * conv_input[b, t+k], each tap cast to float32 where it is read.
+  dw_rows = []
+  for k in range(kernel_size):
+    x_k = conv_input[:, k : k + seq_len, :].astype(jnp.float32)
+    dz_k = dz * tap_masks[k] if tap_masks is not None else dz
+    dw_rows.append(jnp.sum(dz_k * x_k, axis=(0, 1)))
+  dw = jnp.stack(dw_rows, axis=0)
+  dw = dw[:, None, :].astype(conv_weight.dtype) if conv_weight.ndim == 3 else dw.astype(conv_weight.dtype)
+
+  # dx: transposed convolution of dz with the reversed weights (does not read conv_input).
+  dz_pad = jnp.pad(dz, ((0, 0), (0, kernel_size - 1), (0, 0)))
+  w_rev = w3d[::-1]
+  if tap_masks is not None:
+    tap_masks_pad = [jnp.pad(tm, ((0, 0), (0, kernel_size - 1), (0, 0))) for tm in tap_masks]
+    dx = sum(
+        (dz_pad[:, j : j + seq_len, :] * tap_masks_pad[kernel_size - 1 - j][:, j : j + seq_len, :]) * w_rev[j, 0, :]
+        for j in range(kernel_size)
+    )
+  else:
+    dx = sum(dz_pad[:, k : k + seq_len, :] * w_rev[k, 0, :] for k in range(kernel_size))
+  dx = dx.astype(qkv.dtype)
+
+  if return_d_conv_state:
+    d_cs_list = []
+    for m in range(kernel_size - 1):
+      val = jnp.zeros((qkv.shape[0], qkv.shape[2]), dtype=jnp.float32)
+      for t in range(m + 1):
+        k_idx = m - t
+        dz_t = dz[:, t, :] * tap_masks[k_idx][:, t, :] if tap_masks is not None else dz[:, t, :]
+        val = val + dz_t * w3d[k_idx, 0, :]
+      d_cs_list.append(val)
+    d_conv_state = jnp.stack(d_cs_list, axis=1).astype(qkv.dtype)
+    return dx, dw, db, d_conv_state
+  return dx, dw, db
diff --git a/src/maxtext/kernels/gdn/model_runner.py b/src/maxtext/kernels/gdn/model_runner.py
index a1cb7e0b88..24ddaca95f 100644
--- a/src/maxtext/kernels/gdn/model_runner.py
+++ b/src/maxtext/kernels/gdn/model_runner.py
@@ -617,6 +617,8 @@ def run_gdn_kernel_layer(
       cfg, cp_size, layer.num_k_heads, model_mode, use_gdn_kernel=True
   )
   cp_matmul_precision = jax.lax.Precision[getattr(cfg, "gdn_cp_matmul_precision", "highest").upper()]
+  # Conv1D halo input in qkv's dtype instead of a float32 copy (compute_conv1d_bwd.py). Same values.
+  conv_input_in_qkv_dtype = bool(getattr(cfg, "gdn_conv1d_input_in_qkv_dtype", False))
 
   if use_head_sharded_cp:
     qkv = jnp.concatenate([query, key, value_raw], axis=3)
@@ -887,6 +889,7 @@ def shard_mapped_gdn(  # pylint: disable=too-many-positional-arguments
             use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
             compute_dtype=state_dtype,
             segment_ids=seg_val,
+            conv_input_in_qkv_dtype=conv_input_in_qkv_dtype,
         )
 
         if next_cs is not None:
@@ -926,6 +929,7 @@ def shard_mapped_gdn(  # pylint: disable=too-many-positional-arguments
           cp_axis_name=cp_axis_name if use_seq_sharded_cp else None,
           segment_ids=seg_val,
           cp_matmul_precision=cp_matmul_precision,
+          conv_input_in_qkv_dtype=conv_input_in_qkv_dtype,
       )
 
     if not isinstance(qkv, jax.core.Tracer):
@@ -968,6 +972,7 @@ def shard_mapped_gdn(  # pylint: disable=too-many-positional-arguments
         use_qk_norm_in_gdn=cfg.use_qk_norm_in_gdn,
         compute_dtype=state_dtype,
         segment_ids=decoder_segment_ids,
+        conv_input_in_qkv_dtype=conv_input_in_qkv_dtype,
     )
 
   core_attn_out = checkpoint_name(core_attn_out, "gdn_core_attn_out")
diff --git a/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py b/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py
new file mode 100644
index 0000000000..79e59acfbb
--- /dev/null
+++ b/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py
@@ -0,0 +1,289 @@
+# 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.
+
+"""gdn_conv1d_input_in_qkv_dtype: the GDN Conv1D + SiLU with its halo-extended input in qkv's dtype.
+
+The option computes the same float32 values in the same order (it only changes which tensors are materialized), so
+op by op every output is bit-identical. Compiled, the forward outputs and the input / conv-state gradients stay
+bit-identical, but the conv weight gradient is a float32 reduction over batch x sequence whose summation order XLA
+may change when it fuses the tap's bf16 -> float32 cast into it (measured on CPU: a few ulp, in the configuration
+that recomputes the pre-activation without packing or a conv state). The option must also actually change the
+traced program; a test that only compared values would pass if the switch were ignored, so the value checks are
+paired with checks that the two programs differ.
+"""
+
+# pylint: disable=protected-access
+
+import types
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+import jax.numpy as jnp
+from jax.extend import core as jax_core
+import numpy as np
+
+from maxtext.configs import types as config_types
+from maxtext.kernels.gdn import gdn_bwd_pallas
+from maxtext.kernels.gdn.gdn_bwd import compute_conv1d_bwd
+
+_KERNEL_SIZE = 4
+
+
+def _bits(x):
+  """The array's bytes, so equality means bit-identical (NaN-safe, -0.0 != 0.0)."""
+  return np.asarray(jax.device_get(x)).tobytes()
+
+
+def _conv_inputs(seed, *, batch=2, seq_len=24, dim=16, with_state, with_bias, with_segments, dtype=jnp.bfloat16):
+  keys = jax.random.split(jax.random.PRNGKey(seed), 6)
+  qkv = jax.random.normal(keys[0], (batch, seq_len, dim), jnp.float32).astype(dtype)
+  conv_weight = jax.random.normal(keys[1], (_KERNEL_SIZE, 1, dim), jnp.float32)
+  conv_bias = jax.random.normal(keys[2], (dim,), jnp.float32) if with_bias else None
+  conv_state = (
+      jax.random.normal(keys[3], (batch, _KERNEL_SIZE - 1, dim), jnp.float32).astype(dtype) if with_state else None
+  )
+  dy = jax.random.normal(keys[4], (batch, seq_len, dim), jnp.float32)
+  segment_ids = conv_halo_seg = None
+  if with_segments:
+    # Two packed documents and trailing padding (segment 0) per row.
+    row = np.concatenate([np.full(9, 1), np.full(11, 2), np.zeros(seq_len - 20)]).astype(np.int32)
+    segment_ids = jnp.asarray(np.stack([row] * batch))
+    if with_state:
+      conv_halo_seg = jnp.ones((batch, _KERNEL_SIZE - 1), jnp.int32)
+  return qkv, conv_weight, conv_bias, conv_state, dy, segment_ids, conv_halo_seg
+
+
+def _pad_or_concat_out_dtypes(jaxpr, shape):
+  """Dtypes of every pad / concatenate in `jaxpr` (sub-jaxprs included) whose output has `shape`."""
+  found = []
+
+  def visit(jx):
+    for eqn in jx.eqns:
+      if eqn.primitive.name in ("pad", "concatenate"):
+        aval = eqn.outvars[0].aval
+        if tuple(aval.shape) == tuple(shape):
+          found.append(aval.dtype)
+      for sub in jax_core.jaxprs_in_params(eqn.params):
+        visit(sub)
+
+  visit(jaxpr.jaxpr)
+  return found
+
+
+_CASES = [
+    dict(testcase_name=f"state{s}_bias{b}_seg{g}", with_state=s, with_bias=b, with_segments=g)
+    for s in (False, True)
+    for b in (False, True)
+    for g in (False, True)
+]
+
+
+class GdnConv1dInputInQkvDtypeTest(parameterized.TestCase):
+
+  @parameterized.named_parameters(*_CASES)
+  def test_forward_bit_identical(self, with_state, with_bias, with_segments):
+    qkv, w, bias, state, _, seg, halo = _conv_inputs(
+        1, with_state=with_state, with_bias=with_bias, with_segments=with_segments
+    )
+
+    def fwd(narrow):
+      return compute_conv1d_bwd.conv1d_silu_fwd(
+          qkv, w, bias, _KERNEL_SIZE, conv_state=state, segment_ids=seg, conv_halo_seg=halo,
+          conv_input_in_qkv_dtype=narrow,
+      )
+
+    ref_out, ref_act = jax.jit(lambda: fwd(False))()
+    out, act = jax.jit(lambda: fwd(True))()
+    self.assertEqual(ref_out.dtype, out.dtype)
+    self.assertEqual(ref_act.dtype, act.dtype)
+    self.assertEqual(_bits(ref_out), _bits(out))
+    self.assertEqual(_bits(ref_act), _bits(act))
+
+  @parameterized.named_parameters(*_CASES)
+  def test_backward_bit_identical(self, with_state, with_bias, with_segments):
+    qkv, w, bias, state, dy, seg, halo = _conv_inputs(
+        2, with_state=with_state, with_bias=with_bias, with_segments=with_segments
+    )
+    conv_out, _ = compute_conv1d_bwd.conv1d_silu_fwd(
+        qkv, w, bias, _KERNEL_SIZE, conv_state=state, segment_ids=seg, conv_halo_seg=halo
+    )
+    for cached_conv_out in (None, conv_out):
+      for return_d_conv_state in (False, True):
+
+        def bwd(narrow, cached_conv_out=cached_conv_out, return_d_conv_state=return_d_conv_state):
+          return compute_conv1d_bwd.conv1d_silu_bwd(
+              qkv, w, bias, dy, _KERNEL_SIZE, conv_out=cached_conv_out, conv_state=state,
+              return_d_conv_state=return_d_conv_state, segment_ids=seg, conv_halo_seg=halo,
+              conv_input_in_qkv_dtype=narrow,
+          )
+
+        # Op by op: the same values in the same order, so every output is bit-identical.
+        ref, got = bwd(False), bwd(True)
+        self.assertEqual(len(ref), len(got))
+        for r, g in zip(ref, got):
+          if r is None:
+            self.assertIsNone(g)
+            continue
+          self.assertEqual(r.dtype, g.dtype)
+          self.assertEqual(_bits(r), _bits(g))
+        # Compiled: dx, db and d_conv_state stay bit-identical; dw (index 1) is a float32 reduction whose order XLA
+        # may change with the fusion, so it is compared to float32 summation-order tolerance.
+        ref = jax.jit(lambda bwd=bwd: bwd(False))()
+        got = jax.jit(lambda bwd=bwd: bwd(True))()
+        for i, (r, g) in enumerate(zip(ref, got)):
+          if r is None:
+            self.assertIsNone(g)
+            continue
+          self.assertEqual(r.dtype, g.dtype)
+          if i == 1:
+            np.testing.assert_allclose(np.asarray(g), np.asarray(r), rtol=2e-5, atol=1e-6)
+          else:
+            self.assertEqual(_bits(r), _bits(g))
+
+  @parameterized.named_parameters(
+      dict(testcase_name="no_state", with_state=False),
+      dict(testcase_name="state", with_state=True),
+  )
+  def test_conv_input_is_narrow_only_when_enabled(self, with_state):
+    """Positive control: the default path builds the [B, S + K - 1, D] input in float32, the option in bf16."""
+    qkv, w, bias, state, dy, _, _ = _conv_inputs(3, with_state=with_state, with_bias=True, with_segments=False)
+    shape = (qkv.shape[0], qkv.shape[1] + _KERNEL_SIZE - 1, qkv.shape[2])
+    for fn in ("fwd", "bwd"):
+
+      def run(narrow, fn=fn):
+        if fn == "fwd":
+          return compute_conv1d_bwd.conv1d_silu_fwd(
+              qkv, w, bias, _KERNEL_SIZE, conv_state=state, conv_input_in_qkv_dtype=narrow
+          )
+        return compute_conv1d_bwd.conv1d_silu_bwd(
+            qkv, w, bias, dy, _KERNEL_SIZE, conv_state=state, conv_input_in_qkv_dtype=narrow
+        )
+
+      off = _pad_or_concat_out_dtypes(jax.make_jaxpr(lambda run=run: run(False))(), shape)
+      on = _pad_or_concat_out_dtypes(jax.make_jaxpr(lambda run=run: run(True))(), shape)
+      # The backward also pads dz (float32) to the same shape; only the conv input's buffer changes dtype.
+      self.assertNotIn(jnp.dtype(jnp.bfloat16), off, fn)
+      self.assertEqual(on.count(jnp.dtype(jnp.bfloat16)), 1, fn)
+      self.assertEqual(off.count(jnp.dtype(jnp.float32)), on.count(jnp.dtype(jnp.float32)) + 1, fn)
+
+  def test_scope_sets_the_default(self):
+    """conv_input_in_qkv_dtype=None follows conv_input_in_qkv_dtype_scope(); the default is off."""
+    qkv, w, bias, state, _, _, _ = _conv_inputs(4, with_state=True, with_bias=True, with_segments=False)
+
+    def jaxpr_text(narrow):
+      return str(
+          jax.make_jaxpr(
+              lambda: compute_conv1d_bwd.conv1d_silu_fwd(
+                  qkv, w, bias, _KERNEL_SIZE, conv_state=state, conv_input_in_qkv_dtype=narrow
+              )
+          )()
+      )
+
+    off, on = jaxpr_text(False), jaxpr_text(True)
+    self.assertNotEqual(off, on)
+    self.assertEqual(jaxpr_text(None), off)
+    with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(True):
+      self.assertEqual(jaxpr_text(None), on)
+      self.assertEqual(jaxpr_text(False), off)  # an explicit argument wins over the scope
+      with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(False):
+        self.assertEqual(jaxpr_text(None), off)
+      self.assertEqual(jaxpr_text(None), on)
+    self.assertEqual(jaxpr_text(None), off)
+
+  def test_gdn_decoupled_conv1d_grads_bit_identical(self):
+    """Through the custom VJP: the same loss and gradients, and the option reaches the backward rule's trace."""
+    gdn_bwd_pallas.ensure_cpu_interpret_registered()
+    is_tpu = jax.default_backend() == "tpu"
+    batch, chunk_size, num_chunks = 1, (64 if is_tpu else 32), 2
+    seq_len = chunk_size * num_chunks
+    num_k_heads, num_v_heads = 1, 2
+    head_k_dim = head_v_dim = 128 if is_tpu else 64
+    dim = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim
+    keys = jax.random.split(jax.random.PRNGKey(5), 8)
+    qkv = jax.random.normal(keys[0], (batch, seq_len, dim), jnp.float32).astype(jnp.bfloat16)
+    b = jax.random.normal(keys[1], (batch, seq_len, num_v_heads), jnp.float32)
+    a = jax.random.normal(keys[2], (batch, seq_len, num_v_heads), jnp.float32)
+    conv_weight = jax.random.normal(keys[3], (_KERNEL_SIZE, 1, dim), jnp.float32)
+    conv_bias = jax.random.normal(keys[4], (dim,), jnp.float32)
+    a_log = jax.random.normal(keys[5], (num_v_heads,), jnp.float32)
+    dt_bias = jax.random.normal(keys[6], (num_v_heads,), jnp.float32)
+    do = jax.random.normal(keys[7], (batch, seq_len, num_v_heads, head_v_dim), jnp.float32)
+    args = (qkv, b, a, conv_weight, conv_bias, a_log, dt_bias)
+
+    def loss(narrow, qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in):
+      out, _ = gdn_bwd_pallas.gdn_decoupled_conv1d(
+          qkv=qkv_in,
+          b=b_in,
+          a=a_in,
+          conv_weight=cw_in,
+          conv_bias=cb_in,
+          a_log=al_in,
+          dt_bias=dt_in,
+          conv_state=None,
+          recurrent_state=None,
+          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=_KERNEL_SIZE,
+          chunk_size=chunk_size,
+          use_qk_norm_in_gdn=True,
+          compute_dtype=jnp.float32,
+          conv_input_in_qkv_dtype=narrow,
+      )
+      return jnp.sum(out.astype(jnp.float32) * do)
+
+    ref_val, ref_grads = jax.value_and_grad(lambda *x: loss(False, *x), argnums=tuple(range(7)))(*args)
+    val, grads = jax.value_and_grad(lambda *x: loss(True, *x), argnums=tuple(range(7)))(*args)
+    self.assertEqual(_bits(ref_val), _bits(val))
+    for i, (r, g) in enumerate(zip(ref_grads, grads)):
+      self.assertEqual(r.dtype, g.dtype)
+      if i == 3:  # conv_weight: a float32 reduction over batch x sequence (see the module docstring)
+        np.testing.assert_allclose(np.asarray(g), np.asarray(r), rtol=2e-5, atol=1e-6)
+      else:
+        self.assertEqual(_bits(r), _bits(g))
+
+    # The local forward convolves a float32 copy of qkv either way, so a bf16 [B, S + K - 1, D] conv input can only
+    # come from the backward rule: present with the option, absent without it.
+    shape = (batch, seq_len + _KERNEL_SIZE - 1, dim)
+    grad_off = jax.make_jaxpr(jax.grad(lambda *x: loss(False, *x), argnums=tuple(range(7))))(*args)
+    grad_on = jax.make_jaxpr(jax.grad(lambda *x: loss(True, *x), argnums=tuple(range(7))))(*args)
+    self.assertNotIn(jnp.dtype(jnp.bfloat16), _pad_or_concat_out_dtypes(grad_off, shape))
+    self.assertIn(jnp.dtype(jnp.bfloat16), _pad_or_concat_out_dtypes(grad_on, shape))
+
+  def test_requires_gdn_kernel(self):
+    def cfg(**overrides):
+      fields = dict(
+          decoder_block=config_types.DecoderBlockType.QWEN3_5,
+          ici_context_parallelism=1,
+          dcn_context_parallelism=1,
+          ici_context_usp_ulysses_parallelism=1,
+          dcn_context_usp_ulysses_parallelism=1,
+          use_gdn_kernel=False,
+          gdn_conv1d_input_in_qkv_dtype=False,
+      )
+      fields.update(overrides)
+      return types.SimpleNamespace(**fields)
+
+    config_types.MaxTextConfig.validate_gdn_config(cfg(), is_rl_config=False)
+    with self.assertRaisesRegex(ValueError, "gdn_conv1d_input_in_qkv_dtype=True requires use_gdn_kernel=True"):
+      config_types.MaxTextConfig.validate_gdn_config(cfg(gdn_conv1d_input_in_qkv_dtype=True), is_rl_config=False)
+
+  def test_config_default_is_off(self):
+    self.assertIs(config_types.Qwen3Next.model_fields["gdn_conv1d_input_in_qkv_dtype"].default, False)
+
+
+if __name__ == "__main__":
+  absltest.main()

From 8cc954eb2d83da3f53260dc7d76d9b6ff92a0f5f Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 09:10:18 -0700
Subject: [PATCH 09/10] Address review feedback on the MoE and GDN keys

Follow-up to the six commits of this PR (moe_routing_index_math,
moe_unsplit_prefused_wi / moe_ragged_sort_unique_pair_argsort /
moe_prefused_swiglu_vjp, moe_roe_packed_routing_gather,
moe_wi_gmm_fused_swiglu, moe_router_grad_from_wo_input,
gdn_conv1d_input_in_qkv_dtype). No new keys and no behaviour change.

Refactors:
- gmm_v2 kernel: inner_kernel's 2D store is now the helper
  _store_group_rows, which the activation output and the return_pre_act
  gate and up halves all call, instead of the halves using a copy of
  the store; the 3D stores share its row mask and partial-row carry
  (_mask_group_rows, _carry_partial_rows).
- GDN api.py: gdn_decoupled_conv1d and its forward and backward rules keep
  their bodies at the original indentation. A small decorator enters
  conv_input_in_qkv_dtype_scope from the static argument for the primal
  and forward rule; the backward rule runs its argument shim, then its
  body (_gdn_decoupled_conv1d_bwd_impl) inside the scope. The api.py diff
  against the parent shrinks from 848 lines to the real changes.
- moe.py: _use_prefused_swiglu_vjp inlined (it only returned the config
  flag); two pointer comments dropped.
- types.py: the MoE validators' requirement helpers come before the
  validators.

Validators and guards:
- Both chunk loops now refuse moe_router_grad_from_wo_input with chunk
  weight-gradient accumulation at trace time, before any work: the
  pipelined loop had no such check, and the plain loop's came after the
  up projection, which can raise a different error first (the config
  validator rejects the pair in both cases).
- gdn_conv1d_input_in_qkv_dtype is read directly instead of through
  getattr with a default, in its validator and in model_runner
  (gdn_cp_head_sharded_test.py's SimpleNamespace config gains the field).

Tests:
- moe_router_grad_from_wo_input_test.py: chunk_wgrad_guard, both chunk
  loops raise with moe_accumulate_chunk_wgrad past the validator. The op
  precision check also bounds both paths' output and gradients against
  the float32 reference (relative L2 1e-2, measured at most 2.4e-3): the
  two paths share the unsort helpers, so a defect there leaves the
  fused/two-ops ratio at 1, and the forward on the dropping buffer was
  unchecked.
- moe_wi_gmm_fused_swiglu_test.py: GmmSiluVjpTest also with heuristic
  tiling (use_gmm_v2_heuristic_tiling must reach gmm's backward).
- moe_glue_rewrites_test.py and moe_routing_index_math_test.py: validator
  cases for tensor_sequence and autoregressive parallelism, for every
  minimal* remat preset, and an error that names two unmet requirements.
- Test docstrings no longer cite history or TPU runs that did not happen
  (megablox_gmm_silu_test.py is TPU-only and not yet run on TPU); each
  _BENCH_FLAGS is commented as the subset of the benchmark's MoE settings
  that file needs; one-line docstring summaries.
- gdn_conv1d_input_in_qkv_dtype_test.py: dict literals, a docstring for
  _conv_inputs, the repo's license header.

Docs and comments:
- base.yml, types.py and the compute_conv1d_bwd.py comment describe
  gdn_conv1d_input_in_qkv_dtype without model-specific numbers, and say
  the compiled bit identity was measured on CPU.
- nnx_decoders.py and types.py point at each other for the MLP
  checkpoint names of the remat presets (_REMAT_PRESET_MLP_NAMES), so the
  two stay in sync.

Lint (pylint 3.3.8, pyink 24.10.1, codespell 2.4.1, yamllint 1.35.0, repo
config): the changed lines add no messages against the base; the 3 pyink
hunks and 5 pylint messages of the GDN test are gone.

Program identity against the previous head: unchanged, by normalized
TPU v7x lowering, except where intentionally changed (no change was
meant to alter a program).
---
 src/maxtext/configs/base.yml                  |   6 +-
 src/maxtext/configs/types.py                  | 112 +--
 src/maxtext/kernels/gdn/gdn_bwd/api.py        | 870 +++++++++---------
 .../kernels/gdn/gdn_bwd/compute_conv1d_bwd.py |  27 +-
 src/maxtext/kernels/gdn/model_runner.py       |   2 +-
 .../pallas_mosaic_tpu_v2_gmm_kernel.py        | 131 +--
 src/maxtext/layers/moe.py                     |  32 +-
 src/maxtext/layers/nnx_decoders.py            |   2 +
 .../gdn_conv1d_input_in_qkv_dtype_test.py     |  47 +-
 tests/unit/gdn_cp_head_sharded_test.py        |   1 +
 tests/unit/moe_glue_rewrites_test.py          |  25 +-
 .../moe_roe_packed_routing_gather_test.py     |   6 +-
 .../moe_router_grad_from_wo_input_test.py     |  57 +-
 tests/unit/moe_routing_index_math_test.py     |  17 +-
 tests/unit/moe_wi_gmm_fused_swiglu_test.py    |  23 +-
 15 files changed, 750 insertions(+), 608 deletions(-)

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 09f8b45af3..8a3ad1f762 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1639,9 +1639,9 @@ gdn_cp_mode: "auto"
 # sequence-sharded GDN CP (gdn_cp_mode 'seq', or 'auto' when it does not pick head sharding); a no-op otherwise.
 gdn_cp_matmul_precision: "highest"
 # GDN kernel path: build the Conv1D + SiLU halo-extended input [conv_state | qkv] in qkv's dtype and cast each tap
-# to float32 where it is read, instead of materializing a float32 copy (805.5 MB per GDN layer in the backward of
-# the Qwen3.5-397B RL trainer). Same values: outputs and the qkv gradient bit-identical, the conv weight gradient up to
-# float32 summation order. Requires use_gdn_kernel: true.
+# to float32 where it is read, instead of materializing a float32 copy of the whole input. Same values: the outputs
+# and the qkv, conv-state and bias gradients are bit-identical on CPU, the conv weight gradient matches up to float32
+# summation order. Requires use_gdn_kernel: true.
 gdn_conv1d_input_in_qkv_dtype: false
 
 use_tokamax_splash: false
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index db5ed19e63..b95c616ed1 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1113,7 +1113,8 @@ def _reject_unless(setting: str, requirements: list[tuple[bool, str]]) -> None:
     raise ValueError(f"{setting} requires: {'; '.join(unmet)}.")
 
 
-# The MLP checkpoint names that the remat_policy presets save or offload (Decoder.get_remat_policy, nnx_decoders.py).
+# The MLP checkpoint names that the remat_policy presets save or offload (Decoder.minimal_policy and get_remat_policy
+# in nnx_decoders.py, which point back here: keep the two in sync).
 # remat_policy=custom names them per tensor instead (mlpwi_0, mlpwo, moe_mlpwi_0, ... = device / offload).
 _REMAT_PRESET_MLP_NAMES = {
     "minimal": ("mlpwi_0", "mlpwi_1", "mlpwi", "mlpwo"),
@@ -1833,6 +1834,53 @@ def validate_moe_sharding_strategy(self) -> "MoEGeneral":
   # mlp_bias, remat_policy, the parallelism degrees, ...) directly: they run on the combined configs (MaxTextConfig
   # and RLConfig), which have all of them, so a renamed field fails loudly instead of passing the check.
 
+  def _prefused_sparse_matmul_requirements(self) -> list[tuple[bool, str]]:
+    """What the MoE layer needs to run the prefused [w0 | w1] sparse-matmul up projection in training."""
+    return [
+        (self.prefuse_moe_weights, "prefuse_moe_weights=True"),
+        (self.sparse_matmul, "sparse_matmul=True"),
+        (not self.te_moe_block, "te_moe_block=False"),
+        (
+            self.attention not in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"),
+            "an attention other than vllm_rpa* (the fused-MoE inference path)",
+        ),
+    ]
+
+  def _no_weight_quantization_requirement(self) -> tuple[bool, str]:
+    return (not self.quantization and not self.use_qwix_quantization, "no weight quantization")
+
+  def _remat_keeps_moe_tensor_requirement(self, tensor: str, names: tuple[str, ...]) -> tuple[bool, str]:
+    """Requirement: the remat policy saves none of `names`, the checkpoint names of the routed experts' `tensor`.
+
+    The keys that ask for this never form that tensor (it is fused into the op that reads it), so a policy that
+    saves or offloads it would silently stop saving the MoE tensor; the shared expert's dense MLP, which carries the
+    generic names (mlpwi_0, mlpwo) too, would still be saved.
+    """
+    if self.remat_policy == "custom":
+      saved = [name for name in names if getattr(self, name) != RematLocation.REMAT]
+    else:
+      saved = [name for name in names if name in _REMAT_PRESET_MLP_NAMES.get(self.remat_policy, ())]
+    return (
+        not saved,
+        f"a remat policy that does not save the routed experts' {tensor}, which this key never forms"
+        f" (remat_policy={self.remat_policy} saves it as {', '.join(saved)})",
+    )
+
+  def _expert_mlp_dim_shards(self) -> int:
+    """Devices that split the expert weights' MLP dimension inside the expert shard_map.
+
+    By the default logical axis rules: 'mlp_no_fsdp' -> tensor, tensor_sequence, autoregressive.
+    """
+    degrees = (
+        self.ici_tensor_parallelism,
+        self.dcn_tensor_parallelism,
+        self.ici_tensor_sequence_parallelism,
+        self.dcn_tensor_sequence_parallelism,
+        self.ici_autoregressive_parallelism,
+        self.dcn_autoregressive_parallelism,
+    )
+    return math.prod(degrees)
+
   def _ring_of_experts_ragged_sort_requirements(self) -> list[tuple[bool, str]]:
     """What the MoE layer needs to run the ring-of-experts ragged sort."""
     return [
@@ -1892,51 +1940,6 @@ def validate_moe_routing_index_math(self) -> "MoEGeneral":
       )
     return self
 
-  def _prefused_sparse_matmul_requirements(self) -> list[tuple[bool, str]]:
-    """What the MoE layer needs to run the prefused [w0 | w1] sparse-matmul up projection in training."""
-    return [
-        (self.prefuse_moe_weights, "prefuse_moe_weights=True"),
-        (self.sparse_matmul, "sparse_matmul=True"),
-        (not self.te_moe_block, "te_moe_block=False"),
-        (
-            self.attention not in ("vllm_rpa", "vllm_batched_rpa", "vllm_batched_rpa_long_ctx"),
-            "an attention other than vllm_rpa* (the fused-MoE inference path)",
-        ),
-    ]
-
-  def _no_weight_quantization_requirement(self) -> tuple[bool, str]:
-    return (not self.quantization and not self.use_qwix_quantization, "no weight quantization")
-
-  def _remat_keeps_moe_tensor_requirement(self, tensor: str, names: tuple[str, ...]) -> tuple[bool, str]:
-    """Requirement: the remat policy saves none of `names`, the checkpoint names of the routed experts' `tensor`.
-
-    The keys that ask for this never form that tensor (it is fused into the op that reads it), so a policy that
-    saves or offloads it would silently stop saving the MoE tensor; the shared expert's dense MLP, which carries the
-    generic names (mlpwi_0, mlpwo) too, would still be saved.
-    """
-    if self.remat_policy == "custom":
-      saved = [name for name in names if getattr(self, name) != RematLocation.REMAT]
-    else:
-      saved = [name for name in names if name in _REMAT_PRESET_MLP_NAMES.get(self.remat_policy, ())]
-    return (
-        not saved,
-        f"a remat policy that does not save the routed experts' {tensor}, which this key never forms"
-        f" (remat_policy={self.remat_policy} saves it as {', '.join(saved)})",
-    )
-
-  def _expert_mlp_dim_shards(self) -> int:
-    """Devices that split the expert weights' MLP dimension inside the expert shard_map, by the default logical axis
-    rules ('mlp_no_fsdp' -> tensor, tensor_sequence, autoregressive)."""
-    degrees = (
-        self.ici_tensor_parallelism,
-        self.dcn_tensor_parallelism,
-        self.ici_tensor_sequence_parallelism,
-        self.dcn_tensor_sequence_parallelism,
-        self.ici_autoregressive_parallelism,
-        self.dcn_autoregressive_parallelism,
-    )
-    return math.prod(degrees)
-
   @model_validator(mode="after")
   def validate_moe_ragged_sort_unique_pair_argsort(self) -> "MoEGeneral":
     """The pair argsort rewrites the ring-of-experts ragged sort's argsorts, which "count" removes."""
@@ -2241,14 +2244,13 @@ class Qwen3Next(BaseModel):
           "GDN kernel path: build the Conv1D + SiLU input, the halo-extended [conv_state | qkv] (or qkv left-padded"
           " by gdn_conv_kernel_dim - 1 zeros), in qkv's dtype and cast each of the gdn_conv_kernel_dim taps to"
           " float32 where it is read, instead of materializing a float32 copy of the whole input that every tap"
-          " then slices (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py). The copy is f32[1, 16387, 12288] = 805.5 MB"
-          " per GDN layer in the backward pass of the Qwen3.5-397B RL trainer (bf16 activations, sequence-sharded"
-          " context parallelism). bfloat16 -> float32 is lossless and every product and sum is formed in the same"
-          " order, so the forward outputs and the qkv / conv-state / bias gradients are bit-identical; the conv weight"
-          " gradient, a float32 reduction over batch x sequence, may be summed in a different order once XLA fuses"
-          " the cast into it (CPU: a few ulp). Only which tensors live in HBM changes; no optimization barrier is"
-          " involved. Applies to the Conv1D forward and backward that gdn_decoupled_conv1d and its custom-VJP rules"
-          " trace. Requires use_gdn_kernel=True (the only path that runs that Conv1D)."
+          " then slices (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py). bfloat16 -> float32 is lossless and every"
+          " product and sum is formed in the same order, so the forward outputs and the qkv / conv-state / bias"
+          " gradients are bit-identical on CPU; the conv weight gradient, a float32 reduction over batch x sequence,"
+          " may be summed in a different order once XLA fuses the cast into it (CPU: a few ulp). Only which tensors"
+          " live in HBM changes; no optimization barrier is involved. Applies to the Conv1D forward and backward that"
+          " gdn_decoupled_conv1d and its custom-VJP rules trace. Requires use_gdn_kernel=True (the only path that"
+          " runs that Conv1D)."
       ),
   )
 
@@ -5308,7 +5310,7 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None:
         raise ValueError("GatedDeltaNet context parallelism requires context_parallel_load_balance=False.")
 
     if not cfg.use_gdn_kernel:
-      if getattr(cfg, "gdn_conv1d_input_in_qkv_dtype", False):
+      if cfg.gdn_conv1d_input_in_qkv_dtype:
         raise ValueError(
             "gdn_conv1d_input_in_qkv_dtype=True requires use_gdn_kernel=True: only the GDN kernel path runs the"
             " Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py) whose input it keeps in qkv's dtype."
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index ca8fcb8550..c061acbc3e 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -15,6 +15,7 @@
 """Model-facing API and custom VJP boundary for GDN backward kernel."""
 
 import functools
+import inspect
 from typing import Any, Optional, Tuple
 
 import jax
@@ -474,7 +475,22 @@ def _run_cp_gdn_decoupled_fwd_impl(
   return (out, states), t_inv_2, chunk_states, conv_halo, s_in_r, m_local
 
 
+def _conv_input_in_qkv_dtype_scoped(fn):
+  """Traces fn inside compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(<fn's conv_input_in_qkv_dtype argument>)."""
+  signature = inspect.signature(fn)
+
+  @functools.wraps(fn)  # keeps fn's signature, which jax.custom_vjp resolves arguments against
+  def scoped(*args, **kwargs):
+    bound = signature.bind(*args, **kwargs)
+    bound.apply_defaults()
+    with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(bound.arguments["conv_input_in_qkv_dtype"]):
+      return fn(*args, **kwargs)
+
+  return scoped
+
+
 @functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20))
+@_conv_input_in_qkv_dtype_scoped
 def gdn_decoupled_conv1d(
     qkv: jax.Array,
     b: jax.Array,
@@ -508,36 +524,8 @@ def gdn_decoupled_conv1d(
   materializing it in float32 (compute_conv1d_bwd.py). Same values (the conv weight gradient up to float32
   summation order); only HBM traffic changes.
   """
-  with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
-    if _is_cp_active(cp_axis_name):
-      (out, states), _, _, _, _, _ = _run_cp_gdn_decoupled_fwd_impl(
-          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,
-          cp_axis_name=cp_axis_name,
-          segment_ids=segment_ids,
-          cp_matmul_precision=cp_matmul_precision,
-      )
-      return out, states
-
-    segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
-        segment_ids, (conv_state is not None) or (recurrent_state is not None), conv_kernel_size
-    )
-    (out, states), _, _ = _run_local_gdn_decoupled_fwd(
+  if _is_cp_active(cp_axis_name):
+    (out, states), _, _, _, _, _ = _run_cp_gdn_decoupled_fwd_impl(
         qkv,
         b,
         a,
@@ -555,12 +543,39 @@ def gdn_decoupled_conv1d(
         chunk_size=chunk_size,
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
         compute_dtype=compute_dtype,
+        cp_axis_name=cp_axis_name,
         segment_ids=segment_ids,
-        conv_halo_seg=conv_halo_seg,
-        init_seg=init_seg,
+        cp_matmul_precision=cp_matmul_precision,
     )
     return out, states
 
+  segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
+      segment_ids, (conv_state is not None) or (recurrent_state is not None), conv_kernel_size
+  )
+  (out, states), _, _ = _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 out, states
+
 
 def _unwrap_primal(x):
   return jax.tree.map(lambda v: v.value if hasattr(v, "value") else v, x)
@@ -572,6 +587,7 @@ def _unwrap_cotangent(g, like_zero: bool = False):
   return g
 
 
+@_conv_input_in_qkv_dtype_scoped
 def _gdn_decoupled_conv1d_fwd(
     qkv: jax.Array,
     b: jax.Array,
@@ -596,85 +612,40 @@ def _gdn_decoupled_conv1d_fwd(
     conv_input_in_qkv_dtype: bool = False,
 ):
   """Forward rule for custom_vjp registration of gdn_decoupled_conv1d."""
-  with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
-    (
-        qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        segment_ids,
-    ) = _unwrap_primal((
-        qkv,
-        b,
-        a,
-        conv_weight,
-        conv_bias,
-        a_log,
-        dt_bias,
-        conv_state,
-        recurrent_state,
-        segment_ids,
-    ))
-    qkv = checkpoint_name(qkv, "gdn_fwd_conv")
-    qkv = checkpoint_name(qkv, "gdn_conv_out")
-    has_initial_state = (conv_state is not None) or (recurrent_state is not None)
-    if _is_cp_active(cp_axis_name):
-      if segment_ids is not None:
-        seg_metadata = cp_gdn.gather_cp_segment_metadata(
-            segment_ids, cp_axis_name, conv_kernel_size, has_initial_state=has_initial_state
-        )
-      else:
-        seg_metadata = None
-      (out, states), t_inv, chunk_states, conv_halo, s_in_r, m_local = _run_cp_gdn_decoupled_fwd_impl(
-          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,
-          cp_axis_name=cp_axis_name,
-          segment_ids=segment_ids,
-          seg_metadata=seg_metadata,
-          cp_matmul_precision=cp_matmul_precision,
-      )
-      out = checkpoint_name(out, "gdn_core_attn_out")
-      residuals = (
-          checkpoint_name(qkv, "gdn_qkv") if qkv is not None else None,
-          checkpoint_name(b, "gdn_b") if b is not None else None,
-          checkpoint_name(a, "gdn_a") if a is not None else None,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          checkpoint_name(conv_halo, "gdn_conv_state"),
-          checkpoint_name(s_in_r, "gdn_recurrent_state"),
-          checkpoint_name(t_inv, "gdn_t_inv") if t_inv is not None else None,
-          checkpoint_name(chunk_states, "gdn_chunk_states") if chunk_states is not None else None,
-          checkpoint_name(m_local, "gdn_m_local") if m_local is not None else None,
-          conv_state is not None,
-          recurrent_state is not None,
-          seg_metadata if seg_metadata is not None else segment_ids,
+  (
+      qkv,
+      b,
+      a,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      conv_state,
+      recurrent_state,
+      segment_ids,
+  ) = _unwrap_primal((
+      qkv,
+      b,
+      a,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      conv_state,
+      recurrent_state,
+      segment_ids,
+  ))
+  qkv = checkpoint_name(qkv, "gdn_fwd_conv")
+  qkv = checkpoint_name(qkv, "gdn_conv_out")
+  has_initial_state = (conv_state is not None) or (recurrent_state is not None)
+  if _is_cp_active(cp_axis_name):
+    if segment_ids is not None:
+      seg_metadata = cp_gdn.gather_cp_segment_metadata(
+          segment_ids, cp_axis_name, conv_kernel_size, has_initial_state=has_initial_state
       )
-      return (out, states), residuals
-
-    segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(segment_ids, has_initial_state, conv_kernel_size)
-    (out, states), t_inv, chunk_states = _run_local_gdn_decoupled_fwd(
+    else:
+      seg_metadata = None
+    (out, states), t_inv, chunk_states, conv_halo, s_in_r, m_local = _run_cp_gdn_decoupled_fwd_impl(
         qkv,
         b,
         a,
@@ -692,12 +663,12 @@ def _gdn_decoupled_conv1d_fwd(
         chunk_size=chunk_size,
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
         compute_dtype=compute_dtype,
+        cp_axis_name=cp_axis_name,
         segment_ids=segment_ids,
-        conv_halo_seg=conv_halo_seg,
-        init_seg=init_seg,
+        seg_metadata=seg_metadata,
+        cp_matmul_precision=cp_matmul_precision,
     )
     out = checkpoint_name(out, "gdn_core_attn_out")
-    seg_bundle = (segment_ids, conv_halo_seg, init_seg) if segment_ids is not None else None
     residuals = (
         checkpoint_name(qkv, "gdn_qkv") if qkv is not None else None,
         checkpoint_name(b, "gdn_b") if b is not None else None,
@@ -706,14 +677,58 @@ def _gdn_decoupled_conv1d_fwd(
         conv_bias,
         a_log,
         dt_bias,
-        checkpoint_name(conv_state, "gdn_conv_state") if conv_state is not None else None,
-        checkpoint_name(recurrent_state, "gdn_recurrent_state") if recurrent_state is not None else None,
+        checkpoint_name(conv_halo, "gdn_conv_state"),
+        checkpoint_name(s_in_r, "gdn_recurrent_state"),
         checkpoint_name(t_inv, "gdn_t_inv") if t_inv is not None else None,
         checkpoint_name(chunk_states, "gdn_chunk_states") if chunk_states is not None else None,
-        seg_bundle,
+        checkpoint_name(m_local, "gdn_m_local") if m_local is not None else None,
+        conv_state is not None,
+        recurrent_state is not None,
+        seg_metadata if seg_metadata is not None else segment_ids,
     )
     return (out, states), residuals
 
+  segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(segment_ids, has_initial_state, conv_kernel_size)
+  (out, states), t_inv, chunk_states = _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,
+  )
+  out = checkpoint_name(out, "gdn_core_attn_out")
+  seg_bundle = (segment_ids, conv_halo_seg, init_seg) if segment_ids is not None else None
+  residuals = (
+      checkpoint_name(qkv, "gdn_qkv") if qkv is not None else None,
+      checkpoint_name(b, "gdn_b") if b is not None else None,
+      checkpoint_name(a, "gdn_a") if a is not None else None,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      checkpoint_name(conv_state, "gdn_conv_state") if conv_state is not None else None,
+      checkpoint_name(recurrent_state, "gdn_recurrent_state") if recurrent_state is not None else None,
+      checkpoint_name(t_inv, "gdn_t_inv") if t_inv is not None else None,
+      checkpoint_name(chunk_states, "gdn_chunk_states") if chunk_states is not None else None,
+      seg_bundle,
+  )
+  return (out, states), residuals
+
 
 def _gdn_decoupled_conv1d_bwd(
     num_k_heads: int,
@@ -751,295 +766,240 @@ def _gdn_decoupled_conv1d_bwd(
     conv_input_in_qkv_dtype = False
 
   with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
-    m_local_fwd = None
-    seg_slot = None
-    include_segment_ids_grad = False
-    if len(residuals) == 15:
-      (
-          pre_conv_qkv,
-          b,
-          a,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          conv_state,
-          recurrent_state,
-          t_inv_fwd,
-          chunk_states,
-          m_local_fwd,
-          has_user_conv_state,
-          has_user_recurrent_state,
-          seg_slot,
-      ) = residuals
-      include_segment_ids_grad = True
-    elif len(residuals) == 14:
-      (
-          pre_conv_qkv,
-          b,
-          a,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          conv_state,
-          recurrent_state,
-          t_inv_fwd,
-          chunk_states,
-          m_local_fwd,
-          has_user_conv_state,
-          has_user_recurrent_state,
-      ) = residuals
-    elif len(residuals) == 12:
-      (
-          pre_conv_qkv,
-          b,
-          a,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          conv_state,
-          recurrent_state,
-          t_inv_fwd,
-          chunk_states,
-          seg_slot,
-      ) = residuals
-      has_user_conv_state = conv_state is not None
-      has_user_recurrent_state = recurrent_state is not None
-      include_segment_ids_grad = True
-    elif len(residuals) == 11:
-      (
-          pre_conv_qkv,
-          b,
-          a,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          conv_state,
-          recurrent_state,
-          t_inv_fwd,
-          chunk_states,
-      ) = residuals
-      has_user_conv_state = conv_state is not None
-      has_user_recurrent_state = recurrent_state is not None
-    else:
-      (
-          pre_conv_qkv,
-          b,
-          a,
-          conv_weight,
-          conv_bias,
-          a_log,
-          dt_bias,
-          conv_state,
-          recurrent_state,
-          t_inv_fwd,
-      ) = residuals
-      chunk_states = None
-      has_user_conv_state = conv_state is not None
-      has_user_recurrent_state = recurrent_state is not None
-
-    use_cp = _is_cp_active(cp_axis_name)
-    if isinstance(seg_slot, tuple):
-      segment_ids, conv_halo_seg, init_seg = seg_slot
-    elif use_cp:
-      segment_ids, conv_halo_seg, init_seg = seg_slot, None, None
-    else:
-      segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
-          seg_slot, bool(has_user_conv_state) or bool(has_user_recurrent_state), conv_kernel_size
-      )
+    return _gdn_decoupled_conv1d_bwd_impl(
+        num_k_heads,
+        num_v_heads,
+        head_k_dim,
+        head_v_dim,
+        conv_kernel_size,
+        chunk_size,
+        use_qk_norm_in_gdn,
+        compute_dtype,
+        cp_axis_name,
+        cp_matmul_precision,
+        residuals,
+        cotangents,
+    )
 
-    d_out_raw, d_states = cotangents
-    d_conv_state_raw, d_recurrent_state_raw = d_states
-    d_out = _unwrap_cotangent(d_out_raw, like_zero=True)
-    # next_conv_state / next_recurrent_state depend on the inputs even without
-    # caller initial states, so their cotangents are always propagated.
-    d_conv_state = _unwrap_cotangent(d_conv_state_raw, like_zero=False)
-    d_recurrent_state = _unwrap_cotangent(d_recurrent_state_raw, like_zero=False)
-    need_dh0 = bool(has_user_recurrent_state)
-
-    if use_cp and segment_ids is not None and conv_halo_seg is None and init_seg is None:
-      segment_ids, conv_halo_seg, init_seg = cp_gdn.gather_cp_segment_metadata(
-          segment_ids,
-          cp_axis_name,
-          conv_kernel_size,
-          has_initial_state=bool(has_user_conv_state) or bool(has_user_recurrent_state),
-      )
 
-    # When chunk_states is cached in residuals and t_inv was omitted (chunk_size > 64),
-    # the Pallas backward kernel rebuilds T^-1 per chunk in VMEM using xor_dnc instead
-    # of reading it from HBM or recomputing it with XLA.
-    recompute_t_inv_in_kernel = chunk_states is not None and t_inv_fwd is None
-    conv_out_cached = None
-    if not recompute_t_inv_in_kernel and (chunk_states is None or t_inv_fwd is None):
-      qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states(
-          qkv=pre_conv_qkv,
-          b=b,
-          a=a,
-          conv_weight=conv_weight,
-          conv_bias=conv_bias,
-          a_log=a_log,
-          dt_bias=dt_bias,
-          recurrent_state=recurrent_state,
-          conv_state=conv_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,
-          cached_t_inv=t_inv_fwd,
-          segment_ids=segment_ids,
-          conv_halo_seg=conv_halo_seg,
-          init_seg=init_seg,
-      )
-      if chunk_states is None:
-        chunk_states = chunk_states_recomputed
-      if t_inv_fwd is None:
-        t_inv_fwd = t_inv_recomputed
-    else:
-      conv_out_cached, qkv_conv = conv1d_silu_fwd(
-          qkv=pre_conv_qkv,
-          conv_weight=conv_weight,
-          conv_bias=conv_bias,
-          kernel_size=conv_kernel_size,
-          conv_state=conv_state,
-          segment_ids=segment_ids,
-          conv_halo_seg=conv_halo_seg,
-      )
-    t_inv = t_inv_fwd
-
-    if use_cp:
-      # Under jax.shard_map(..., check_vma=False), _shard_map_transpose (shard_map.py:1942)
-      # divides cotangents of unmapped out_specs (next_conv_state, next_recurrent_state) by D,
-      # and (shard_map.py:1959) emits a single lax.psum across cp_axis_name for cotangents of
-      # unmapped in_specs (conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state),
-      # avoiding redundant All-Reduce collectives inside the custom_vjp rule.
-      dht_final = (
-          jax.lax.psum(d_recurrent_state.astype(jnp.float32), cp_axis_name)
-          if d_recurrent_state is not None
-          else jnp.zeros(
-              (pre_conv_qkv.shape[0], num_v_heads, head_k_dim, head_v_dim),
-              dtype=jnp.float32,
-          )
-      )
-      d_cs_ext_unscaled = (
-          jax.lax.psum(d_conv_state.astype(jnp.float32), cp_axis_name).astype(pre_conv_qkv.dtype)
-          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,
-      )
-      dht_local, _ = cp_gdn.incoming_grad_state(
-          dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision
-      )
-      bwd_out = pallas_gdn_bwd_kernel(
-          qkv_conv=qkv_conv,
-          b=b,
-          a=a,
-          a_log=a_log,
-          dt_bias=dt_bias,
-          do=d_out,
-          chunk_states=chunk_states,
-          t_inv=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,
-          d_recurrent_state=dht_local,
-          return_dh0=need_dh0,
-          recompute_t_inv=recompute_t_inv_in_kernel,
-      )
-      if need_dh0:
-        dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
-      else:
-        dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out  # pylint: disable=unbalanced-tuple-unpacking
-        dh0_local = None
-
-      d_pre_conv_qkv, d_conv_weight, d_conv_bias, d_cs_local = conv1d_silu_bwd(
-          qkv=pre_conv_qkv,
-          conv_weight=conv_weight,
-          conv_bias=conv_bias,
-          dy=dy_conv,
-          kernel_size=conv_kernel_size,
-          conv_out=conv_out_cached,
-          conv_state=conv_state,
-          return_d_conv_state=True,
-          segment_ids=segment_ids,
-          conv_halo_seg=conv_halo_seg,
-      )
-      d_cs_ext_for_halo = d_cs_ext_unscaled
-      if d_cs_ext_unscaled is not None and segment_ids is not None:
-        # With packing, next_conv_state is the window ending at the last valid
-        # token of the rank that owns it (see broadcast_end_conv_state). Scatter
-        # its cotangent into that rank's tokens / halo instead of the last rank's tail.
-        cs_has_valid = jnp.any(segment_ids > 0, axis=1)
-        is_owner = cp_gdn.select_end_conv_state_rank(cs_has_valid, cp_axis_name)
-        d_cs_owned = jnp.where(is_owner[:, None, None], d_cs_ext_unscaled, jnp.zeros_like(d_cs_ext_unscaled))
-        d_pre_conv_qkv, d_cs_local = local_compute_conv1d.extract_segment_conv_state_split_adjoint(
-            d_cs_owned,
-            d_pre_conv_qkv,
-            d_cs_local,
-            segment_ids,
-            conv_kernel_size,
-            conv_halo_seg,
-        )
-        d_cs_ext_for_halo = None
-      d_pre_conv_qkv, _ = cp_gdn.halo_exchange_for_conv_bwd(
-          dx=d_pre_conv_qkv,
-          d_conv_state=d_cs_local,
-          d_conv_state_ext=d_cs_ext_for_halo,
-          kernel_size=conv_kernel_size,
-          cp_axis=cp_axis_name,
-          segment_ids=segment_ids,
-      )
+def _gdn_decoupled_conv1d_bwd_impl(
+    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,
+    cp_axis_name: str | tuple[str, ...] | None,
+    cp_matmul_precision: Any,
+    residuals: tuple[Any, ...],
+    cotangents: tuple[Any, ...],
+):
+  """_gdn_decoupled_conv1d_bwd after its argument shim, traced inside its conv_input_in_qkv_dtype scope."""
+  m_local_fwd = None
+  seg_slot = None
+  include_segment_ids_grad = False
+  if len(residuals) == 15:
+    (
+        pre_conv_qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        t_inv_fwd,
+        chunk_states,
+        m_local_fwd,
+        has_user_conv_state,
+        has_user_recurrent_state,
+        seg_slot,
+    ) = residuals
+    include_segment_ids_grad = True
+  elif len(residuals) == 14:
+    (
+        pre_conv_qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        t_inv_fwd,
+        chunk_states,
+        m_local_fwd,
+        has_user_conv_state,
+        has_user_recurrent_state,
+    ) = residuals
+  elif len(residuals) == 12:
+    (
+        pre_conv_qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        t_inv_fwd,
+        chunk_states,
+        seg_slot,
+    ) = residuals
+    has_user_conv_state = conv_state is not None
+    has_user_recurrent_state = recurrent_state is not None
+    include_segment_ids_grad = True
+  elif len(residuals) == 11:
+    (
+        pre_conv_qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        t_inv_fwd,
+        chunk_states,
+    ) = residuals
+    has_user_conv_state = conv_state is not None
+    has_user_recurrent_state = recurrent_state is not None
+  else:
+    (
+        pre_conv_qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        t_inv_fwd,
+    ) = residuals
+    chunk_states = None
+    has_user_conv_state = conv_state is not None
+    has_user_recurrent_state = recurrent_state is not None
 
-      idx = jax.lax.axis_index(cp_axis_name)
-      d_init_cs_rank0 = jnp.where(idx == 0, d_cs_local, jnp.zeros_like(d_cs_local))
-      d_conv_state_out = d_init_cs_rank0 if has_user_conv_state else None
-      if need_dh0 and dh0_local is not None:
-        d_recurrent_state_out = jnp.where(idx == 0, dh0_local, jnp.zeros_like(dh0_local))
-      else:
-        d_recurrent_state_out = None
-      out_grads = (
-          d_pre_conv_qkv,
-          d_b,
-          d_a,
-          d_conv_weight,
-          d_conv_bias,
-          d_a_log,
-          d_dt_bias,
-          d_conv_state_out,
-          d_recurrent_state_out,
-      )
-      if include_segment_ids_grad:
-        return out_grads + (None,)
-      return out_grads
+  use_cp = _is_cp_active(cp_axis_name)
+  if isinstance(seg_slot, tuple):
+    segment_ids, conv_halo_seg, init_seg = seg_slot
+  elif use_cp:
+    segment_ids, conv_halo_seg, init_seg = seg_slot, None, None
+  else:
+    segment_ids, conv_halo_seg, init_seg = _local_segment_metadata(
+        seg_slot, bool(has_user_conv_state) or bool(has_user_recurrent_state), conv_kernel_size
+    )
+
+  d_out_raw, d_states = cotangents
+  d_conv_state_raw, d_recurrent_state_raw = d_states
+  d_out = _unwrap_cotangent(d_out_raw, like_zero=True)
+  # next_conv_state / next_recurrent_state depend on the inputs even without
+  # caller initial states, so their cotangents are always propagated.
+  d_conv_state = _unwrap_cotangent(d_conv_state_raw, like_zero=False)
+  d_recurrent_state = _unwrap_cotangent(d_recurrent_state_raw, like_zero=False)
+  need_dh0 = bool(has_user_recurrent_state)
+
+  if use_cp and segment_ids is not None and conv_halo_seg is None and init_seg is None:
+    segment_ids, conv_halo_seg, init_seg = cp_gdn.gather_cp_segment_metadata(
+        segment_ids,
+        cp_axis_name,
+        conv_kernel_size,
+        has_initial_state=bool(has_user_conv_state) or bool(has_user_recurrent_state),
+    )
 
+  # When chunk_states is cached in residuals and t_inv was omitted (chunk_size > 64),
+  # the Pallas backward kernel rebuilds T^-1 per chunk in VMEM using xor_dnc instead
+  # of reading it from HBM or recomputing it with XLA.
+  recompute_t_inv_in_kernel = chunk_states is not None and t_inv_fwd is None
+  conv_out_cached = None
+  if not recompute_t_inv_in_kernel and (chunk_states is None or t_inv_fwd is None):
+    qkv_conv, chunk_states_recomputed, t_inv_recomputed = _compute_forward_conv_and_states(
+        qkv=pre_conv_qkv,
+        b=b,
+        a=a,
+        conv_weight=conv_weight,
+        conv_bias=conv_bias,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        recurrent_state=recurrent_state,
+        conv_state=conv_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,
+        cached_t_inv=t_inv_fwd,
+        segment_ids=segment_ids,
+        conv_halo_seg=conv_halo_seg,
+        init_seg=init_seg,
+    )
+    if chunk_states is None:
+      chunk_states = chunk_states_recomputed
+    if t_inv_fwd is None:
+      t_inv_fwd = t_inv_recomputed
+  else:
+    conv_out_cached, qkv_conv = conv1d_silu_fwd(
+        qkv=pre_conv_qkv,
+        conv_weight=conv_weight,
+        conv_bias=conv_bias,
+        kernel_size=conv_kernel_size,
+        conv_state=conv_state,
+        segment_ids=segment_ids,
+        conv_halo_seg=conv_halo_seg,
+    )
+  t_inv = t_inv_fwd
+
+  if use_cp:
+    # Under jax.shard_map(..., check_vma=False), _shard_map_transpose (shard_map.py:1942)
+    # divides cotangents of unmapped out_specs (next_conv_state, next_recurrent_state) by D,
+    # and (shard_map.py:1959) emits a single lax.psum across cp_axis_name for cotangents of
+    # unmapped in_specs (conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state),
+    # avoiding redundant All-Reduce collectives inside the custom_vjp rule.
+    dht_final = (
+        jax.lax.psum(d_recurrent_state.astype(jnp.float32), cp_axis_name)
+        if d_recurrent_state is not None
+        else jnp.zeros(
+            (pre_conv_qkv.shape[0], num_v_heads, head_k_dim, head_v_dim),
+            dtype=jnp.float32,
+        )
+    )
+    d_cs_ext_unscaled = (
+        jax.lax.psum(d_conv_state.astype(jnp.float32), cp_axis_name).astype(pre_conv_qkv.dtype)
+        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,
+    )
+    dht_local, _ = cp_gdn.incoming_grad_state(
+        dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision
+    )
     bwd_out = pallas_gdn_bwd_kernel(
         qkv_conv=qkv_conv,
         b=b,
@@ -1056,15 +1016,14 @@ def _gdn_decoupled_conv1d_bwd(
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
         segment_ids=segment_ids,
         init_seg=init_seg,
-        d_recurrent_state=d_recurrent_state,
+        d_recurrent_state=dht_local,
         return_dh0=need_dh0,
         recompute_t_inv=recompute_t_inv_in_kernel,
     )
-    # pylint: disable=unbalanced-tuple-unpacking
     if need_dh0:
       dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
     else:
-      dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out
+      dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out  # pylint: disable=unbalanced-tuple-unpacking
       dh0_local = None
 
     d_pre_conv_qkv, d_conv_weight, d_conv_bias, d_cs_local = conv1d_silu_bwd(
@@ -1079,21 +1038,39 @@ def _gdn_decoupled_conv1d_bwd(
         segment_ids=segment_ids,
         conv_halo_seg=conv_halo_seg,
     )
-    # pylint: enable=unbalanced-tuple-unpacking
-    if d_conv_state is not None and segment_ids is not None:
+    d_cs_ext_for_halo = d_cs_ext_unscaled
+    if d_cs_ext_unscaled is not None and segment_ids is not None:
+      # With packing, next_conv_state is the window ending at the last valid
+      # token of the rank that owns it (see broadcast_end_conv_state). Scatter
+      # its cotangent into that rank's tokens / halo instead of the last rank's tail.
+      cs_has_valid = jnp.any(segment_ids > 0, axis=1)
+      is_owner = cp_gdn.select_end_conv_state_rank(cs_has_valid, cp_axis_name)
+      d_cs_owned = jnp.where(is_owner[:, None, None], d_cs_ext_unscaled, jnp.zeros_like(d_cs_ext_unscaled))
       d_pre_conv_qkv, d_cs_local = local_compute_conv1d.extract_segment_conv_state_split_adjoint(
-          d_conv_state,
+          d_cs_owned,
           d_pre_conv_qkv,
-          d_cs_local if has_user_conv_state else None,
+          d_cs_local,
           segment_ids,
           conv_kernel_size,
           conv_halo_seg,
       )
-    elif d_conv_state is not None and pre_conv_qkv.shape[1] >= conv_kernel_size - 1:
-      d_pre_conv_qkv = d_pre_conv_qkv.at[:, -(conv_kernel_size - 1) :, :].add(d_conv_state.astype(d_pre_conv_qkv.dtype))
+      d_cs_ext_for_halo = None
+    d_pre_conv_qkv, _ = cp_gdn.halo_exchange_for_conv_bwd(
+        dx=d_pre_conv_qkv,
+        d_conv_state=d_cs_local,
+        d_conv_state_ext=d_cs_ext_for_halo,
+        kernel_size=conv_kernel_size,
+        cp_axis=cp_axis_name,
+        segment_ids=segment_ids,
+    )
 
-    d_conv_state_out = d_cs_local if has_user_conv_state else None
-    d_recurrent_state_out = dh0_local if need_dh0 else None
+    idx = jax.lax.axis_index(cp_axis_name)
+    d_init_cs_rank0 = jnp.where(idx == 0, d_cs_local, jnp.zeros_like(d_cs_local))
+    d_conv_state_out = d_init_cs_rank0 if has_user_conv_state else None
+    if need_dh0 and dh0_local is not None:
+      d_recurrent_state_out = jnp.where(idx == 0, dh0_local, jnp.zeros_like(dh0_local))
+    else:
+      d_recurrent_state_out = None
     out_grads = (
         d_pre_conv_qkv,
         d_b,
@@ -1109,6 +1086,75 @@ def _gdn_decoupled_conv1d_bwd(
       return out_grads + (None,)
     return out_grads
 
+  bwd_out = pallas_gdn_bwd_kernel(
+      qkv_conv=qkv_conv,
+      b=b,
+      a=a,
+      a_log=a_log,
+      dt_bias=dt_bias,
+      do=d_out,
+      chunk_states=chunk_states,
+      t_inv=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,
+      d_recurrent_state=d_recurrent_state,
+      return_dh0=need_dh0,
+      recompute_t_inv=recompute_t_inv_in_kernel,
+  )
+  # pylint: disable=unbalanced-tuple-unpacking
+  if need_dh0:
+    dy_conv, d_b, d_a, d_a_log, d_dt_bias, dh0_local = bwd_out
+  else:
+    dy_conv, d_b, d_a, d_a_log, d_dt_bias = bwd_out
+    dh0_local = None
+
+  d_pre_conv_qkv, d_conv_weight, d_conv_bias, d_cs_local = conv1d_silu_bwd(
+      qkv=pre_conv_qkv,
+      conv_weight=conv_weight,
+      conv_bias=conv_bias,
+      dy=dy_conv,
+      kernel_size=conv_kernel_size,
+      conv_out=conv_out_cached,
+      conv_state=conv_state,
+      return_d_conv_state=True,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+  )
+  # pylint: enable=unbalanced-tuple-unpacking
+  if d_conv_state is not None and segment_ids is not None:
+    d_pre_conv_qkv, d_cs_local = local_compute_conv1d.extract_segment_conv_state_split_adjoint(
+        d_conv_state,
+        d_pre_conv_qkv,
+        d_cs_local if has_user_conv_state else None,
+        segment_ids,
+        conv_kernel_size,
+        conv_halo_seg,
+    )
+  elif d_conv_state is not None and pre_conv_qkv.shape[1] >= conv_kernel_size - 1:
+    d_pre_conv_qkv = d_pre_conv_qkv.at[:, -(conv_kernel_size - 1) :, :].add(d_conv_state.astype(d_pre_conv_qkv.dtype))
+
+  d_conv_state_out = d_cs_local if has_user_conv_state else None
+  d_recurrent_state_out = dh0_local if need_dh0 else None
+  out_grads = (
+      d_pre_conv_qkv,
+      d_b,
+      d_a,
+      d_conv_weight,
+      d_conv_bias,
+      d_a_log,
+      d_dt_bias,
+      d_conv_state_out,
+      d_recurrent_state_out,
+  )
+  if include_segment_ids_grad:
+    return out_grads + (None,)
+  return out_grads
+
 
 gdn_decoupled_conv1d.defvjp(
     _gdn_decoupled_conv1d_fwd,
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py
index 7628934cbe..a17489e967 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/compute_conv1d_bwd.py
@@ -139,8 +139,16 @@ def conv1d_silu_bwd(
   """
   if _resolve_conv_input_in_qkv_dtype(conv_input_in_qkv_dtype):
     return _conv1d_silu_bwd_narrow_input(
-        qkv, conv_weight, conv_bias, dy, kernel_size, conv_out, conv_state, return_d_conv_state, segment_ids,
-        conv_halo_seg
+        qkv,
+        conv_weight,
+        conv_bias,
+        dy,
+        kernel_size,
+        conv_out,
+        conv_state,
+        return_d_conv_state,
+        segment_ids,
+        conv_halo_seg,
     )
   _, seq_len, _ = qkv.shape
   if conv_weight.ndim == 3:
@@ -233,18 +241,19 @@ def conv1d_silu_bwd(
 # gdn_conv1d_input_in_qkv_dtype: the halo-extended conv input in qkv's dtype.
 #
 # The default path above builds conv_input = [conv_state | qkv] (or qkv left-padded by kernel_size - 1 zeros) in
-# float32 and slices the kernel_size taps out of it. That materializes a float32 copy of the whole input: on the
-# Qwen3.5-397B RL trainer (bf16 qkv, sequence-sharded context parallelism) one f32[1, 16387, 12288] = 805.5 MB buffer
-# per GDN layer in the backward pass, written once and read kernel_size times.
+# float32 and slices the kernel_size taps out of it. That materializes a float32 copy of the whole input, one buffer
+# of batch x (sequence + kernel_size - 1) x channels per GDN layer in the backward pass, written once and read
+# kernel_size times.
 #
 # The functions below build conv_input in qkv's dtype (bf16 for bf16 qkv and conv_state; the promoted dtype if they
 # differ) and cast each tap slice to float32 where it is read, so XLA fuses the cast into the consumer and the
 # float32 copy is never written. Every product and sum is formed in the same order, on the same float32 values, as in
 # the default path (bf16 -> float32 is exact), so op by op every output is bit-identical; only which tensors are
-# materialized in HBM changes. Compiled, the forward outputs and dx / d_conv_state / db stay bit-identical, while the
-# weight gradient dw, a float32 reduction over batch x sequence, can come out in a different summation order when XLA
-# fuses the tap's cast into the reduction (CPU: a few ulp, only when the pre-activation is recomputed without packing
-# or a conv state). No jax.lax.optimization_barrier is used: the narrow buffer alone is what removes the copy.
+# materialized in HBM changes. Compiled (measured on CPU), the forward outputs and dx / d_conv_state / db stay
+# bit-identical, while the weight gradient dw, a float32 reduction over batch x sequence, can come out in a different
+# summation order when XLA fuses the tap's cast into the reduction (CPU: a few ulp, only when the pre-activation is
+# recomputed without packing or a conv state). No jax.lax.optimization_barrier is used: the narrow buffer alone is
+# what removes the copy.
 # Tests: tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py.
 
 
diff --git a/src/maxtext/kernels/gdn/model_runner.py b/src/maxtext/kernels/gdn/model_runner.py
index 24ddaca95f..a6cce0e185 100644
--- a/src/maxtext/kernels/gdn/model_runner.py
+++ b/src/maxtext/kernels/gdn/model_runner.py
@@ -618,7 +618,7 @@ def run_gdn_kernel_layer(
   )
   cp_matmul_precision = jax.lax.Precision[getattr(cfg, "gdn_cp_matmul_precision", "highest").upper()]
   # Conv1D halo input in qkv's dtype instead of a float32 copy (compute_conv1d_bwd.py). Same values.
-  conv_input_in_qkv_dtype = bool(getattr(cfg, "gdn_conv1d_input_in_qkv_dtype", False))
+  conv_input_in_qkv_dtype = cfg.gdn_conv1d_input_in_qkv_dtype
 
   if use_head_sharded_cp:
     qkv = jnp.concatenate([query, key, value_raw], axis=3)
diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
index f02f6452d6..8168559c99 100644
--- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
+++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py
@@ -794,21 +794,11 @@ def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int):
 
       acc = apply_act_fn(acc, cfgs.fuse_act)
 
-      gm_id = pl.program_id(cfgs.gm_axis)
-
-      # Mask out rows that does not belong to the current group.
-      m_start = metadata_ref.gm_id_to_m_offset[gm_id]
-      m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1]
-      m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane
-
-      m_start_local = m_start - m_offset
-      m_end_local = m_end - m_offset
-
-      iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0)
-      mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local)
-      acc_masked = jnp.where(mask, acc, 0)
-
       # Write the final output to the output ref.
+      if not cfgs.out_is_3d:
+        _store_group_rows(tiled_out_ref, partial_out_ref, acc, metadata_ref, cfgs)
+        return
+      acc_masked, gm_id, m_end_local = _mask_group_rows(acc, metadata_ref, cfgs)
       if cfgs.out_window:
         # Full-D0 out slab resident across the N tiles of this gm tile. The
         # cross-gm partial-row carry is kept per N tile, in the 2D domain.
@@ -825,39 +815,12 @@ def rhs_block(start_k: int, end_k: int, start_n: int, end_n: int):
         partial_out_ref[n_id] = jnp.where(m_end_local % sub == 0, jnp.zeros_like(last), last)
         relayout.store_2d_as_3d_window(tiled_out_ref, out2d, out2d.shape[0], cfgs.out_tile_d0, n_id)
         return
-      if cfgs.out_is_3d:
-        # Block is [tile_m // sublane, sublane, tile_d0, 128]; scatter the
-        # [tile_m, tile_n] rows back into per-token [tile_d0, 128] slabs.
-        relayout.store_2d_as_3d(
-            tiled_out_ref, acc_masked.astype(tiled_out_ref.dtype), acc_masked.shape[0], cfgs.out_tile_d0
-        )
-      else:
-        tiled_out_ref[...] = acc_masked.reshape(tiled_out_ref.shape).astype(tiled_out_ref.dtype)
-
-      # If this is the first tile for grid[n_id, :, :], we initialize the
-      # partial out to zeros. Otherwise, partial out from last tile of
-      # grid[n_id-1, :, :] can be used and cause numeric issues.
-      partial_out_zeros = jnp.zeros_like(partial_out_ref)
-
-      # Accumulate the partial output from the previous step.
-      tiled_out_ref[0] += jnp.where(gm_id == 0, partial_out_zeros, partial_out_ref[...])
-
-      # Consider following case where size_lhs_sublane = 4, number denotes group
-      # id and | denotes boundaries between sublanes:
-      # | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 |
-      #
-      # Assuming group id of current step is 1, current step will not completely
-      # fill size_lhs_sublane rows and will be revisited at the next step. By
-      # storing the partial rows into the partial_out_ref, the next step can
-      # read them and accumulate to them.  Additionally, for group id of 2,
-      # since it completely fills the size_lhs_sublane rows, we need to zero out
-      # partial_out_ref to avoid numeric error for group 3.
-      last_row = m_end_local // cfgs.dims.size_lhs_sublane
-      partial_out_ref[...] = jnp.where(
-          m_end_local % cfgs.dims.size_lhs_sublane == 0,
-          partial_out_zeros,
-          tiled_out_ref[last_row],
+      # Block is [tile_m // sublane, sublane, tile_d0, 128]; scatter the
+      # [tile_m, tile_n] rows back into per-token [tile_d0, 128] slabs.
+      relayout.store_2d_as_3d(
+          tiled_out_ref, acc_masked.astype(tiled_out_ref.dtype), acc_masked.shape[0], cfgs.out_tile_d0
       )
+      _carry_partial_rows(tiled_out_ref, partial_out_ref, gm_id, m_end_local, cfgs)
     else:
       acc_ref[...] = acc
 
@@ -900,42 +863,80 @@ def matmul_last():
   )
 
 
-def _store_group_rows(
-    tiled_ref: jax.Array,  # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n]
-    partial_ref: jax.Array,  # [size_lhs_sublane, tile_n]
-    acc: jax.Array,  # [tile_m, tile_n]
-    metadata_ref: MetadataRef,
-    cfgs: GmmConfigs,
-):
-  """Stores the current group's rows of `acc` into a 2D output tile, with the partial-row carry.
+def _mask_group_rows(acc: jax.Array, metadata_ref: MetadataRef, cfgs: GmmConfigs):
+  """Zeroes the rows of `acc` that do not belong to the current gm tile's group.
 
-  The same steps as the 2D store at the end of inner_kernel's last k step: zero
-  the rows of other groups, round to the output dtype, add the sublane rows the
-  previous gm tile left in `partial_ref`, and leave this tile's unfinished last
-  sublane rows there for the next one. Used for the extra (return_pre_act)
-  outputs written from the same accumulator.
+  Returns the masked `acc`, the gm tile id and the group's end row relative to
+  the tile's first sublane row (what _carry_partial_rows needs).
   """
   gm_id = pl.program_id(cfgs.gm_axis)
+
+  # Mask out rows that does not belong to the current group.
   m_start = metadata_ref.gm_id_to_m_offset[gm_id]
   m_end = metadata_ref.gm_id_to_m_offset[gm_id + 1]
   m_offset = m_start - m_start % cfgs.dims.size_lhs_sublane
+
   m_start_local = m_start - m_offset
   m_end_local = m_end - m_offset
 
   iota = lax.broadcasted_iota(jnp.int32, acc.shape, 0)
   mask = jnp.logical_and(m_start_local <= iota, iota < m_end_local)
-  tiled_ref[...] = jnp.where(mask, acc, 0).reshape(tiled_ref.shape).astype(tiled_ref.dtype)
+  return jnp.where(mask, acc, 0), gm_id, m_end_local
 
-  partial_zeros = jnp.zeros_like(partial_ref)
-  tiled_ref[0] += jnp.where(gm_id == 0, partial_zeros, partial_ref[...])
+
+def _carry_partial_rows(
+    tiled_out_ref: jax.Array,
+    partial_out_ref: jax.Array,
+    gm_id: jax.Array,
+    m_end_local: jax.Array,
+    cfgs: GmmConfigs,
+):
+  """Adds the sublane rows the previous gm tile left in `partial_out_ref`; leaves this tile's unfinished ones there."""
+  # If this is the first tile for grid[n_id, :, :], we initialize the
+  # partial out to zeros. Otherwise, partial out from last tile of
+  # grid[n_id-1, :, :] can be used and cause numeric issues.
+  partial_out_zeros = jnp.zeros_like(partial_out_ref)
+
+  # Accumulate the partial output from the previous step.
+  tiled_out_ref[0] += jnp.where(gm_id == 0, partial_out_zeros, partial_out_ref[...])
+
+  # Consider following case where size_lhs_sublane = 4, number denotes group
+  # id and | denotes boundaries between sublanes:
+  # | 0 0 1 2 | 2 2 2 2 | 3 3 4 4 |
+  #
+  # Assuming group id of current step is 1, current step will not completely
+  # fill size_lhs_sublane rows and will be revisited at the next step. By
+  # storing the partial rows into the partial_out_ref, the next step can
+  # read them and accumulate to them.  Additionally, for group id of 2,
+  # since it completely fills the size_lhs_sublane rows, we need to zero out
+  # partial_out_ref to avoid numeric error for group 3.
   last_row = m_end_local // cfgs.dims.size_lhs_sublane
-  partial_ref[...] = jnp.where(
+  partial_out_ref[...] = jnp.where(
       m_end_local % cfgs.dims.size_lhs_sublane == 0,
-      partial_zeros,
-      tiled_ref[last_row],
+      partial_out_zeros,
+      tiled_out_ref[last_row],
   )
 
 
+def _store_group_rows(
+    tiled_ref: jax.Array,  # [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n]
+    partial_ref: jax.Array,  # [size_lhs_sublane, tile_n]
+    acc: jax.Array,  # [tile_m, tile_n]
+    metadata_ref: MetadataRef,
+    cfgs: GmmConfigs,
+):
+  """The 2D store of inner_kernel's last k step: the current group's rows of `acc`, with the partial-row carry.
+
+  Zeroes the rows of other groups, rounds to the output dtype, stores the tile
+  and carries its unfinished last sublane rows to the next gm tile. Used for the
+  plain 2D output and for the extra (return_pre_act) outputs written from the
+  same accumulator.
+  """
+  acc_masked, gm_id, m_end_local = _mask_group_rows(acc, metadata_ref, cfgs)
+  tiled_ref[...] = acc_masked.reshape(tiled_ref.shape).astype(tiled_ref.dtype)
+  _carry_partial_rows(tiled_ref, partial_ref, gm_id, m_end_local, cfgs)
+
+
 def fill_metadata(
     lhs_group_sizes_ref: jax.Array,  # int32[size_lhs_group]
     group_offset_ref: jax.Array,  # int32[1]
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index cee0570dc3..176dea4880 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -1771,15 +1771,6 @@ def _keep_prefused_wi_whole(self):
       raise ValueError(f"moe_unsplit_prefused_wi=True cannot apply: {', '.join(blockers)}.")
     return True
 
-  def _use_prefused_swiglu_vjp(self):
-    """moe_prefused_swiglu_vjp: whether gmm_up hands [gate | up] to the activation as one array.
-
-    The per-half outputs are never formed (the fused activation keeps [gate | up] as its residual), so the config
-    validator admits the key only where no remat policy saves them (mlpwi_0 / moe_mlpwi_0 / moe_mlpwi_1), without
-    MLP biases (added per half) and without embedding chunks (partial sums per half).
-    """
-    return self.config.moe_prefused_swiglu_vjp
-
   def permute(
       self,
       inputs,
@@ -3518,7 +3509,7 @@ 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 partial_accum0 is None and partial_accum1 is None and self._use_prefused_swiglu_vjp():
+        if partial_accum0 is None and partial_accum1 is None and self.config.moe_prefused_swiglu_vjp:
           # moe_prefused_swiglu_vjp: the whole [gate | up] output goes to the activation as one array
           # (apply_ffn_activation_maybe_fused, which layer_w1=None selects), whose backward then writes d[gate | up]
           # in one pass. There are no per-half outputs to tag; the config validator rejects a remat policy that
@@ -3880,7 +3871,7 @@ def _cn(t):
         w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias)
 
       gmm_operands = get_gmm_operands_for_local_experts(x, routing, route_metadata)
-      gmm_fn = functools.partial(gmm, **gmm_operands)  # get_gmm_for_local_experts
+      gmm_fn = functools.partial(gmm, **gmm_operands)
       up_res = gmm_up_or_silu(
           x,
           w0,
@@ -4068,8 +4059,10 @@ def _ring_reduce_scatter(output):
       return output
 
     def _moe_body_outputs(output, routing, weights):
-      """What _moe_body returns: the layer output, the routing's side outputs, and the weights `(w0, w1, wo)`
-      (as returned by the GMMs under moe_accumulate_chunk_wgrad)."""
+      """What _moe_body returns: the layer output, the routing's side outputs and the weights `(w0, w1, wo)`.
+
+      The weights are as the GMMs return them under moe_accumulate_chunk_wgrad.
+      """
       return (
           output,
           routing.lb_loss,
@@ -4098,6 +4091,9 @@ def _moe_body(
         gather_weights=True,
         return_weights=False,
     ):
+      if router_grad_from_wo_input and (return_weights or not gather_weights):
+        # The fused op does not return the weights (the config validator rejects the two keys together).
+        raise NotImplementedError("moe_router_grad_from_wo_input does not support moe_accumulate_chunk_wgrad.")
       batch_size, sequence_length, embed_dim = x.shape
       if self.config.num_moe_emb_chunks > 0:
         output0, output1, gmm_fn, routing, route_metadata, wo_bias = moe_emb_chunking(
@@ -4145,7 +4141,7 @@ def _cn(t):
           w0_bias, w1_bias, wo_bias = self.transform_bias(routing.selected_experts, w0_bias, w1_bias, wo_bias)
 
         gmm_operands = get_gmm_operands_for_local_experts(x, routing, route_metadata)
-        gmm_fn = functools.partial(gmm, **gmm_operands)  # get_gmm_for_local_experts
+        gmm_fn = functools.partial(gmm, **gmm_operands)
         up_res = gmm_up_or_silu(
             x,
             w0,
@@ -4166,10 +4162,7 @@ def _cn(t):
 
       if router_grad_from_wo_input:
         # moe_router_grad_from_wo_input: the activation, the down projection and the unsort as one op whose
-        # backward does not read the down projection's output; then the same reduce-scatter as below. The op does
-        # not return the weights, so it cannot serve moe_accumulate_chunk_wgrad (the config validator rejects it).
-        if return_weights or not gather_weights:
-          raise NotImplementedError("moe_router_grad_from_wo_input does not support moe_accumulate_chunk_wgrad.")
+        # backward does not read the down projection's output; then the same reduce-scatter as below.
         output, _ = _ring_down_project_and_unpermute(
             output0, output1, wo, routing, route_metadata, batch_size, sequence_length
         )
@@ -4490,6 +4483,9 @@ def _with_routed(dispatched_c, routed):
         for c in range(n_chunks):
           gather_c = not accum_chunk_wgrad or c == 0
           return_c = accum_chunk_wgrad and c < n_chunks - 1
+          if router_grad_from_wo_input and (return_c or not gather_c):
+            # The fused op does not return the weights (the config validator rejects the two keys together).
+            raise NotImplementedError("moe_router_grad_from_wo_input does not support moe_accumulate_chunk_wgrad.")
           if c + 1 < n_chunks:
             # 1. The next chunk's token all-gather waits for this chunk's routed tokens.
             x_next, routed = barrier_tie(x[:, (c + 1) * chunk : (c + 2) * chunk, :], dispatched[c][0])
diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py
index 82450502ca..e6e2e40331 100644
--- a/src/maxtext/layers/nnx_decoders.py
+++ b/src/maxtext/layers/nnx_decoders.py
@@ -1249,6 +1249,8 @@ def get_deepseek():
 
     return layer_map[cfg.decoder_block]
 
+  # configs/types.py _REMAT_PRESET_MLP_NAMES lists the MLP checkpoint names that minimal_policy and the presets of
+  # get_remat_policy save (the MoE key validators read it): keep the two in sync.
   def minimal_policy(self, with_context=False, with_quantization=False):
     """Helper for creating minimal checkpoint policies."""
     names = [
diff --git a/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py b/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py
index 79e59acfbb..0735d75ffa 100644
--- a/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py
+++ b/tests/unit/gdn_conv1d_input_in_qkv_dtype_test.py
@@ -4,7 +4,7 @@
 # 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
+#       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,
@@ -47,6 +47,7 @@ def _bits(x):
 
 
 def _conv_inputs(seed, *, batch=2, seq_len=24, dim=16, with_state, with_bias, with_segments, dtype=jnp.bfloat16):
+  """Seeded Conv1D + SiLU inputs: (qkv, conv_weight, conv_bias, conv_state, dy, segment_ids, conv_halo_seg)."""
   keys = jax.random.split(jax.random.PRNGKey(seed), 6)
   qkv = jax.random.normal(keys[0], (batch, seq_len, dim), jnp.float32).astype(dtype)
   conv_weight = jax.random.normal(keys[1], (_KERNEL_SIZE, 1, dim), jnp.float32)
@@ -83,7 +84,7 @@ def visit(jx):
 
 
 _CASES = [
-    dict(testcase_name=f"state{s}_bias{b}_seg{g}", with_state=s, with_bias=b, with_segments=g)
+    {"testcase_name": f"state{s}_bias{b}_seg{g}", "with_state": s, "with_bias": b, "with_segments": g}
     for s in (False, True)
     for b in (False, True)
     for g in (False, True)
@@ -100,7 +101,13 @@ def test_forward_bit_identical(self, with_state, with_bias, with_segments):
 
     def fwd(narrow):
       return compute_conv1d_bwd.conv1d_silu_fwd(
-          qkv, w, bias, _KERNEL_SIZE, conv_state=state, segment_ids=seg, conv_halo_seg=halo,
+          qkv,
+          w,
+          bias,
+          _KERNEL_SIZE,
+          conv_state=state,
+          segment_ids=seg,
+          conv_halo_seg=halo,
           conv_input_in_qkv_dtype=narrow,
       )
 
@@ -124,8 +131,16 @@ def test_backward_bit_identical(self, with_state, with_bias, with_segments):
 
         def bwd(narrow, cached_conv_out=cached_conv_out, return_d_conv_state=return_d_conv_state):
           return compute_conv1d_bwd.conv1d_silu_bwd(
-              qkv, w, bias, dy, _KERNEL_SIZE, conv_out=cached_conv_out, conv_state=state,
-              return_d_conv_state=return_d_conv_state, segment_ids=seg, conv_halo_seg=halo,
+              qkv,
+              w,
+              bias,
+              dy,
+              _KERNEL_SIZE,
+              conv_out=cached_conv_out,
+              conv_state=state,
+              return_d_conv_state=return_d_conv_state,
+              segment_ids=seg,
+              conv_halo_seg=halo,
               conv_input_in_qkv_dtype=narrow,
           )
 
@@ -153,8 +168,8 @@ def bwd(narrow, cached_conv_out=cached_conv_out, return_d_conv_state=return_d_co
             self.assertEqual(_bits(r), _bits(g))
 
   @parameterized.named_parameters(
-      dict(testcase_name="no_state", with_state=False),
-      dict(testcase_name="state", with_state=True),
+      {"testcase_name": "no_state", "with_state": False},
+      {"testcase_name": "state", "with_state": True},
   )
   def test_conv_input_is_narrow_only_when_enabled(self, with_state):
     """Positive control: the default path builds the [B, S + K - 1, D] input in float32, the option in bf16."""
@@ -265,15 +280,15 @@ def loss(narrow, qkv_in, b_in, a_in, cw_in, cb_in, al_in, dt_in):
 
   def test_requires_gdn_kernel(self):
     def cfg(**overrides):
-      fields = dict(
-          decoder_block=config_types.DecoderBlockType.QWEN3_5,
-          ici_context_parallelism=1,
-          dcn_context_parallelism=1,
-          ici_context_usp_ulysses_parallelism=1,
-          dcn_context_usp_ulysses_parallelism=1,
-          use_gdn_kernel=False,
-          gdn_conv1d_input_in_qkv_dtype=False,
-      )
+      fields = {
+          "decoder_block": config_types.DecoderBlockType.QWEN3_5,
+          "ici_context_parallelism": 1,
+          "dcn_context_parallelism": 1,
+          "ici_context_usp_ulysses_parallelism": 1,
+          "dcn_context_usp_ulysses_parallelism": 1,
+          "use_gdn_kernel": False,
+          "gdn_conv1d_input_in_qkv_dtype": False,
+      }
       fields.update(overrides)
       return types.SimpleNamespace(**fields)
 
diff --git a/tests/unit/gdn_cp_head_sharded_test.py b/tests/unit/gdn_cp_head_sharded_test.py
index c6d19b09d8..10681602f6 100644
--- a/tests/unit/gdn_cp_head_sharded_test.py
+++ b/tests/unit/gdn_cp_head_sharded_test.py
@@ -69,6 +69,7 @@ def create_gdn_config(
       use_gdn_kernel=True,
       enable_gdn_sequence_packing=enable_gdn_sequence_packing,
       gdn_cp_mode=gdn_cp_mode,
+      gdn_conv1d_input_in_qkv_dtype=False,
       load_balance_loss_weight=0.0,
       scan_layers=False,
       using_pipeline_parallelism=False,
diff --git a/tests/unit/moe_glue_rewrites_test.py b/tests/unit/moe_glue_rewrites_test.py
index ae98a6c562..93ae001d70 100644
--- a/tests/unit/moe_glue_rewrites_test.py
+++ b/tests/unit/moe_glue_rewrites_test.py
@@ -233,6 +233,16 @@ def both(h, c):
             ("chunk_wgrad", {"moe_accumulate_chunk_wgrad": True}, "moe_accumulate_chunk_wgrad=False"),
             ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"),
             ("tensor_parallel", _TP_MESH, "an unsharded expert MLP dimension"),
+            (
+                "tensor_sequence_parallel",
+                {**_TP_MESH, "ici_tensor_parallelism": 1, "ici_tensor_sequence_parallelism": 2},
+                "an unsharded expert MLP dimension",
+            ),
+            (
+                "autoregressive_parallel",
+                {**_TP_MESH, "ici_tensor_parallelism": 1, "ici_autoregressive_parallelism": 2},
+                "an unsharded expert MLP dimension",
+            ),
         )
     ]
     + [(f"swiglu_vjp_{name}", "moe_prefused_swiglu_vjp", o, m) for name, o, m in _OFF_PREFUSED_PATH]
@@ -243,6 +253,18 @@ def both(h, c):
             ("emb_chunks", {"num_moe_emb_chunks": 2}, "num_moe_emb_chunks=0"),
             ("minimal_remat", {"remat_policy": "minimal"}, "remat_policy=minimal saves it as mlpwi_0"),
             ("offloaded_remat", {"remat_policy": "minimal_offloaded"}, "remat_policy=minimal_offloaded saves it as"),
+            ("context_remat", {"remat_policy": "minimal_with_context"}, "remat_policy=minimal_with_context saves it as"),
+            ("flash_remat", {"remat_policy": "minimal_flash"}, "remat_policy=minimal_flash saves it as"),
+            (
+                "quantization_remat",
+                {"remat_policy": "minimal_with_quantization"},
+                "remat_policy=minimal_with_quantization saves it as",
+            ),
+            (
+                "context_quantization_remat",
+                {"remat_policy": "minimal_with_context_and_quantization"},
+                "remat_policy=minimal_with_context_and_quantization saves it as",
+            ),
             ("saved_gate", {"remat_policy": "custom", "mlpwi_0": "device"}, "remat_policy=custom saves it as mlpwi_0"),
             ("saved_moe_gate", {"remat_policy": "custom", "moe_mlpwi_0": "device"}, "saves it as moe_mlpwi_0"),
             ("offloaded_moe_up", {"remat_policy": "custom", "moe_mlpwi_1": "offload"}, "saves it as moe_mlpwi_1"),
@@ -520,7 +542,8 @@ def _unsplit_wi_tensor_parallel_guard():
   return ok
 
 
-# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not
+# all set them).
 _BENCH_FLAGS = {
     "moe_chunk_pipeline": True,
     "moe_expert_weight_prefetch": "forward",
diff --git a/tests/unit/moe_roe_packed_routing_gather_test.py b/tests/unit/moe_roe_packed_routing_gather_test.py
index b5ece5ac61..da3be34079 100644
--- a/tests/unit/moe_roe_packed_routing_gather_test.py
+++ b/tests/unit/moe_roe_packed_routing_gather_test.py
@@ -46,8 +46,7 @@
     but exactly one [..., 2 * top_k] float32 array; under moe_routing_maps=device that array is tagged
     moe_routing_maps, under remat it is untagged, and nothing is ever tagged moe_route.
   - Deliberately wrong versions must FAIL: a packed gather whose result is shifted along the sequence (each chunk
-    gets another chunk's routing), and the packed path with its load-balance loss dropped (what an earlier version
-    of this change did).
+    gets another chunk's routing), and the packed path with its load-balance loss dropped.
 """
 
 import contextlib
@@ -325,7 +324,8 @@ def _dropped_load_balance_loss():
   return mock.patch.object(jax.lax, "pmean", lambda x, axis_name, **kwargs: jnp.zeros_like(x))
 
 
-# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not
+# all set them).
 _BENCH_FLAGS = {
     "moe_expert_weight_prefetch": "forward",
     "ring_of_experts_row_major_reduce_scatter": True,
diff --git a/tests/unit/moe_router_grad_from_wo_input_test.py b/tests/unit/moe_router_grad_from_wo_input_test.py
index 329e5816b4..17b0ac4d6f 100644
--- a/tests/unit/moe_router_grad_from_wo_input_test.py
+++ b/tests/unit/moe_router_grad_from_wo_input_test.py
@@ -40,11 +40,14 @@
     runs: dh 0.996-1.005, dwo 0.992-1.003, dw 0.941-1.065 (dw, one value per slot, is the noisiest), so 1.10 sits
     above the spread. A backward that rounds the routing weight to bf16 measures 1.31-1.34 on dh and dwo on every
     seed, and one that shrinks the routing-weight gradient by 0.5% measures 3.0-3.3 on dw; both must fail.
+    Both paths share the unsort helpers, so a defect there moves both errors together and leaves the ratio at 1;
+    each path's error, and the forward output's, must therefore also stay within _PRECISION_MAX_ERR = 1e-2 of the
+    float32 reference (measured at most 2.4e-3).
   - The shrunk Qwen3.5 routed MoE of moe_roe_local_routing_test.py (16 experts, top 4, cp-as-ep, ragged sort, local
     routing, two token chunks): output and loss bit-identical to the key off, every gradient within relative L2 1e-2
     (measured <= 4.2e-3): plain and pipelined chunk loop, dropless, padding excluded with "count", saved routing
-    maps, the prefused-weight keys, and the benchmark's MoE settings with every compatible key of this series on
-    against all off (moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp,
+    maps, the prefused-weight keys, and the benchmark's MoE settings with every compatible MoE key on against all
+    off (moe_routing_index_math="count", moe_unsplit_prefused_wi, moe_prefused_swiglu_vjp,
     moe_roe_packed_routing_gather; expert-weight prefetch, row-major reduce-scatter, routing maps on device,
     ragged_buffer_factor 2.75, fp8 fake-quant, padding excluded). A 0.5% error in the routing-weight gradient sits
     inside the layer's 1e-2 tolerance; the precision check above is what catches it. Structure: the compiled
@@ -52,7 +55,8 @@
     longer recomputed in the backward), the pipelined forward keeps the same number of optimization barriers (tie 2
     moves into the op), and every chunk but the last reads the next chunk's routed tokens from the op's tie output.
   - On an expert 4 x tensor 2 mesh the layer runs without the key and raises at trace time when the key reaches it
-    past the validator.
+    past the validator. With moe_accumulate_chunk_wgrad past the validator, the plain and the pipelined chunk loop
+    both raise (the fused op does not return the weights).
   - Deliberately wrong versions must FAIL: the fused backward with the routing weight dropped from the input and
     weight gradients (at the op and in the layer).
   - The validator rejects the key for each requirement it does not meet, naming it.
@@ -92,6 +96,8 @@
 _GRAD_REL_L2 = 1e-2
 # The fused backward's error against a float32 run, over the two-ops backward's error (see the module docstring).
 _PRECISION_RATIO = 1.10
+# The largest relative L2 error either path's output or gradients may have against the float32 run.
+_PRECISION_MAX_ERR = 1e-2
 _PRECISION_SEEDS = range(6)
 
 
@@ -371,7 +377,8 @@ def _op_backward_run(seed, buffer_size, index_math, path):
   """The op (path "fused"), the two ops apart ("two_ops", both bf16) or the dense float32 function ("reference"),
   differentiated at (h, w, wo) with the cotangent g on a 4-device EP mesh. h is an input, so the only roundings are
   the backward's own (and, in the two-ops path, the forward GEMM output's). Returns (dh, dw, dwo), summed or gathered
-  over the shards. The routing weights are in [0.25, 0.3), the bottom of a binade, where bf16 rounds them worst."""
+  over the shards, after each shard's forward output. The routing weights are in [0.25, 0.3), the bottom of a binade,
+  where bf16 rounds them worst."""
   x, ids, _, _, wo, cot, h = _op_inputs(seed, _PRECISION_TOKENS)
   w = jax.random.uniform(jax.random.PRNGKey(1000 + seed), (_PRECISION_TOKENS * _OP_TOPK,), jnp.float32, 0.25, 0.3)
   ep, local = _OP_EP, _OP_EXPERTS // _OP_EP
@@ -416,11 +423,11 @@ def f(h, w, wo):
 
     out, vjp = jax.vjp(f, h[:rows], w, wo)
     dh, dw, dwo = vjp(cot.astype(out.dtype))
-    return dh[None], dw[None], dwo[None]
+    return out[None], dh[None], dw[None], dwo[None]
 
-  sm = jax.shard_map(body, mesh=mesh, in_specs=(P("ep"), P(), P("ep"), P()), out_specs=(P("ep"),) * 3, check_vma=False)
-  dh, dw, dwo = [np.asarray(jnp.asarray(a).astype(jnp.float32)) for a in jax.jit(sm)(h_s, w, wo_s, cot)]
-  return dh, dw.sum(0), dwo.reshape(_OP_EXPERTS, _OP_F, _OP_D)
+  sm = jax.shard_map(body, mesh=mesh, in_specs=(P("ep"), P(), P("ep"), P()), out_specs=(P("ep"),) * 4, check_vma=False)
+  out, dh, dw, dwo = [np.asarray(jnp.asarray(a).astype(jnp.float32)) for a in jax.jit(sm)(h_s, w, wo_s, cot)]
+  return out, dh, dw.sum(0), dwo.reshape(_OP_EXPERTS, _OP_F, _OP_D)
 
 
 def _op_precision_case(label, buffer_size, index_math, control=None):
@@ -432,17 +439,18 @@ def _op_precision_case(label, buffer_size, index_math, control=None):
     with control if control is not None else contextlib.nullcontext():
       fused = _op_backward_run(seed, buffer_size, index_math, "fused")
     parts = []
-    for name, r, a, b in zip(("dh", "dw", "dwo"), reference, two_ops, fused):
+    for name, r, a, b in zip(("out", "dh", "dw", "dwo"), reference, two_ops, fused):
       err_two_ops, err_fused = _rel(a, r), _rel(b, r)
       ratio = err_fused / err_two_ops
       ok = ratio <= _PRECISION_RATIO and bool(np.isfinite(b).all()) and err_two_ops > 0
+      ok &= max(err_two_ops, err_fused) <= _PRECISION_MAX_ERR
       passed &= ok
       ratios.append(ratio)
       parts.append(f"{name} {err_two_ops:.2e}/{err_fused:.2e}={ratio:.3f}{'' if ok else ' FAIL'}")
     lines.append(f"  seed {seed}: error vs float32, two ops / fused = ratio: {'; '.join(parts)}")
   print(
-      f"{label}: {'PASS' if passed else 'FAIL'} (ratio <= {_PRECISION_RATIO}; over the seeds min"
-      f" {min(ratios):.3f} max {max(ratios):.3f})",
+      f"{label}: {'PASS' if passed else 'FAIL'} (ratio <= {_PRECISION_RATIO}, errors <= {_PRECISION_MAX_ERR};"
+      f" over the seeds min {min(ratios):.3f} max {max(ratios):.3f})",
       flush=True,
   )
   print("\n".join(lines), flush=True)
@@ -721,7 +729,8 @@ def _tensor_parallel_guard():
   return ok
 
 
-# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not
+# all set them).
 _BENCH_FLAGS = {
     "moe_chunk_pipeline": True,
     "moe_expert_weight_prefetch": "forward",
@@ -732,7 +741,7 @@ def _tensor_parallel_guard():
     "fp8_moe_fake_quant": True,
     "moe_exclude_padding_from_routing": True,
 }
-# Every key of this series that composes with the others on the megablox path the CPU runs (the fused SiLU GMM is
+# Every MoE performance key that composes with the others on the megablox path the CPU runs (the fused SiLU GMM is
 # tokamax-gmm_v2-only; moe_wi_gmm_fused_swiglu_test.py stacks it with the others using kernel stand-ins).
 _ALL_KEYS = {
     "moe_routing_index_math": "count",
@@ -743,6 +752,27 @@ def _tensor_parallel_guard():
 }
 
 
+def _chunk_wgrad_guard():
+  """With moe_accumulate_chunk_wgrad past the validator, both chunk loops refuse the key at trace time."""
+  results = []
+  for loop, overrides in (("plain", {}), ("pipelined", {"moe_chunk_pipeline": True})):
+    cfg = _cfg(**overrides)
+    mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+    x, forced, token_mask = _inputs(cfg, "real")
+    with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules):
+      _, params, _ = nnx.split(_build(cfg, mesh), nnx.Param, ...)
+    override = _ConfigOverride(cfg, moe_router_grad_from_wo_input=True, moe_accumulate_chunk_wgrad=True)
+    try:
+      _run(override, mesh, params, x, forced, token_mask)
+    except NotImplementedError as e:
+      raised = "does not support moe_accumulate_chunk_wgrad" in str(e)
+    else:
+      raised = False
+    results.append(raised)
+    print(f"chunk_wgrad_guard: [{'ok' if raised else 'FAIL'}] {loop} loop raises at trace time: {raised}", flush=True)
+  return all(results)
+
+
 def main():
   assert jax.device_count() == _NUM_DEVICES, jax.devices()
   pipelined = {"moe_chunk_pipeline": True}
@@ -775,6 +805,7 @@ def main():
           "layer_bench_flags_all_keys", _BENCH_FLAGS, mode="padmask", keys=_ALL_KEYS, same_barriers=False
       ),
       "tensor_parallel_guard": _tensor_parallel_guard(),
+      "chunk_wgrad_guard": _chunk_wgrad_guard(),
   }
   controls = {
       "control_op_weight_dropped": _op_case(
diff --git a/tests/unit/moe_routing_index_math_test.py b/tests/unit/moe_routing_index_math_test.py
index 1d79314ba4..80fa35b372 100644
--- a/tests/unit/moe_routing_index_math_test.py
+++ b/tests/unit/moe_routing_index_math_test.py
@@ -201,6 +201,13 @@ def test_rejected_off_the_ring_ragged_sort(self, overrides, requirement):
   def test_accepted_on_the_ring_ragged_sort(self):
     self.assertEqual(_cfg(moe_routing_index_math="count").moe_routing_index_math, "count")
 
+  def test_names_every_unmet_requirement(self):
+    """The error names all unmet requirements, not only the first."""
+    with self.assertRaises(ValueError) as raised:
+      _cfg(moe_routing_index_math="count", te_moe_block=True, use_ragged_sort=False)
+    self.assertIn("te_moe_block=False", str(raised.exception))
+    self.assertIn("use_ragged_sort=True", str(raised.exception))
+
 
 def _cfg(**overrides):
   """The shrunk Qwen3.5 routed-MoE config on the cp-as-ep mesh (moe_roe_local_routing_test.py's, local routing)."""
@@ -324,8 +331,11 @@ def _count_argsorts(jaxpr):
 
 
 def _run(cfg, mesh, params, x, forced, token_mask):
-  """Loss, output, gradients, the overflow flag, the sort / top_k equation counts of the gradient, and the
-  packed_preprocess flag of every ragged_gather_reduce call traced."""
+  """Runs the layer's value_and_grad; returns (loss, output, gradients), the overflow flag and the counts.
+
+  The counts are the gradient's sort / top_k equations and the packed_preprocess flag of every
+  ragged_gather_reduce call traced.
+  """
   preprocess_flags = []
   real_gather_reduce = ragged_sort.ragged_gather_reduce
 
@@ -454,7 +464,8 @@ def body(tokens, ids):
   return raised
 
 
-# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not
+# all set them).
 _BENCH_FLAGS = {
     "moe_chunk_pipeline": True,
     "moe_expert_weight_prefetch": "forward",
diff --git a/tests/unit/moe_wi_gmm_fused_swiglu_test.py b/tests/unit/moe_wi_gmm_fused_swiglu_test.py
index 9d75356fb6..09076e8853 100644
--- a/tests/unit/moe_wi_gmm_fused_swiglu_test.py
+++ b/tests/unit/moe_wi_gmm_fused_swiglu_test.py
@@ -14,8 +14,8 @@
 """moe_wi_gmm_fused_swiglu on CPU: megablox ops.gmm_silu's own code, its wiring in the MoE layer, and the validator.
 
 With the key, the up projection calls megablox `gmm_silu` (the tokamax gmm_v2 kernel with the SwiGLU in its epilogue)
-instead of `gmm` followed by the activation. The kernels are Pallas TPU kernels and do not run on CPU; their numerics
-are tested on TPU in megablox_gmm_silu_test.py. Two CPU checks cover everything around them:
+instead of `gmm` followed by the activation. The kernels are Pallas TPU kernels and do not run on CPU; their own test,
+megablox_gmm_silu_test.py, is TPU-only (not yet run on TPU). Two CPU checks cover everything around them:
 
 ops.gmm_silu (in-process, `GmmSiluVjpTest`). Only the Pallas entry `gmm_v2.gmm_v2` and gmm's TPU backward
 `ops._gmm_bwd` are replaced: the first by a jnp grouped matmul that honours fuse_act and return_pre_act and records the
@@ -34,7 +34,8 @@
 removal, the whole [w0 | w1] weight, and the skipped activation), not the kernel. The fused kernel's real numerics
 differ by rounding (silu reads the f32 accumulator). Cases: the plain and the pipelined chunk loop, a dropless buffer
 (where each shard's group offset is not 0), padding excluded from routing, the whole prefused wi
-(moe_unsplit_prefused_wi), with the other MoE keys of this series, and with those and the benchmark's MoE settings
+(moe_unsplit_prefused_wi), with it and the other MoE keys (moe_routing_index_math="count",
+moe_roe_packed_routing_gather), and with those and the benchmark's MoE settings
 (expert-weight prefetch, row-major reduce-scatter, routing maps on device, ragged_buffer_factor 2.75, fp8 fake-quant,
 moe_gmm_v2_dlhs_transpose_rhs). Two-sided: the fused path must trace exactly one gmm_silu call per chunk and the plain
 one none; gmm_silu must get the routing operands, the wi tiling the plain path gives gmm, and the DLHS options of the
@@ -186,8 +187,9 @@ class GmmSiluVjpTest(parameterized.TestCase):
       group_offset=(0, 1),
       # (use_dlhs_transpose_rhs, fuse_dlhs_scale, fuse_3d_dlhs_ps): each alone, so a swapped binding shows.
       dlhs_options=((False, False, False), (True, False, False), (False, True, False), (False, False, True)),
+      heuristic=(False, True),
   )
-  def test_custom_vjp(self, group_offset, dlhs_options):
+  def test_custom_vjp(self, group_offset, dlhs_options, heuristic):
     dlhs_transpose, fuse_scale, fuse_3d = dlhs_options
     seen = {"tile_n": [], "return_pre_act": [], "bwd": []}
 
@@ -204,7 +206,7 @@ def kernel_stand_in(
     ):
       del kwargs
       self.assertEqual(fuse_act, "silu")
-      seen["tile_n"].append(tile_info.tile_n)
+      seen["tile_n"].append(tile_info if callable(tile_info) else tile_info.tile_n)
       seen["return_pre_act"].append(return_pre_act)
       acc = _grouped_f32(lhs, rhs, group_sizes, 0 if group_offset is None else group_offset)
       n = acc.shape[-1] // 2
@@ -232,6 +234,7 @@ def gmm_silu(l, r):
           preferred_element_type=jnp.bfloat16,
           tiling=_OP_TILING,
           group_offset=offset,
+          use_gmm_v2_heuristic_tiling=heuristic,
           use_dlhs_transpose_rhs=dlhs_transpose,
           fuse_dlhs_scale=fuse_scale,
           fuse_3d_dlhs_ps=fuse_3d,
@@ -245,8 +248,9 @@ def gmm_silu(l, r):
       _, vjp = jax.vjp(gmm_silu, lhs, rhs)
       dlhs, _ = vjp(cotangent)
 
-    # The forward tile: tiling[2] halved, every call; act alone undifferentiated, act + gate + up differentiated.
-    self.assertEqual(seen["tile_n"], [_OP_TILING[2] // 2] * 2)
+    # The forward tile: tiling[2] halved (or gmm_v2's heuristic), every call; act alone undifferentiated, act + gate +
+    # up differentiated.
+    self.assertEqual(seen["tile_n"], [gmm_v2.calculate_tiling if heuristic else _OP_TILING[2] // 2] * 2)
     self.assertEqual(seen["return_pre_act"], [False, True])
     # The backward: gmm's backward gets [d gate | d up] of silu(gate) * up on the bf16 halves, bit for bit.
     acc = _grouped_f32(lhs, rhs, group_sizes, offset)
@@ -272,7 +276,7 @@ def gmm_silu(l, r):
         "quantization_rule": None,
         "use_tokamax_backend": True,
         "use_gmm_v2": True,
-        "use_gmm_v2_heuristic_tiling": False,
+        "use_gmm_v2_heuristic_tiling": heuristic,
         "return_rhs": False,
         "return_lhs": False,
         "fuse_dlhs_scale": fuse_scale,
@@ -494,7 +498,8 @@ def _ragged_dropless_matches_unsorted(label, ragged_gmm=None):
   return passed
 
 
-# The MoE settings of the Qwen3.5-397B RL trainer benchmark that the cases above do not all set.
+# The subset of the Qwen3.5-397B RL trainer benchmark's MoE settings that this file needs (the cases above do not
+# all set them).
 _BENCH_FLAGS = {
     "moe_chunk_pipeline": True,
     "moe_expert_weight_prefetch": "forward",

From 20c877aa154d2f21d97e62e871169ed8fb430b1f Mon Sep 17 00:00:00 2001
From: Param Bole <parambole@google.com>
Date: Thu, 15 Oct 2026 06:51:47 -0700
Subject: [PATCH 10/10] GDN: build the forward kernel's Conv1D output from
 dense per-token views (gdn_conv1d_dense_token_stack)

The Pallas GDN forward kernel (kernels/gdn/wrapper.py fused_conv1d_gdn)
convolves a chunk one token at a time and stacks the per-token
[1, 1, dim] float32 rows with jnp.stack(axis=1)
(kernels/gdn/compute_conv1d.py causal_conv1d). The stacked
[1, chunk, 1, dim] value has a second-minor dimension of 1, and the TPU
compiler of the October 2026 Pathways build (and of the libtpu 0.0.51
nightlies) lays it out in (8, 128) tiles padded from 1 to 8 sublanes.
The SiLU applied next inherits that layout, so its exp and reciprocal
run on 8x as many vector registers as the data fills. Older compilers
keep the value dense.

With gdn_conv1d_dense_token_stack=True each row is viewed as
[dim / 1024, 8, 128], which is free because the row already fills
dim / 1024 float32 registers; the views are concatenated along the
leading axis and the result is reshaped to [1, chunk, 1, dim]
(compute_conv1d.stack_conv1d_tokens). The same values in the same
order; the compiler keeps the value dense on old and new builds alike.

The option reaches the kernel through gdn_decoupled_conv1d's new static
argument conv_dense_token_stack. Its primal and forward rule pass it to
every forward-kernel call: the local call, both sequence-sharded
context-parallel passes, and the forward rule's replay under remat (the
backward rule never runs the forward kernel). fused_conv1d_gdn hands it
to the per-sequence kernel only: its batched (decode) kernel, which
convolves one token per sequence, keeps jnp.stack for those one-token
rows, and the tokamax Op wrapper (kernels/gdn/pallas_mosaic_tpu.py)
does not take the option. MaxText's GDN layer runs neither, since
gdn_decoupled_conv1d calls fused_conv1d_gdn prefill-only. The option is
left out of the kernel name and metadata, so outside the kernel body
the lowered program is unchanged and profiles keep the kernel name.

Numerics: every element lands where jnp.stack puts it, so the kernel's
outputs do not change. On CPU, causal_conv1d and the SiLU after it give
the same bits with and without the option (the kernel itself does not
run on CPU). The TPU numerics A/B of the measured build below passed
all six gates (nxs6gr).

The validator requires a decoder with Gated DeltaNet layers (qwen3_next
or qwen3_5), use_gdn_kernel=True, and a width that each forward-kernel
call convolves (2 * key heads * key head dim + value heads * value head
dim; one rank's heads under head-sharded GDN context parallelism) that
is a multiple of 1024. Like the existing sequence-length check, it
counts a context parallelism left at -1 (filled in when the mesh is
built) as 1; fused_conv1d_gdn checks the per-call width again when the
layer is traced, which catches such a config. Where the width is not a
multiple of 1024 the measured build below silently kept jnp.stack; this
change rejects the config instead. Default off.

Measured on Qwen3.5-397B-A17B, 128 TPU v7x, on the RL trainer benchmark
(#5701 driver, mesh fsdp 32 x context 4 x expert 2, 9 micro-batches per
step, ragged_buffer_factor 2.75, staged step boundary), with the same
change behind an environment switch (MAXTEXT_GDN_CONV_DENSE_STACK=1) on
the stack-6 base: 85.731 s per step, regular micro-batch 9.3552 s
(o275x-s6gr-l3-prof), against 86.926 s / 9.4910 s for the same-pod arm
o275x-s6n2-l3-prof: -1.195 s per step. That arm's only change, a 2D GDN
out_proj, is a no-op (XLA already canonicalises the 4D out_proj into the
same 2D dot; TensorCore busy time -2 ms per micro-batch). The stack-6
control in another pod measured 86.970 s / 9.494 s (o275x-s6c-l3).

Equivalence to the measured build, on CPU: lowered for a TPU7x target
(no TPU needed), the GDN layer's forward and backward program with the
key on is the measured build's program with its switch on, and with the
key off is the base program (StableHLO and every Mosaic kernel body,
compared without source locations). 18 configurations: no context
parallelism, head-sharded context parallelism 2, sequence-sharded
context parallelism 4 (also at the 397B's 16 key and 64 value heads),
with and without packing and remat, and gdn_conv1d_input_in_qkv_dtype
off and on. In each, the key changes only the forward-kernel bodies.
Outside that layer the two differ: the measured build's switch also
regrouped the batched kernel's one-token rows (at a decode tile of 1)
and reached the tokamax Op wrapper.

Tests (gdn_conv1d_dense_token_stack_test.py, new), on CPU:
- stack_conv1d_tokens(dense=True) equals jnp.stack(axis=1) bit for bit,
  and dense=False is jnp.stack itself (the same jaxpr); causal_conv1d
  and the SiLU after it give the same bits with the option at widths
  1024 and 12288, packed or not, with and without conv bias, for full
  and partial chunks; a token-minor regrouping is caught.
- Traced and lowered for a TPU7x target on the CPU host, every
  forward-kernel call assembles its conv output densely exactly when
  the option is on: fused_conv1d_gdn; the custom VJP's primal, forward
  rule and remat replay; both sequence-sharded CP passes, under the
  forward rule and, forward only, under the primal; head-sharded CP;
  and model_runner reading the config key. The batched (decode) kernel
  is the same program with and without the option. The Mosaic lowering
  keeps the kernel name and changes only the kernel body. The option
  composes with gdn_conv1d_input_in_qkv_dtype, and the backward rule
  accepts every positional form. An option the conv ignores, one the
  forward rule drops and one the primal drops each fail every check
  they break.
- The width check and the validator reject the option where it would
  do nothing, including one rank's width under head-sharded CP. The
  validator's copy of model_runner's head-vs-sequence CP rule agrees
  with model_runner at every GDN CP mode, degrees 1-8 on each context
  axis, nine head-count pairs and head dims 128 and 256. The kernel
  name and metadata do not change.
---
 src/maxtext/configs/base.yml                  |   6 +
 src/maxtext/configs/types.py                  |  63 ++
 src/maxtext/kernels/gdn/compute_conv1d.py     |  34 +-
 src/maxtext/kernels/gdn/config.py             |  13 +-
 src/maxtext/kernels/gdn/gdn_bwd/api.py        |  50 +-
 src/maxtext/kernels/gdn/model_runner.py       |   5 +
 src/maxtext/kernels/gdn/wrapper.py            |  22 +-
 .../unit/gdn_conv1d_dense_token_stack_test.py | 922 ++++++++++++++++++
 8 files changed, 1101 insertions(+), 14 deletions(-)
 create mode 100644 tests/unit/gdn_conv1d_dense_token_stack_test.py

diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 8a3ad1f762..bba0085bab 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1643,6 +1643,12 @@ gdn_cp_matmul_precision: "highest"
 # and the qkv, conv-state and bias gradients are bit-identical on CPU, the conv weight gradient matches up to float32
 # summation order. Requires use_gdn_kernel: true.
 gdn_conv1d_input_in_qkv_dtype: false
+# GDN kernel path: build the Pallas GDN forward kernel's Conv1D output from dense per-token [conv_dim / 1024, 8, 128]
+# views instead of jnp.stack (kernels/gdn/compute_conv1d.py). Same values; the October 2026 TPU compiler lays the
+# stacked value out with 7 of 8 sublanes as padding, so the SiLU after it runs on 8x the vector registers. Requires
+# use_gdn_kernel: true and a GDN Conv1D width (2 * key heads * key head dim + value heads * value head dim) that is
+# a multiple of 1024.
+gdn_conv1d_dense_token_stack: false
 
 use_tokamax_splash: false
 # Setting this flag will use a non-pallas implementation.
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index b95c616ed1..8f9afcd803 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2253,6 +2253,24 @@ class Qwen3Next(BaseModel):
           " runs that Conv1D)."
       ),
   )
+  gdn_conv1d_dense_token_stack: bool = Field(
+      False,
+      description=(
+          "GDN kernel path: in the Pallas GDN forward kernel, build each chunk's Conv1D output from per-token"
+          " [conv_dim / 1024, 8, 128] float32 views concatenated along the leading axis, instead of jnp.stack(axis=1)"
+          " of the per-token [1, 1, conv_dim] rows (kernels/gdn/compute_conv1d.py stack_conv1d_tokens). The value is"
+          " the same, element for element; only the layout the TPU compiler gives it changes. The stacked value's"
+          " second-minor dimension is 1, and the compiler of the October 2026 Pathways build lays it out in (8, 128)"
+          " tiles padded from 1 to 8 sublanes, so the SiLU that follows runs its exp and reciprocal on 8x as many"
+          " vector registers as the data fills; the dense views keep the value in dense tiles, as older compilers"
+          " already do for the stack. Applies to every forward kernel call of gdn_decoupled_conv1d and its forward"
+          " rule (both sequence-sharded context-parallel passes, and the forward rule's replay under remat); the"
+          " backward kernel is unchanged. Requires use_gdn_kernel=True (the only path that runs that kernel) and a"
+          " Conv1D width, 2 * gdn_num_key_heads * gdn_key_head_dim + gdn_num_value_heads * gdn_value_head_dim (one"
+          " rank's heads under head-sharded GDN context parallelism), that is a multiple of 1024: one float32 vector"
+          " register of 8 x 128."
+      ),
+  )
 
 
 # ----------------------------------------------------------------------------
@@ -5294,7 +5312,13 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None:
     """
     if is_rl_config:
       MaxTextConfig.validate_gdn_remat_requires_kernel(cfg)
+    dense_token_stack = getattr(cfg, "gdn_conv1d_dense_token_stack", False)
     if cfg.decoder_block not in (DecoderBlockType.QWEN3_NEXT, DecoderBlockType.QWEN3_5):
+      if dense_token_stack:
+        raise ValueError(
+            "gdn_conv1d_dense_token_stack=True requires a decoder with Gated DeltaNet layers (qwen3_next or"
+            f" qwen3_5), got decoder_block={cfg.decoder_block.value}."
+        )
       return
     cp_size = math.prod(
         max(1, getattr(cfg, f"{mesh}_{axis}_parallelism"))
@@ -5315,6 +5339,11 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None:
             "gdn_conv1d_input_in_qkv_dtype=True requires use_gdn_kernel=True: only the GDN kernel path runs the"
             " Conv1D + SiLU (kernels/gdn/gdn_bwd/compute_conv1d_bwd.py) whose input it keeps in qkv's dtype."
         )
+      if dense_token_stack:
+        raise ValueError(
+            "gdn_conv1d_dense_token_stack=True requires use_gdn_kernel=True: only the Pallas GDN forward kernel"
+            " runs the in-VMEM Conv1D (kernels/gdn/compute_conv1d.py) whose output it assembles."
+        )
       return
     if cfg.ici_tensor_parallelism > 1 or cfg.dcn_tensor_parallelism > 1:
       raise ValueError(
@@ -5341,6 +5370,40 @@ def validate_gdn_config(cfg: Any, *, is_rl_config: bool) -> None:
             f"degree ({cp_size}) times gdn_chunk_size ({cfg.gdn_chunk_size}), so that every per-rank sequence "
             "is a whole number of kernel chunks."
         )
+    if dense_token_stack:
+      MaxTextConfig.validate_gdn_dense_token_stack_width(cfg, cp_size)
+
+  @staticmethod
+  def validate_gdn_dense_token_stack_width(cfg: Any, cp_size: int) -> None:
+    """Fails if gdn_conv1d_dense_token_stack cannot regroup the Conv1D rows the GDN forward kernel convolves.
+
+    The option regroups each token's float32 Conv1D row in whole vector registers of 8 x 128
+    (kernels/gdn/compute_conv1d.py DENSE_TOKEN_STACK_ALIGNMENT), so the width each forward-kernel call convolves must
+    be a multiple of 1024: the whole Conv1D, or one rank's heads under head-sharded GDN context parallelism.
+
+    `cp_size` counts a context parallelism left at -1 (filled in when the mesh is built) as 1, as validate_gdn_config's
+    sequence-length check does. If the filled-in degree head-shards the heads to a width that is not a multiple of
+    1024, wrapper.fused_conv1d_gdn's own width check raises when the layer is traced.
+    """
+    num_k_heads, num_v_heads = cfg.gdn_num_key_heads, cfg.gdn_num_value_heads
+    # Head-sharded as model_runner._resolve_gdn_cp_mode decides it for training: gdn_cp_mode 'head', or 'auto' at a
+    # degree of 2 that divides the key heads. (Head counts the degree does not divide fail there on their own.)
+    cp_mode = getattr(cfg, "gdn_cp_mode", "auto")
+    head_sharded = cp_size > 1 and (
+        cp_mode == "head" or (cp_mode == "auto" and cp_size <= min(2, num_k_heads) and num_k_heads % cp_size == 0)
+    )
+    ranks = cp_size if head_sharded and num_k_heads % cp_size == 0 and num_v_heads % cp_size == 0 else 1
+    width = 2 * (num_k_heads // ranks) * cfg.gdn_key_head_dim + (num_v_heads // ranks) * cfg.gdn_value_head_dim
+    if width % 1024:
+      if ranks > 1:
+        where = f"one rank's heads under head-sharded GDN context parallelism of degree {cp_size}"
+      else:
+        where = "the whole Conv1D"
+      raise ValueError(
+          f"gdn_conv1d_dense_token_stack=True requires the width each GDN forward kernel convolves ({where}: 2 * key"
+          f" heads * gdn_key_head_dim + value heads * gdn_value_head_dim = {width}) to be a multiple of 1024: it"
+          " regroups each token's float32 row in whole vector registers of 8 x 128."
+      )
 
   @model_validator(mode="after")
   def set_derived_and_validate_values(self) -> "MaxTextConfig":
diff --git a/src/maxtext/kernels/gdn/compute_conv1d.py b/src/maxtext/kernels/gdn/compute_conv1d.py
index 05e40e3c69..bdcba5efb6 100644
--- a/src/maxtext/kernels/gdn/compute_conv1d.py
+++ b/src/maxtext/kernels/gdn/compute_conv1d.py
@@ -20,6 +20,38 @@
 
 from . import config
 
+# One float32 vector register is 8 sublanes x 128 lanes. gdn_conv1d_dense_token_stack regroups causal_conv1d's
+# per-token rows in units of one register (see stack_conv1d_tokens), so it needs rows of a multiple of this size.
+_F32_VREG_SUBLANES = 8
+_VREG_LANES = 128
+DENSE_TOKEN_STACK_ALIGNMENT = _F32_VREG_SUBLANES * _VREG_LANES
+
+
+def stack_conv1d_tokens(out_list: list[jax.Array], dense: bool) -> jax.Array:
+  """Stacks causal_conv1d's per-token rows, each [seq, 1, dim], into its [seq, chunk, 1, dim] output.
+
+  dense=False is jnp.stack(out_list, axis=1). The stacked value's second-minor dimension is 1, and the TPU compiler
+  of the October 2026 Pathways build (and of the libtpu 0.0.51 nightlies) lays it out in (8, 128) tiles padded from
+  1 to 8 sublanes. The SiLU that wrapper.inner_kernel applies next inherits that layout, so its exp and reciprocal
+  run on 8x as many vector registers as the data fills. Older compilers keep the value in dense (1, 128) tiles.
+
+  dense=True (config gdn_conv1d_dense_token_stack) views each [1, 1, dim] float32 row as [dim / 1024, 8, 128], which
+  is free because the row already fills dim / 1024 registers, concatenates the views along the leading axis and
+  reshapes the result to [1, chunk, 1, dim]. The values and their order are the same; only the layout the compiler
+  chooses changes, and it stays dense on old and new compilers alike. Requires seq == 1 (GDNMode.PER_SEQ), float32
+  rows and dim a multiple of DENSE_TOKEN_STACK_ALIGNMENT.
+  """
+  if not dense:
+    return jnp.stack(out_list, axis=1)
+  row = out_list[0]
+  if row.shape[0] != 1 or row.size % DENSE_TOKEN_STACK_ALIGNMENT or jnp.dtype(row.dtype).itemsize != 4:
+    raise ValueError(
+        "gdn_conv1d_dense_token_stack needs one sequence per tile and float32 rows whose size is a multiple of"
+        f" {DENSE_TOKEN_STACK_ALIGNMENT}, got a {jnp.dtype(row.dtype)} row of shape {row.shape}."
+    )
+  dense_rows = jnp.concatenate([o.reshape(-1, _F32_VREG_SUBLANES, _VREG_LANES) for o in out_list], axis=0)
+  return dense_rows.reshape(row.shape[0], len(out_list), *row.shape[1:])
+
 
 def encode_segment_ids(
     segment_ids: jax.Array,
@@ -296,4 +328,4 @@ def causal_conv1d(
         new_conv_state,
     )
 
-  return jnp.stack(out_list, axis=1), new_conv_state
+  return stack_conv1d_tokens(out_list, cfg.conv_dense_token_stack), new_conv_state
diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py
index 72c15644e7..a61ccf964c 100644
--- a/src/maxtext/kernels/gdn/config.py
+++ b/src/maxtext/kernels/gdn/config.py
@@ -33,6 +33,12 @@
     "emit_t_inv": True,
 }
 
+# GDNConfig fields that change how the kernel lays out a value it computes either way, not what it computes or
+# returns. They are left out of the kernel name and metadata, so the lowered program outside the kernel body is the
+# same with and without them, and a profile shows the same kernel name for both variants (the run's config tells
+# them apart).
+_LAYOUT_ONLY_FIELDS = ("conv_dense_token_stack",)
+
 
 class GDNMode(enum.StrEnum):
   """Execution mode for GDN kernel."""
@@ -92,6 +98,9 @@ class GDNConfig:
   tri_inv_algo: str = "xor_dnc"
   head_group_size: int | None = None
   emit_t_inv: bool = True
+  # Build the Conv1D output from dense per-token views (compute_conv1d.stack_conv1d_tokens; config
+  # gdn_conv1d_dense_token_stack). GDNMode.PER_SEQ only.
+  conv_dense_token_stack: bool = False
 
   @property
   def chunk_size(self) -> int:
@@ -136,11 +145,13 @@ def get_kernel_name(self) -> str:
     return name
 
   def get_metadata(self) -> dict[str, str | int | float]:
-    """Flattened config fields for the kernel's metadata (knobs at their defaults are left out)."""
+    """Flattened config fields for the kernel's metadata (knobs at their defaults and layout-only fields are left out)."""
     cfgs_dict = dataclasses.asdict(self)
     for field_name, default in _TRI_INV_FIELD_DEFAULTS.items():
       if cfgs_dict.get(field_name) == default:
         cfgs_dict.pop(field_name)
+    for field_name in _LAYOUT_ONLY_FIELDS:
+      cfgs_dict.pop(field_name)
     ret = {}
     for path, val in jax.tree_util.tree_leaves_with_path(cfgs_dict):
       key = jax.tree_util.keystr(path, simple=True, separator=".")
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index c061acbc3e..15733ef79a 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -148,8 +148,13 @@ def _run_local_gdn_decoupled_fwd(
     segment_ids: Optional[jax.Array] = None,
     conv_halo_seg: Optional[jax.Array] = None,
     init_seg: Optional[jax.Array] = None,
+    conv_dense_token_stack: bool = False,
 ) -> Tuple[Tuple[jax.Array, Tuple[jax.Array, jax.Array]], Optional[jax.Array], Optional[jax.Array],]:
-  """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU."""
+  """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU.
+
+  `conv_dense_token_stack` goes to the Pallas forward kernel (wrapper.fused_conv1d_gdn); the CPU fallback has no
+  in-VMEM Conv1D and does not use it.
+  """
   batch_size, seq_len, dim_size = qkv.shape
   if seq_len % chunk_size != 0:
     raise ValueError(
@@ -273,6 +278,7 @@ def _run_local_gdn_decoupled_fwd(
       tri_inv_algo=local_compute_gdn.DEFAULT_TRI_INV_ALGO,
       head_group_size=head_group_size,
       emit_t_inv=emit_t_inv,
+      conv_dense_token_stack=conv_dense_token_stack,
   )
 
   core_attn_out = core_attn_out_flat.reshape(batch_size, seq_len, num_v_heads, head_v_dim)
@@ -348,8 +354,9 @@ def _run_cp_gdn_decoupled_fwd_impl(
     segment_ids: Optional[jax.Array] = None,
     seg_metadata: Optional[Tuple[Optional[jax.Array], Optional[jax.Array], Optional[jax.Array]]] = None,
     cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION,
+    conv_dense_token_stack: bool = False,
 ):
-  """Runs 2-pass sequence-sharded CP forward for GDN."""
+  """Runs 2-pass sequence-sharded CP forward for GDN. Both passes get `conv_dense_token_stack`."""
   batch_size = qkv.shape[0]
   if seg_metadata is not None:
     s_enc_local, conv_halo_seg, init_seg = seg_metadata
@@ -401,6 +408,7 @@ def _run_cp_gdn_decoupled_fwd_impl(
       segment_ids=s_enc_local,
       conv_halo_seg=conv_halo_seg,
       init_seg=init_seg,
+      conv_dense_token_stack=conv_dense_token_stack,
   )
 
   # Convolve only the K channel slice (1/6th of qkv) needed by compose_local_from_t_inv
@@ -462,6 +470,7 @@ def _run_cp_gdn_decoupled_fwd_impl(
       segment_ids=s_enc_local,
       conv_halo_seg=conv_halo_seg,
       init_seg=init_seg,
+      conv_dense_token_stack=conv_dense_token_stack,
   )
   if s_enc_local is not None:
     # Owner = last rank with a valid local token (rank 0 if none). A rank whose
@@ -489,7 +498,7 @@ def scoped(*args, **kwargs):
   return scoped
 
 
-@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20))
+@functools.partial(jax.custom_vjp, nondiff_argnums=(9, 10, 11, 12, 13, 14, 15, 16, 17, 19, 20, 21))
 @_conv_input_in_qkv_dtype_scoped
 def gdn_decoupled_conv1d(
     qkv: jax.Array,
@@ -513,6 +522,7 @@ def gdn_decoupled_conv1d(
     segment_ids: Optional[jax.Array] = None,
     cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION,
     conv_input_in_qkv_dtype: bool = False,
+    conv_dense_token_stack: bool = False,
 ) -> Tuple[jax.Array, Tuple[jax.Array, jax.Array]]:
   """Decoupled Conv1D + GDN with Pallas backward pass and optional sequence-sharded CP.
 
@@ -523,6 +533,11 @@ def gdn_decoupled_conv1d(
   traced by this op and its forward and backward rules keep the halo-extended conv input in qkv's dtype instead of
   materializing it in float32 (compute_conv1d_bwd.py). Same values (the conv weight gradient up to float32
   summation order); only HBM traffic changes.
+
+  `conv_dense_token_stack` (config gdn_conv1d_dense_token_stack) makes every Pallas forward kernel this op and its
+  forward rule run (both sequence-sharded CP passes, and the forward rule's replay under remat) build its Conv1D
+  output from dense per-token views (compute_conv1d.stack_conv1d_tokens). Same values; only the layout the TPU
+  compiler gives that value changes. The backward rule never runs the forward kernel.
   """
   if _is_cp_active(cp_axis_name):
     (out, states), _, _, _, _, _ = _run_cp_gdn_decoupled_fwd_impl(
@@ -546,6 +561,7 @@ def gdn_decoupled_conv1d(
         cp_axis_name=cp_axis_name,
         segment_ids=segment_ids,
         cp_matmul_precision=cp_matmul_precision,
+        conv_dense_token_stack=conv_dense_token_stack,
     )
     return out, states
 
@@ -573,6 +589,7 @@ def gdn_decoupled_conv1d(
       segment_ids=segment_ids,
       conv_halo_seg=conv_halo_seg,
       init_seg=init_seg,
+      conv_dense_token_stack=conv_dense_token_stack,
   )
   return out, states
 
@@ -610,6 +627,7 @@ def _gdn_decoupled_conv1d_fwd(
     segment_ids: Optional[jax.Array] = None,
     cp_matmul_precision: jax.lax.Precision = cp_gdn.DEFAULT_PRECISION,
     conv_input_in_qkv_dtype: bool = False,
+    conv_dense_token_stack: bool = False,
 ):
   """Forward rule for custom_vjp registration of gdn_decoupled_conv1d."""
   (
@@ -667,6 +685,7 @@ def _gdn_decoupled_conv1d_fwd(
         segment_ids=segment_ids,
         seg_metadata=seg_metadata,
         cp_matmul_precision=cp_matmul_precision,
+        conv_dense_token_stack=conv_dense_token_stack,
     )
     out = checkpoint_name(out, "gdn_core_attn_out")
     residuals = (
@@ -710,6 +729,7 @@ def _gdn_decoupled_conv1d_fwd(
       segment_ids=segment_ids,
       conv_halo_seg=conv_halo_seg,
       init_seg=init_seg,
+      conv_dense_token_stack=conv_dense_token_stack,
   )
   out = checkpoint_name(out, "gdn_core_attn_out")
   seg_bundle = (segment_ids, conv_halo_seg, init_seg) if segment_ids is not None else None
@@ -742,14 +762,16 @@ def _gdn_decoupled_conv1d_bwd(
     cp_axis_name: str | tuple[str, ...] | None = None,
     cp_matmul_precision: Any = cp_gdn.DEFAULT_PRECISION,
     conv_input_in_qkv_dtype: bool = False,
+    conv_dense_token_stack: bool = False,
     residuals: tuple[Any, ...] | None = None,
     cotangents: tuple[Any, ...] | None = None,
 ):
   """Backward rule for custom_vjp registration of gdn_decoupled_conv1d."""
   # Support positional calls that omit the trailing nondiff arguments: 10 args (no cp_axis_name,
-  # no cp_matmul_precision, no conv_input_in_qkv_dtype), 11 args (no cp_matmul_precision, no
-  # conv_input_in_qkv_dtype) or 12 args (no conv_input_in_qkv_dtype). jax.custom_vjp itself always
-  # passes all 11 nondiff arguments.
+  # no cp_matmul_precision, no conv_input_in_qkv_dtype, no conv_dense_token_stack), 11 args (no
+  # cp_matmul_precision, no conv_input_in_qkv_dtype, no conv_dense_token_stack), 12 args (no
+  # conv_input_in_qkv_dtype, no conv_dense_token_stack) or 13 args (no conv_dense_token_stack).
+  # jax.custom_vjp itself always passes all 12 nondiff arguments.
   if residuals is None and cotangents is None and isinstance(cp_axis_name, tuple) and len(cp_axis_name) >= 10:
     residuals, cotangents = cp_axis_name, cp_matmul_precision
     cp_axis_name, cp_matmul_precision = None, cp_gdn.DEFAULT_PRECISION
@@ -761,9 +783,19 @@ def _gdn_decoupled_conv1d_bwd(
   ):
     residuals, cotangents = cp_matmul_precision, conv_input_in_qkv_dtype
     cp_matmul_precision, conv_input_in_qkv_dtype = cp_gdn.DEFAULT_PRECISION, False
-  elif cotangents is None and isinstance(conv_input_in_qkv_dtype, tuple) and len(conv_input_in_qkv_dtype) >= 10:
-    residuals, cotangents = conv_input_in_qkv_dtype, residuals
-    conv_input_in_qkv_dtype = False
+  elif (
+      residuals is None
+      and cotangents is None
+      and isinstance(conv_input_in_qkv_dtype, tuple)
+      and len(conv_input_in_qkv_dtype) >= 10
+  ):
+    residuals, cotangents = conv_input_in_qkv_dtype, conv_dense_token_stack
+    conv_input_in_qkv_dtype, conv_dense_token_stack = False, False
+  elif cotangents is None and isinstance(conv_dense_token_stack, tuple) and len(conv_dense_token_stack) >= 10:
+    residuals, cotangents = conv_dense_token_stack, residuals
+    conv_dense_token_stack = False
+  # The backward rule recomputes the Conv1D with JAX ops; it never runs the Pallas forward kernel.
+  del conv_dense_token_stack
 
   with compute_conv1d_bwd.conv_input_in_qkv_dtype_scope(conv_input_in_qkv_dtype):
     return _gdn_decoupled_conv1d_bwd_impl(
diff --git a/src/maxtext/kernels/gdn/model_runner.py b/src/maxtext/kernels/gdn/model_runner.py
index a6cce0e185..384916933f 100644
--- a/src/maxtext/kernels/gdn/model_runner.py
+++ b/src/maxtext/kernels/gdn/model_runner.py
@@ -619,6 +619,8 @@ def run_gdn_kernel_layer(
   cp_matmul_precision = jax.lax.Precision[getattr(cfg, "gdn_cp_matmul_precision", "highest").upper()]
   # Conv1D halo input in qkv's dtype instead of a float32 copy (compute_conv1d_bwd.py). Same values.
   conv_input_in_qkv_dtype = cfg.gdn_conv1d_input_in_qkv_dtype
+  # Forward kernel's Conv1D output built from dense per-token views (compute_conv1d.py). Same values.
+  conv_dense_token_stack = bool(getattr(cfg, "gdn_conv1d_dense_token_stack", False))
 
   if use_head_sharded_cp:
     qkv = jnp.concatenate([query, key, value_raw], axis=3)
@@ -890,6 +892,7 @@ def shard_mapped_gdn(  # pylint: disable=too-many-positional-arguments
             compute_dtype=state_dtype,
             segment_ids=seg_val,
             conv_input_in_qkv_dtype=conv_input_in_qkv_dtype,
+            conv_dense_token_stack=conv_dense_token_stack,
         )
 
         if next_cs is not None:
@@ -930,6 +933,7 @@ def shard_mapped_gdn(  # pylint: disable=too-many-positional-arguments
           segment_ids=seg_val,
           cp_matmul_precision=cp_matmul_precision,
           conv_input_in_qkv_dtype=conv_input_in_qkv_dtype,
+          conv_dense_token_stack=conv_dense_token_stack,
       )
 
     if not isinstance(qkv, jax.core.Tracer):
@@ -973,6 +977,7 @@ def shard_mapped_gdn(  # pylint: disable=too-many-positional-arguments
         compute_dtype=state_dtype,
         segment_ids=decoder_segment_ids,
         conv_input_in_qkv_dtype=conv_input_in_qkv_dtype,
+        conv_dense_token_stack=conv_dense_token_stack,
     )
 
   core_attn_out = checkpoint_name(core_attn_out, "gdn_core_attn_out")
diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py
index 686eac6a4d..952aac211a 100644
--- a/src/maxtext/kernels/gdn/wrapper.py
+++ b/src/maxtext/kernels/gdn/wrapper.py
@@ -371,6 +371,7 @@ def _run(allocations):
         "tri_inv_algo",
         "head_group_size",
         "emit_t_inv",
+        "conv_dense_token_stack",
     ),
 )
 def fused_conv1d_gdn(
@@ -405,12 +406,20 @@ def fused_conv1d_gdn(
     tri_inv_algo: str = compute_gdn.DEFAULT_TRI_INV_ALGO,
     head_group_size: int | None = None,
     emit_t_inv: bool = True,
+    conv_dense_token_stack: bool = False,
 ) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array | None, jax.Array]:
   """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states).
 
   When `head_group_size` is None, `compute_gdn.default_head_group_size` selects an
   appropriate KQ-head group size for `chunk_size > 64`. With `emit_t_inv=False`,
   the returned `t_inv` is None.
+
+  `conv_dense_token_stack` (config gdn_conv1d_dense_token_stack) makes the per-sequence kernel build its Conv1D
+  output from dense per-token views instead of jnp.stack (compute_conv1d.stack_conv1d_tokens): the same values, in
+  a layout the TPU compiler keeps dense. It requires `dim` to be a multiple of
+  compute_conv1d.DENSE_TOKEN_STACK_ALIGNMENT. The batched (decode) kernel, which convolves one token per sequence,
+  does not take the option and keeps jnp.stack for its one-token rows; MaxText's GDN layer never runs that kernel,
+  since gdn_bwd/api.py calls this function prefill-only.
   """
   act_in_dtype = qkv.dtype
   act_out_dtype = qkv.dtype
@@ -435,6 +444,12 @@ def fused_conv1d_gdn(
   assert query_start_loc.shape == (num_seqs + 1,)
   assert state_indices.shape == (num_seqs,)
   assert distribution.shape == (3,)
+  if conv_dense_token_stack and dim % compute_conv1d.DENSE_TOKEN_STACK_ALIGNMENT:
+    raise ValueError(
+        f"conv_dense_token_stack (gdn_conv1d_dense_token_stack) requires the Conv1D width ({dim}) to be a multiple of"
+        f" {compute_conv1d.DENSE_TOKEN_STACK_ALIGNMENT}; with head-sharded GDN context parallelism this is the width"
+        " of one rank's heads."
+    )
 
   num_lanes = pltpu.get_tpu_info().num_lanes
   packing = 4 // act_in_dtype.itemsize
@@ -530,7 +545,7 @@ def call_kernel(
   ) -> tuple[jax.Array, ...]:
     if mode == config.GDNMode.BATCHED:
       tile_size = decode_tile_size
-      tri_inv_kwargs = {}
+      per_seq_kwargs = {}
     else:
       tile_size = mixed_tile_size
       resolved_head_group_size = (
@@ -538,10 +553,11 @@ def call_kernel(
           if head_group_size is None
           else head_group_size
       )
-      tri_inv_kwargs = {
+      per_seq_kwargs = {
           "tri_inv_algo": compute_gdn.validate_tri_inv_algo(tri_inv_algo),
           "head_group_size": resolved_head_group_size,
           "emit_t_inv": emit_t_inv,
+          "conv_dense_token_stack": conv_dense_token_stack,
       }
 
     cfg = config.GDNConfig(
@@ -563,7 +579,7 @@ def call_kernel(
             recurrent_state=in_recurrent_state.dtype,
             conv_state=in_conv_state.dtype,
         ),
-        **tri_inv_kwargs,
+        **per_seq_kwargs,
     )
 
     if mode == config.GDNMode.BATCHED:
diff --git a/tests/unit/gdn_conv1d_dense_token_stack_test.py b/tests/unit/gdn_conv1d_dense_token_stack_test.py
new file mode 100644
index 0000000000..da223e6194
--- /dev/null
+++ b/tests/unit/gdn_conv1d_dense_token_stack_test.py
@@ -0,0 +1,922 @@
+# 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.
+
+"""gdn_conv1d_dense_token_stack: the Pallas GDN forward kernel's Conv1D output built from dense per-token views.
+
+The option changes how compute_conv1d.causal_conv1d assembles a value the kernel computes either way: each token's
+float32 row is viewed as [dim / 1024, 8, 128] and the views are concatenated, instead of jnp.stack(axis=1) of the
+rows. Every element lands where jnp.stack puts it; only the layout the TPU compiler gives the value changes.
+
+The forward kernel cannot run on CPU (neither Pallas interpreter supports its emit_pipeline DMAs), so:
+  - Values (in this process): causal_conv1d, the one function whose code changes, gives the same bits with the
+    option as without it at the kernel's in-VMEM shapes (widths 1024 and 12288, packed segment ids on/off, conv
+    bias on/off, full and partial chunks), and so does the SiLU the kernel applies next. A token-minor regrouping
+    is caught, so the comparison can fail. Without the option the rows go through jnp.stack itself (the same
+    jaxpr), so the default program is the base program.
+  - Programs, traced for a TPU target on this CPU host with no TPU or libtpu (runtime_utils.target_platform follows
+    the abstract mesh). In this process: fused_conv1d_gdn hands the option to its per-sequence kernel only, and its
+    batched (decode) kernel is the same program with and without it. In a child process (see main): every
+    forward-kernel call carries the dense regrouping exactly when the option is on: wrapper.fused_conv1d_gdn;
+    gdn_decoupled_conv1d's primal, its forward rule and the rule's replay under remat; both sequence-sharded
+    context-parallel passes, under the forward rule and (forward only) under the primal; head-sharded context
+    parallelism; and model_runner.run_gdn_kernel_layer reading the config key. The option composes with
+    gdn_conv1d_input_in_qkv_dtype, and the backward rule accepts every positional form. Deliberately wrong
+    versions must FAIL every check they break: an option the conv ignores, one the forward rule drops, and one the
+    primal drops.
+  - Guards: the per-call width check and the config validator reject the option where it would do nothing, and the
+    validator's copy of model_runner's head-vs-sequence context-parallel rule agrees with model_runner.
+"""
+
+# pylint: disable=protected-access
+
+import collections
+import functools
+import itertools
+import os
+import re
+import subprocess
+import sys
+import types
+
+from absl.testing import absltest
+from absl.testing import parameterized
+import jax
+from jax.extend import core as jax_core
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.common.common_types import ShardMode
+from maxtext.configs import pyconfig
+from maxtext.configs import types as config_types
+from maxtext.kernels.gdn import compute_conv1d
+from maxtext.kernels.gdn import config as gdn_config
+from maxtext.kernels.gdn import model_runner
+from maxtext.kernels.gdn import wrapper
+from maxtext.kernels.gdn.gdn_bwd import api as gdn_api
+from maxtext.kernels.gdn.gdn_bwd import cp_gdn
+from tests.utils.test_helpers import get_test_config_path
+
+_PASS_TOKEN = "GDN_CONV1D_DENSE_TOKEN_STACK_CHECKS_PASSED"
+_CHUNK = 64
+_KERNEL = 4
+_HEAD_DIM = 128
+
+
+def _bits(x):
+  """The array's bytes, so equality means bit-identical (NaN-safe, -0.0 != 0.0)."""
+  return np.asarray(jax.device_get(x)).tobytes()
+
+
+def _gdn_config(dim, num_v_heads, has_seg_ids, dense):
+  return gdn_config.GDNConfig(
+      mode=gdn_config.GDNMode.PER_SEQ,
+      dtypes=gdn_config.Dtypes(jnp.bfloat16, jnp.bfloat16, jnp.float32, jnp.float32, jnp.float32),
+      batch_size=_CHUNK,
+      dim_size=dim,
+      kernel_size=_KERNEL,
+      tile_size=_CHUNK,
+      num_kq_heads=num_v_heads // 2,
+      num_v_heads=num_v_heads,
+      kq_head_dim=_HEAD_DIM,
+      v_head_dim=_HEAD_DIM,
+      has_seg_ids=has_seg_ids,
+      conv_dense_token_stack=dense,
+  )
+
+
+def _conv_inputs(dim, num_v_heads, seed):
+  """causal_conv1d's in-VMEM operands for one PER_SEQ tile: lhs = [halo | chunk tokens], weights, bias, b lanes."""
+  rng = np.random.default_rng(seed)
+  lhs = jnp.asarray(rng.standard_normal((1, _CHUNK + _KERNEL - 1, 1, dim)).astype(np.float32))
+  weight = jnp.asarray((rng.standard_normal((_KERNEL, 1, dim)) * 0.5).astype(np.float32))
+  bias = jnp.asarray(rng.standard_normal((dim,)).astype(np.float32))
+  # b's lanes num_v_heads (signed active segment of each token: two documents, then padding) and num_v_heads + 1
+  # (chunk header: the halo's segment ids, then the previous chunk's segment).
+  b_vreg = np.zeros((1, _CHUNK, 1, 128), np.float32)
+  b_vreg[0, :, 0, num_v_heads] = np.concatenate([np.full(21, 1.0), np.full(30, 2.0), np.full(_CHUNK - 51, -2.0)])
+  b_vreg[0, :_KERNEL, 0, num_v_heads + 1] = 1.0
+  return lhs, weight, bias, jnp.asarray(b_vreg)
+
+
+def _stack_counts(jaxpr, dim):
+  """(dense, stacked): how many times `jaxpr` (sub-jaxprs included) assembles a [1, chunk, 1, dim] Conv1D output from
+  dense views (a concatenate to [chunk * dim / 1024, 8, 128]) or by stacking (a stack, or a concatenate, whose result
+  is [1, chunk, 1, dim]; jax 0.11 traces jnp.stack as its own primitive)."""
+  dense_shape, stack_shape = (_CHUNK * dim // 1024, 8, 128), (1, _CHUNK, 1, dim)
+  dense = stacked = 0
+  for eqn in jaxpr.eqns:
+    if eqn.primitive.name in ("concatenate", "stack"):
+      shape = tuple(eqn.outvars[0].aval.shape)
+      dense += eqn.primitive.name == "concatenate" and shape == dense_shape
+      stacked += shape == stack_shape
+    for sub in jax_core.jaxprs_in_params(eqn.params):
+      sub_dense, sub_stacked = _stack_counts(sub, dim)
+      dense, stacked = dense + sub_dense, stacked + sub_stacked
+  return dense, stacked
+
+
+def _forward_kernels(closed_jaxpr):
+  """(name, kernel jaxpr) of every forward-kernel call (a pallas_call named fused_conv1d_gdn*), in trace order."""
+  found = []
+
+  def visit(jaxpr):
+    for eqn in jaxpr.eqns:
+      if eqn.primitive.name == "pallas_call" and str(eqn.params.get("name", "")).startswith("fused_conv1d_gdn"):
+        found.append((str(eqn.params["name"]), eqn.params["jaxpr"]))
+        continue
+      for sub in jax_core.jaxprs_in_params(eqn.params):
+        visit(sub)
+
+  visit(closed_jaxpr.jaxpr)
+  return found
+
+
+def _token_minor_stack(out_list):
+  """Deliberately wrong regrouping: interleaves the tokens inside each group of 8 x 128 registers."""
+  row = out_list[0]
+  views = jnp.stack([o.reshape(-1, 8, 128) for o in out_list], axis=1)
+  return views.reshape(row.shape[0], len(out_list), *row.shape[1:])
+
+
+class StackConv1dTokensTest(parameterized.TestCase):
+  """stack_conv1d_tokens(dense=True) equals jnp.stack(axis=1) bit for bit, and rejects rows it cannot regroup."""
+
+  @parameterized.named_parameters(
+      ("chunk64_width1024", 64, 1024),
+      ("chunk64_width12288", 64, 12288),
+      ("chunk3_width2048", 3, 2048),
+  )
+  def test_dense_equals_stack(self, num_tokens, dim):
+    rng = np.random.default_rng(num_tokens + dim)
+    rows = [jnp.asarray(rng.standard_normal((1, 1, dim)).astype(np.float32)) for _ in range(num_tokens)]
+    ref = jnp.stack(rows, axis=1)
+    for got in (
+        compute_conv1d.stack_conv1d_tokens(rows, True),
+        jax.jit(compute_conv1d.stack_conv1d_tokens, static_argnums=1)(rows, True),
+    ):
+      self.assertEqual(got.shape, ref.shape)
+      self.assertEqual(got.dtype, ref.dtype)
+      self.assertEqual(_bits(got), _bits(ref))
+    self.assertEqual(_bits(compute_conv1d.stack_conv1d_tokens(rows, False)), _bits(ref))
+    # Positive control: a regrouping that moves elements is caught by the same comparison. (With one register per
+    # token, width 1024, the token-minor interleave is the same order, so it only moves elements at wider rows.)
+    if dim > 1024:
+      self.assertNotEqual(_bits(_token_minor_stack(rows)), _bits(ref))
+
+  def test_default_is_jnp_stack(self):
+    """dense=False traces to jnp.stack(axis=1) itself, so with the option off causal_conv1d is the base program.
+
+    Equal values or shapes would not show this: a rewrite of the default path that keeps the values (a concatenate of
+    expanded rows, say) is a different program.
+    """
+    rows = [jnp.zeros((1, 1, 2048), jnp.float32) for _ in range(_CHUNK)]
+    ref = str(jax.make_jaxpr(lambda r: jnp.stack(r, axis=1))(rows))
+    self.assertEqual(str(jax.make_jaxpr(lambda r: compute_conv1d.stack_conv1d_tokens(r, False))(rows)), ref)
+    # Positive control: the comparison sees a program change that keeps the values.
+    self.assertNotEqual(str(jax.make_jaxpr(lambda r: compute_conv1d.stack_conv1d_tokens(r, True))(rows)), ref)
+
+  @parameterized.named_parameters(
+      ("two_sequences", (2, 1, 1024), jnp.float32),
+      ("partial_register", (1, 1, 1536), jnp.float32),
+      ("bfloat16_rows", (1, 1, 2048), jnp.bfloat16),
+  )
+  def test_rejects_rows_it_cannot_regroup(self, shape, dtype):
+    rows = [jnp.zeros(shape, dtype)] * 4
+    with self.assertRaisesRegex(ValueError, "one sequence per tile and float32 rows"):
+      compute_conv1d.stack_conv1d_tokens(rows, True)
+    self.assertEqual(compute_conv1d.stack_conv1d_tokens(rows, False).shape, (shape[0], 4) + shape[1:])
+
+
+_CONV_CASES = [
+    {
+        "testcase_name": f"width{dim}_seg{int(seg)}_bias{int(bias)}_real{real}",
+        "dim": dim,
+        "num_v_heads": nv,
+        "has_seg_ids": seg,
+        "with_bias": bias,
+        "real": real,
+    }
+    for dim, nv in ((1024, 4), (12288, 64))
+    for seg in (False, True)
+    for bias in (False, True)
+    for real in (_CHUNK, 37)
+]
+
+
+class CausalConv1dTest(parameterized.TestCase):
+  """causal_conv1d with the option gives the same bits; and the option changes the traced conv."""
+
+  @parameterized.named_parameters(*_CONV_CASES)
+  def test_bit_identical(self, dim, num_v_heads, has_seg_ids, with_bias, real):
+    lhs, weight, bias, b_vreg = _conv_inputs(dim, num_v_heads, seed=dim + 2 * has_seg_ids + with_bias)
+
+    def run(dense):
+      out, state = compute_conv1d.causal_conv1d(
+          real_sizes=jnp.array([real], jnp.int32),
+          lhs=lhs,
+          conv_weight=weight,
+          conv_bias=bias if with_bias else None,
+          cfg=_gdn_config(dim, num_v_heads, has_seg_ids, dense),
+          b_vreg=b_vreg if has_seg_ids else None,
+      )
+      return out, state, jax.nn.silu(out)  # inner_kernel applies the SiLU to the conv output next
+
+    for ref, got in zip(jax.jit(lambda: run(False))(), jax.jit(lambda: run(True))()):
+      self.assertEqual(ref.shape, got.shape)
+      self.assertEqual(ref.dtype, got.dtype)
+      self.assertEqual(_bits(ref), _bits(got))
+
+  def test_option_changes_the_traced_conv(self):
+    """Positive control for the value test: the option must change the program, not only keep the values."""
+    dim, num_v_heads = 2048, 8
+    lhs, weight, bias, _ = _conv_inputs(dim, num_v_heads, seed=0)
+
+    def stacks(dense):
+      cfg = _gdn_config(dim, num_v_heads, False, dense)
+      jaxpr = jax.make_jaxpr(lambda: compute_conv1d.causal_conv1d(jnp.array([_CHUNK]), lhs, weight, bias, cfg))()
+      return _stack_counts(jaxpr.jaxpr, dim)
+
+    self.assertEqual(stacks(True), (1, 0))
+    self.assertEqual(stacks(False), (0, 1))
+
+
+class KernelIdentityTest(absltest.TestCase):
+  """The option is left out of the kernel name and metadata, so the program outside the kernel body is unchanged."""
+
+  def test_name_and_metadata_unchanged(self):
+    off, on = _gdn_config(1024, 4, True, False), _gdn_config(1024, 4, True, True)
+    self.assertEqual(off.get_kernel_name(), on.get_kernel_name())
+    self.assertEqual(off.get_metadata(), on.get_metadata())
+    self.assertNotIn("conv_dense_token_stack", off.get_metadata())
+    self.assertFalse(gdn_config.GDNConfig.__dataclass_fields__["conv_dense_token_stack"].default)
+
+
+def _tpu_target(shape=(), axes=()):
+  """A context in which traces target TPU7x as in an AOT compile: no TPU and no libtpu needed."""
+  device = jax.sharding.AbstractDevice(device_kind="TPU7x", num_cores=1, platform="tpu")
+  mesh = jax.sharding.AbstractMesh(
+      shape, axes, axis_types=(jax.sharding.AxisType.Auto,) * len(axes), abstract_device=device
+  )
+  return jax.sharding.use_abstract_mesh(mesh), mesh
+
+
+def _fused_args(dim, num_v_heads, seq=128):
+  a = jax.ShapeDtypeStruct
+  return (
+      a((seq, dim), jnp.bfloat16),
+      a((seq, num_v_heads), jnp.bfloat16),
+      a((seq, num_v_heads), jnp.bfloat16),
+      a((2, _KERNEL - 1, dim), jnp.bfloat16),
+      a((2, num_v_heads, _HEAD_DIM, _HEAD_DIM), jnp.float32),
+      a((dim, 1, _KERNEL), jnp.float32),
+      None,
+      a((num_v_heads,), jnp.float32),
+      a((num_v_heads,), jnp.float32),
+      a((2,), jnp.int32),
+      a((1,), jnp.int32),
+      a((3,), jnp.int32),
+      a((1,), jnp.int32),
+  )
+
+
+def _trace_fused(num_k_heads, num_v_heads, dense, *, prefill_only=True):
+  """wrapper.fused_conv1d_gdn traced for a TPU7x target.
+
+  Prefill-only (as gdn_decoupled_conv1d calls it) runs the per-sequence kernel alone; prefill_only=False (mixed
+  mode) runs the batched (decode) kernel first, at a decode tile of 1 because _fused_args has one sequence.
+  """
+  dim = (2 * num_k_heads + num_v_heads) * _HEAD_DIM
+  ctx, _ = _tpu_target()
+  fn = functools.partial(
+      wrapper.fused_conv1d_gdn,
+      n_kq=num_k_heads,
+      n_v=num_v_heads,
+      d_k=_HEAD_DIM,
+      d_v=_HEAD_DIM,
+      kernel_size=_KERNEL,
+      compute_precision=jnp.dtype(jnp.float32),
+      mixed_tile_size=_CHUNK,
+      is_prefill_only=prefill_only,
+      conv_dense_token_stack=dense,
+  )
+  with ctx:
+    return jax.make_jaxpr(fn)(*_fused_args(dim, num_v_heads))
+
+
+class WidthCheckTest(absltest.TestCase):
+  """fused_conv1d_gdn rejects the option for a Conv1D width that is not a whole number of float32 registers."""
+
+  def test_rejects_width_512(self):
+    _trace_fused(1, 2, False)  # width 512; the default path is unchanged
+    with self.assertRaisesRegex(ValueError, r"requires the Conv1D width \(512\) to be a multiple of 1024"):
+      _trace_fused(1, 2, True)
+
+  def test_accepts_width_1024(self):
+    _trace_fused(2, 4, True)
+
+
+class BatchedKernelTest(absltest.TestCase):
+  """fused_conv1d_gdn hands the option to its per-sequence kernel only; the batched (decode) kernel keeps jnp.stack."""
+
+  def test_batched_kernel_does_not_take_the_option(self):
+    # Mixed mode runs the batched kernel, then the per-sequence kernel. With one sequence the batched kernel's tile
+    # holds one sequence, so its one-token rows are the only batched rows the dense regrouping could take.
+    kernels = {dense: _forward_kernels(_trace_fused(2, 4, dense, prefill_only=False)) for dense in (False, True)}
+    names = [name for name, _ in kernels[True]]
+    self.assertEqual(names, [name for name, _ in kernels[False]])
+    self.assertLen(names, 2)
+    self.assertTrue(names[0].startswith("fused_conv1d_gdn_batched_b1_"), names)
+    self.assertTrue(names[1].startswith("fused_conv1d_gdn_per_seq_"), names)
+    (_, batched_off), (_, per_seq_off) = kernels[False]
+    (_, batched_on), (_, per_seq_on) = kernels[True]
+    self.assertEqual(str(batched_on), str(batched_off))
+    self.assertEqual(_stack_counts(per_seq_on, 1024), (1, 0))
+    self.assertEqual(_stack_counts(per_seq_off, 1024), (0, 1))
+    # Positive control: the same comparison tells the per-sequence kernel's two programs apart.
+    self.assertNotEqual(str(per_seq_on), str(per_seq_off))
+
+
+def _validator_cfg(**overrides):
+  """The fields validate_gdn_config reads, for a Qwen3.5-397B-shaped GDN kernel path that passes it."""
+  fields = {
+      "decoder_block": config_types.DecoderBlockType.QWEN3_5,
+      "ici_context_parallelism": 1,
+      "dcn_context_parallelism": 1,
+      "ici_context_usp_ulysses_parallelism": 1,
+      "dcn_context_usp_ulysses_parallelism": 1,
+      "use_gdn_kernel": True,
+      "ici_tensor_parallelism": 1,
+      "dcn_tensor_parallelism": 1,
+      "gdn_chunk_size": 64,
+      "gdn_key_head_dim": 128,
+      "gdn_value_head_dim": 128,
+      "gdn_num_key_heads": 16,
+      "gdn_num_value_heads": 64,
+      "use_qk_norm_in_gdn": True,
+      "max_target_length": 4096,
+      "gdn_conv1d_input_in_qkv_dtype": False,
+      "gdn_conv1d_dense_token_stack": False,
+      "gdn_cp_mode": "auto",
+  }
+  fields.update(overrides)
+  return types.SimpleNamespace(**fields)
+
+
+class ValidatorTest(parameterized.TestCase):
+  """The validator rejects the key where it would do nothing; the default config is untouched."""
+
+  @parameterized.named_parameters(
+      ("no_cp", {}),
+      ("head_cp4", {"ici_context_parallelism": 4, "gdn_cp_mode": "head"}),  # 3072 per rank
+      ("auto_cp2_head_sharded", {"ici_context_parallelism": 2}),  # 6144 per rank
+      ("auto_cp4_seq_sharded", {"ici_context_parallelism": 4, "gdn_num_key_heads": 2, "gdn_num_value_heads": 4}),
+      # 2048: model_runner runs auto CP 4 sequence-sharded even with 4 key heads (head-sharded would be 512 per rank).
+      (
+          "auto_cp4_seq_sharded_four_key_heads",
+          {"ici_context_parallelism": 4, "gdn_num_key_heads": 4, "gdn_num_value_heads": 8},
+      ),
+  )
+  def test_accepted(self, overrides):
+    cfg = _validator_cfg(gdn_conv1d_dense_token_stack=True, **overrides)
+    config_types.MaxTextConfig.validate_gdn_config(cfg, is_rl_config=False)
+
+  def test_head_sharding_agrees_with_model_runner(self):
+    """The validator's copy of model_runner._resolve_gdn_cp_mode's head-vs-sequence rule agrees with it.
+
+    configs/types.py cannot import the kernel module, so the validator repeats the rule that decides whether one
+    forward-kernel call convolves every head or one rank's heads. Over every GDN context-parallel mode, degrees 1-8 on
+    each context axis, and head counts and head dims that model_runner runs, the validator must accept the key exactly
+    when the width one forward-kernel call convolves, as model_runner shards it, is a multiple of 1024.
+    """
+    seen = collections.Counter()
+    for cp_mode, cp, axis, (nk, nv), head_dim in itertools.product(
+        ("auto", "head", "seq"),
+        range(1, 9),
+        ("ici_context_parallelism", "ici_context_usp_ulysses_parallelism", "dcn_context_parallelism"),
+        ((1, 2), (2, 4), (3, 6), (4, 8), (6, 12), (8, 16), (16, 32), (16, 64), (32, 64)),
+        (128, 256),
+    ):
+      fields = {
+          axis: cp,
+          "gdn_cp_mode": cp_mode,
+          "gdn_num_key_heads": nk,
+          "gdn_num_value_heads": nv,
+          "gdn_key_head_dim": head_dim,
+          "gdn_value_head_dim": head_dim,
+          "max_target_length": 840 * _CHUNK,  # a whole number of chunks per rank at every degree 1-8
+      }
+      try:
+        head_sharded, _ = model_runner._resolve_gdn_cp_mode(_validator_cfg(**fields), cp, nk)
+      except ValueError:
+        continue  # model_runner refuses this config itself
+      ranks = cp if head_sharded else 1
+      width = 2 * (nk // ranks) * head_dim + (nv // ranks) * head_dim
+      try:
+        config_types.MaxTextConfig.validate_gdn_config(
+            _validator_cfg(gdn_conv1d_dense_token_stack=True, **fields), is_rl_config=False
+        )
+        accepted = True
+      except ValueError as e:
+        self.assertIn("gdn_conv1d_dense_token_stack=True requires the width", str(e))
+        accepted = False
+      self.assertEqual(accepted, width % 1024 == 0, f"{fields}: model_runner head_sharded={head_sharded}")
+      seen[(head_sharded, accepted)] += 1
+    # Head- and sequence-sharded configs, each both accepted and rejected: the comparison can fail either way.
+    self.assertLen(seen, 4, seen)
+
+  @parameterized.named_parameters(
+      ("no_gdn_layers", {"decoder_block": config_types.DecoderBlockType.LLAMA2}, "a decoder with Gated DeltaNet"),
+      ("no_gdn_kernel", {"use_gdn_kernel": False}, "use_gdn_kernel=True"),
+      ("width_512", {"gdn_num_key_heads": 1, "gdn_num_value_heads": 2}, r"the width .* \(the whole Conv1D: .* = 512\)"),
+      (
+          "head_cp8_rank_width_1536",
+          {"ici_context_parallelism": 8, "gdn_cp_mode": "head"},
+          r"the width .* \(one rank's heads under head-sharded GDN context parallelism of degree 8: .* = 1536\)",
+      ),
+      (
+          "auto_cp2_rank_width_512",
+          {"ici_context_parallelism": 2, "gdn_num_key_heads": 2, "gdn_num_value_heads": 4},
+          r"the width .* \(one rank's heads under head-sharded GDN context parallelism of degree 2: .* = 512\)",
+      ),
+  )
+  def test_rejected(self, overrides, message):
+    cfg = _validator_cfg(gdn_conv1d_dense_token_stack=True, **overrides)
+    with self.assertRaisesRegex(ValueError, f"gdn_conv1d_dense_token_stack=True requires {message}"):
+      config_types.MaxTextConfig.validate_gdn_config(cfg, is_rl_config=False)
+    # The same config with the key off passes.
+    config_types.MaxTextConfig.validate_gdn_config(_validator_cfg(**overrides), is_rl_config=False)
+
+  def test_pyconfig(self):
+    def cfg(**kwargs):
+      return pyconfig.initialize(
+          [sys.argv[0], get_test_config_path()],
+          run_name="gdn_conv1d_dense_token_stack",
+          enable_checkpointing=False,
+          log_config=False,
+          skip_jax_distributed_system=True,
+          **kwargs,
+      )
+
+    self.assertFalse(cfg().gdn_conv1d_dense_token_stack)
+    self.assertFalse(config_types.Qwen3Next().gdn_conv1d_dense_token_stack)
+    on = cfg(model_name="qwen3.5-397b-a17b", use_gdn_kernel=True, gdn_conv1d_dense_token_stack=True)
+    self.assertTrue(on.gdn_conv1d_dense_token_stack)
+    with self.assertRaisesRegex(ValueError, "requires a decoder with Gated DeltaNet layers"):
+      cfg(gdn_conv1d_dense_token_stack=True)
+
+
+@pytest.mark.cpu_only
+def test_programs_in_a_child_process():
+  """Runs `main` below in a fresh CPU process (the program checks trace TPU kernels on this host)."""
+  env = os.environ.copy()
+  env["JAX_PLATFORMS"] = "cpu"
+  repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
+  env["PYTHONPATH"] = os.pathsep.join([repo_root, env["PYTHONPATH"]]) if env.get("PYTHONPATH") else repo_root
+  result = subprocess.run([sys.executable, __file__], env=env, capture_output=True, text=True, check=False)
+  print(result.stdout)
+  assert result.returncode == 0, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+  assert _PASS_TOKEN in result.stdout, f"stdout:\n{result.stdout}\nstderr:\n{result.stderr}"
+
+
+# ----------------------------------------------------------------------------------------------------------------------
+# Program checks (main). Each forward-kernel call is a pallas_call named fused_conv1d_gdn*; _stack_counts reads which
+# way its kernel jaxpr assembles the Conv1D output.
+
+
+def _forward_kernel_stacks(closed_jaxpr, dim):
+  """[_stack_counts] of every fused_conv1d_gdn kernel in the jaxpr, in trace order."""
+  return [_stack_counts(kernel, dim) for _, kernel in _forward_kernels(closed_jaxpr)]
+
+
+def _bf16_conv_inputs(closed_jaxpr, shape):
+  """bf16 pad / concatenate results of `shape` outside the Pallas kernels (gdn_conv1d_input_in_qkv_dtype's marker)."""
+  found = 0
+
+  def visit(jaxpr):
+    nonlocal found
+    for eqn in jaxpr.eqns:
+      if eqn.primitive.name == "pallas_call":
+        continue
+      if eqn.primitive.name in ("pad", "concatenate"):
+        aval = eqn.outvars[0].aval
+        found += tuple(aval.shape) == tuple(shape) and aval.dtype == jnp.bfloat16
+      for sub in jax_core.jaxprs_in_params(eqn.params):
+        visit(sub)
+
+  visit(closed_jaxpr.jaxpr)
+  return found
+
+
+class _Value:
+  """A parameter stand-in: .value and [...] like an nnx Param."""
+
+  def __init__(self, value):
+    self.value = value
+
+  def __getitem__(self, idx):
+    return self.value[idx]
+
+
+def _trace_model_runner(mode, nk, nv, *, dense, qkv_dtype=False, remat=True, seg=True, seq=512, grad=True):
+  """model_runner.run_gdn_kernel_layer on a stand-in GDN layer, traced for a TPU7x target.
+
+  grad=True traces value_and_grad, so the custom VJP's forward rule makes the forward-kernel calls; grad=False traces
+  the forward pass alone (as a forward-only RL pass runs it), so its primal makes them.
+  """
+  cp = {"none": 1, "seq4": 4, "head2": 2}[mode]
+  if mode == "none":
+    ctx, mesh = _tpu_target()
+    layer_mesh = None
+  else:
+    ctx, mesh = _tpu_target((cp,), ("context",))
+    layer_mesh = mesh
+  cfg = types.SimpleNamespace(
+      ici_context_parallelism=cp,
+      dcn_context_parallelism=1,
+      ici_context_usp_ulysses_parallelism=1,
+      dcn_context_usp_ulysses_parallelism=1,
+      gdn_cp_mode={"head2": "head", "seq4": "seq"}.get(mode, "auto"),
+      gdn_cp_matmul_precision="highest",
+      shard_mode=ShardMode.AUTO,
+      logical_axis_rules=[],
+      gdn_conv_kernel_dim=_KERNEL,
+      gdn_chunk_size=_CHUNK,
+      use_qk_norm_in_gdn=True,
+      context_parallel_load_balance=False,
+      gdn_conv1d_input_in_qkv_dtype=qkv_dtype,
+      gdn_conv1d_dense_token_stack=dense,
+  )
+  key_dim, conv_dim = nk * _HEAD_DIM, (2 * nk + nv) * _HEAD_DIM
+
+  def loss(params, x, aux):
+    def body(params, x):
+      layer = types.SimpleNamespace(
+          config=cfg,
+          mesh=layer_mesh,
+          num_k_heads=nk,
+          num_v_heads=nv,
+          head_k_dim=_HEAD_DIM,
+          head_v_dim=_HEAD_DIM,
+          key_dim=key_dim,
+          v_heads_per_k_head=nv // nk,
+          conv1d=types.SimpleNamespace(kernel=_Value(params["conv_w"]), bias=None),
+          A_log=_Value(params["a_log"]),
+          dt_bias=_Value(params["dt_bias"]),
+      )
+      out, next_cs, next_rs = model_runner.run_gdn_kernel_layer(
+          layer=layer,
+          query=x["query"],
+          key=x["key"],
+          value=x["value_raw"].reshape(1, seq, nv, _HEAD_DIM),
+          value_raw=x["value_raw"],
+          b=x["b"],
+          a=x["a"],
+          active_cache=None,
+          decoder_segment_ids=aux["seg"] if seg else None,
+          model_mode="train",
+          state_dtype=jnp.float32,
+      )
+      return (
+          jnp.sum(out.astype(jnp.float32) * aux["do"])
+          + 0.5 * jnp.sum(next_cs.astype(jnp.float32))
+          + 0.25 * jnp.sum(next_rs.astype(jnp.float32))
+      )
+
+    if remat:
+      body = jax.checkpoint(body, policy=jax.checkpoint_policies.nothing_saveable)
+    return body(params, x)
+
+  a = jax.ShapeDtypeStruct
+  params = {
+      "conv_w": a((_KERNEL, 1, conv_dim), jnp.float32),
+      "a_log": a((nv,), jnp.float32),
+      "dt_bias": a((nv,), jnp.float32),
+  }
+  x = {
+      "query": a((1, seq, nk, _HEAD_DIM), jnp.bfloat16),
+      "key": a((1, seq, nk, _HEAD_DIM), jnp.bfloat16),
+      "value_raw": a((1, seq, nk, (nv // nk) * _HEAD_DIM), jnp.bfloat16),
+      "b": a((1, seq, nv), jnp.bfloat16),
+      "a": a((1, seq, nv), jnp.bfloat16),
+  }
+  aux = {"seg": a((1, seq), jnp.int32), "do": a((1, seq, nv, _HEAD_DIM), jnp.float32)}
+  fn = jax.value_and_grad(loss, argnums=(0, 1)) if grad else loss
+  with ctx:
+    return jax.make_jaxpr(fn)(params, x, aux)
+
+
+def _trace_decoupled(dense, *, grad, remat, qkv_dtype=False, nk=2, nv=4, seq=128):
+  """gdn_decoupled_conv1d for a TPU7x target: the primal alone, or value_and_grad (forward rule), optionally remat."""
+  dim = (2 * nk + nv) * _HEAD_DIM
+  op = functools.partial(
+      gdn_api.gdn_decoupled_conv1d,
+      num_k_heads=nk,
+      num_v_heads=nv,
+      head_k_dim=_HEAD_DIM,
+      head_v_dim=_HEAD_DIM,
+      conv_kernel_size=_KERNEL,
+      chunk_size=_CHUNK,
+      use_qk_norm_in_gdn=True,
+      compute_dtype=jnp.float32,
+      conv_input_in_qkv_dtype=qkv_dtype,
+      conv_dense_token_stack=dense,
+  )
+
+  def loss(qkv, b, a, cw, alog, dtb, do):
+    def body(qkv, b, a, cw, alog, dtb):
+      out, _ = op(
+          qkv=qkv,
+          b=b,
+          a=a,
+          conv_weight=cw,
+          conv_bias=None,
+          a_log=alog,
+          dt_bias=dtb,
+          conv_state=None,
+          recurrent_state=None,
+      )
+      return jnp.sum(out.astype(jnp.float32) * do)
+
+    if remat:
+      body = jax.checkpoint(body, policy=jax.checkpoint_policies.nothing_saveable)
+    return body(qkv, b, a, cw, alog, dtb)
+
+  a = jax.ShapeDtypeStruct
+  args = (
+      a((1, seq, dim), jnp.bfloat16),
+      a((1, seq, nv), jnp.bfloat16),
+      a((1, seq, nv), jnp.bfloat16),
+      a((_KERNEL, 1, dim), jnp.float32),
+      a((nv,), jnp.float32),
+      a((nv,), jnp.float32),
+      a((1, seq, nv, _HEAD_DIM), jnp.float32),
+  )
+  fn = jax.value_and_grad(loss, argnums=tuple(range(6))) if grad else loss
+  ctx, _ = _tpu_target()
+  with ctx:
+    return jax.make_jaxpr(fn)(*args), dim
+
+
+def _check(label, ok, detail=""):
+  print(f"{label}: [{'ok' if ok else 'FAIL'}] {detail}", flush=True)
+  return ok
+
+
+def _all_dense(stacks, expected_calls):
+  return len(stacks) == expected_calls and all(s == (1, 0) for s in stacks)
+
+
+def _all_stacked(stacks, expected_calls):
+  return len(stacks) == expected_calls and all(s == (0, 1) for s in stacks)
+
+
+# The cases of _option_reaches_every_forward_kernel whose forward-kernel calls go through the custom VJP's forward
+# rule (traced under value_and_grad) or through its primal (traced forward only). A deliberately wrong version in main
+# that drops the option in one of the two must break every case in its group.
+_FORWARD_RULE_CASES = (
+    "decoupled_forward_rule",
+    "decoupled_forward_rule_remat",
+    "runner_no_cp",
+    "runner_seq_cp4",
+    "runner_head_cp2",
+    "runner_seq_cp4_qwen35_397b_heads",
+)
+_PRIMAL_CASES = ("decoupled_primal", "runner_seq_cp4_forward_only")
+
+
+def _option_reaches_every_forward_kernel(names=None):
+  """Every forward-kernel call carries the dense regrouping iff the option is on; returns (ok, {case: (ok, detail)}).
+
+  `names` runs only those cases (the deliberately wrong versions in main use it).
+  """
+  results = {}
+
+  def want(name):
+    return names is None or name in names
+
+  # wrapper.fused_conv1d_gdn: the per-sequence kernel (the only one a prefill-only call runs).
+  for nk, nv in ((2, 4), (16, 64)):
+    dim = (2 * nk + nv) * _HEAD_DIM
+    name = f"fused_conv1d_gdn_width{dim}"
+    if want(name):
+      on = _forward_kernel_stacks(_trace_fused(nk, nv, True), dim)
+      off = _forward_kernel_stacks(_trace_fused(nk, nv, False), dim)
+      results[name] = (_all_dense(on, 1) and _all_stacked(off, 1), f"on={on} off={off}")
+  # gdn_decoupled_conv1d: the primal (1 call), the forward rule (1), the forward rule and its remat replay (2).
+  for name, grad, remat, calls in (
+      ("decoupled_primal", False, False, 1),
+      ("decoupled_forward_rule", True, False, 1),
+      ("decoupled_forward_rule_remat", True, True, 2),
+  ):
+    if want(name):
+      jaxpr_on, dim = _trace_decoupled(True, grad=grad, remat=remat)
+      jaxpr_off, _ = _trace_decoupled(False, grad=grad, remat=remat)
+      on, off = _forward_kernel_stacks(jaxpr_on, dim), _forward_kernel_stacks(jaxpr_off, dim)
+      results[name] = (_all_dense(on, calls) and _all_stacked(off, calls), f"on={on} off={off}")
+  # model_runner reading the config key, under remat. With value_and_grad (the custom VJP's forward rule): no context
+  # parallelism (1 call + replay), sequence-sharded CP 4 (2 passes + their replays), head-sharded CP 2 (1 call per
+  # rank + replay), and the production head counts (16 key, 64 value heads) under sequence-sharded CP 4. Forward
+  # only (the custom VJP's primal): sequence-sharded CP 4 (2 passes, no replay).
+  for name, mode, nk, nv, seq, grad, calls in (
+      ("runner_no_cp", "none", 4, 8, 512, True, 2),
+      ("runner_seq_cp4", "seq4", 4, 8, 512, True, 4),
+      ("runner_head_cp2", "head2", 4, 8, 512, True, 2),
+      ("runner_seq_cp4_qwen35_397b_heads", "seq4", 16, 64, 256, True, 4),
+      ("runner_seq_cp4_forward_only", "seq4", 4, 8, 512, False, 2),
+  ):
+    if want(name):
+      dim = (2 * nk + nv) * _HEAD_DIM // (2 if mode == "head2" else 1)
+      on = _forward_kernel_stacks(_trace_model_runner(mode, nk, nv, dense=True, seq=seq, grad=grad), dim)
+      off = _forward_kernel_stacks(_trace_model_runner(mode, nk, nv, dense=False, seq=seq, grad=grad), dim)
+      results[name] = (_all_dense(on, calls) and _all_stacked(off, calls), f"on={on} off={off}")
+  return all(ok for ok, _ in results.values()), results
+
+
+def _composes_with_input_in_qkv_dtype():
+  """All four combinations: each key's own marker is present iff that key is on."""
+  ok = True
+  for qkv_dtype in (False, True):
+    for dense in (False, True):
+      jaxpr, dim = _trace_decoupled(dense, grad=True, remat=False, qkv_dtype=qkv_dtype)
+      stacks = _forward_kernel_stacks(jaxpr, dim)
+      narrow = _bf16_conv_inputs(jaxpr, (1, 128 + _KERNEL - 1, dim))
+      good = (_all_dense(stacks, 1) if dense else _all_stacked(stacks, 1)) and (narrow > 0) == qkv_dtype
+      ok &= _check(
+          f"compose qkv_dtype={qkv_dtype} dense={dense}", good, f"forward kernels={stacks} bf16 conv inputs={narrow}"
+      )
+  return ok
+
+
+def _backward_rule_positional_forms():
+  """The backward rule traces the same program from every positional form custom_vjp or a caller may use."""
+  nk, nv, seq = 2, 4, 128
+  dim = (2 * nk + nv) * _HEAD_DIM
+  nondiff = (nk, nv, _HEAD_DIM, _HEAD_DIM, _KERNEL, _CHUNK, True, jnp.float32)
+  prec = cp_gdn.DEFAULT_PRECISION
+  forms = {
+      10: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, res, cot),
+      11: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, res, cot),
+      12: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, prec, res, cot),
+      13: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, prec, False, res, cot),
+      14: lambda res, cot: gdn_api._gdn_decoupled_conv1d_bwd(*nondiff, None, prec, False, True, res, cot),
+  }
+
+  def program(form):
+    def f(qkv, b, a, cw, alog, dtb, do):
+      _, res = gdn_api._gdn_decoupled_conv1d_fwd(qkv, b, a, cw, None, alog, dtb, None, None, *nondiff)
+      return forms[form](res, (do, (None, None)))
+
+    a = jax.ShapeDtypeStruct
+    args = (
+        a((1, seq, dim), jnp.bfloat16),
+        a((1, seq, nv), jnp.bfloat16),
+        a((1, seq, nv), jnp.bfloat16),
+        a((_KERNEL, 1, dim), jnp.float32),
+        a((nv,), jnp.float32),
+        a((nv,), jnp.float32),
+        a((1, seq, nv, _HEAD_DIM), jnp.float32),
+    )
+    ctx, _ = _tpu_target()
+    with ctx:
+      return str(jax.make_jaxpr(f)(*args))
+
+  ref = program(14)
+  ok = True
+  for form in (10, 11, 12, 13):
+    try:
+      same = program(form) == ref
+    except Exception as e:  # pylint: disable=broad-exception-caught
+      same = False
+      print(f"  form {form} raised {type(e).__name__}: {e}", flush=True)
+    ok &= _check(f"backward rule, {form} positional arguments", same)
+  return ok
+
+
+_BODY = re.compile(r"\\22body\\22: \\22[A-Za-z0-9+/=]+\\22")
+
+
+def _lowers_to_mosaic():
+  """Lowered for a TPU7x target (Pallas -> Mosaic, as an AOT compile does), the option changes the kernel body only."""
+  ok = True
+  for nk, nv in ((2, 4), (16, 64)):
+    texts = {}
+    for dense in (False, True):
+      ctx, _ = _tpu_target()
+      fn = functools.partial(
+          wrapper.fused_conv1d_gdn,
+          n_kq=nk,
+          n_v=nv,
+          d_k=_HEAD_DIM,
+          d_v=_HEAD_DIM,
+          kernel_size=_KERNEL,
+          compute_precision=jnp.dtype(jnp.float32),
+          mixed_tile_size=_CHUNK,
+          is_prefill_only=True,
+          conv_dense_token_stack=dense,
+      )
+      with ctx:
+        lowered = jax.jit(fn).trace(*_fused_args((2 * nk + nv) * _HEAD_DIM, nv)).lower(lowering_platforms=("tpu",))
+      texts[dense] = lowered.as_text(debug_info=False)
+    names = {d: re.findall(r'kernel_name = "([^"]+)"', t) for d, t in texts.items()}
+    good = (
+        len(names[False]) == 1
+        and names[False] == names[True]
+        and names[False][0].startswith("fused_conv1d_gdn")
+        and len(_BODY.findall(texts[True])) == 1
+        and texts[False] != texts[True]
+        and _BODY.sub("BODY", texts[False]) == _BODY.sub("BODY", texts[True])
+    )
+    ok &= _check(f"mosaic lowering width {(2 * nk + nv) * _HEAD_DIM}", good, f"kernels={names[True]}")
+  return ok
+
+
+def _with_wrong_version(obj, attr, wrong, names):
+  """The results of the forward-kernel cases `names` (None: all) with `obj.attr` replaced by a wrong version."""
+  real = getattr(obj, attr)
+  # jax.jit caches traces (fused_conv1d_gdn is jitted): clear them so the patched function is traced, and again
+  # afterwards so no later check reuses a trace of it.
+  jax.clear_caches()
+  setattr(obj, attr, wrong)
+  try:
+    return _option_reaches_every_forward_kernel(names)[1]
+  finally:
+    setattr(obj, attr, real)
+    jax.clear_caches()
+
+
+def _ignored_option(names):
+  """Deliberately wrong: the conv ignores the option (always jnp.stack)."""
+  real = compute_conv1d.stack_conv1d_tokens
+
+  def ignore(out_list, dense):
+    del dense
+    return real(out_list, False)
+
+  return _with_wrong_version(compute_conv1d, "stack_conv1d_tokens", ignore, names)
+
+
+def _option_dropped_by_forward_rule(names):
+  """Deliberately wrong: the custom VJP's forward rule (and so its remat replay) drops the option."""
+  real_fwd = gdn_api.gdn_decoupled_conv1d.fwd
+
+  def drop(*args):
+    args = list(args)
+    args[21] = False  # conv_dense_token_stack
+    return real_fwd(*args)
+
+  return _with_wrong_version(gdn_api.gdn_decoupled_conv1d, "fwd", drop, names)
+
+
+def _option_dropped_by_primal(names):
+  """Deliberately wrong: the custom VJP's primal (every forward-only call, with or without CP) drops the option."""
+  real_fun = gdn_api.gdn_decoupled_conv1d.fun
+
+  @functools.wraps(real_fun)  # custom_vjp binds keyword arguments against the primal's signature
+  def drop(*args):
+    args = list(args)
+    args[21] = False  # conv_dense_token_stack
+    return real_fun(*args)
+
+  return _with_wrong_version(gdn_api.gdn_decoupled_conv1d, "fun", drop, names)
+
+
+def main():
+  ok = True
+  reached, details = _option_reaches_every_forward_kernel()
+  for name, (good, detail) in details.items():
+    ok &= _check(name, good, detail)
+  ok &= _check("option reaches every forward kernel", reached)
+  ok &= _lowers_to_mosaic()
+  ok &= _composes_with_input_in_qkv_dtype()
+  ok &= _backward_rule_positional_forms()
+  # Each deliberately wrong version must make every case it breaks FAIL. A case that passes, or that does not exist
+  # (a misspelt name), counts as a miss.
+  controls = {
+      "control_ignored_option": (_ignored_option, tuple(details)),
+      "control_option_dropped_by_forward_rule": (_option_dropped_by_forward_rule, _FORWARD_RULE_CASES),
+      "control_option_dropped_by_primal": (_option_dropped_by_primal, _PRIMAL_CASES),
+  }
+  for name, (wrong_version, must_break) in controls.items():
+    results = wrong_version(must_break)
+    missed = [case for case in must_break if results.get(case, (True, ""))[0]]
+    detail = f"detected by all {len(must_break)} cases it breaks" if not missed else f"NOT detected by {missed}"
+    ok &= _check(name, not missed, detail)
+  if not ok:
+    sys.exit(1)
+  print(_PASS_TOKEN, flush=True)
+
+
+if __name__ == "__main__":
+  main()