MoE: fused TensorCore Pallas permute/combine kernels and split-expert wi layout for TPU v6e

[csgoogle]

Draft. Part 2 of 3 of the Gemma4-26B-A4B / TPU v6e optimization stack. Supersedes draft PR #5474 (split by area and rebuilt without the all-gather scheduling fence, with every environment-variable gate replaced by a config key).

Context

Workload: Gemma4-26B-A4B (128 experts, top-8), v6e-128, seq 16k, per-device batch 2, FSDP-128, full remat, bf16, sparse_matmul + Tokamax GMM v2, moe_use_direct_token_gather=true. Step time went from 7.58 s to 3.82 s across the three PRs.

On v6e:

1. The MoE token gathers (the permute into expert-sorted order, the unpermute + Top-K weighted sum "combine", and their backward scatters) are offloaded by XLA to SparseCore, where they serialize behind the expert-weight all-gathers and show up as exposed SparseCore / collective wait around the grouped matmuls. Commits 1 and 2 add two fused TensorCore Pallas kernels (src/maxtext/kernels/moe_combine_tc.py) that keep both on the TensorCore and never materialize the [tokens * top_k, emb] unsorted copy.
2. With prefuse_moe_weights=true, wi is stored as [E, emb, 2*mlp] (~1 GB across FSDP), which either incurs ~400 ms/step of concat/split copies around gmm_v2 or stalls SparseCore behind a single ~1 GB all-gather. Commit 3 adds moe_split_expert_weight_layout=true, storing wi as [2, E, emb/2, 2*mlp] so its two leading slices (W[:E/2] and W[E/2:], ~507 MB each) are all-gathered separately and consumed by two chained gmm_v2 calls (src/maxtext/kernels/megablox/split_gmm.py) into one shared output buffer.

Base: 4d291bf02 (the commit the stack was developed and measured on); independent of Part 1 except for the measured numbers.

Commits

