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