diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index ba872c9a4b..fca14624ed 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -1523,12 +1523,13 @@ class Qwen3Next(BaseModel):
       "highest",
       description=(
           "Precision of the f32 matmuls that compose the GDN recurrent state across sequence-sharded context"
-          " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3."
-          " Only the XLA composition math changes; the Pallas kernels keep their own precision. It applies only"
-          " with use_gdn_kernel=True and sequence-sharded GDN context parallelism: gdn_cp_mode='seq', or 'auto',"
-          " which shards the sequence unless the GDN context parallelism is 2 and divides the number of key"
-          " heads. It has no effect with use_gdn_kernel=False (that path's context parallelism always uses"
-          " HIGHEST) or with head-sharded context parallelism."
+          " parallelism (kernels/gdn/gdn_bwd/cp_gdn.py). On TPU 'highest' runs 6 bf16 passes and 'high' runs 3;"
+          " with 'high' the rank-local forward transition product and backward state gradient run as Pallas scans"
+          " (kernels/gdn/gdn_bwd/dh_scan.py) at the same 3 passes. Only the composition math changes; the Pallas GDN"
+          " kernels keep their own precision. It applies only with use_gdn_kernel=True and sequence-sharded GDN"
+          " context parallelism: gdn_cp_mode='seq', or 'auto', which shards the sequence unless the GDN context"
+          " parallelism is 2 and divides the number of key heads. It has no effect with use_gdn_kernel=False (that"
+          " path's context parallelism always uses HIGHEST) or with head-sharded context parallelism."
       ),
   )
 
