397B v7x trainer: MoE wo FSDP sharding and valid-rows-only MoE kernels

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
                       # ==========================================

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",
+  )

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

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))