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