diff --git a/src/maxtext/kernels/gdn/compute_gdn.py b/src/maxtext/kernels/gdn/compute_gdn.py
index 36d8c7bf62..45a09d5eff 100644
--- a/src/maxtext/kernels/gdn/compute_gdn.py
+++ b/src/maxtext/kernels/gdn/compute_gdn.py
@@ -51,18 +51,24 @@ def invert_triangular_matrix(
   num_blocks = chunk // block_size
 
   def local_forward_sub(t_mat: jax.Array, b_mat: jax.Array) -> jax.Array:
-    x_list = []
-    for i in range(block_size):
-      b_i = b_mat[:, i, :]
-      if i == 0:
-        x_i = b_i
+    # Right-looking substitution: once row j is final, subtract t[i, j] * x_j from every later row i.
+    # The rows stay in place as two sublane-aligned halves, so no step re-stacks the solved rows or
+    # moves a row of t from lanes to sublanes. Each row subtracts the same products as a left-looking
+    # solve, one at a time in column order.
+    half = block_size // 2
+    x_lo, x_hi = b_mat[:, :half, :], b_mat[:, half:, :]
+    t_lo, t_hi = t_mat[:, :half, :], t_mat[:, half:, :]
+    rows_lo = jax.lax.broadcasted_iota(jnp.int32, (1, half, 1), 1)
+    rows_hi = jax.lax.broadcasted_iota(jnp.int32, (1, block_size - half, 1), 1)
+    for j in range(block_size - 1):
+      if j < half:
+        x_j = x_lo[:, j : j + 1, :]
+        x_lo = jnp.where(rows_lo > j, x_lo - t_lo[:, :, j : j + 1] * x_j, x_lo)
+        x_hi = x_hi - t_hi[:, :, j : j + 1] * x_j
       else:
-        stacked_x = jnp.stack(x_list, axis=1)
-        all_prev_t = t_mat[:, i, :i]
-        prev_sum = jnp.sum(all_prev_t[..., None] * stacked_x, axis=1)
-        x_i = b_i - prev_sum
-      x_list.append(x_i)
-    return jnp.stack(x_list, axis=1)
+        x_j = x_hi[:, j - half : j - half + 1, :]
+        x_hi = jnp.where(rows_hi > j - half, x_hi - t_hi[:, :, j : j + 1] * x_j, x_hi)
+    return jnp.concatenate([x_lo, x_hi], axis=1)
 
   x_blocks = []
   iota_r = jax.lax.broadcasted_iota(jnp.int32, t.shape, 1)
@@ -114,8 +120,8 @@ def chunked_gdn_per_seq(
     cfg: config.GDNConfig,
     seg_c_col: jax.Array | None = None,  # [1, chunk, 1]
     seg_aux_col: jax.Array | None = None,  # [1, chunk, 1]
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Perform chunked GDN over input [num_heads, chunk, head_dim]."""
+) -> tuple[jax.Array | None, jax.Array, jax.Array]:
+  """Perform chunked GDN over input [num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode."""
   # bf16 activations keep single-pass MXU matmuls even with an f32 compute dtype
   # (HIGHEST would triple the forward MXU cost); HIGHEST needs f32 activations and compute.
   dot_prec = (
@@ -236,23 +242,36 @@ def chunked_gdn_per_seq(
   # [num_v_heads, chunk, v_head_dim]
   u, w = jnp.split(merged_uw, [cfg.v_head_dim], axis=-1)
 
-  # [num_v_heads, chunk, kq_head_dim]
-  q_large_gating = q_repeat * gating_forward
-  # NOTE: Concatenate lhs with same rhs to leverage weight
-  # stationary architecture.
-  # [num_v_heads, 2 * chunk, kq_head_dim]
-  merged_w_q = jnp.concat([w, q_large_gating], axis=1)
-  # [num_v_heads, 2 * chunk, v_head_dim]
-  merged_ws_out_updated = jax.lax.dot(
-      merged_w_q,
-      state_prev,
-      dimension_numbers=(((2,), (1,)), ((0,), (0,))),
-      precision=dot_prec,
-      preferred_element_type=jnp.float32,
-  )
+  states_only = cfg.mode == config.GDNMode.PER_SEQ_STATES
+  if states_only:
+    # No output: only the w rows of the merged matmul below. Each row's dot product is the same,
+    # so ws is bit-identical.
+    # [num_v_heads, chunk, v_head_dim]
+    ws = jax.lax.dot(
+        w,
+        state_prev,
+        dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+        precision=dot_prec,
+        preferred_element_type=jnp.float32,
+    )
+  else:
+    # [num_v_heads, chunk, kq_head_dim]
+    q_large_gating = q_repeat * gating_forward
+    # NOTE: Concatenate lhs with same rhs to leverage weight
+    # stationary architecture.
+    # [num_v_heads, 2 * chunk, kq_head_dim]
+    merged_w_q = jnp.concat([w, q_large_gating], axis=1)
+    # [num_v_heads, 2 * chunk, v_head_dim]
+    merged_ws_out_updated = jax.lax.dot(
+        merged_w_q,
+        state_prev,
+        dimension_numbers=(((2,), (1,)), ((0,), (0,))),
+        precision=dot_prec,
+        preferred_element_type=jnp.float32,
+    )
 
-  # NOTE: Splitting along non sublane/lane dim is free.
-  ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1)
+    # NOTE: Splitting along non sublane/lane dim is free.
+    ws, out_updated = jnp.split(merged_ws_out_updated, 2, axis=1)
   ws = ws.astype(cfg.dtypes.compute)
 
   # [num_v_heads, chunk, v_head_dim]
@@ -273,6 +292,8 @@ def chunked_gdn_per_seq(
   # [num_v_heads, kq_head_dim, v_head_dim]
   state_updated = state_prev * gating_last
   state = state_updated + state_new
+  if states_only:
+    return None, state, t_inv
 
   # [num_kq_heads, chunk, chunk]
   out_qk = jax.lax.dot(
@@ -314,8 +335,8 @@ def chunked_gdn(
     a_log: jax.Array,
     dt_bias: jax.Array,
     cfg: config.GDNConfig,
-) -> tuple[jax.Array, jax.Array, jax.Array]:
-  """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]."""
+) -> tuple[jax.Array | None, jax.Array, jax.Array]:
+  """Perform chunked GDN over input [seq, num_heads, chunk, head_dim]; out is None in PER_SEQ_STATES mode."""
 
   mask_dtype = get_mask_dtype(cfg.dtypes.compute)
   iota = jax.lax.broadcasted_iota(mask_dtype, (cfg.seq_tile_size, 1, cfg.chunk_size, 1), 2)
@@ -373,10 +394,11 @@ def chunked_gdn(
         seg_c_col=seg_c_per_seq[idx] if cfg.has_seg_ids else None,
         seg_aux_col=seg_aux_per_seq[idx] if cfg.has_seg_ids else None,
     )
-    out_list.append(out.swapaxes(0, 1))
+    if out is not None:
+      out_list.append(out.swapaxes(0, 1))
     state_list.append(state)
     t_inv_list.append(t_inv)
-  out = jnp.stack(out_list, axis=0)
+  out = jnp.stack(out_list, axis=0) if out_list else None
   state = jnp.stack(state_list, axis=0)
   t_inv = jnp.stack(t_inv_list, axis=0)
   return out, state, t_inv
diff --git a/src/maxtext/kernels/gdn/config.py b/src/maxtext/kernels/gdn/config.py
index 8ad3933295..b06926d24c 100644
--- a/src/maxtext/kernels/gdn/config.py
+++ b/src/maxtext/kernels/gdn/config.py
@@ -24,7 +24,8 @@
 from jax.experimental.pallas import tpu as pltpu
 import jax.numpy as jnp
 
-DEFAULT_VMEM_LIMIT_FACTOR: float = 0.90
+# The 64-v-head prefill kernel needs about 58 MiB of the 64 MiB VMEM on TPU7x.
+DEFAULT_VMEM_LIMIT_FACTOR: float = 0.95
 
 
 class GDNMode(enum.StrEnum):
@@ -32,6 +33,8 @@ class GDNMode(enum.StrEnum):
 
   BATCHED = enum.auto()
   PER_SEQ = enum.auto()
+  # PER_SEQ that only carries the conv/recurrent states and t_inv: no output, no chunk states.
+  PER_SEQ_STATES = enum.auto()
 
   def get_seq_tile_size(self, tile_size: int) -> int:
     if self == GDNMode.BATCHED:
@@ -82,6 +85,8 @@ class GDNConfig:
   num_buffers: int = 2
   has_seg_ids: bool = False
   use_qk_norm_in_gdn: bool = True
+  # Prefill-only tiles are whole chunks: qkv is DMA'd in act_in dtype and upcast in VMEM.
+  is_prefill_only: bool = False
 
   @property
   def chunk_size(self) -> int:
@@ -150,7 +155,13 @@ def get_scratch_shape_dict(self) -> dict[str, Any]:
       carry_conv_scratch = pltpu.VMEM(conv_shape, jnp.float32)
       carry_recurrent_scratch = pltpu.VMEM(recurrent_shape, jnp.float32)
 
+    # Compact fp32 copy of a prefill qkv chunk, read by Conv1D and the q/k/v loads.
+    qkv_scratch = None
+    if self.is_prefill_only:
+      qkv_scratch = pltpu.VMEM((self.seq_tile_size, self.chunk_size, 1, self.dim_size), jnp.float32)
+
     return {
         "carry_conv_scratch_ref": carry_conv_scratch,
         "carry_recurrent_scratch_ref": carry_recurrent_scratch,
+        "qkv_scratch_ref": qkv_scratch,
     }
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/api.py b/src/maxtext/kernels/gdn/gdn_bwd/api.py
index 4ea358228b..e56aa53cd8 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/api.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/api.py
@@ -25,6 +25,7 @@
 from .. import compute_conv1d as local_compute_conv1d
 from .. import wrapper as local_gdn_wrapper
 from . import cp_gdn
+from . import dh_scan
 from .compute_conv1d_bwd import conv1d_silu_bwd
 from .compute_conv1d_bwd import conv1d_silu_fwd
 from .jax_compute_gdn_states import _compute_forward_conv_and_states
@@ -130,6 +131,50 @@ def decoupled_conv1d_gdn_bwd_kernel(
   )
 
 
+def _fused_kernel_inputs(
+    qkv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    conv_weight: jax.Array,
+    conv_state: Optional[jax.Array],
+    recurrent_state: Optional[jax.Array],
+    num_v_heads: int,
+    head_k_dim: int,
+    head_v_dim: int,
+    conv_kernel_size: int,
+) -> dict[str, jax.Array]:
+  """Lays out [batch, seq, ...] inputs for the fused Conv1D-GDN kernel: one equal-length prefill sequence per row."""
+  num_seqs, seq_len, dim_size = qkv.shape
+  has_init = (conv_state is not None) or (recurrent_state is not None)
+
+  if conv_state is None:
+    tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype)
+  elif conv_state.shape[0] == num_seqs:
+    tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0)))
+  else:
+    tokamax_conv_state = conv_state
+
+  if recurrent_state is None:
+    tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32)
+  elif recurrent_state.shape[0] == num_seqs:
+    tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0)))
+  else:
+    tokamax_recurrent_state = recurrent_state.astype(jnp.float32)
+
+  return {
+      "qkv": qkv.reshape(-1, dim_size),
+      "b": b.reshape(-1, b.shape[-1]),
+      "a": a.reshape(-1, a.shape[-1]),
+      "conv_state": tokamax_conv_state,
+      "recurrent_state": tokamax_recurrent_state,
+      "conv_weight": jnp.swapaxes(conv_weight, 0, 2),
+      "query_start_loc": jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32),
+      "state_indices": jnp.arange(1, num_seqs + 1, dtype=jnp.int32),
+      "distribution": jnp.array([0, 0, num_seqs], dtype=jnp.int32),
+      "seq_lens": jnp.full((num_seqs,), seq_len + (1 if has_init else 0), dtype=jnp.int32),
+  }
+
+
 def _run_local_gdn_decoupled_fwd(
     qkv: jax.Array,
     b: jax.Array,
@@ -158,7 +203,7 @@ def _run_local_gdn_decoupled_fwd(
     Optional[jax.Array],
 ]:
   """Runs local GDN forward pass on TPU returning (t_inv, chunk_states), or pure JAX on CPU."""
-  batch_size, seq_len, dim_size = qkv.shape
+  batch_size, seq_len, _ = qkv.shape
   if seq_len % chunk_size != 0:
     raise ValueError(
         f"GDN forward kernel requires the local sequence length ({seq_len}) to be a multiple of"
@@ -215,58 +260,27 @@ def _run_local_gdn_decoupled_fwd(
     )
     return (out, states), t_inv, chunk_states
 
-  num_seqs = batch_size
   num_chunks = seq_len // chunk_size
-
-  qkv_flat = qkv.reshape(-1, dim_size)
-  b_flat = b.reshape(-1, b.shape[-1])
-  a_flat = a.reshape(-1, a.shape[-1])
-  tokamax_conv_weight = jnp.swapaxes(conv_weight, 0, 2)
-
-  has_init = (conv_state is not None) or (recurrent_state is not None)
-  query_start_loc = jnp.arange(0, (num_seqs + 1) * seq_len, seq_len, dtype=jnp.int32)
-  state_indices = jnp.arange(1, num_seqs + 1, dtype=jnp.int32)
-  seq_lens = jnp.full((num_seqs,), seq_len + (1 if has_init else 0), dtype=jnp.int32)
-  distribution = jnp.array([0, 0, num_seqs], dtype=jnp.int32)
-
-  if conv_state is None:
-    tokamax_conv_state = jnp.zeros((num_seqs + 1, conv_kernel_size - 1, dim_size), dtype=qkv.dtype)
-  elif conv_state.shape[0] == num_seqs:
-    tokamax_conv_state = jnp.pad(conv_state, ((1, 0), (0, 0), (0, 0)))
-  else:
-    tokamax_conv_state = conv_state
-
-  if recurrent_state is None:
-    tokamax_recurrent_state = jnp.zeros((num_seqs + 1, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32)
-  elif recurrent_state.shape[0] == num_seqs:
-    tokamax_recurrent_state = jnp.pad(recurrent_state.astype(jnp.float32), ((1, 0), (0, 0), (0, 0), (0, 0)))
-  else:
-    tokamax_recurrent_state = recurrent_state.astype(jnp.float32)
-
   (
       core_attn_out_flat,
       (new_conv_state, new_recurrent_state),
       t_inv_raw,
       chunk_states_raw,
   ) = local_gdn_wrapper.fused_conv1d_gdn(
-      qkv=qkv_flat,
-      b=b_flat,
-      a=a_flat,
-      conv_state=tokamax_conv_state,
-      recurrent_state=tokamax_recurrent_state,
-      conv_weight=tokamax_conv_weight,
+      **_fused_kernel_inputs(
+          qkv, b, a, conv_weight, conv_state, recurrent_state, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size
+      ),
       conv_bias=conv_bias,
       a_log=a_log,
       dt_bias=dt_bias,
-      query_start_loc=query_start_loc,
-      state_indices=state_indices,
-      distribution=distribution,
-      seq_lens=seq_lens,
       n_kq=num_k_heads,
       n_v=num_v_heads,
       d_k=head_k_dim,
       d_v=head_v_dim,
       kernel_size=conv_kernel_size,
+      # seq_len is a multiple of chunk_size (checked above), so the prefill kernel writes every
+      # output row and the output buffer needs no zero fill.
+      zero_initialize_out=False,
       compute_precision=jnp.dtype(compute_dtype),
       mixed_tile_size=chunk_size,
       is_prefill_only=True,
@@ -292,6 +306,90 @@ def _run_local_gdn_decoupled_fwd(
   )
 
 
+def _run_local_gdn_states_fwd(
+    qkv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    conv_weight: jax.Array,
+    conv_bias: Optional[jax.Array],
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    conv_state: Optional[jax.Array],
+    recurrent_state: Optional[jax.Array],
+    *,
+    num_k_heads: int,
+    num_v_heads: int,
+    head_k_dim: int,
+    head_v_dim: int,
+    conv_kernel_size: int,
+    chunk_size: int,
+    use_qk_norm_in_gdn: bool,
+    compute_dtype: jnp.dtype,
+    segment_ids: Optional[jax.Array] = None,
+    conv_halo_seg: Optional[jax.Array] = None,
+    init_seg: Optional[jax.Array] = None,
+) -> Tuple[Tuple[jax.Array, jax.Array], jax.Array]:
+  """_run_local_gdn_decoupled_fwd for callers that need only (states, t_inv), bit-identical to it.
+
+  On TPU the kernel skips the output and the chunk states; on CPU XLA drops the unused pure-JAX outputs.
+  """
+  if target_platform(qkv) == "cpu":
+    (_, states), t_inv, _ = _run_local_gdn_decoupled_fwd(
+        qkv,
+        b,
+        a,
+        conv_weight,
+        conv_bias,
+        a_log,
+        dt_bias,
+        conv_state,
+        recurrent_state,
+        num_k_heads=num_k_heads,
+        num_v_heads=num_v_heads,
+        head_k_dim=head_k_dim,
+        head_v_dim=head_v_dim,
+        conv_kernel_size=conv_kernel_size,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        compute_dtype=compute_dtype,
+        segment_ids=segment_ids,
+        conv_halo_seg=conv_halo_seg,
+        init_seg=init_seg,
+    )
+    return states, t_inv
+  batch_size, seq_len, _ = qkv.shape
+  if seq_len % chunk_size != 0:
+    raise ValueError(
+        f"GDN forward kernel requires the local sequence length ({seq_len}) to be a multiple of"
+        f" chunk_size ({chunk_size}); with sequence-sharded context parallelism this is seq_len / cp."
+    )
+  fallback_reason = pallas_unsupported_reason(head_k_dim=head_k_dim, head_v_dim=head_v_dim, chunk_size=chunk_size)
+  if fallback_reason is not None:
+    raise ValueError(f"GDN Pallas TPU forward kernel does not support {fallback_reason}.")
+
+  (new_conv_state, new_recurrent_state), t_inv_raw = local_gdn_wrapper.fused_conv1d_gdn_states(
+      **_fused_kernel_inputs(
+          qkv, b, a, conv_weight, conv_state, recurrent_state, num_v_heads, head_k_dim, head_v_dim, conv_kernel_size
+      ),
+      conv_bias=conv_bias,
+      a_log=a_log,
+      dt_bias=dt_bias,
+      n_kq=num_k_heads,
+      n_v=num_v_heads,
+      d_k=head_k_dim,
+      d_v=head_v_dim,
+      kernel_size=conv_kernel_size,
+      compute_precision=jnp.dtype(compute_dtype),
+      mixed_tile_size=chunk_size,
+      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+      init_seg=init_seg,
+  )
+  t_inv = t_inv_raw.astype(jnp.float32).reshape(batch_size, seq_len // chunk_size, num_v_heads, chunk_size, chunk_size)
+  return (new_conv_state[1:].astype(qkv.dtype), new_recurrent_state[1:].astype(jnp.float32)), t_inv
+
+
 def _is_cp_active(cp_axis_name: str | tuple[str, ...] | None) -> bool:
   if cp_axis_name is None:
     return False
@@ -377,8 +475,9 @@ def _run_cp_gdn_decoupled_fwd_impl(
   conv_halo = checkpoint_name(conv_halo, "gdn_cp_conv_halo")
   zero_rs = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), dtype=jnp.float32)
 
-  # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local
-  (_, states_pass1), t_inv, _ = _run_local_gdn_decoupled_fwd(
+  # Pass 1: Local GDN with zero initial recurrent state -> yields t_inv and S_ext_local. Its output
+  # is not needed, so the kernel neither computes it nor writes it or the chunk states.
+  states_pass1, t_inv = _run_local_gdn_states_fwd(
       qkv,
       b,
       a,
@@ -414,24 +513,43 @@ def _run_cp_gdn_decoupled_fwd_impl(
       conv_halo_seg=conv_halo_seg,
       output_dtype=jnp.float32,
   )
-  m_local, s_ext_local = cp_gdn.compose_local_from_t_inv(
-      qkv_conv=k_conv,
-      b=b,
-      a=a,
-      a_log=a_log,
-      dt_bias=dt_bias,
-      t_inv=t_inv,
-      num_k_heads=num_k_heads,
-      num_v_heads=num_v_heads,
-      head_k_dim=head_k_dim,
-      head_v_dim=head_v_dim,
-      chunk_size=chunk_size,
-      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
-      s_ext_pass1=states_pass1[1],
-      segment_ids=s_enc_local,
-      init_seg=init_seg,
-      precision=cp_matmul_precision,
-  )
+  if target_platform(qkv) == "tpu" and cp_matmul_precision == jax.lax.Precision.HIGH:
+    # The same chunk transitions as compose_local_from_t_inv at the same 3-pass precision, multiplied in
+    # sequence in a Pallas scan instead of XLA's HBM-bound tree of f32 [chunks, heads, d_k, d_k] products.
+    m_local = dh_scan.gdn_cp_m_scan(
+        k_conv,
+        b,
+        a,
+        a_log,
+        dt_bias,
+        t_inv,
+        num_v_heads=num_v_heads,
+        kq_head_dim=head_k_dim,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        segment_ids=s_enc_local,
+        init_seg=init_seg,
+    )
+    s_ext_local = states_pass1[1].astype(jnp.float32)
+  else:
+    m_local, s_ext_local = cp_gdn.compose_local_from_t_inv(
+        qkv_conv=k_conv,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        t_inv=t_inv,
+        num_k_heads=num_k_heads,
+        num_v_heads=num_v_heads,
+        head_k_dim=head_k_dim,
+        head_v_dim=head_v_dim,
+        chunk_size=chunk_size,
+        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        s_ext_pass1=states_pass1[1],
+        segment_ids=s_enc_local,
+        init_seg=init_seg,
+        precision=cp_matmul_precision,
+    )
   m_local = checkpoint_name(m_local, "gdn_cp_m_local")
 
   h_init = recurrent_state.astype(jnp.float32) if recurrent_state is not None else zero_rs
@@ -909,25 +1027,47 @@ def _gdn_decoupled_conv1d_bwd(
         if d_conv_state is not None
         else None
     )
-    dm_local, ds_ext_local = cp_gdn.compose_bwd_local_from_t_inv(
-        qkv_conv=qkv_conv,
-        b=b,
-        a=a,
-        a_log=a_log,
-        dt_bias=dt_bias,
-        do=d_out,
-        t_inv=t_inv,
-        num_k_heads=num_k_heads,
-        num_v_heads=num_v_heads,
-        head_k_dim=head_k_dim,
-        head_v_dim=head_v_dim,
-        chunk_size=chunk_size,
-        use_qk_norm_in_gdn=use_qk_norm_in_gdn,
-        m_local_cached=m_local_fwd,
-        segment_ids=segment_ids,
-        init_seg=init_seg,
-        precision=cp_matmul_precision,
-    )
+    if target_platform(qkv_conv) == "tpu" and m_local_fwd is not None and cp_matmul_precision == jax.lax.Precision.HIGH:
+      # ds_ext_local of compose_bwd_local_from_t_inv at the same 3-pass precision, as a reverse Pallas scan
+      # (Horner order) instead of XLA's HBM-bound tree. dm_local is the transposed forward transition, which
+      # compose_bwd_local_from_t_inv also returns when given m_local_cached.
+      ds_ext_local = dh_scan.gdn_cp_dh_scan(
+          qkv_conv,
+          b,
+          a,
+          a_log,
+          dt_bias,
+          d_out,
+          t_inv,
+          num_v_heads=num_v_heads,
+          kq_head_dim=head_k_dim,
+          v_head_dim=head_v_dim,
+          chunk_size=chunk_size,
+          use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+          segment_ids=segment_ids,
+          init_seg=init_seg,
+      )
+      dm_local = jnp.swapaxes(m_local_fwd.astype(jnp.float32), -1, -2)
+    else:
+      dm_local, ds_ext_local = cp_gdn.compose_bwd_local_from_t_inv(
+          qkv_conv=qkv_conv,
+          b=b,
+          a=a,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          do=d_out,
+          t_inv=t_inv,
+          num_k_heads=num_k_heads,
+          num_v_heads=num_v_heads,
+          head_k_dim=head_k_dim,
+          head_v_dim=head_v_dim,
+          chunk_size=chunk_size,
+          use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+          m_local_cached=m_local_fwd,
+          segment_ids=segment_ids,
+          init_seg=init_seg,
+          precision=cp_matmul_precision,
+      )
     dht_local, _ = cp_gdn.incoming_grad_state(
         dm_local, ds_ext_local, dht_final, cp_axis_name, precision=cp_matmul_precision
     )
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
index 3d379e5cf7..f52565aa50 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/bwd_memory_ref.py
@@ -22,15 +22,23 @@
 def make_bwd_block_specs(
     num_chunks: int,
     chunk_size: int,
-    dim_size: int,
+    qk_tile_size: int,
+    v_tile_size: int,
     num_v_heads: int,
     kq_head_dim: int,
     v_head_dim: int,
+    k_block_offset: int,
+    v_block_offset: int,
     padded_num_v_heads: int | None = None,
     has_dht: bool = False,
     has_dh0: bool = False,
 ) -> Tuple[list[pl.BlockSpec], list[pl.BlockSpec], int, int]:
-  """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward."""
+  """Constructs reverse-scan Pallas emit_pipeline in_specs and out_specs for GDN backward.
+
+  Head group g of q/k/v and dq/dk/dv is a column block of a token-major array (the k and v blocks of qkv_conv start
+  at block k_block_offset / v_block_offset), and do / chunk_states / t_inv are indexed by group, so the kernel reads
+  and writes them in place instead of through group-major copies.
+  """
   if padded_num_v_heads is None:
     padded_num_v_heads = ((num_v_heads + 127) // 128) * 128
 
@@ -38,47 +46,57 @@ def rc(c):
     return num_chunks - 1 - c
 
   in_specs = [
-      # 0: qkv_conv [B, G, C, chunk_size, dim_size]
+      # 0: q [B, C, chunk_size, dim_size], columns of group g
       pl.BlockSpec(
-          (None, None, None, chunk_size, dim_size),
-          lambda b, g, c: (b, g, rc(c), 0, 0),
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
+      ),
+      # 1: k [B, C, chunk_size, dim_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, k_block_offset + g),
+      ),
+      # 2: v [B, C, chunk_size, dim_size or v_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, v_tile_size),
+          lambda b, g, c: (b, rc(c), 0, v_block_offset + g),
       ),
-      # 1: b [B, G, C, chunk_size, padded_num_v_heads]
+      # 3: b [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 2: a [B, G, C, chunk_size, padded_num_v_heads]
+      # 4: a [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 3: do [B, G, C, chunk_size, num_v_heads, v_head_dim]
+      # 5: do [B, C, chunk_size, G, num_v_heads, v_head_dim]
       pl.BlockSpec(
-          (None, None, None, chunk_size, num_v_heads, v_head_dim),
-          lambda b, g, c: (b, g, rc(c), 0, 0, 0),
+          (None, None, chunk_size, None, num_v_heads, v_head_dim),
+          lambda b, g, c: (b, rc(c), 0, g, 0, 0),
       ),
-      # 4: chunk_states [B, G, C, num_v_heads, kq_head_dim, v_head_dim]
+      # 6: chunk_states [B, C, G, num_v_heads, kq_head_dim, v_head_dim]
       pl.BlockSpec(
           (None, None, None, num_v_heads, kq_head_dim, v_head_dim),
-          lambda b, g, c: (b, g, rc(c), 0, 0, 0),
+          lambda b, g, c: (b, rc(c), g, 0, 0, 0),
       ),
-      # 5: t_inv [B, G, C, num_v_heads, chunk_size, chunk_size]
+      # 7: t_inv [B, C, G, num_v_heads, chunk_size, chunk_size]
       pl.BlockSpec(
           (None, None, None, num_v_heads, chunk_size, chunk_size),
-          lambda b, g, c: (b, g, rc(c), 0, 0, 0),
+          lambda b, g, c: (b, rc(c), g, 0, 0, 0),
       ),
-      # 6: a_log [B, G, 1, padded_num_v_heads]
+      # 8: a_log [B, G, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, 0, 0),
       ),
-      # 7: dt_bias [B, G, 1, padded_num_v_heads]
+      # 9: dt_bias [B, G, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, 0, 0),
       ),
-      # 8: reset [B, G, C, 1, 128]
+      # 10: reset [B, G, C, 1, 128]
       pl.BlockSpec(
           (None, None, None, 1, 128),
           lambda b, g, c: (b, g, rc(c), 0, 0),
@@ -86,7 +104,7 @@ def rc(c):
   ]
   if has_dht:
     in_specs.append(
-        # 9: dht [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
+        # 11: dht [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
         pl.BlockSpec(
             (None, None, 1, num_v_heads, kq_head_dim, v_head_dim),
             lambda b, g, c: (b, g, 0, 0, 0, 0),
@@ -94,27 +112,37 @@ def rc(c):
     )
 
   out_specs = [
-      # 0: dy_conv [B, G, C, chunk_size, dim_size]
+      # 0: dq [B, C, chunk_size, q_size], columns of group g
       pl.BlockSpec(
-          (None, None, None, chunk_size, dim_size),
-          lambda b, g, c: (b, g, rc(c), 0, 0),
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
+      ),
+      # 1: dk [B, C, chunk_size, k_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, qk_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
+      ),
+      # 2: dv [B, C, chunk_size, v_size], columns of group g
+      pl.BlockSpec(
+          (None, None, chunk_size, v_tile_size),
+          lambda b, g, c: (b, rc(c), 0, g),
       ),
-      # 1: db [B, G, C, chunk_size, padded_num_v_heads]
+      # 3: db [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 2: da [B, G, C, chunk_size, padded_num_v_heads]
+      # 4: da [B, G, C, chunk_size, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, chunk_size, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 3: dal [B, G, C, 1, padded_num_v_heads]
+      # 5: dal [B, G, C, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
       ),
-      # 4: ddt [B, G, C, 1, padded_num_v_heads]
+      # 6: ddt [B, G, C, 1, padded_num_v_heads]
       pl.BlockSpec(
           (None, None, None, 1, padded_num_v_heads),
           lambda b, g, c: (b, g, rc(c), 0, 0),
@@ -122,7 +150,7 @@ def rc(c):
   ]
   if has_dh0:
     out_specs.append(
-        # 5: dh0 [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
+        # 7: dh0 [B, G, 1, num_v_heads, kq_head_dim, v_head_dim]
         pl.BlockSpec(
             (None, None, 1, num_v_heads, kq_head_dim, v_head_dim),
             lambda b, g, c: (b, g, 0, 0, 0, 0),
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
new file mode 100644
index 0000000000..b4dba19cf4
--- /dev/null
+++ b/src/maxtext/kernels/gdn/gdn_bwd/dh_scan.py
@@ -0,0 +1,543 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+r"""Pallas scans that fold one sequence-sharded CP rank's GDN chunk transitions on Mosaic TPU.
+
+Per chunk $c$ of the rank, with $w_c = T_c^{-1} k^{\beta g}_c$, the forward state map is
+$h \mapsto M_c h + S_c$ with $M_c = g^{last}_c I - (k^{bwd}_c)^\top w_c$, and the backward one is
+$\mathrm{d}h \mapsto M_c^\top \mathrm{d}h + \mathrm{d}S_c$ with
+$\mathrm{d}S_c = q_g^\top \mathrm{d}o - (k^{\beta g}_c)^\top T_c^{-\top} (\mathrm{attn}_c^\top \mathrm{d}o)$.
+
+- `gdn_cp_m_scan` returns $M_{N-1} \cdots M_0$, the `m_local` of `cp_gdn.compose_local_from_t_inv`
+  (its `s_ext_pass1` branch).
+- `gdn_cp_dh_scan` returns $\mathrm{d}S_0 + M_0^\top(\mathrm{d}S_1 + M_1^\top(\cdots))$, the `ds_ext_local` of
+  `cp_gdn.compose_bwd_local_from_t_inv`, as a reverse scan in the form of the backward kernel's state recurrence.
+
+Both are the same terms as the `cp_gdn` balanced trees in sequential order, with every matmul at 3-pass bf16
+(`Precision.HIGH`, as `gdn_cp_matmul_precision=high` gives the trees) and the segment-masked gating cumsum at
+`HIGHEST`. They read the head groups of qkv / do / t_inv in place, with the block layout of
+`pallas_gdn_bwd_kernel`, and keep the running product in VMEM, so no [chunks, heads, d_k, d_k] f32 intermediate
+reaches HBM.
+"""
+
+import functools
+from typing import Optional
+
+import jax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+import jax.numpy as jnp
+
+from .. import compute_conv1d as local_compute_conv1d
+from .pallas_mosaic_tpu_bwd import _gdn_matmul
+from .runtime_utils import ensure_cpu_interpret_registered
+from .runtime_utils import target_platform
+
+# Value heads per grid step. The dh-scan's per-chunk working set grows with the head tile; the M-scan is lighter
+# and runs all 64 heads of the Qwen3.5-397B layer in one group.
+_DH_SCAN_HEAD_TILE = 32
+_M_SCAN_HEAD_TILE = 64
+
+_precise_matmul = functools.partial(_gdn_matmul, compute_dtype=jnp.bfloat16, precise=True)
+
+
+def _pick_head_tile(num_v_heads: int, repeats: int, target: int) -> int:
+  """Largest divisor of num_v_heads up to target that holds whole GVA groups, as pallas_gdn_bwd_kernel picks it."""
+  for candidate in range(min(num_v_heads, target), 0, -1):
+    if num_v_heads % candidate == 0 and candidate % repeats == 0:
+      return candidate
+  return repeats if num_v_heads % repeats == 0 else num_v_heads
+
+
+def _grouped_gate_inputs(
+    b: jax.Array,
+    a: jax.Array,
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    segment_ids: Optional[jax.Array],
+    init_seg: Optional[jax.Array],
+    *,
+    num_chunks: int,
+    chunk_size: int,
+    num_groups: int,
+    tile_v_heads: int,
+):
+  """Lays out b, a, a_log, dt_bias and the segment metadata per head group, as pallas_gdn_bwd_kernel does.
+
+  Returns b / a as [B, G, C, chunk_size, padded], a_log / dt_bias as [B, G, 1, padded] and, per chunk, the active
+  segment entering it (lane 0 of [B, G, C, 1, 128]). With segment IDs, their signed encoding rides in b's spare lane
+  `tile_v_heads`, and b stays f32 then: bf16 rounds integers above 256, so distinct documents would share an ID.
+  """
+  batch_size = b.shape[0]
+  seq_len = num_chunks * chunk_size
+  has_seg_ids = segment_ids is not None
+  padded = ((tile_v_heads + int(has_seg_ids) + 127) // 128) * 128
+  lane_pad = ((0, 0),) * 4 + ((0, padded - tile_v_heads),)
+
+  def per_group(x):
+    x = x.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_heads).transpose((0, 3, 1, 2, 4))
+    return jnp.pad(x, lane_pad)
+
+  def per_group_param(x):
+    x = jnp.pad(x.reshape(num_groups, tile_v_heads), ((0, 0), (0, padded - tile_v_heads)))
+    return jnp.broadcast_to(x[None, :, None, :], (batch_size, num_groups, 1, padded))
+
+  b_5d = per_group(b.astype(jnp.float32) if has_seg_ids else b)
+  a_5d = per_group(a)
+  if has_seg_ids:
+    s_enc = local_compute_conv1d.encode_segment_ids(segment_ids.reshape(batch_size, seq_len), init_seg=init_seg)
+    s_enc_3d = s_enc.reshape(batch_size, num_chunks, chunk_size)
+    b_5d = b_5d.at[:, :, :, :, tile_v_heads].set(s_enc_3d[:, None, :, :].astype(b_5d.dtype))
+    if init_seg is not None:
+      init_active = jnp.abs(init_seg.astype(jnp.float32).reshape(batch_size, 1))
+    else:
+      init_active = jnp.zeros((batch_size, 1), dtype=jnp.float32)
+    seg_prev = jnp.concatenate([init_active, jnp.abs(s_enc_3d[:, :-1, -1])], axis=1)
+    seg_prev_hbm = jnp.pad(seg_prev[:, :, None, None], ((0, 0), (0, 0), (0, 0), (0, 127)))
+  else:
+    seg_prev_hbm = jnp.zeros((batch_size, num_chunks, 1, 128), dtype=jnp.float32)
+  seg_prev_5d = jnp.broadcast_to(seg_prev_hbm[:, None], (batch_size, num_groups, num_chunks, 1, 128))
+  return b_5d, a_5d, per_group_param(a_log), per_group_param(dt_bias), seg_prev_5d, padded
+
+
+def _chunk_gates(b_ref, a_ref, a_log_ref, dt_bias_ref, seg_prev_ref, *, chunk_size, num_v_heads, has_seg_ids):
+  """Segment masks and decay gates of one chunk and head group, as cp_gdn computes them.
+
+  Returns (valid [C, 1] or None, beta_h [H, C], cumsum_h [H, C], gating_forward [H, C, 1], gating_last [H, 1, 1],
+  gating_backward [H, C, 1], mask_causal [C, C]).
+  """
+  b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32)
+  a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32)
+  a_log_val = a_log_ref[...][0, :num_v_heads].astype(jnp.float32)
+  dt_bias_val = dt_bias_ref[...][0, :num_v_heads].astype(jnp.float32)
+  beta = jax.nn.sigmoid(b_val)
+  log_g = -jnp.exp(a_log_val) * jax.nn.softplus(a_val + dt_bias_val)
+  mask_tril = jnp.tril(jnp.ones((chunk_size, chunk_size), dtype=jnp.float32))
+
+  valid_c = None
+  mask_cumsum = mask_tril
+  mask_causal = mask_tril
+  if has_seg_ids:
+    seg_c = b_ref[...][:, num_v_heads : num_v_heads + 1].astype(jnp.float32)
+    seg_prev = seg_prev_ref[...][0, 0].astype(jnp.float32)
+    valid_c = seg_c > 0.5
+    active_c = jnp.abs(seg_c)
+    beta = jnp.where(valid_c, beta, 0.0)
+    log_g = jnp.where(valid_c, log_g, 0.0)
+    same_active = (jnp.abs(active_c - active_c.T) < 0.5) & (active_c > 0.5)
+    same_valid = (jnp.abs(seg_c - seg_c.T) < 0.5) & valid_c
+    mask_cumsum = mask_tril * same_active.astype(jnp.float32)
+    mask_causal = mask_tril * same_valid.astype(jnp.float32)
+
+  # The cumsum feeds exp(); like cp_gdn it stays at HIGHEST (0/1 mask times f32: exact up to summation order).
+  cumsum_h = jnp.transpose(jnp.dot(mask_cumsum, log_g, precision=jax.lax.Precision.HIGHEST), (1, 0))
+  gating_forward = jnp.exp(cumsum_h)[:, :, None]
+  gating_last = jnp.exp(cumsum_h[:, -1])[:, None, None]
+  if has_seg_ids:
+    m_in = ((jnp.abs(seg_c - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(jnp.float32)
+    m_out = ((jnp.abs(seg_c - active_c[-1:]) < 0.5) & (active_c[-1:] > 0.5)).astype(jnp.float32)
+    m_keep = ((jnp.abs(active_c[-1:] - seg_prev) < 0.5) & (seg_prev > 0.5)).astype(jnp.float32)
+    gating_forward = gating_forward * m_in[None, :, :]
+    gating_last = gating_last * m_keep[None, :, :]
+    bwd_diff = jnp.where(m_out.T > 0.5, cumsum_h[:, -1:] - cumsum_h, -1e4)
+    gating_backward = jnp.exp(bwd_diff)[:, :, None] * m_out[None, :, :]
+  else:
+    gating_backward = jnp.exp(cumsum_h[:, -1:] - cumsum_h)[:, :, None]
+  beta_h = jnp.transpose(beta, (1, 0))
+  return valid_c, beta_h, cumsum_h, gating_forward, gating_last, gating_backward, mask_causal
+
+
+def _unit_rows(x: jax.Array, use_qk_norm: bool) -> jax.Array:
+  """The q / k l2 normalization of cp_gdn: x * rsqrt(sum(x**2) + 1e-6) over the head dim."""
+  if not use_qk_norm:
+    return x
+  return x * jax.lax.rsqrt(jnp.sum(x**2, axis=-1, keepdims=True) + 1e-6)
+
+
+def _m_scan_body(
+    k_ref,
+    b_ref,
+    a_ref,
+    t_inv_ref,
+    a_log_ref,
+    dt_bias_ref,
+    seg_prev_ref,
+    m_ref,
+    acc_scr,
+    *,
+    chunk_size: int,
+    num_kq_heads: int,
+    num_v_heads: int,
+    kq_head_dim: int,
+    num_chunks: int,
+    has_seg_ids: bool,
+    use_qk_norm: bool,
+):
+  """Forward step per (batch, group, chunk): acc <- M_c acc = g_last acc - k_bwd^T (w (acc)), from acc = I."""
+  c = pl.program_id(2)
+
+  @pl.when(c == 0)
+  def _init():
+    rows = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 0)
+    cols = jax.lax.broadcasted_iota(jnp.int32, (kq_head_dim, kq_head_dim), 1)
+    acc_scr[...] = jnp.broadcast_to((rows == cols).astype(jnp.float32)[None], acc_scr.shape)
+
+  valid_c, beta_h, _, gating_forward, gating_last, gating_backward, _ = _chunk_gates(
+      b_ref,
+      a_ref,
+      a_log_ref,
+      dt_bias_ref,
+      seg_prev_ref,
+      chunk_size=chunk_size,
+      num_v_heads=num_v_heads,
+      has_seg_ids=has_seg_ids,
+  )
+  k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  if valid_c is not None:
+    k_orig = jnp.where(valid_c[:, :, None], k_orig, 0.0)
+  k_h = jnp.transpose(jnp.repeat(_unit_rows(k_orig, use_qk_norm), num_v_heads // num_kq_heads, axis=1), (1, 0, 2))
+  k_beta_g = k_h * beta_h[:, :, None] * gating_forward
+  k_bwd = k_h * gating_backward
+
+  acc = acc_scr[...]
+  w = _precise_matmul(t_inv_ref[...].astype(jnp.float32), k_beta_g)
+  acc_new = acc * gating_last - _precise_matmul(jnp.swapaxes(k_bwd, -1, -2), _precise_matmul(w, acc))
+  acc_scr[...] = acc_new
+
+  @pl.when(c == num_chunks - 1)
+  def _store():
+    m_ref[0, ...] = acc_new
+
+
+def _dh_scan_body(
+    q_ref,
+    k_ref,
+    b_ref,
+    a_ref,
+    do_ref,
+    t_inv_ref,
+    a_log_ref,
+    dt_bias_ref,
+    seg_prev_ref,
+    ds_ref,
+    acc_scr,
+    *,
+    chunk_size: int,
+    num_kq_heads: int,
+    num_v_heads: int,
+    kq_head_dim: int,
+    num_chunks: int,
+    has_seg_ids: bool,
+    use_qk_norm: bool,
+):
+  """Reverse step per (batch, group, chunk): acc <- dS_c + M_c^T acc, from acc = 0.
+
+  M_c^T acc + dS_c is evaluated as the backward kernel's state recurrence does it:
+  g_last acc + q_g^T do - k_beta_g^T (t_inv^T (attn^T do + k_bwd acc)).
+  """
+  c = pl.program_id(2)
+  repeats = num_v_heads // num_kq_heads
+
+  @pl.when(c == 0)
+  def _init():
+    acc_scr[...] = jnp.zeros_like(acc_scr)
+
+  valid_c, beta_h, cumsum_h, gating_forward, gating_last, gating_backward, mask_causal = _chunk_gates(
+      b_ref,
+      a_ref,
+      a_log_ref,
+      dt_bias_ref,
+      seg_prev_ref,
+      chunk_size=chunk_size,
+      num_v_heads=num_v_heads,
+      has_seg_ids=has_seg_ids,
+  )
+  q_orig = q_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  do_val = do_ref[...].astype(jnp.float32)
+  if valid_c is not None:
+    q_orig = jnp.where(valid_c[:, :, None], q_orig, 0.0)
+    k_orig = jnp.where(valid_c[:, :, None], k_orig, 0.0)
+    do_val = jnp.where(valid_c[:, :, None], do_val, 0.0)
+  q_scaled = _unit_rows(q_orig, use_qk_norm) * (1.0 / jnp.sqrt(kq_head_dim))
+  k_scaled = _unit_rows(k_orig, use_qk_norm)
+
+  q_h = jnp.transpose(jnp.repeat(q_scaled, repeats, axis=1), (1, 0, 2))
+  k_h = jnp.transpose(jnp.repeat(k_scaled, repeats, axis=1), (1, 0, 2))
+  do_h = jnp.transpose(do_val, (1, 0, 2))
+  diff = cumsum_h[:, :, None] - cumsum_h[:, None, :]
+  g_mat_causal = jnp.exp(jnp.where(mask_causal[None] > 0.5, diff, -1e4)) * mask_causal[None]
+  q_g = q_h * gating_forward
+  k_beta_g = k_h * beta_h[:, :, None] * gating_forward
+  k_bwd = k_h * gating_backward
+
+  # q k^T once per kq head, then repeated to its value heads: the same dot products as with repeated q and k.
+  qk = _precise_matmul(jnp.transpose(q_scaled, (1, 0, 2)), jnp.swapaxes(jnp.transpose(k_scaled, (1, 0, 2)), -1, -2))
+  attn = jnp.repeat(qk, repeats, axis=0) * g_mat_causal
+  dv_attn = _precise_matmul(jnp.swapaxes(attn, -1, -2), do_h)
+
+  acc = acc_scr[...]
+  t_inv_t = jnp.swapaxes(t_inv_ref[...].astype(jnp.float32), -1, -2)
+  d_v_beta = _precise_matmul(t_inv_t, dv_attn + _precise_matmul(k_bwd, acc))
+  # q_g^T do - k_beta_g^T d_v_beta as one matmul over the stacked chunk rows.
+  acc_new = acc * gating_last + _precise_matmul(
+      jnp.swapaxes(jnp.concatenate([q_g, -k_beta_g], axis=1), -1, -2),
+      jnp.concatenate([do_h, d_v_beta], axis=1),
+  )
+  acc_scr[...] = acc_new
+
+  @pl.when(c == num_chunks - 1)
+  def _store():
+    ds_ref[0, ...] = acc_new
+
+
+def _run_scan(body, args, in_specs, out_shape, out_spec, acc_shape, *, grid, interpret, name):
+  """One pallas_call that pipelines `body` over grid=(B, G, C) with the running product in a VMEM scratch."""
+
+  def outer(*refs):
+    pltpu.emit_pipeline(body, grid=grid, in_specs=in_specs, out_specs=[out_spec])(
+        *refs[: len(args) + 1], scratches=tuple(refs[len(args) + 1 :])
+    )
+
+  hbm = pltpu.MemorySpace.HBM
+  return pl.pallas_call(
+      outer,
+      grid=(),
+      out_shape=out_shape,
+      in_specs=[pl.BlockSpec(memory_space=hbm)] * len(args),
+      out_specs=pl.BlockSpec(memory_space=hbm),
+      scratch_shapes=[pltpu.VMEM(acc_shape, jnp.float32)],
+      compiler_params=pltpu.CompilerParams(
+          vmem_limit_bytes=int(0.9 * pltpu.get_tpu_info().vmem_capacity_bytes),
+          disable_bounds_checks=True,
+      ),
+      interpret=interpret,
+      name=name,
+  )(*args)
+
+
+def _resolve_interpret(x: jax.Array, interpret: bool | pltpu.InterpretParams | None):
+  if interpret is None:
+    interpret = target_platform(x) == "cpu"
+  if interpret:
+    ensure_cpu_interpret_registered()
+  return interpret
+
+
+def gdn_cp_m_scan(
+    k_conv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    t_inv: jax.Array,
+    *,
+    num_v_heads: int,
+    kq_head_dim: int,
+    chunk_size: int = 64,
+    use_qk_norm_in_gdn: bool = False,
+    segment_ids: Optional[jax.Array] = None,
+    init_seg: Optional[jax.Array] = None,
+    interpret: bool | pltpu.InterpretParams | None = None,
+) -> jax.Array:
+  """M_local of `cp_gdn.compose_local_from_t_inv` (s_ext_pass1 branch), [B, num_v_heads, d_k, d_k] f32.
+
+  Args:
+    k_conv: [B, S, num_kq_heads * d_k] conv output of the K channels, as passed to compose_local_from_t_inv.
+    b, a: [B, S, num_v_heads] beta and decay pre-activations.
+    a_log, dt_bias: [num_v_heads] decay parameters.
+    t_inv: [B, S / chunk_size, num_v_heads, chunk_size, chunk_size] per-chunk T^{-1} of the forward kernel.
+    num_v_heads, kq_head_dim, chunk_size, use_qk_norm_in_gdn: GDN layer geometry.
+    segment_ids, init_seg: Rank-local segment metadata, as passed to compose_local_from_t_inv.
+    interpret: Pallas interpret mode; None interprets on CPU.
+  """
+  interpret = _resolve_interpret(k_conv, interpret)
+  batch_size, seq_len, k_size = k_conv.shape
+  num_chunks = seq_len // chunk_size
+  num_kq_heads = k_size // kq_head_dim
+  repeats = num_v_heads // num_kq_heads
+  tile_v_heads = _pick_head_tile(num_v_heads, repeats, _M_SCAN_HEAD_TILE)
+  num_groups = num_v_heads // tile_v_heads
+  tile_k_size = (tile_v_heads // repeats) * kq_head_dim
+  has_seg_ids = segment_ids is not None
+
+  b_5d, a_5d, a_log_4d, dt_bias_4d, seg_prev_5d, padded = _grouped_gate_inputs(
+      b,
+      a,
+      a_log,
+      dt_bias,
+      segment_ids,
+      init_seg,
+      num_chunks=num_chunks,
+      chunk_size=chunk_size,
+      num_groups=num_groups,
+      tile_v_heads=tile_v_heads,
+  )
+  args = [
+      k_conv.reshape(batch_size, num_chunks, chunk_size, k_size),
+      b_5d,
+      a_5d,
+      t_inv.reshape(batch_size, num_chunks, num_groups, tile_v_heads, chunk_size, chunk_size).astype(jnp.float32),
+      a_log_4d,
+      dt_bias_4d,
+      seg_prev_5d,
+  ]
+  in_specs = [
+      # k [B, C, chunk_size, k_size], columns of group g
+      pl.BlockSpec((None, None, chunk_size, tile_k_size), lambda b_, g, c: (b_, c, 0, g)),
+      # b, a [B, G, C, chunk_size, padded]
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, c, 0, 0)),
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, c, 0, 0)),
+      # t_inv [B, C, G, tile_v_heads, chunk_size, chunk_size]
+      pl.BlockSpec((None, None, None, tile_v_heads, chunk_size, chunk_size), lambda b_, g, c: (b_, c, g, 0, 0, 0)),
+      # a_log, dt_bias [B, G, 1, padded]
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      # segment entering the chunk [B, G, C, 1, 128]
+      pl.BlockSpec((None, None, None, 1, 128), lambda b_, g, c: (b_, g, c, 0, 0)),
+  ]
+  acc_shape = (tile_v_heads, kq_head_dim, kq_head_dim)
+  body = functools.partial(
+      _m_scan_body,
+      chunk_size=chunk_size,
+      num_kq_heads=tile_v_heads // repeats,
+      num_v_heads=tile_v_heads,
+      kq_head_dim=kq_head_dim,
+      num_chunks=num_chunks,
+      has_seg_ids=has_seg_ids,
+      use_qk_norm=use_qk_norm_in_gdn,
+  )
+  m_6d = _run_scan(
+      body,
+      args,
+      in_specs,
+      jax.ShapeDtypeStruct((batch_size, num_groups, 1) + acc_shape, jnp.float32),
+      pl.BlockSpec((None, None, 1) + acc_shape, lambda b_, g, c: (b_, g, 0, 0, 0, 0)),
+      acc_shape,
+      grid=(batch_size, num_groups, num_chunks),
+      interpret=interpret,
+      name="gdn_cp_m_scan",
+  )
+  return m_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, kq_head_dim)
+
+
+def gdn_cp_dh_scan(
+    qkv_conv: jax.Array,
+    b: jax.Array,
+    a: jax.Array,
+    a_log: jax.Array,
+    dt_bias: jax.Array,
+    do: jax.Array,
+    t_inv: jax.Array,
+    *,
+    num_v_heads: int,
+    kq_head_dim: int,
+    v_head_dim: int,
+    chunk_size: int = 64,
+    use_qk_norm_in_gdn: bool = False,
+    segment_ids: Optional[jax.Array] = None,
+    init_seg: Optional[jax.Array] = None,
+    interpret: bool | pltpu.InterpretParams | None = None,
+) -> jax.Array:
+  """dS_ext_local of `cp_gdn.compose_bwd_local_from_t_inv`, [B, num_v_heads, d_k, d_v] f32.
+
+  Args:
+    qkv_conv: [B, S, 2 * num_kq_heads * d_k + num_v_heads * d_v] conv output (q, k, v channels).
+    b, a: [B, S, num_v_heads] beta and decay pre-activations.
+    a_log, dt_bias: [num_v_heads] decay parameters.
+    do: [B, S, num_v_heads, d_v] cotangent of the GDN output.
+    t_inv: [B, S / chunk_size, num_v_heads, chunk_size, chunk_size] per-chunk T^{-1} of the forward kernel.
+    num_v_heads, kq_head_dim, v_head_dim, chunk_size, use_qk_norm_in_gdn: GDN layer geometry.
+    segment_ids, init_seg: Rank-local segment metadata, as passed to compose_bwd_local_from_t_inv.
+    interpret: Pallas interpret mode; None interprets on CPU.
+  """
+  interpret = _resolve_interpret(qkv_conv, interpret)
+  batch_size, seq_len, dim_size = qkv_conv.shape
+  num_chunks = seq_len // chunk_size
+  num_kq_heads = (dim_size - num_v_heads * v_head_dim) // (kq_head_dim * 2)
+  repeats = num_v_heads // num_kq_heads
+  tile_v_heads = _pick_head_tile(num_v_heads, repeats, _DH_SCAN_HEAD_TILE)
+  num_groups = num_v_heads // tile_v_heads
+  tile_q_size = (tile_v_heads // repeats) * kq_head_dim
+  k_block_offset = num_kq_heads * kq_head_dim // tile_q_size
+  has_seg_ids = segment_ids is not None
+
+  b_5d, a_5d, a_log_4d, dt_bias_4d, seg_prev_5d, padded = _grouped_gate_inputs(
+      b,
+      a,
+      a_log,
+      dt_bias,
+      segment_ids,
+      init_seg,
+      num_chunks=num_chunks,
+      chunk_size=chunk_size,
+      num_groups=num_groups,
+      tile_v_heads=tile_v_heads,
+  )
+  qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, dim_size)
+  args = [
+      qkv_conv_4d,
+      qkv_conv_4d,
+      b_5d,
+      a_5d,
+      do.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_heads, v_head_dim),
+      t_inv.reshape(batch_size, num_chunks, num_groups, tile_v_heads, chunk_size, chunk_size).astype(jnp.float32),
+      a_log_4d,
+      dt_bias_4d,
+      seg_prev_5d,
+  ]
+
+  def rc(c):
+    return num_chunks - 1 - c
+
+  in_specs = [
+      # q, k [B, C, chunk_size, dim_size], columns of group g (k after the q blocks)
+      pl.BlockSpec((None, None, chunk_size, tile_q_size), lambda b_, g, c: (b_, rc(c), 0, g)),
+      pl.BlockSpec((None, None, chunk_size, tile_q_size), lambda b_, g, c: (b_, rc(c), 0, k_block_offset + g)),
+      # b, a [B, G, C, chunk_size, padded]
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, rc(c), 0, 0)),
+      pl.BlockSpec((None, None, None, chunk_size, padded), lambda b_, g, c: (b_, g, rc(c), 0, 0)),
+      # do [B, C, chunk_size, G, tile_v_heads, d_v]
+      pl.BlockSpec((None, None, chunk_size, None, tile_v_heads, v_head_dim), lambda b_, g, c: (b_, rc(c), 0, g, 0, 0)),
+      # t_inv [B, C, G, tile_v_heads, chunk_size, chunk_size]
+      pl.BlockSpec((None, None, None, tile_v_heads, chunk_size, chunk_size), lambda b_, g, c: (b_, rc(c), g, 0, 0, 0)),
+      # a_log, dt_bias [B, G, 1, padded]
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      pl.BlockSpec((None, None, 1, padded), lambda b_, g, c: (b_, g, 0, 0)),
+      # segment entering the chunk [B, G, C, 1, 128]
+      pl.BlockSpec((None, None, None, 1, 128), lambda b_, g, c: (b_, g, rc(c), 0, 0)),
+  ]
+  acc_shape = (tile_v_heads, kq_head_dim, v_head_dim)
+  body = functools.partial(
+      _dh_scan_body,
+      chunk_size=chunk_size,
+      num_kq_heads=tile_v_heads // repeats,
+      num_v_heads=tile_v_heads,
+      kq_head_dim=kq_head_dim,
+      num_chunks=num_chunks,
+      has_seg_ids=has_seg_ids,
+      use_qk_norm=use_qk_norm_in_gdn,
+  )
+  ds_6d = _run_scan(
+      body,
+      args,
+      in_specs,
+      jax.ShapeDtypeStruct((batch_size, num_groups, 1) + acc_shape, jnp.float32),
+      pl.BlockSpec((None, None, 1) + acc_shape, lambda b_, g, c: (b_, g, 0, 0, 0, 0)),
+      acc_shape,
+      grid=(batch_size, num_groups, num_chunks),
+      interpret=interpret,
+      name="gdn_cp_dh_scan",
+  )
+  return ds_6d[:, :, 0].reshape(batch_size, num_v_heads, kq_head_dim, v_head_dim)
diff --git a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py
index f4c490af8b..d57de46af2 100644
--- a/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py
+++ b/src/maxtext/kernels/gdn/gdn_bwd/pallas_mosaic_tpu_bwd.py
@@ -139,13 +139,15 @@ def _bwd_gdn_pipeline_body(
     has_dh0: bool = False,
 ) -> None:
   """Inner kernel executed per (batch, group, chunk) by emit_pipeline with manual GDN backward."""
-  if len(refs) == 17:
+  if len(refs) == 21:
     has_dht = True
     has_dh0 = True
   idx = 0
   # pylint: disable=unbalanced-tuple-unpacking
   (
-      qkv_conv_ref,
+      q_ref,
+      k_ref,
+      v_ref,
       b_ref,
       a_ref,
       do_ref,
@@ -154,22 +156,24 @@ def _bwd_gdn_pipeline_body(
       a_log_ref,
       dt_bias_ref,
       reset_ref,
-  ) = refs[idx : idx + 9]
-  idx += 9
+  ) = refs[idx : idx + 11]
+  idx += 11
   if has_dht:
     dht_ref = refs[idx]
     idx += 1
   else:
     dht_ref = None
   (
-      dy_conv_ref,
+      dq_ref,
+      dk_ref,
+      dv_ref,
       d_b_ref,
       d_a_ref,
       d_a_log_ref,
       d_dt_bias_ref,
-  ) = refs[idx : idx + 5]
+  ) = refs[idx : idx + 7]
   # pylint: enable=unbalanced-tuple-unpacking
-  idx += 5
+  idx += 7
   if has_dh0:
     dh0_ref = refs[idx]
     idx += 1
@@ -198,17 +202,12 @@ def _init():
 
   is_reset = reset_ref[...][0, 0] > 0.5
   d_state = jnp.where(is_reset, 0.0, d_state_scr[...])
-  compute_dtype = qkv_conv_ref.dtype
-  y_c = qkv_conv_ref[...]
-
-  # Slice chunk inputs for this head group
-  q_orig = y_c[:, :q_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
-  k_orig = y_c[:, q_size : q_size + k_size].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
-  v = (
-      y_c[:, q_size + k_size : q_size + k_size + v_size]
-      .reshape((chunk_size, num_v_heads, v_head_dim))
-      .astype(jnp.float32)
-  )
+  compute_dtype = q_ref.dtype
+
+  # Chunk inputs for this head group
+  q_orig = q_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  k_orig = k_ref[...].reshape((chunk_size, num_kq_heads, kq_head_dim)).astype(jnp.float32)
+  v = v_ref[...].reshape((chunk_size, num_v_heads, v_head_dim)).astype(jnp.float32)
 
   b_val = b_ref[...][:, :num_v_heads].astype(jnp.float32)
   a_val = a_ref[...][:, :num_v_heads].astype(jnp.float32)
@@ -446,11 +445,10 @@ def _init():
     d_q = jnp.where(valid_c[:, :, None], d_q, 0.0)
     d_k = jnp.where(valid_c[:, :, None], d_k, 0.0)
 
-  # Flatten gradients to chunk_size x dim_size and write to refs
-  d_q_flat = d_q.reshape(chunk_size, q_size).astype(dy_conv_ref.dtype)
-  d_k_flat = d_k.reshape(chunk_size, k_size).astype(dy_conv_ref.dtype)
-  d_v_flat = d_v_val.reshape(chunk_size, v_size).astype(dy_conv_ref.dtype)
-  dy_conv_ref[...] = jnp.concatenate([d_q_flat, d_k_flat, d_v_flat], axis=-1)
+  # Flatten gradients to chunk_size x q/k/v size and write to refs
+  dq_ref[...] = d_q.reshape(chunk_size, q_size).astype(dq_ref.dtype)
+  dk_ref[...] = d_k.reshape(chunk_size, k_size).astype(dk_ref.dtype)
+  dv_ref[...] = d_v_val.reshape(chunk_size, v_size).astype(dv_ref.dtype)
 
   if padded_num_v_heads > num_v_heads:
     d_b_ref[...] = jnp.pad(d_b_val, ((0, 0), (0, padded_num_v_heads - num_v_heads))).astype(d_b_ref.dtype)
@@ -619,14 +617,15 @@ def pallas_gdn_bwd_kernel(
 
   padded_tile_v_heads = cfg.padded_num_v_heads
 
-  # 1. Prepare qkv_conv into (B, G, C, chunk_size, group_dim_size)
-  q = qkv_conv[:, :, :q_size]
-  k = qkv_conv[:, :, q_size : q_size + k_size]
-  v = qkv_conv[:, :, q_size + k_size :]
-  q_5d = q.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_q_size).transpose((0, 3, 1, 2, 4))
-  k_5d = k.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_k_size).transpose((0, 3, 1, 2, 4))
-  v_5d = v.reshape(batch_size, num_chunks, chunk_size, num_groups, tile_v_size).transpose((0, 3, 1, 2, 4))
-  qkv_conv_5d = jnp.concatenate([q_5d, k_5d, v_5d], axis=-1)
+  # 1. View qkv_conv as (B, C, chunk_size, dim_size); the kernel reads group g's q, k and v as column blocks
+  # (k after the num_groups q blocks, v after q + k). If q + k is not a whole number of v blocks, v is sliced out.
+  qkv_conv_4d = qkv_conv.reshape(batch_size, num_chunks, chunk_size, dim_size)
+  if (q_size + k_size) % tile_v_size == 0:
+    v_4d = qkv_conv_4d
+    v_block_offset = (q_size + k_size) // tile_v_size
+  else:
+    v_4d = qkv_conv_4d[..., q_size + k_size :]
+    v_block_offset = 0
 
   # 2. Prepare b and a into (B, G, C, chunk_size, padded_tile_v_heads)
   # With sequence packing, the signed segment IDs are packed into the spare lane
@@ -659,7 +658,7 @@ def pallas_gdn_bwd_kernel(
         ),
     )
 
-  # 3. Prepare do into (B, G, C, chunk_size, tile_v_heads, v_head_dim)
+  # 3. View do as (B, C, chunk_size, G, tile_v_heads, v_head_dim); the kernel reads group g in place
   do_6d = do.reshape(
       batch_size,
       num_chunks,
@@ -667,9 +666,9 @@ def pallas_gdn_bwd_kernel(
       num_groups,
       tile_v_heads,
       cfg.v_head_dim,
-  ).transpose((0, 3, 1, 2, 4, 5))
+  )
 
-  # 4. Prepare chunk_states into (B, G, C, tile_v_heads, kq_head_dim, v_head_dim)
+  # 4. View chunk_states as (B, C, G, tile_v_heads, kq_head_dim, v_head_dim)
   chunk_states_6d = chunk_states.reshape(
       batch_size,
       num_chunks,
@@ -677,22 +676,18 @@ def pallas_gdn_bwd_kernel(
       tile_v_heads,
       cfg.kq_head_dim,
       cfg.v_head_dim,
-  ).swapaxes(1, 2)
-
-  # 5. Prepare t_inv into (B, G, C, tile_v_heads, chunk_size, chunk_size)
-  t_inv_6d = (
-      t_inv.reshape(
-          batch_size,
-          num_chunks,
-          num_groups,
-          tile_v_heads,
-          chunk_size,
-          chunk_size,
-      )
-      .swapaxes(1, 2)
-      .astype(jnp.float32)
   )
 
+  # 5. View t_inv as (B, C, G, tile_v_heads, chunk_size, chunk_size)
+  t_inv_6d = t_inv.reshape(
+      batch_size,
+      num_chunks,
+      num_groups,
+      tile_v_heads,
+      chunk_size,
+      chunk_size,
+  ).astype(jnp.float32)
+
   # 6. Prepare a_log and dt_bias into (B, G, 1, padded_tile_v_heads)
   a_log_2d = a_log.reshape(num_groups, tile_v_heads)
   if padded_tile_v_heads > tile_v_heads:
@@ -739,7 +734,9 @@ def pallas_gdn_bwd_kernel(
   reset_hbm_5d = jnp.broadcast_to(reset_hbm[:, None, :, :, :], (batch_size, num_groups, num_chunks, 1, 128))
 
   in_args = [
-      qkv_conv_5d,
+      qkv_conv_4d,
+      qkv_conv_4d,
+      v_4d,
       b_5d,
       a_5d,
       do_6d,
@@ -767,20 +764,22 @@ def pallas_gdn_bwd_kernel(
   in_specs, out_specs, num_in, num_out = make_bwd_block_specs(
       num_chunks=num_chunks,
       chunk_size=chunk_size,
-      dim_size=group_dim_size,
+      qk_tile_size=tile_q_size,
+      v_tile_size=tile_v_size,
       num_v_heads=tile_v_heads,
       kq_head_dim=cfg.kq_head_dim,
       v_head_dim=cfg.v_head_dim,
+      k_block_offset=q_size // tile_k_size,
+      v_block_offset=v_block_offset,
       padded_num_v_heads=padded_tile_v_heads,
       has_dht=has_dht,
       has_dh0=has_dh0,
   )
 
   out_shapes_list = [
-      jax.ShapeDtypeStruct(
-          (batch_size, num_groups, num_chunks, chunk_size, group_dim_size),
-          qkv_conv.dtype,
-      ),
+      jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, q_size), qkv_conv.dtype),
+      jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, k_size), qkv_conv.dtype),
+      jax.ShapeDtypeStruct((batch_size, num_chunks, chunk_size, v_size), qkv_conv.dtype),
       jax.ShapeDtypeStruct(b_5d.shape, b_5d.dtype),
       jax.ShapeDtypeStruct(a_5d.shape, a_5d.dtype),
       jax.ShapeDtypeStruct(
@@ -849,7 +848,9 @@ def outer(*refs):
 
   if has_dh0:
     (
-        dy_conv_chunks,
+        dq_4d,
+        dk_4d,
+        dv_4d,
         d_b_chunks,
         d_a_chunks,
         d_a_log_chunks,
@@ -858,7 +859,9 @@ def outer(*refs):
     ) = pallas_out
   else:
     (
-        dy_conv_chunks,
+        dq_4d,
+        dk_4d,
+        dv_4d,
         d_b_chunks,
         d_a_chunks,
         d_a_log_chunks,
@@ -866,13 +869,10 @@ def outer(*refs):
     ) = pallas_out
     dh0_6d = None
 
-  # Reconstruct dy_conv: (B, G, C, S, group_dim_size) -> (B, seq_len, dim_size)
-  dq_5d = dy_conv_chunks[..., :tile_q_size]
-  dk_5d = dy_conv_chunks[..., tile_q_size : tile_q_size + tile_k_size]
-  dv_5d = dy_conv_chunks[..., tile_q_size + tile_k_size :]
-  dq = dq_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, q_size)
-  dk = dk_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, k_size)
-  dv = dv_5d.transpose((0, 2, 3, 1, 4)).reshape(batch_size, seq_len, v_size)
+  # Reconstruct dy_conv: (B, C, S, q/k/v size) -> (B, seq_len, dim_size)
+  dq = dq_4d.reshape(batch_size, seq_len, q_size)
+  dk = dk_4d.reshape(batch_size, seq_len, k_size)
+  dv = dv_4d.reshape(batch_size, seq_len, v_size)
   dy_conv_flat = jnp.concatenate([dq, dk, dv], axis=-1).astype(qkv_conv.dtype)
 
   # Reconstruct d_b and d_a: (B, G, C, S, padded_tile_v_heads) -> (B, seq_len, num_v_heads)
diff --git a/src/maxtext/kernels/gdn/memory_ref.py b/src/maxtext/kernels/gdn/memory_ref.py
index f29d79d81d..a1c4545ef5 100644
--- a/src/maxtext/kernels/gdn/memory_ref.py
+++ b/src/maxtext/kernels/gdn/memory_ref.py
@@ -281,6 +281,36 @@ def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
     ).wait()
 
 
+@jax.tree_util.register_dataclass
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class ChunkInBufferedRef(BaseBufferedRef):
+  """Input DMA reference that copies whole chunks of a [num_chunks, chunk, dim] array."""
+
+  def copy_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
+    """Initiates asynchronous copy of each tile's chunk into current input VMEM slot."""
+    assert self.sem_recvs is not None
+    assert self.window_ref is not None
+    slot = self.current_copy_in_slot
+    sem = self.sem_recvs.at[slot]
+    vmem_ref = self.window_ref.at[slot]
+    p_id = grid_indices[0]
+
+    for idx in range(self.cfg.seq_tile_size):
+      chunk_idx = self.metadata_ref.get_record(p_id, idx).r_base // self.cfg.chunk_size
+      pltpu.make_async_copy(src_ref.at[chunk_idx], vmem_ref.at[idx], sem).start()
+
+  def wait_in(self, src_ref: jax.Array, grid_indices: tuple[int | jax.Array]):
+    """Waits for asynchronous chunk copy to complete."""
+    assert self.sem_recvs is not None
+    assert self.window_ref is not None
+    slot = self.current_wait_in_slot
+    sem = self.sem_recvs.at[slot]
+    vmem_ref = self.window_ref.at[slot]
+
+    for idx in range(self.cfg.seq_tile_size):
+      pltpu.make_async_copy(vmem_ref.at[idx], vmem_ref.at[idx], sem).wait()
+
+
 @jax.tree_util.register_dataclass
 @dataclasses.dataclass(frozen=True, kw_only=True)
 class OutBufferedRef(BaseBufferedRef):
@@ -496,15 +526,18 @@ def create_allocs(
     qkv_ref: jax.Array,
     b_ref: jax.Array,
     a_ref: jax.Array,
-    out_ref: jax.Array,
+    out_ref: jax.Array | None,
     conv_state_ref: jax.Array,
     recurrent_state_ref: jax.Array,
     cfg: config.GDNConfig,
     t_inv_ref: jax.Array | None = None,
     chunk_states_ref: jax.Array | None = None,
 ) -> tuple[Any, ...]:
-  """Creates all Pallas double-buffered allocations for GDN kernel."""
-  qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size)
+  """Creates all Pallas double-buffered allocations for GDN kernel (no output buffers without out_ref)."""
+  if cfg.is_prefill_only:
+    qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, cfg.dim_size)
+  else:
+    qkv_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.dim_size)
   ba_shape = (cfg.seq_tile_size, cfg.chunk_size, 1, cfg.aligned_num_v_heads)
 
   out_shape = (
@@ -539,25 +572,31 @@ def create_allocs(
       cfg=cfg,
       metadata_ref=metadata_ref,
   )
-  qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref)
+  if cfg.is_prefill_only:
+    qkv_alloc = ChunkInBufferedRef.input(
+        spec=qkv_spec,
+        dtype_or_type=qkv_ref,
+        buffer_count=pipeline_mode.buffer_count,
+        use_lookahead=pipeline_mode.use_lookahead,
+        cfg=cfg,
+        metadata_ref=metadata_ref,
+    )
+  else:
+    qkv_alloc = in_buffered_partial(spec=qkv_spec, dtype_or_type=qkv_ref)
   b_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=b_ref)
   a_alloc = in_buffered_partial(spec=ba_spec, dtype_or_type=a_ref)
 