#	Commit	What / why	s/step before → after
1	MoE: fused TensorCore Pallas combine kernel (moe_combine_kernel="tc")	Fused unpermute + Top-K weighted sum; backward is a window-DMA permutation for dx and a masked lane reduction for dw	4.527 → 4.233 (−294 ms standalone, works with or without #3)
2	MoE: fused TensorCore Pallas permute kernel (moe_permute_kernel="tc")	Token gather into expert-sorted order with unit weights; backward = combine forward with unit weights. Moves the group_size histogram above the gather (no behaviour change)	4.233 → 4.309 standalone, but −492 ms of exposed collective wait once #3 is on
3	MoE: split-expert wi layout with chained gmm_v2 (moe_split_expert_weight_layout)	Stores the prefused wi as two expert halves ([2, E, emb/2, 2*mlp]) that are all-gathered separately and consumed by two chained GMM v2 calls (kernels/megablox/split_gmm.py), so the ~1 GB fused weight is never concatenated or split	4.285 → 4.036 (−249 ms)

Only gated code is added to layers/moe.py; with the defaults (moe_combine_kernel="xla", moe_permute_kernel="xla", moe_split_expert_weight_layout=false) the MoE path is unchanged.

How to enable

# Recipe 1 (Standard checkpoint layout, works at any FSDP size including FSDP=256; saves ~294 ms/step):
moe_combine_kernel: "tc"               # requires sparse_matmul=true, use_ring_of_experts=false, use_ragged_sort=false
moe_permute_kernel: "xla"
moe_split_expert_weight_layout: false

# Recipe 2 (Maximum throughput with split-expert layout, saves an extra ~197 ms/step; requires fsdp * cp <= num_experts):
moe_combine_kernel: "tc"
moe_permute_kernel: "tc"               # requires moe_use_direct_token_gather=true, use_ring_of_experts=false
moe_split_expert_weight_layout: true   # requires prefuse_moe_weights=true, shard_exp_on_fsdp=true, use_gmm_v2=true

When using Recipe 2, combine moe_permute_kernel="tc" + moe_split_expert_weight_layout=true with the XLA flags
--xla_tpu_offload_gather_to_sparsecore=false --xla_tpu_offload_all_supported_gathers_to_sparsecore=false
so that no MoE gather is left serialized on SparseCore behind the expert-weight all-gathers.

Note

With moe_split_expert_weight_layout=true (Recipe 2), the checkpoint shape of wi changes to [2, *scan_dims, num_experts, emb/2, 2*mlp] (a reshape of the fused weight). Use split_gmm.convert_checkpoint_tree(state, to_split=True/False) (included in src/maxtext/kernels/megablox/split_gmm.py) to convert checkpoints in either direction. With moe_split_expert_weight_layout=false (Recipe 1), checkpoint shapes are unchanged.

Tests

• tests/unit/moe_combine_tc_test.py (new, Pallas interpret mode on CPU with race detection, compiled on TPU): combine forward (≤ 1 bf16 ulp vs reference) and backward (dx exact, dw to 1e-2) over block sizes / top_k / expert counts / skewed routing / shared boundary granules / align=16 / f32 weights; jax.checkpoint + lax.scan gradient check; permute forward (exact) and backward; fallbacks. Also checks the per-block scatter staging table (which boundary granules are read from HBM vs carried from the previous block's output buffer, never twice). 18 tests, ~55 s on CPU.
• tests/unit/split_gmm_test.py (new): forward + gradients of the chained split GMM against a dense reference, using a pure-JAX model of gmm_v2 that reproduces the boundary-tile masking (CPU).
• tests/unit/pyconfig_test.py: validation of moe_combine_kernel, moe_permute_kernel, and moe_split_expert_weight_layout.
• src/maxtext/kernels/moe_combine_bench.py: single-chip correctness + benchmark script.
• pyink 24.10.1 (--pyink-indentation=2 --line-length=122) clean.

Related PRs

• Part 1: MaxText layer / model optimizations — Gemma4 26B MoE on TPU v6e: MaxText layer/model optimizations (one commit per change) #5572
• Part 3 (Optional): transposed-RHS gmm_v2 kernel for the backward dlhs — Megablox: transposed-RHS gmm_v2 kernel for the backward dlhs #5574

[gemini-code-assist]

Code Review

This pull request introduces fused TensorCore Pallas kernels for MoE "combine" and "permute" operations to optimize TPU performance by avoiding large unsorted memory copies. It integrates these kernels into the MoE layer, exposes them via new configuration options, and adds corresponding benchmarks, unit tests, and validation checks. The feedback suggests adding a safety check in kernel_supported to ensure that the number of experts per token (k) does not exceed _LANES (128), preventing potential layout failures in the Pallas kernel.

[gemini-code-assist]

The Pallas kernel implementation relies on _LANES (128) for the layout of top-k weights and row info in VMEM. If k (the number of experts per token) exceeds _LANES, operations like _padded_weights (which uses _LANES - w.shape[1]) will fail or produce incorrect results. Adding a check for k <= _LANES in kernel_supported ensures that the kernel safely falls back to the XLA path for unsupported shapes.

Suggested change

	def kernel_supported(x_shape, x_dtype, num_tokens, k, block_tokens, align=8) -> bool:
	"""True iff the Pallas kernels can handle this problem (see the module docstring)."""
	if len(x_shape) != 2:
	return False
	r, e = x_shape
	return (
	jnp.dtype(x_dtype) == jnp.bfloat16
	and e % _LANES == 0
	and r == num_tokens * k
	and r % align == 0
	and block_tokens % _LANES == 0
	and num_tokens % block_tokens == 0
	and num_tokens <= _MAX_TOKENS
	)
	def kernel_supported(x_shape, x_dtype, num_tokens, k, block_tokens, align=8) -> bool:
	"""True iff the Pallas kernels can handle this problem (see the module docstring)."""
	if len(x_shape) != 2:
	return False
	r, e = x_shape
	return (
	jnp.dtype(x_dtype) == jnp.bfloat16
	and e % _LANES == 0
	and r == num_tokens * k
	and r % align == 0
	and block_tokens % _LANES == 0
	and num_tokens % block_tokens == 0
	and num_tokens <= _MAX_TOKENS
	and k <= _LANES
	)