-  out_alloc = OutBufferedRef.output(
-      spec=block_spec_partial(block_shape=out_shape),
-      dtype_or_type=out_ref,
-      buffer_count=pipeline_mode.buffer_count,
-      use_lookahead=pipeline_mode.use_lookahead,
-      cfg=cfg,
-      metadata_ref=metadata_ref,
-  )
-
-  conv_spec = block_spec_partial(block_shape=conv_shape)
-  recurrent_spec = block_spec_partial(block_shape=recurrent_shape)
+  # A prefill call with one sequence copies the states in only before its first tile and out only after
+  # its last, so no copy is in flight while a tile uses the slot: one buffer suffices and frees VMEM.
+  state_pipeline_mode = pipeline_mode
+  if cfg.is_prefill_only and metadata_ref.s_idx_has_initial_state.shape[0] == 1:
+    state_pipeline_mode = pl.Buffered(buffer_count=1, use_lookahead=False)
+  conv_spec = block_spec_partial(block_shape=conv_shape, pipeline_mode=state_pipeline_mode)
+  recurrent_spec = block_spec_partial(block_shape=recurrent_shape, pipeline_mode=state_pipeline_mode)
   state_buffered_partial = functools.partial(
       StateBufferedRef.input_output,
-      buffer_count=pipeline_mode.buffer_count,
-      use_lookahead=pipeline_mode.use_lookahead,
+      buffer_count=state_pipeline_mode.buffer_count,
+      use_lookahead=state_pipeline_mode.use_lookahead,
       cfg=cfg,
       metadata_ref=metadata_ref,
   )
@@ -570,9 +609,19 @@ def create_allocs(
       a_alloc,
       conv_alloc,
       recurrent_alloc,
-      out_alloc,
   ]
 
+  if out_ref is not None:
+    out_alloc = OutBufferedRef.output(
+        spec=block_spec_partial(block_shape=out_shape),
+        dtype_or_type=out_ref,
+        buffer_count=pipeline_mode.buffer_count,
+        use_lookahead=pipeline_mode.use_lookahead,
+        cfg=cfg,
+        metadata_ref=metadata_ref,
+    )
+    allocs.append(out_alloc)
+
   if t_inv_ref is not None:
     t_inv_shape = (
         cfg.seq_tile_size,
diff --git a/src/maxtext/kernels/gdn/vmem_ldst.py b/src/maxtext/kernels/gdn/vmem_ldst.py
index 7d604a69b1..3cc5225c48 100644
--- a/src/maxtext/kernels/gdn/vmem_ldst.py
+++ b/src/maxtext/kernels/gdn/vmem_ldst.py
@@ -134,6 +134,28 @@ def load_compact_to_large(vmem_ref: jax.Ref) -> jax.Array:
   return jnp.concat(vreg_list, axis=-1).reshape(new_shape)
 
 
+def upcast_large_to_compact(large_ref: jax.Ref, compact_ref: jax.Ref) -> None:
+  """Upcast activations from large layout into a compact fp32 reference.
+
+  The inverse of load_as_qkv_large: one strided store per 128-lane column puts
+  each token's lanes into its own row of the compact reference.
+
+  Args:
+    large_ref: Activations in VMEM of shape [seq_tile_size, chunk_size,
+      dim_size].
+    compact_ref: fp32 reference in VMEM of shape [seq_tile_size, chunk_size, 1,
+      dim_size].
+  """
+
+  num_lanes = pltpu.get_tpu_info().num_lanes
+  lanes_per_col = compact_ref.shape[-1] // num_lanes
+  for idx in range(compact_ref.shape[0]):
+    compact_flat_ref = compact_ref.at[idx].reshape(-1, num_lanes)  # pyrefly: ignore[missing-attribute]
+    for lane in range(lanes_per_col):
+      col = large_ref[idx, :, lane * num_lanes : (lane + 1) * num_lanes]
+      compact_flat_ref[lane::lanes_per_col] = col.astype(compact_ref.dtype)
+
+
 def load_and_select_states(
     metadata_ref: memory_ref.MetadataRef,
     p_id: jax.Array,
diff --git a/src/maxtext/kernels/gdn/wrapper.py b/src/maxtext/kernels/gdn/wrapper.py
index 7b951c7cee..65b52a8a89 100644
--- a/src/maxtext/kernels/gdn/wrapper.py
+++ b/src/maxtext/kernels/gdn/wrapper.py
@@ -95,14 +95,14 @@ def _pack_fwd_segment_metadata(
 
 def inner_kernel(
     # Inputs.
-    qkv_slot_ref: jax.Ref,  # [seq, chunk, 1, dim_size]
+    qkv_slot_ref: jax.Ref,  # [seq, chunk, 1, dim_size] ([seq, chunk, dim_size] for prefill-only)
     b_slot_ref: jax.Ref,  # [seq, chunk, 1, num_v_heads]
     a_slot_ref: jax.Ref,  # [seq, chunk, 1, num_v_heads]
     conv_state_slot_ref: jax.Ref,  # [seq, prev_kernel_size, 1, dim_size]
     recurrent_slot_ref: jax.Ref,  # [seq, num_v_heads, kq_head, v_head]
-    # Outputs.
-    out_slot_ref: jax.Array,  # [seq * chunk, num_v_heads, v_head]
-    t_inv_slot_ref: jax.Array,  # [seq, num_v_heads, chunk, chunk]
+    # Outputs, then scratches:
+    # out [seq * chunk, num_v_heads, v_head] (not in PER_SEQ_STATES), t_inv [seq, num_v_heads, chunk, chunk],
+    # chunk_states [seq, num_v_heads, kq_head, v_head] (PER_SEQ only).
     *args,
     cfg: config.GDNConfig,
     **kwargs,
@@ -114,21 +114,28 @@ def inner_kernel(
   kernel invokes vmem_ldst to pre-processes data needed for compute and
   invokes compute_conv1d and compute_gdn for actual compute.
   """
-  if cfg.mode == config.GDNMode.PER_SEQ:
-    chunk_states_slot_ref = args[0]
-    metadata_ref = args[1]
-    weights_ref = args[2]
-    carry_conv_scratch_ref = args[3] if len(args) > 3 else None
-    carry_recurrent_scratch_ref = args[4] if len(args) > 4 else None
+  if cfg.mode == config.GDNMode.PER_SEQ_STATES:
+    out_slot_ref = chunk_states_slot_ref = None
+    t_inv_slot_ref, *args = args
+  elif cfg.mode == config.GDNMode.PER_SEQ:
+    out_slot_ref, t_inv_slot_ref, chunk_states_slot_ref, *args = args
   else:
+    out_slot_ref, t_inv_slot_ref, *args = args
     chunk_states_slot_ref = None
-    metadata_ref = args[0]
-    weights_ref = args[1]
-    carry_conv_scratch_ref = args[2] if len(args) > 2 else None
-    carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None
+  metadata_ref = args[0]
+  weights_ref = args[1]
+  carry_conv_scratch_ref = args[2] if len(args) > 2 else None
+  carry_recurrent_scratch_ref = args[3] if len(args) > 3 else None
+  qkv_scratch_ref = args[4] if len(args) > 4 else None
 
   p_id = pl.program_id(0)
 
+  # Prefill-only qkv arrives in act_in dtype and large layout. Upcast it into the
+  # compact fp32 scratch, which then serves as the qkv slot below.
+  if qkv_scratch_ref is not None:
+    vmem_ldst.upcast_large_to_compact(qkv_slot_ref, qkv_scratch_ref)
+    qkv_slot_ref = qkv_scratch_ref
+
   # Prepare states.
   real_sizes, prev_conv, prev_recurrent = vmem_ldst.load_and_select_states(
       metadata_ref=metadata_ref,
@@ -216,7 +223,8 @@ def inner_kernel(
     )
 
   # Store output, recurrent, and t_inv to vmem.
-  out_slot_ref[...] = out.astype(out_slot_ref.dtype)
+  if out_slot_ref is not None:
+    out_slot_ref[...] = out.astype(out_slot_ref.dtype)
   recurrent_slot_ref[...] = new_recurrent_state.astype(recurrent_slot_ref.dtype)
   t_inv_slot_ref[...] = t_inv.astype(t_inv_slot_ref.dtype)
   if chunk_states_slot_ref is not None:
@@ -236,19 +244,21 @@ def outer_kernel(
     recurrent_state_ref: jax.Array,
     _: jax.Array,
     weights_ref: memory_ref.WeightRefs,
-    # Outputs.
-    out_ref: jax.Array,
-    conv_state_out_ref: jax.Array,
-    recurrent_state_out_ref: jax.Array,
-    t_inv_ref: jax.Array,
+    # Outputs: out (not in PER_SEQ_STATES), conv_state, recurrent_state, t_inv, chunk_states (PER_SEQ only).
+    # The state outputs alias the state inputs, which the pipeline updates in place.
     *args,
     carry_conv_scratch_ref: jax.Array | None = None,
     carry_recurrent_scratch_ref: jax.Array | None = None,
+    qkv_scratch_ref: jax.Array | None = None,
     cfg: config.GDNConfig,
     **kwargs,
 ) -> None:
   """Setup memory allocations and emit pipeline for running inner_kernel."""
-  del conv_state_out_ref, recurrent_state_out_ref
+  if cfg.mode == config.GDNMode.PER_SEQ_STATES:
+    out_ref = None
+    _, _, t_inv_ref, *args = args
+  else:
+    out_ref, _, _, t_inv_ref, *args = args
 
   chunk_states_ref = args[0] if (len(args) > 0 and cfg.mode == config.GDNMode.PER_SEQ) else None
 
@@ -269,16 +279,11 @@ def outer_kernel(
   a_alloc = allocs[2]
   conv_alloc = allocs[3]
   recurrent_alloc = allocs[4]
-  out_alloc = allocs[5]
-  t_inv_alloc = allocs[6]
-  chunk_states_alloc = allocs[7] if len(allocs) > 7 else None
+  # [out,] t_inv[, chunk_states], in the order of the pipelined output refs below.
+  out_specs = [alloc.spec for alloc in allocs[5:]]
 
   num_tiles = metadata_ref.num_tiles[...]
 
-  out_specs = [out_alloc.spec, t_inv_alloc.spec]
-  if chunk_states_alloc is not None:
-    out_specs.append(chunk_states_alloc.spec)
-
   pipeline_func = pltpu.emit_pipeline(
       body=functools.partial(
           inner_kernel,
@@ -297,9 +302,7 @@ def outer_kernel(
 
   @pl.with_scoped(allocations=allocs)
   def _run(allocations):
-    out_args = [out_ref, t_inv_ref]
-    if chunk_states_ref is not None:
-      out_args.append(chunk_states_ref)
+    out_args = [ref for ref in (out_ref, t_inv_ref, chunk_states_ref) if ref is not None]
     pipeline_func(
         qkv_ref,
         b_ref,
@@ -312,6 +315,7 @@ def _run(allocations):
             weights_ref,
             carry_conv_scratch_ref,
             carry_recurrent_scratch_ref,
+            qkv_scratch_ref,
         ),
         allocations=allocations,
     )
@@ -320,23 +324,7 @@ def _run(allocations):
   _run()
 
 
-@jax.jit(
-    donate_argnames=("conv_state", "recurrent_state"),
-    static_argnames=(
-        "n_kq",
-        "n_v",
-        "d_k",
-        "d_v",
-        "kernel_size",
-        "decode_tile_size",
-        "mixed_tile_size",
-        "zero_initialize_out",
-        "compute_precision",
-        "is_prefill_only",
-        "use_qk_norm_in_gdn",
-    ),
-)
-def fused_conv1d_gdn(
+def _fused_conv1d_gdn(
     qkv: jax.Array,  # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size]
     b: jax.Array,  # [batch_size, n_v]
     a: jax.Array,  # [batch_size, n_v]
@@ -365,15 +353,15 @@ def fused_conv1d_gdn(
     segment_ids: jax.Array | None = None,
     conv_halo_seg: jax.Array | None = None,
     init_seg: jax.Array | None = None,
-) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]:
-  """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states)."""
+    prefill_mode: config.GDNMode = config.GDNMode.PER_SEQ,
+) -> tuple[jax.Array | None, tuple[jax.Array, jax.Array], jax.Array, jax.Array | None]:
+  """Implements fused_conv1d_gdn; prefill_mode=PER_SEQ_STATES returns no out and no chunk_states."""
   act_in_dtype = qkv.dtype
   act_out_dtype = qkv.dtype
   conv_out_dtype = conv_state.dtype
   recurrent_out_dtype = recurrent_state.dtype
   assert a.dtype == b.dtype == qkv.dtype == act_in_dtype
 
-  qkv = qkv.astype(jnp.float32)
   b = b.astype(jnp.float32)
   a = a.astype(jnp.float32)
   conv_state = conv_state.astype(jnp.float32)
@@ -441,7 +429,6 @@ def fused_conv1d_gdn(
   batch_padding_size = padded_batch_size - batch_size
   aligned_num_v_heads = tiling.align_to(n_v + extra_lanes, num_lanes)
   num_v_padding_size = aligned_num_v_heads - n_v
-  qkv = jnp.pad(qkv, ((0, batch_padding_size), (0, 0)))
   b = jnp.pad(b, ((0, batch_padding_size), (0, num_v_padding_size)))
   a = jnp.pad(a, ((0, batch_padding_size), (0, num_v_padding_size)))
   if has_seg_ids:
@@ -456,7 +443,12 @@ def fused_conv1d_gdn(
     b = b.at[:batch_size, n_v].set(s_enc_flat)
     b = b.at[:batch_size, n_v + 1].set(seg_aux_flat)
 
-  qkv = qkv.reshape(padded_batch_size, 1, -1)
+  if is_prefill_only:
+    # Whole chunks, DMA'd in act_in dtype and upcast to fp32 in VMEM by the kernel.
+    qkv = qkv.reshape(batch_size // mixed_tile_size, mixed_tile_size, dim)
+  else:
+    qkv = jnp.pad(qkv.astype(jnp.float32), ((0, batch_padding_size), (0, 0)))
+    qkv = qkv.reshape(padded_batch_size, 1, -1)
   b = b.reshape(padded_batch_size, 1, -1)
   a = a.reshape(padded_batch_size, 1, -1)
 
@@ -500,6 +492,7 @@ def call_kernel(
         v_head_dim=d_v,
         has_seg_ids=has_seg_ids,
         use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+        is_prefill_only=is_prefill_only,
         dtypes=config.Dtypes(
             act_in=act_in_dtype,
             act_out=act_out_dtype,
@@ -534,7 +527,7 @@ def call_kernel(
     input_output_aliases = {len(metadata_obj) + 3: 1, len(metadata_obj) + 4: 2}
     out_shape = cfg.get_out_shape()
 
-    if in_act is None and zero_initialize_out:
+    if in_act is None and zero_initialize_out and mode != config.GDNMode.PER_SEQ_STATES:
       in_act = jnp.zeros_like(out_shape)
     if in_act is not None:
       out_shape = in_act
@@ -547,7 +540,12 @@ def call_kernel(
         cfg.dtypes.compute,
     )
 
-    if mode == config.GDNMode.PER_SEQ:
+    if mode == config.GDNMode.PER_SEQ_STATES:
+      # No output, so the states move up to outputs 0 and 1.
+      input_output_aliases = {len(metadata_obj) + 3: 0, len(metadata_obj) + 4: 1}
+      out_shape_tuple = (in_conv_state, in_recurrent_state, t_inv_shape)
+      out_specs_tuple = (hbm_spec, hbm_spec, hbm_spec)
+    elif mode == config.GDNMode.PER_SEQ:
       chunk_states_shape = jax.ShapeDtypeStruct(
           (num_chunks, cfg.num_v_heads, cfg.kq_head_dim, cfg.v_head_dim),
           cfg.dtypes.compute,
@@ -611,11 +609,15 @@ def call_kernel(
     out_conv_state = conv_state
     out_recurrent_state = recurrent_state
 
-  out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel(
-      out_conv_state, out_recurrent_state, out_act, config.GDNMode.PER_SEQ
-  )
+  if prefill_mode == config.GDNMode.PER_SEQ_STATES:
+    out_conv_state, out_recurrent_state, t_inv = call_kernel(out_conv_state, out_recurrent_state, None, prefill_mode)
+    out_act = chunk_states = None
+  else:
+    out_act, out_conv_state, out_recurrent_state, t_inv, chunk_states = call_kernel(
+        out_conv_state, out_recurrent_state, out_act, prefill_mode
+    )
+    out_act = out_act.reshape(padded_batch_size, -1)[:batch_size]
 
-  out_act = out_act.reshape(padded_batch_size, -1)[:batch_size]
   out_conv_state = out_conv_state.astype(conv_out_dtype)
   out_conv_state = out_conv_state.reshape(conv_state_shape)
   if has_seg_ids and segment_ids is not None:
@@ -630,3 +632,158 @@ def call_kernel(
   out_recurrent_state = out_recurrent_state.astype(recurrent_out_dtype)
 
   return out_act, (out_conv_state, out_recurrent_state), t_inv, chunk_states
+
+
+@jax.jit(
+    donate_argnames=("conv_state", "recurrent_state"),
+    static_argnames=(
+        "n_kq",
+        "n_v",
+        "d_k",
+        "d_v",
+        "kernel_size",
+        "decode_tile_size",
+        "mixed_tile_size",
+        "zero_initialize_out",
+        "compute_precision",
+        "is_prefill_only",
+        "use_qk_norm_in_gdn",
+    ),
+)
+def fused_conv1d_gdn(
+    qkv: jax.Array,  # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size]
+    b: jax.Array,  # [batch_size, n_v]
+    a: jax.Array,  # [batch_size, n_v]
+    conv_state: jax.Array,  # [num_seqs + 1, kernel_size - 1, dim_size]
+    recurrent_state: jax.Array,  # [num_seqs + 1, nv, dk, dv]
+    conv_weight: jax.Array,  # [kernel_size - 1, dim_size]
+    conv_bias: jax.Array | None,  # [dim_size]
+    a_log: jax.Array,  # [n_v]
+    dt_bias: jax.Array,  # [n_v]
+    query_start_loc: jax.Array,  # [num_seqs + 1]
+    state_indices: jax.Array,  # [num_seqs]
+    distribution: jax.Array,  # [3]
+    seq_lens: jax.Array,  # [num_seqs]
+    *,
+    n_kq: int,
+    n_v: int,
+    d_k: int,
+    d_v: int,
+    kernel_size: int,
+    zero_initialize_out: bool = True,
+    compute_precision: jnp.dtype = jnp.float32.dtype,
+    decode_tile_size: int | None = None,
+    mixed_tile_size: int | None = None,
+    is_prefill_only: bool = False,
+    use_qk_norm_in_gdn: bool = True,
+    segment_ids: jax.Array | None = None,
+    conv_halo_seg: jax.Array | None = None,
+    init_seg: jax.Array | None = None,
+) -> tuple[jax.Array, tuple[jax.Array, jax.Array], jax.Array, jax.Array]:
+  """Perform conv1d and gdn in a single fused kernel, returning (out, states, t_inv, chunk_states)."""
+  return _fused_conv1d_gdn(
+      qkv,
+      b,
+      a,
+      conv_state,
+      recurrent_state,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      query_start_loc,
+      state_indices,
+      distribution,
+      seq_lens,
+      n_kq=n_kq,
+      n_v=n_v,
+      d_k=d_k,
+      d_v=d_v,
+      kernel_size=kernel_size,
+      zero_initialize_out=zero_initialize_out,
+      compute_precision=compute_precision,
+      decode_tile_size=decode_tile_size,
+      mixed_tile_size=mixed_tile_size,
+      is_prefill_only=is_prefill_only,
+      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+      init_seg=init_seg,
+  )
+
+
+@jax.jit(
+    donate_argnames=("conv_state", "recurrent_state"),
+    static_argnames=(
+        "n_kq",
+        "n_v",
+        "d_k",
+        "d_v",
+        "kernel_size",
+        "mixed_tile_size",
+        "compute_precision",
+        "use_qk_norm_in_gdn",
+    ),
+)
+def fused_conv1d_gdn_states(
+    qkv: jax.Array,  # [batch_size, n_kq * d_k * 2 + n_v * d_v = dim_size]
+    b: jax.Array,  # [batch_size, n_v]
+    a: jax.Array,  # [batch_size, n_v]
+    conv_state: jax.Array,  # [num_seqs + 1, kernel_size - 1, dim_size]
+    recurrent_state: jax.Array,  # [num_seqs + 1, nv, dk, dv]
+    conv_weight: jax.Array,  # [kernel_size - 1, dim_size]
+    conv_bias: jax.Array | None,  # [dim_size]
+    a_log: jax.Array,  # [n_v]
+    dt_bias: jax.Array,  # [n_v]
+    query_start_loc: jax.Array,  # [num_seqs + 1]
+    state_indices: jax.Array,  # [num_seqs]
+    distribution: jax.Array,  # [3]
+    seq_lens: jax.Array,  # [num_seqs]
+    *,
+    n_kq: int,
+    n_v: int,
+    d_k: int,
+    d_v: int,
+    kernel_size: int,
+    mixed_tile_size: int,
+    compute_precision: jnp.dtype = jnp.float32.dtype,
+    use_qk_norm_in_gdn: bool = True,
+    segment_ids: jax.Array | None = None,
+    conv_halo_seg: jax.Array | None = None,
+    init_seg: jax.Array | None = None,
+) -> tuple[tuple[jax.Array, jax.Array], jax.Array]:
+  """Prefill-only fused_conv1d_gdn that skips the output and chunk states, returning (states, t_inv).
+
+  For callers that only carry the state forward, such as pass 1 of sequence-sharded context
+  parallelism: the kernel neither computes the output nor writes it or the chunk states to HBM. The
+  states and t_inv are bit-identical to fused_conv1d_gdn's.
+  """
+  _, states, t_inv, _ = _fused_conv1d_gdn(
+      qkv,
+      b,
+      a,
+      conv_state,
+      recurrent_state,
+      conv_weight,
+      conv_bias,
+      a_log,
+      dt_bias,
+      query_start_loc,
+      state_indices,
+      distribution,
+      seq_lens,
+      n_kq=n_kq,
+      n_v=n_v,
+      d_k=d_k,
+      d_v=d_v,
+      kernel_size=kernel_size,
+      compute_precision=compute_precision,
+      mixed_tile_size=mixed_tile_size,
+      is_prefill_only=True,
+      use_qk_norm_in_gdn=use_qk_norm_in_gdn,
+      segment_ids=segment_ids,
+      conv_halo_seg=conv_halo_seg,
+      init_seg=init_seg,
+      prefill_mode=config.GDNMode.PER_SEQ_STATES,
+  )
+  return states, t_inv
diff --git a/src/maxtext/kernels/ragged/ragged_sort.py b/src/maxtext/kernels/ragged/ragged_sort.py
index dc563cef68..35ac56d09c 100644
--- a/src/maxtext/kernels/ragged/ragged_sort.py
+++ b/src/maxtext/kernels/ragged/ragged_sort.py
@@ -72,6 +72,8 @@ def ring_ragged_sort(
         only the tokens destined for local experts, padded with zeros elsewhere.
       - 1D tensor ``group_sizes_local`` tracking expert token counts.
       - 1D tensor ``topk_argsort_revert_indices`` for inverse routing.
+      - 1D tensor ``topk_argsort_indices``, the sort permutation (the inverse of
+        ``topk_argsort_revert_indices``), for the backward of :func:`ring_ragged_unsort`.
   """
 
   @jax.custom_vjp
@@ -88,17 +90,16 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
     topk_indices_flat = topk_indices_local.flatten()  # num_tokens_local x topk
     topk_argsort_indices = jnp.argsort(topk_indices_flat)  # num_tokens_local x topk
 
-    token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk)  # num_tokens_local x topk
-    token_indices_sorted = token_indices[topk_argsort_indices]  # num_tokens_local x topk
-
     group_sizes_local = jax.nn.one_hot(topk_indices_flat, num_experts, dtype=jnp.int32).sum(axis=0)  # GLOBAL_NUM_EXPERTS
 
     topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices)  # num_tokens_local x topk
     # Named for the `moe_route` remat key. Saving these integer results lets the backward replay
     # the token gather, whose output the expert GMMs need again, without replaying both argsorts.
-    token_indices_sorted = checkpoint_name(token_indices_sorted, "moe_route")
+    topk_argsort_indices = checkpoint_name(topk_argsort_indices, "moe_route")
     group_sizes_local = checkpoint_name(group_sizes_local, "moe_route")
     topk_argsort_revert_indices = checkpoint_name(topk_argsort_revert_indices, "moe_route")
+    # Flat slot j belongs to token j // topk.
+    token_indices_sorted = topk_argsort_indices // topk  # num_tokens_local x topk
     shard_idx = jax.lax.axis_index(ep_name)
 
     local_num_experts = num_experts // ep_size
@@ -145,7 +146,7 @@ def _ring_ragged_sort_fwd(hidden_states_local, topk_indices_local):
           use_single_sparsecore=use_single_sparsecore,
       )
 
-    out = (x, group_sizes_local, topk_argsort_revert_indices)
+    out = (x, group_sizes_local, topk_argsort_revert_indices, topk_argsort_indices)
 
     res = (
         topk_argsort_revert_indices,
@@ -173,7 +174,7 @@ def _ring_ragged_sort_bwd(res, g_out):
         local_buffer_size,
         _,
     ) = res
-    g_x, _, _ = g_out
+    g_x, _, _, _ = g_out
     # Restrict to the [start, end) source range via a validity bitmask. The
     # ragged kernel packs valid rows to the front of each row-partition and
     # only iterates over the populated prefix, so we hand it the mask directly
@@ -236,6 +237,7 @@ def ring_ragged_unsort(
     sorted_tokens_local,
     group_sizes_local,
     topk_argsort_revert_indices,
+    topk_argsort_indices,
     topk,
     local_num_experts,
     ep_name,
@@ -268,6 +270,8 @@ def ring_ragged_unsort(
     sorted_tokens_local: 2D ``[buffer_size, hidden]`` output tensor from local experts.
     group_sizes_local: 1D tensor tracking the token loads per expert.
     topk_argsort_revert_indices: 1D permutation restoring flat token positions.
+    topk_argsort_indices: 1D sort permutation from :func:`ring_ragged_sort`, the inverse of
+      ``topk_argsort_revert_indices``. Read by the backward only.
     topk: scalar ``int`` representing the routing top-k factor.
     local_num_experts: scalar ``int`` representing the count of experts hosted on this shard.
     ep_name: ``str`` identifying the expert parallel axis name.
@@ -284,6 +288,7 @@ def _ring_ragged_unsort(
       sorted_tokens_local,
       group_sizes_local,
       topk_argsort_revert_indices,
+      topk_argsort_indices,
       topk_weights_flat,
   ):
     """Unsort and scatter activations."""
@@ -291,6 +296,7 @@ def _ring_ragged_unsort(
         sorted_tokens_local,
         group_sizes_local,
         topk_argsort_revert_indices,
+        topk_argsort_indices,
         topk_weights_flat,
     )[0]
 
@@ -299,6 +305,7 @@ def _ring_ragged_unsort_fwd(
       sorted_tokens_local,
       group_sizes_local,
       topk_argsort_revert_indices,
+      topk_argsort_indices,
       topk_weights_flat,
   ):
     """Executes unsorting sending tokens back."""
@@ -361,6 +368,7 @@ def _ring_ragged_unsort_fwd(
     res = (
         sorted_tokens_local,
         topk_argsort_revert_indices,
+        topk_argsort_indices,
         topk_weights_flat,
         shard_output_start,
         shard_output_end,
@@ -388,6 +396,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
     (
         sorted_tokens_local,
         topk_argsort_revert_indices,
+        topk_argsort_indices,
         topk_weights_flat,
         shard_output_start,
         shard_output_end,
@@ -396,9 +405,9 @@ def _ring_ragged_unsort_bwd(res, g_out):
     g_hidden_states_local = g_out
 
     n = topk_argsort_revert_indices.shape[0]
-    # Build the inverse permutation idx_inv such that idx_inv[j] = i
+    # The forward sort's permutation is the inverse permutation: idx_inv[j] = i
     # where revert[i] = j.
-    idx_inv = jnp.argsort(topk_argsort_revert_indices)
+    idx_inv = topk_argsort_indices
 
     # Handle the same two buffering modes for backward pass.
     # ragged_gather does the fan-out, by indexing into the un-expanded
@@ -451,7 +460,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
       dot_local = jnp.sum(gathered.astype(jnp.float32) * sorted_tokens_local.astype(jnp.float32), axis=-1)
       slots = jnp.where(mask, sliced_idx_inv, n)
       grad_topk_weights = jnp.zeros((n,), jnp.float32).at[slots].set(dot_local, mode="drop")
-    return grad_sorted_tokens, None, None, grad_topk_weights
+    return grad_sorted_tokens, None, None, None, grad_topk_weights
 
   _ring_ragged_unsort.defvjp(_ring_ragged_unsort_fwd, _ring_ragged_unsort_bwd)
 
@@ -462,6 +471,7 @@ def _ring_ragged_unsort_bwd(res, g_out):
       sorted_tokens_local,
       group_sizes_local,
       topk_argsort_revert_indices,
+      topk_argsort_indices,
       topk_weights_flat,
   )
 
diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py
index dda2f599eb..8899fc12e9 100644
--- a/src/maxtext/layers/moe.py
+++ b/src/maxtext/layers/moe.py
@@ -109,6 +109,9 @@ class RouteOutput:
   # log_moe_drops=True on the ring-of-experts ragged path only (None otherwise): the drop audit
   # `(dropped_assignments, max_load_ratio)`, see `RoutedMoE.permute`. Logging only.
   moe_audit: Optional[tuple[jax.Array, jax.Array]] = None
+  # Ring-of-experts ragged path only (None otherwise): the sort permutation, the inverse of
+  # `sorted_selected_experts` there. The unsort backward reads it instead of argsorting that.
+  topk_argsort_indices: Optional[jax.Array] = None
 
 
 def combine_moe_audits(audits):
@@ -301,6 +304,39 @@ def forward_row_major(x: jax.Array) -> jax.Array:
   return _forward_layout_constraint(x, tuple(range(x.ndim)))
 
 
+def _split_gate_up(fused):
+  """The gate and up halves of the prefused gate/up GEMM output, with the remat names of `gmm_up`."""
+  n = fused.shape[-1] // 2
+  gate = adc.checkpoint_name(adc.checkpoint_name(fused[..., :n], "mlpwi_0"), "moe_mlpwi_0")
+  up = adc.checkpoint_name(fused[..., n:], "moe_mlpwi_1")
+  return gate, up
+
+
+@functools.partial(jax.custom_vjp, nondiff_argnums=(1,))
+def fused_glu(fused, glu):
+  """`glu(gate, up)` of the two halves of `fused`, the [..., 2 * mlp] output of the prefused gate/up GEMM.
+
+  Autodiff transposes each half's slice into a zero-padded [..., 2 * mlp] gradient and adds the two, one more
+  full pass over the gradient. This backward computes the same two half gradients, with `jax.vjp` of `glu`, and
+  concatenates them, which XLA writes straight from the activation's backward fusion.
+  """
+  return glu(*_split_gate_up(fused))
+
+
+def _fused_glu_fwd(fused, glu):
+  gate, up = _split_gate_up(fused)
+  return glu(gate, up), (gate, up)
+
+
+def _fused_glu_bwd(glu, residuals, cotangent):
+  _, glu_vjp = jax.vjp(glu, *residuals)
+  # The zero-padded add turns -0 into +0; so does this, so the gradient is bit-identical to autodiff's.
+  return (jnp.concatenate([jnp.where(d == 0, 0, d) for d in glu_vjp(cotangent)], axis=-1),)
+
+
+fused_glu.defvjp(_fused_glu_fwd, _fused_glu_bwd)
+
+
 def get_batchsplit_init_kernel_axes():
   return (
       ("expert_only", "embed_moe", None),
@@ -384,6 +420,37 @@ def fold_padding_into_forced_routing(forced_routed_experts, token_mask, num_expe
   return jnp.where(is_pad, jnp.asarray(PAD_TOKEN_ROUTED_EXPERT, base.dtype), base)
 
 
+def _top_k_indices_where_needed(probs, needed, k):
+  """`jax.lax.top_k(probs, k)[1]` on the tokens where `needed` is set; other tokens may get any ids.
+
+  Under router replay almost every token is forced or padding. When at most 256 tokens need the
+  router's own choice, top_k sorts only those rows (gathered, then scattered back) instead of every
+  row; otherwise it sorts every row. top_k is per row, so the needed rows are exact either way.
+
+  Args:
+    probs: `[..., num_experts]` router probabilities.
+    needed: `[..., 1]` True for tokens whose ids are read.
+    k: top_k.
+  """
+  max_rows = 256
+  flat = probs.reshape(-1, probs.shape[-1])
+  if flat.shape[0] <= max_rows:
+    return jax.lax.top_k(probs, k)[1]
+  needed = needed.reshape(-1)
+
+  def compact(flat, needed):
+    # Fill entries repeat row 0, which only writes row 0's own top_k to row 0.
+    rows = jnp.nonzero(needed, size=max_rows)[0]
+    return jnp.zeros((flat.shape[0], k), jnp.int32).at[rows].set(jax.lax.top_k(flat[rows], k)[1])
+
+  def full(flat, needed):
+    del needed
+    return jax.lax.top_k(flat, k)[1]
+
+  indices = jax.lax.cond(jnp.sum(needed) <= max_rows, compact, full, flat, needed)
+  return indices.reshape(probs.shape[:-1] + (k,))
+
+
 def _top_2_in_group_sum(scores_grouped: jax.Array) -> jax.Array:
   """Sums the two largest scores along the last axis, accumulating in float32.
 
@@ -1136,7 +1203,6 @@ def get_topk(
     valid_token_mask = None
     if forced_routed_experts is not None:
       router_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1)
-      _, auto_indices = jax.lax.top_k(router_probs, self.num_experts_per_tok)
       slot_valid = valid_expert_mask(forced_routed_experts, self.num_experts)
       # A token has real forced routing only if its slots are within [0, num_experts).
       # Tokens with UNSET_ROUTED_EXPERT (-1) slots must fall back to auto_indices
@@ -1144,7 +1210,6 @@ def get_topk(
       sorted_experts = jnp.sort(forced_routed_experts, axis=-1)
       has_duplicate = jnp.any(sorted_experts[..., 1:] == sorted_experts[..., :-1], axis=-1, keepdims=True)
       token_is_forced = jnp.logical_and(jnp.all(slot_valid, axis=-1, keepdims=True), jnp.logical_not(has_duplicate))
-      top_k_indices = jnp.where(token_is_forced, forced_routed_experts, auto_indices)
       if self.config.moe_exclude_padding_from_routing:
         # A padding token (every slot PAD_TOKEN_ROUTED_EXPERT) is never forced (its slots are not
         # valid ids) and gets no valid slot, so its weights are zeroed below and its ids become -1.
@@ -1152,6 +1217,11 @@ def get_topk(
         valid_token_mask = jnp.where(token_is_forced, slot_valid, jnp.logical_not(pad_token))
       else:
         valid_token_mask = jnp.where(token_is_forced, slot_valid, jnp.ones_like(slot_valid, dtype=jnp.bool_))
+      # Only tokens that are neither forced nor padding read auto_indices: every other id is replaced
+      # or masked below.
+      uses_auto = jnp.logical_and(jnp.logical_not(token_is_forced), jnp.any(valid_token_mask, axis=-1, keepdims=True))
+      auto_indices = _top_k_indices_where_needed(router_probs, uses_auto, self.num_experts_per_tok)
+      top_k_indices = jnp.where(token_is_forced, forced_routed_experts, auto_indices)
       gather_indices = jnp.where(valid_token_mask, top_k_indices, 0)
       if self.config.decoder_block == ctypes.DecoderBlockType.GEMMA4:
         top_k_weights = jnp.take_along_axis(router_probs, gather_indices, axis=-1).astype(self.dtype)
@@ -1420,7 +1490,7 @@ def permute(
       else:
         buffer_size = None
 
-      sorted_inputs, group_size, sorted_selected_experts = ring_ragged_sort(
+      sorted_inputs, group_size, sorted_selected_experts, topk_argsort_indices = ring_ragged_sort(
           inputs_2d,
           topk_indices_2d,
           self.config.num_experts,
@@ -1437,6 +1507,7 @@ def permute(
           use_single_sparsecore=self.config.ragged_sort_use_single_sparsecore,
       )
     else:
+      topk_argsort_indices = None
       flatten_selected_experts = jnp.ravel(selected_experts)
 
       # Must precede roll_to_expert_id: `(-1 - roll) % num_experts` wraps padding
@@ -1542,6 +1613,7 @@ def permute(
         local_group_size,
         has_overflow,
         moe_audit,
+        topk_argsort_indices,
     )
 
   def unpermute(
@@ -1553,6 +1625,7 @@ def unpermute(
       sequence_length,
       use_custom_sort_vjp=True,
       group_sizes=None,
+      topk_argsort_indices=None,
   ):
     """Unpermute tokens to original order and combine weights."""
 
@@ -1566,6 +1639,7 @@ def unpermute(
           intermediate,
           group_sizes,
           sorted_selected_experts,
+          topk_argsort_indices,
           self.num_experts_per_tok,
           local_num_experts,
           self._expert_parallelism_name,
@@ -2414,6 +2488,7 @@ def _ep_all_gather(z):
           local_group_sizes,
           has_overflow,
           moe_audit,
+          topk_argsort_indices,
       ) = self.permute(
           x,
           logits,
@@ -2439,6 +2514,7 @@ def _ep_all_gather(z):
               local_group_sizes=local_group_sizes,
               has_overflow=has_overflow,
               moe_audit=moe_audit,
+              topk_argsort_indices=topk_argsort_indices,
           ),
           RouteMetadata(
               expert_shard_id=expert_shard_id,
@@ -2472,6 +2548,7 @@ def ra2a_and_route(
           local_group_sizes,
           _,
           _,
+          _,
       ) = self.permute(
           x,
           logits,
@@ -2678,6 +2755,10 @@ def gmm_up(
         # Weights are stored as (G,K,2N); w0/w1 are adjacent slices so XLA elides this concat.
         w_fused = jnp.concatenate([w0, w1], axis=-1)
         out = gmm_fn(x, w_fused, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes)
+        if not self.config.mlp_bias:
+          # Returned whole, as (out, None): `_down_projection` splits it in `fused_glu`, whose backward concatenates
+          # the halves' gradients, and the chunk pipeline's tie 4 holds the GEMM output, not a copy of its first half.
+          return out, None
         n = out.shape[-1] // 2
         layer_w0, layer_w1 = out[:, :n], out[:, n:]
         if self.config.mlp_bias and w0_bias is not None and w1_bias is not None:
@@ -2931,7 +3012,10 @@ def _up_projection(x, routing, route_metadata, w0, w1, w0_bias, w1_bias, wo_bias
 
     def _down_projection(output0, output1, gmm_fn, wo, wo_bias, routing, route_metadata):
       """The activation and the down projection, per routed token."""
-      intermediate_layer = self.apply_ffn_activation(output0, output1)
+      if output1 is None:  # output0 is the prefused gate/up GEMM output (see `gmm_up`)
+        intermediate_layer = fused_glu(output0, self.apply_ffn_activation)
+      else:
+        intermediate_layer = self.apply_ffn_activation(output0, output1)
       wo_gather_axes, wo_tile_size = get_wo_gmm_params()
       intermediate_output = gmm_fn(
           intermediate_layer,
@@ -2963,6 +3047,7 @@ def _ring_unpermute(intermediate_output, routing, batch_size, sequence_length):
           sequence_length=sequence_length,
           use_custom_sort_vjp=self.config.use_custom_sort_vjp,
           group_sizes=routing.group_sizes,
+          topk_argsort_indices=routing.topk_argsort_indices,
       )
       return jnp.reshape(
           output,
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index 0f1d0ddc1b..e0fc9e6f18 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -1546,11 +1546,16 @@ def _fwd_bwd_kernel(self, params, rest, batch, graphdef=None):
       keep the parameters' dtype; and the loss denominator to normalize them by. Under
       `retry_when_tokens_dropped`, followed by `has_moe_overflow`: whether any MoE layer's ragged
       buffer dropped tokens, in which case the gradients and denominator are zero and `fwd_bwd`
-      reruns the micro-batch once the flag reaches the host.
+      reruns the micro-batch once the flag reaches the host. Frozen parameters' gradients are zero.
     """
     loss_callable = self._loss_fn if self._loss_fn is not None else maxtext_train.loss_fn
 
     def diff_wrapper(p, r, b):
+      if self._freeze_mask_fn is not None:
+        # Parameters `trainable_parameters_mask` freezes enter as constants. `_update_kernel` zeroes their
+        # gradients anyway, so the update is unchanged, but the backward pass no longer computes their weight
+        # gradients or reduce-scatters them (the MoE router gate's, every layer, every micro-batch).
+        p = jax.tree.map(lambda x, frozen: jax.lax.stop_gradient(x) if frozen else x, p, self._freeze_mask_fn(p))
       mdl = nnx.merge(graphdef if graphdef is not None else self._model_graphdef, p, r, copy=True)
       if self._gen_model_input_fn is not None:
         # A gen_model_input_fn maps a payload to the loss fn's *keyword arguments* -- see
diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py
index 674b7dc280..b562444b17 100644
--- a/tests/post_training/unit/maxtext_engine_test.py
+++ b/tests/post_training/unit/maxtext_engine_test.py
@@ -336,6 +336,8 @@ def __init__(self):
       payload = DummyPayload()
       t.compile(payload)
       t.fwd_bwd(payload)
+      # Not differentiated at all, rather than differentiated and zeroed in `update()`.
+      np.testing.assert_array_equal(np.asarray(t._accumulated_grads["router_gate"]), [0.0, 0.0])
       t.update()
 
       np.testing.assert_allclose(np.asarray(t.model.router_gate.get_value()), [10.0, 10.0])
diff --git a/tests/unit/gdn_bwd_pallas_test.py b/tests/unit/gdn_bwd_pallas_test.py
index ee91c2b33e..d44e8f44b7 100644
--- a/tests/unit/gdn_bwd_pallas_test.py
+++ b/tests/unit/gdn_bwd_pallas_test.py
@@ -24,9 +24,12 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import pytest
 
 from maxtext.configs import types as config_types
+from maxtext.kernels.gdn import compute_gdn
 from maxtext.kernels.gdn import gdn_bwd_pallas
+from maxtext.kernels.gdn.gdn_bwd import api as gdn_api
 from maxtext.models import qwen3
 from maxtext.utils import maxtext_utils
 
@@ -138,6 +141,28 @@ def test_chunk_forward_matches_jax_reference(self):
     np.testing.assert_allclose(out_emit, expected_out[0], rtol=1e-3, atol=1e-3)
     np.testing.assert_allclose(state_emit, expected_state[0], rtol=5e-3, atol=5e-3)
 
+  def test_invert_triangular_matrix_matches_float64_inverse(self):
+    """The forward kernels' blocked inverse of T = I + strictly-lower(decay * beta k k^T) is exact to f32 rounding.
+
+    Chunks of 16 and less are one substitution block; 32 and 64 also substitute into rows that the
+    inter-block matmul has already updated.
+    """
+    rng = np.random.default_rng(0)
+    num_heads, head_dim = 4, 128
+    for chunk_size in (8, 16, 32, 64):
+      k = rng.standard_normal((num_heads, chunk_size, head_dim))
+      k /= np.linalg.norm(k, axis=-1, keepdims=True)
+      beta = 1 / (1 + np.exp(-rng.standard_normal((num_heads, chunk_size, 1))))
+      log_g = np.cumsum(-0.05 * np.abs(rng.standard_normal((num_heads, chunk_size))), axis=1)
+      decay = np.exp(log_g[:, :, None] - log_g[:, None, :]) * np.tril(np.ones((chunk_size, chunk_size)), -1)
+      t = (np.eye(chunk_size) + decay * np.einsum("hik,hjk->hij", k * beta, k)).astype(np.float32)
+
+      t_inv = compute_gdn.invert_triangular_matrix(jnp.asarray(t), precision=jax.lax.Precision.HIGHEST)
+
+      np.testing.assert_allclose(
+          np.asarray(t_inv, np.float64), np.linalg.inv(t.astype(np.float64)), rtol=0, atol=1e-6, err_msg=f"{chunk_size=}"
+      )
+
   def test_compute_forward_conv_and_states(self):
     batch_size = 1
     chunk_size = 16
@@ -726,6 +751,107 @@ def test_run_local_gdn_decoupled_fwd_returns_cached_chunk_states(self):
     np.testing.assert_allclose(chunk_states, exp_chunk_states, rtol=tol, atol=tol)
     np.testing.assert_allclose(t_inv, exp_t_inv, rtol=tol, atol=tol)
 
+  def test_run_local_gdn_states_fwd_matches_full_forward(self):
+    """The states-only forward (CP pass 1) returns the full forward's states and t_inv bit for bit.
+
+    On TPU this compares the PER_SEQ_STATES kernel, which skips the output, with the PER_SEQ kernel.
+    """
+    is_tpu = jax.default_backend() == "tpu"
+    batch_size = 2
+    chunk_size = 64 if is_tpu else 32
+    seq_len = 4 * chunk_size
+    num_k_heads = 2
+    num_v_heads = 4
+    head_k_dim = 128 if is_tpu else 64
+    head_v_dim = 128 if is_tpu else 64
+    conv_kernel_size = 4
+    dim_size = num_k_heads * head_k_dim * 2 + num_v_heads * head_v_dim
+
+    qkv, b, a, conv_weight, conv_bias, a_log, dt_bias = _init_fwd_inputs(
+        jax.random.PRNGKey(7), batch_size, seq_len, dim_size, num_v_heads, conv_kernel_size
+    )
+    qkv, b, a = (x.astype(jnp.bfloat16) for x in (qkv, b, a))
+    # Pass-1 inputs: the conv halo, a zero recurrent state, and packed rows with trailing padding.
+    conv_state = jax.random.normal(jax.random.PRNGKey(8), (batch_size, conv_kernel_size - 1, dim_size), jnp.bfloat16)
+    recurrent_state = jnp.zeros((batch_size, num_v_heads, head_k_dim, head_v_dim), jnp.float32)
+    seg = np.zeros((batch_size, seq_len), np.int32)
+    seg[0, : seq_len // 3] = 1
+    seg[0, seq_len // 3 : seq_len - 10] = 2
+    seg[1, : seq_len - 40] = 1
+    segment_ids, conv_halo_seg, init_seg = gdn_api._local_segment_metadata(jnp.asarray(seg), True, conv_kernel_size)
+    args = (qkv, b, a, conv_weight, conv_bias, a_log, dt_bias, conv_state, recurrent_state)
+    kwargs = {
+        "num_k_heads": num_k_heads,
+        "num_v_heads": num_v_heads,
+        "head_k_dim": head_k_dim,
+        "head_v_dim": head_v_dim,
+        "conv_kernel_size": conv_kernel_size,
+        "chunk_size": chunk_size,
+        "use_qk_norm_in_gdn": True,
+        "compute_dtype": jnp.float32,
+        "segment_ids": segment_ids,
+        "conv_halo_seg": conv_halo_seg,
+        "init_seg": init_seg,
+    }
+
+    (_, exp_states), exp_t_inv, _ = gdn_api._run_local_gdn_decoupled_fwd(*args, **kwargs)
+    states, t_inv = gdn_api._run_local_gdn_states_fwd(*args, **kwargs)
+
+    for got, exp in zip((*states, t_inv), (*exp_states, exp_t_inv)):
+      self.assertEqual(got.dtype, exp.dtype)
+      np.testing.assert_array_equal(np.asarray(got, np.float32), np.asarray(exp, np.float32))
+
+  @pytest.mark.tpu_only
+  def test_prefill_bf16_qkv_matches_fp32_qkv_path(self):
+    """Prefill-only kernels upcast bf16 qkv chunks in VMEM; the general path upcasts in HBM. Same bits.
+
+    Covers the full prefill kernel (PER_SEQ) and the states-only kernel (PER_SEQ_STATES).
+    """
+    from maxtext.kernels.gdn import wrapper  # pylint: disable=import-outside-toplevel
+
+    batch_size, chunk_size, num_k_heads, num_v_heads, head_dim, conv_kernel_size = 2, 64, 2, 8, 128, 4
+    seq_len = 4 * chunk_size
+    dim_size = num_k_heads * head_dim * 2 + num_v_heads * head_dim
+    qkv, b, a, conv_weight, conv_bias, a_log, dt_bias = _init_fwd_inputs(
+        jax.random.PRNGKey(5), batch_size, seq_len, dim_size, num_v_heads, conv_kernel_size
+    )
+
+    def run(fn, **kwargs):
+      # Fresh states for every call: the wrapper donates them.
+      k1, k2 = jax.random.split(jax.random.PRNGKey(6))
+      conv_state = jax.random.normal(k1, (batch_size + 1, conv_kernel_size - 1, dim_size), jnp.bfloat16)
+      recurrent_state = 0.1 * jax.random.normal(k2, (batch_size + 1, num_v_heads, head_dim, head_dim), jnp.float32)
+      return fn(
+          qkv=qkv.astype(jnp.bfloat16).reshape(-1, dim_size),
+          b=b.astype(jnp.bfloat16).reshape(-1, num_v_heads),
+          a=a.astype(jnp.bfloat16).reshape(-1, num_v_heads),
+          conv_state=conv_state,
+          recurrent_state=recurrent_state,
+          conv_weight=jnp.swapaxes(conv_weight, 0, 2),
+          conv_bias=conv_bias,
+          a_log=a_log,
+          dt_bias=dt_bias,
+          query_start_loc=jnp.arange(0, (batch_size + 1) * seq_len, seq_len, dtype=jnp.int32),
+          state_indices=jnp.arange(1, batch_size + 1, dtype=jnp.int32),
+          distribution=jnp.array([0, 0, batch_size], dtype=jnp.int32),
+          seq_lens=jnp.full((batch_size,), seq_len + 1, dtype=jnp.int32),
+          n_kq=num_k_heads,
+          n_v=num_v_heads,
+          d_k=head_dim,
+          d_v=head_dim,
+          kernel_size=conv_kernel_size,
+          mixed_tile_size=chunk_size,
+          **kwargs,
+      )
+
+    out, (conv_out, rec_out), t_inv, chunk_states = run(wrapper.fused_conv1d_gdn, is_prefill_only=True)
+    (states_conv, states_rec), states_t_inv = run(wrapper.fused_conv1d_gdn_states)
+    ref_out, (ref_conv, ref_rec), ref_t_inv, ref_chunk_states = run(wrapper.fused_conv1d_gdn, is_prefill_only=False)
+    got = (out, conv_out, rec_out, t_inv, chunk_states, states_conv, states_rec, states_t_inv)
+    exp = (ref_out, ref_conv, ref_rec, ref_t_inv, ref_chunk_states, ref_conv, ref_rec, ref_t_inv)
+    for g, e in zip(got, exp):
+      np.testing.assert_array_equal(np.asarray(g), np.asarray(e))
+
   def test_decoupled_conv1d_gdn_kernel_gradient_with_initial_states(self):
     """Verifies custom VJP gradients when initial conv_state and recurrent_state are provided."""
     is_tpu = jax.default_backend() == "tpu"
@@ -879,6 +1005,101 @@ def test_gdn_kernel_bwd_multi_group_head_parallel(self):
     np.testing.assert_allclose(dal1, dal2, rtol=1e-3, atol=1e-3)
     np.testing.assert_allclose(ddt1, ddt2, rtol=1e-3, atol=1e-3)
 
+  @pytest.mark.tpu_only
+  def test_gdn_kernel_bwd_group_blocks_match_per_group_calls(self):
+    """Verifies the in-place group reads (q/k/v column blocks, do, chunk_states, t_inv) and dq/dk/dv writes.
+
+    Every head group runs the same kernel body as a single-group call on that group's heads, so the 2-group call at
+    the Qwen3.5-397B per-device layer geometry (64 v / 16 kq heads, bf16) must match per-group calls bit for bit.
+    """
+    batch_size = 1
+    chunk_size = 64
+    num_chunks = 8
+    seq_len = num_chunks * chunk_size
+    num_k_heads = 16
+    num_v_heads = 64
+    head_tile = 32
+    head_dim = 128
+    q_size = num_k_heads * head_dim
+    dim_size = 2 * q_size + num_v_heads * head_dim
+    tile_q = q_size * head_tile // num_v_heads
+    tile_v = head_tile * head_dim
+
+    k1, k2, k3, k4, k5, k6, k7, k8, k9 = jax.random.split(jax.random.PRNGKey(2468), 9)
+    qkv = jax.random.normal(k1, (batch_size, seq_len, dim_size), dtype=jnp.float32).astype(jnp.bfloat16)
+    b = jax.random.normal(k2, (batch_size, seq_len, num_v_heads), dtype=jnp.float32).astype(jnp.bfloat16)
+    a = jax.random.normal(k3, (batch_size, seq_len, num_v_heads), dtype=jnp.float32).astype(jnp.bfloat16)
+    a_log = jax.random.normal(k4, (num_v_heads,), dtype=jnp.float32)
+    dt_bias = jax.random.normal(k5, (num_v_heads,), dtype=jnp.float32)
+    do = jax.random.normal(k6, (batch_size, seq_len, num_v_heads, head_dim), dtype=jnp.float32).astype(jnp.bfloat16)
+    chunk_states = 0.1 * jax.random.normal(k7, (batch_size, num_chunks, num_v_heads, head_dim, head_dim))
+    t_inv = 0.1 * jax.random.normal(k8, (batch_size, num_chunks, num_v_heads, chunk_size, chunk_size))
+    dht = 0.1 * jax.random.normal(k9, (batch_size, num_v_heads, head_dim, head_dim))
+    # Two packed documents, the second starting mid-chunk, then padding.
+    segment_ids = jnp.zeros((batch_size, seq_len), jnp.int32)
+    segment_ids = segment_ids.at[:, : seq_len // 2 + 5].set(1).at[:, seq_len // 2 + 5 : seq_len - 16].set(2)
+
+    run = functools.partial(
+        gdn_bwd_pallas.pallas_gdn_bwd_kernel,
+        kq_head_dim=head_dim,
+        v_head_dim=head_dim,
+        chunk_size=chunk_size,
+        head_tile=head_tile,
+        segment_ids=segment_ids,
+        return_dh0=True,
+    )
+    full = run(
+        qkv_conv=qkv,
+        b=b,
+        a=a,
+        a_log=a_log,
+        dt_bias=dt_bias,
+        do=do,
+        chunk_states=chunk_states,
+        t_inv=t_inv,
+        num_v_heads=num_v_heads,
+        d_recurrent_state=dht,
+    )
+
+    per_group = []
+    for g in range(num_v_heads // head_tile):
+      heads = slice(g * head_tile, (g + 1) * head_tile)
+      qkv_g = jnp.concatenate(
+          [
+              qkv[..., g * tile_q : (g + 1) * tile_q],
+              qkv[..., q_size + g * tile_q : q_size + (g + 1) * tile_q],
+              qkv[..., 2 * q_size + g * tile_v : 2 * q_size + (g + 1) * tile_v],
+          ],
+          axis=-1,
+      )
+      per_group.append(
+          run(
+              qkv_conv=qkv_g,
+              b=b[..., heads],
+              a=a[..., heads],
+              a_log=a_log[heads],
+              dt_bias=dt_bias[heads],
+              do=do[:, :, heads],
+              chunk_states=chunk_states[:, :, heads],
+              t_inv=t_inv[:, :, heads],
+              num_v_heads=head_tile,
+              d_recurrent_state=dht[:, heads],
+          )
+      )
+
+    dy = [out[0] for out in per_group]
+    ref_dy = jnp.concatenate(
+        [d[..., :tile_q] for d in dy] + [d[..., tile_q : 2 * tile_q] for d in dy] + [d[..., 2 * tile_q :] for d in dy],
+        axis=-1,
+    )
+    np.testing.assert_array_equal(full[0], ref_dy)
+    np.testing.assert_array_equal(full[1], jnp.concatenate([out[1] for out in per_group], axis=-1))
+    np.testing.assert_array_equal(full[2], jnp.concatenate([out[2] for out in per_group], axis=-1))
+    np.testing.assert_array_equal(full[5], jnp.concatenate([out[5] for out in per_group], axis=1))
+    # d_a_log / d_dt_bias also sum per-chunk partials in XLA, whose order may depend on the number of groups.
+    np.testing.assert_allclose(full[3], jnp.concatenate([out[3] for out in per_group]), rtol=1e-6, atol=0)
+    np.testing.assert_allclose(full[4], jnp.concatenate([out[4] for out in per_group]), rtol=1e-6, atol=0)
+
   def test_gdn_kernel_bwd_variable_length_segment_ids_reset(self):
     """Verifies segment_ids document boundaries reset carried state gradient to prevent leakage."""
     is_tpu = jax.default_backend() == "tpu"
@@ -1511,53 +1732,65 @@ def test_bwd_block_specs_and_symbolic_zero_dht_dh0_omission(self):
     in_00, out_00, n_in_00, n_out_00 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=False,
         has_dh0=False,
     )
-    self.assertEqual(n_in_00, 9)
-    self.assertEqual(n_out_00, 5)
-    self.assertLen(in_00, 9)
-    self.assertLen(out_00, 5)
+    self.assertEqual(n_in_00, 11)
+    self.assertEqual(n_out_00, 7)
+    self.assertLen(in_00, 11)
+    self.assertLen(out_00, 7)
 
     _, _, n_in_10, n_out_10 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=True,
         has_dh0=False,
     )
-    self.assertEqual((n_in_10, n_out_10), (10, 5))
+    self.assertEqual((n_in_10, n_out_10), (12, 7))
 
     _, _, n_in_01, n_out_01 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=False,
         has_dh0=True,
     )
-    self.assertEqual((n_in_01, n_out_01), (9, 6))
+    self.assertEqual((n_in_01, n_out_01), (11, 8))
 
     _, _, n_in_11, n_out_11 = gdn_bwd_pallas.make_bwd_block_specs(
         num_chunks=2,
         chunk_size=16,
-        dim_size=128,
+        qk_tile_size=32,
+        v_tile_size=64,
         num_v_heads=4,
         kq_head_dim=16,
         v_head_dim=16,
+        k_block_offset=1,
+        v_block_offset=1,
         has_dht=True,
         has_dh0=True,
     )
-    self.assertEqual((n_in_11, n_out_11), (10, 6))
+    self.assertEqual((n_in_11, n_out_11), (12, 8))
 
     is_tpu = jax.default_backend() == "tpu"
     batch_size = 1
diff --git a/tests/unit/gdn_cp_dh_scan_test.py b/tests/unit/gdn_cp_dh_scan_test.py
new file mode 100644
index 0000000000..28857fa05b
--- /dev/null
+++ b/tests/unit/gdn_cp_dh_scan_test.py
@@ -0,0 +1,189 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""The Pallas CP composition scans (`dh_scan`) against the XLA `cp_gdn` trees they replace on TPU.
+
+Runs the kernel in Pallas interpret mode. XLA:CPU computes the tree's f32 matmuls in plain f32, so the comparison
+bounds the scans' 3-pass bf16 matmul error and the reduction order, not TPU numerics.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from maxtext.kernels.gdn.gdn_bwd import cp_gdn
+from maxtext.kernels.gdn.gdn_bwd import dh_scan
+
+_NK, _NV, _D, _C, _NC = 2, 8, 128, 64, 4
+
+
+def _inputs(seed=0):
+  """One rank's K-slice conv output, gates and per-chunk T^{-1}, with weak decay."""
+  keys = jax.random.split(jax.random.PRNGKey(seed), 6)
+  s = _NC * _C
+  k_conv = jax.random.normal(keys[0], (1, s, _NK * _D), jnp.float32)
+  b = jax.random.normal(keys[1], (1, s, _NV), jnp.float32).astype(jnp.bfloat16)
+  # Weak decay keeps the product of the 4 chunk transitions far from 0.
+  a = (jax.random.normal(keys[2], (1, s, _NV), jnp.float32) - 6.0).astype(jnp.bfloat16)
+  a_log = jax.random.normal(keys[3], (_NV,), jnp.float32) * 0.3
+  dt_bias = jax.random.normal(keys[4], (_NV,), jnp.float32) * 0.3
+  t_inv = jnp.tril(jax.random.normal(keys[5], (1, _NC, _NV, _C, _C), jnp.float32) * 0.05, -1) + jnp.eye(_C)
+  return k_conv, b, a, a_log, dt_bias, t_inv
+
+
+def _tree_m(k_conv, b, a, a_log, dt_bias, t_inv, segment_ids, init_seg):
+  """m_local of the XLA tree (f32 on CPU)."""
+  m, _ = cp_gdn.compose_local_from_t_inv(
+      k_conv,
+      b,
+      a,
+      a_log,
+      dt_bias,
+      t_inv,
+      num_k_heads=_NK,
+      num_v_heads=_NV,
+      head_k_dim=_D,
+      head_v_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      s_ext_pass1=jnp.zeros((1, _NV, _D, _D), jnp.float32),
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+  )
+  return m
+
+
+def _scan_m(k_conv, b, a, a_log, dt_bias, t_inv, segment_ids, init_seg):
+  """m_local of the Pallas scan, interpreted."""
+  return dh_scan.gdn_cp_m_scan(
+      k_conv,
+      b,
+      a,
+      a_log,
+      dt_bias,
+      t_inv,
+      num_v_heads=_NV,
+      kq_head_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+      interpret=True,
+  )
+
+
+def _rel_err(got, ref):
+  got, ref = np.asarray(got, np.float64), np.asarray(ref, np.float64)
+  return np.abs(got - ref).max() / np.abs(ref).max()
+
+
+# A document continuing from the previous rank and spanning the whole rank (the long single-document case), and a
+# packed row: that document ends inside chunk 1, two more follow with boundaries inside chunks, then padding.
+_SINGLE = np.ones((1, _NC * _C), np.int32)
+_PACKED = np.zeros((1, _NC * _C), np.int32)
+_PACKED[0, :70], _PACKED[0, 70:200], _PACKED[0, 200:240] = 1, 2, 3
+
+
+@pytest.mark.parametrize("head_tile", [64, 4])  # one head group, two head groups
+@pytest.mark.parametrize("layout", ["none", "single", "packed"])
+def test_m_scan_matches_compose_local_tree(layout, head_tile, monkeypatch):
+  monkeypatch.setattr(dh_scan, "_M_SCAN_HEAD_TILE", head_tile)
+  segment_ids = {"none": None, "single": jnp.asarray(_SINGLE), "packed": jnp.asarray(_PACKED)}[layout]
+  init_seg = None if segment_ids is None else jnp.ones((1,), jnp.float32)
+  args = _inputs()
+  ref = _tree_m(*args, segment_ids, init_seg)
+  got = _scan_m(*args, segment_ids, init_seg)
+  assert float(jnp.abs(ref).max()) > 1e-4, "degenerate reference: the transitions vanish"
+  assert _rel_err(got, ref) < 5e-5
+
+
+def test_m_scan_sees_chunk_order_and_segment_starts():
+  """Negative controls: the comparison above can fail."""
+  args = _inputs()
+  segment_ids = jnp.asarray(_SINGLE)
+  ref = _tree_m(*args, segment_ids, jnp.ones((1,), jnp.float32))
+  k_conv, b, a, a_log, dt_bias, t_inv = args
+  reversed_chunks = _scan_m(k_conv, b, a, a_log, dt_bias, t_inv[:, ::-1], segment_ids, jnp.ones((1,), jnp.float32))
+  assert _rel_err(reversed_chunks, ref) > 1e-3
+  # Without the incoming document, the rank's first token starts a new one, and nothing from the incoming state
+  # reaches the rank's end state.
+  fresh = _scan_m(*args, segment_ids, jnp.zeros((1,), jnp.float32))
+  assert float(jnp.abs(fresh).max()) == 0.0
+
+
+def _bwd_inputs(seed=1):
+  """One rank's full conv output (q, k, v), gates, output cotangent and per-chunk T^{-1}."""
+  keys = jax.random.split(jax.random.PRNGKey(seed), 7)
+  s = _NC * _C
+  qkv_conv = jax.random.normal(keys[0], (1, s, 2 * _NK * _D + _NV * _D), jnp.float32).astype(jnp.bfloat16)
+  b = jax.random.normal(keys[1], (1, s, _NV), jnp.float32).astype(jnp.bfloat16)
+  a = (jax.random.normal(keys[2], (1, s, _NV), jnp.float32) - 2.0).astype(jnp.bfloat16)
+  a_log = jax.random.normal(keys[3], (_NV,), jnp.float32) * 0.3
+  dt_bias = jax.random.normal(keys[4], (_NV,), jnp.float32) * 0.3
+  do = jax.random.normal(keys[5], (1, s, _NV, _D), jnp.float32).astype(jnp.bfloat16)
+  t_inv = jnp.tril(jax.random.normal(keys[6], (1, _NC, _NV, _C, _C), jnp.float32) * 0.05, -1) + jnp.eye(_C)
+  return qkv_conv, b, a, a_log, dt_bias, do, t_inv
+
+
+@pytest.mark.parametrize("head_tile", [32, 4])  # one head group, two head groups
+@pytest.mark.parametrize("layout", ["none", "single", "packed"])
+def test_dh_scan_matches_compose_bwd_local_tree(layout, head_tile, monkeypatch):
+  monkeypatch.setattr(dh_scan, "_DH_SCAN_HEAD_TILE", head_tile)
+  segment_ids = {"none": None, "single": jnp.asarray(_SINGLE), "packed": jnp.asarray(_PACKED)}[layout]
+  init_seg = None if segment_ids is None else jnp.ones((1,), jnp.float32)
+  args = _bwd_inputs()
+  _, ref = cp_gdn.compose_bwd_local_from_t_inv(
+      *args,
+      num_k_heads=_NK,
+      num_v_heads=_NV,
+      head_k_dim=_D,
+      head_v_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+  )
+  got = dh_scan.gdn_cp_dh_scan(
+      *args,
+      num_v_heads=_NV,
+      kq_head_dim=_D,
+      v_head_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+      interpret=True,
+  )
+  assert float(jnp.abs(ref).max()) > 1e-4, "degenerate reference"
+  assert _rel_err(got, ref) < 5e-5
+  # Negative control: a perturbed t_inv must show.
+  qkv_conv, b, a, a_log, dt_bias, do, t_inv = args
+  bad = dh_scan.gdn_cp_dh_scan(
+      qkv_conv,
+      b,
+      a,
+      a_log,
+      dt_bias,
+      do,
+      t_inv * 1.01,
+      num_v_heads=_NV,
+      kq_head_dim=_D,
+      v_head_dim=_D,
+      chunk_size=_C,
+      use_qk_norm_in_gdn=True,
+      segment_ids=segment_ids,
+      init_seg=init_seg,
+      interpret=True,
+  )
+  assert _rel_err(bad, ref) > 1e-3
diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py
index fa4025072d..fa8f72cadc 100644
--- a/tests/unit/moe_test.py
+++ b/tests/unit/moe_test.py
@@ -14,6 +14,7 @@
 """Mixture of Experts (MoE) tests."""
 
 import functools
+import re
 from types import SimpleNamespace
 import unittest
 from unittest import mock
@@ -30,6 +31,7 @@
 from jax.sharding import Mesh, PartitionSpec as P
 from maxtext.common.common_types import Config, DType
 from maxtext.configs import pyconfig
+from maxtext.kernels.ragged import ragged_sort
 from maxtext.layers import linears
 from maxtext.layers import moe
 from maxtext.layers import nnx_wrappers
@@ -267,6 +269,28 @@ def test_init(self):
     self.model.init({"params": self.rng, "dropout": self.rng}, x)
 
 
+class TopKIndicesWhereNeededTest(unittest.TestCase):
+
+  def test_matches_top_k_on_needed_tokens(self):
+    """Both branches give jax.lax.top_k's ids on the needed tokens: 256 compact, 257 sorts every row."""
+    # Coarse logits make ties common; top_k breaks them by the lower index.
+    logits = jnp.round(jax.random.normal(jax.random.PRNGKey(0), (2, 512, 32)) * 4) / 4
+    probs = jax.nn.softmax(logits, axis=-1)
+    expected = np.asarray(jax.lax.top_k(probs, 4)[1])
+    top_k_where_needed = jax.jit(moe._top_k_indices_where_needed, static_argnums=2)  # pylint: disable=protected-access
+    order = jax.random.permutation(jax.random.PRNGKey(1), 2 * 512)
+    for num_needed in (256, 257):
+      needed = jnp.zeros(2 * 512, jnp.bool_).at[order[:num_needed]].set(True)
+      got = np.asarray(top_k_where_needed(probs, needed.reshape(2, 512, 1), 4)).reshape(-1, 4)
+      needed = np.asarray(needed)
+      np.testing.assert_array_equal(got[needed], expected.reshape(-1, 4)[needed])
+      if num_needed <= 256:
+        # The compact branch leaves the other tokens (but row 0) at 0.
+        np.testing.assert_array_equal(got[1:][~needed[1:]], 0)
+      else:
+        np.testing.assert_array_equal(got, expected.reshape(-1, 4))
+
+
 class DeepSeekRoutingTest(unittest.TestCase):
 
   def setUp(self):
@@ -3507,5 +3531,83 @@ def test_fp8_moe_with_weight_quant_config(self):
     self.assertEqual(model.wi_scale.shape, expected_wi_shape)
 
 
+class FusedGluTest(unittest.TestCase):
+  """`fused_glu` matches autodiff of the split prefused gate/up output, bit for bit."""
+
+  def test_matches_autodiff_of_split_bitwise(self):
+    def glu(gate, up):
+      return (jax.nn.silu(gate) * up).astype(jnp.bfloat16)
+
+    def split_then_glu(fused):
+      n = fused.shape[-1] // 2
+      return glu(fused[:, :n], fused[:, n:])
+
+    k_fused, k_cotangent = jax.random.split(jax.random.key(0))
+    fused = (3 * jax.random.normal(k_fused, (256, 2 * 128))).astype(jnp.bfloat16)
+    # Zero cotangent rows give -0 half gradients, which autodiff's zero-padded add turns into +0.
+    cotangent = jax.random.normal(k_cotangent, (256, 128)).astype(jnp.bfloat16).at[64:128].set(0)
+    bits = lambda a: np.asarray(a).view(np.uint16)
+    for wrap in (lambda f: f, jax.checkpoint):
+      want, want_vjp = jax.vjp(jax.jit(wrap(split_then_glu)), fused)
+      got, got_vjp = jax.vjp(jax.jit(wrap(lambda f: moe.fused_glu(f, glu))), fused)
+      np.testing.assert_array_equal(bits(got), bits(want))
+      np.testing.assert_array_equal(bits(got_vjp(cotangent)[0]), bits(want_vjp(cotangent)[0]))
+
+
+class RingRaggedUnsortGradTest(parameterized.TestCase):
+  """ring_ragged_unsort's custom backward against autodiff of the same forward written with plain gathers."""
+
+  @parameterized.named_parameters(("unbuffered", None), ("buffered", 8))
+  def test_grad_matches_autodiff_without_argsort(self, buffer_size):
+    ep, num_tokens, hidden, topk, num_experts = 2, 8, 128, 2, 4
+    local_num_experts = num_experts // ep
+    rng = np.random.default_rng(0)
+    # Small integers keep every product and sum exact, so the comparison is bitwise.
+    hidden_states = jnp.asarray(rng.integers(-4, 5, (ep, num_tokens, hidden)), jnp.bfloat16)
+    experts = jnp.asarray(
+        [[rng.choice(num_experts, topk, replace=False) for _ in range(num_tokens)] for _ in range(ep)], jnp.int32
+    )
+    weights = jnp.asarray(rng.integers(1, 4, (ep, num_tokens * topk)), jnp.float32)
+    cotangent = jnp.asarray(rng.integers(-4, 5, (ep, num_tokens, hidden)), jnp.float32)
+    fallbacks = {"enforce_gather_fallback": True, "enforce_gather_reduce_fallback": True}
+
+    def sort(x, e):
+      return ragged_sort.ring_ragged_sort(x, e, num_experts, topk, "expert", ep, buffer_size=buffer_size, **fallbacks)
+
+    # Each vmap lane is one expert shard. buffer_size=8 is the balanced load: the shards receive 9 and 6 slots,
+    # so one drops a slot and the other leaves buffer rows unused.
+    tokens, group_sizes, revert, perm = jax.vmap(sort, axis_name="expert")(hidden_states, experts)
+
+    def unsort(tokens, weights, group_sizes, revert, perm):
+      return ragged_sort.ring_ragged_unsort(
+          tokens, group_sizes, revert, perm, topk, local_num_experts, "expert", weights, **fallbacks
+      )
+
+    def reference(tokens, weights, group_sizes, revert, perm):
+      del perm
+      offsets = jnp.cumulative_sum(group_sizes, include_initial=True)
+      shard = jax.lax.axis_index("expert")
+      start, end = offsets[shard * local_num_experts], offsets[(shard + 1) * local_num_experts]
+      if tokens.shape[0] < revert.shape[0]:  # packed buffer: this shard's rows start at 0
+        pos, valid = revert - start, (revert >= start) & (revert - start < jnp.minimum(end - start, tokens.shape[0]))
+      else:
+        pos, valid = revert, (revert >= start) & (revert < end)
+      rows = jnp.where(valid[:, None], tokens[jnp.where(valid, pos, 0)].astype(jnp.float32) * weights[:, None], 0.0)
+      return rows.reshape(num_tokens, topk, hidden).sum(axis=1)
+
+    def grads(fn):
+      def loss(tokens, weights):
+        out = jax.vmap(fn, axis_name="expert")(tokens, weights, group_sizes, revert, perm)
+        return jnp.sum(out.astype(jnp.float32) * cotangent)
+
+      return jax.grad(loss, argnums=(0, 1))
+
+    for actual, expected in zip(grads(unsort)(tokens, weights), grads(reference)(tokens, weights)):
+      np.testing.assert_array_equal(actual, expected)
+    # The backward reuses the sort's permutation instead of argsorting the revert indices.
+    jaxpr = str(jax.make_jaxpr(grads(unsort))(tokens, weights))
+    self.assertIsNone(re.search(r"\bsort\[", jaxpr), "the unsort backward argsorts")
+
+
 if __name__ == "__main__":
   absltest.main()