Gemma4 26B MoE on TPU v6e: MaxText layer/model optimizations (one commit per change)

[csgoogle]

Draft. Part 1 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).

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. Step time went from 7.58 s (baseline) to 3.82 s across the three PRs; this PR holds the 7 pure MaxText layer / model changes — one commit per optimization, requiring no new config flags, no custom Pallas kernels, and no checkpoint layout changes.

Base: 4d291bf02 (the commit the stack was developed and measured on).

Commits

#	Commit	What / why	s/step before → after
1	Gemma4: replicate pre_forward_scale_2 like the other norm scales	Removes a per-layer all-gather of a norm scale that every other scale avoided	4.805 → 4.793
2	RMSNorm / Gemma4 router: multiply norm scales in fp32	Removes bf16 round trips (converts + copies) around the scale multiply	→ 4.773
3	Embed: shard the embedding table on the vocabulary dimension	Avoids the per-step all-gather of the [vocab, emb] table under FSDP	→ 4.762
4	RotaryEmbedding: compute split-half RoPE directly	One rotate + multiply instead of concat / slice copies	→ 4.716
5	Vocabulary tiling: tile along unsharded seq_len to avoid FSDP all-to-all	The previous tiling axis forced an all-to-all per tile	→ 4.587 (after use_iota_embed=False, a config change measured at 4.716 → 4.613)
6	MoE: TensorCore top_k_weights, per-expert scale and group_size; single rsqrt in RMSNorm	Keeps small routing ops off SparseCore, where they serialize behind the weight all-gathers	→ 4.527
7	MoE: gate/up split custom VJP	Replaces two zero-pads + fp32 add in backward of the (gate, up) slice with a single concatenate	4.309 → 4.285

Commit messages carry the per-commit details and measurements. All 7 commits are active by default and need no configuration flags.

Tests

• Existing unit tests for the touched layers; pyink 24.10.1 (--pyink-indentation=2 --line-length=122) clean.

Related PRs

• Part 2: fused TensorCore permute / combine kernels & split-expert wi layout — MoE: fused TensorCore Pallas permute/combine kernels and split-expert wi layout for TPU v6e #5573
• 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 the moe_split_expert_weight_layout configuration, which stores the prefused MoE wi weight in a split layout to eliminate concatenation and splitting overhead during forward and backward passes. It implements a new split_gmm kernel to chain grouped matrix multiplications and optimizes several MoE activation routing, sorting, and normalization paths to improve XLA fusion and performance. Feedback on the changes suggests handling None gradients defensively in the custom VJP _split_gate_up_bwd to prevent potential runtime errors during partial differentiation.

[gemini-code-assist]

In JAX custom VJPs for multi-output functions, if one of the outputs is not differentiated (e.g., during testing, debugging, or partial differentiation), its corresponding cotangent in grads can be None. To prevent runtime errors when jnp.concatenate receives None, it is safer to handle None gradients defensively by replacing them with zeros of the same shape and dtype as the other gradient.

Suggested change

	def _split_gate_up_bwd(_residuals: None, grads: tuple[jax.Array, jax.Array]) -> tuple[jax.Array]:
	grad_gate, grad_up = grads
	return (jnp.concatenate([grad_gate, grad_up], axis=-1),)
	def _split_gate_up_bwd(_residuals: None, grads: tuple[jax.Array, jax.Array]) -> tuple[jax.Array]:
	grad_gate, grad_up = grads
	if grad_gate is None and grad_up is None:
	return (None,)
	if grad_gate is None:
	grad_gate = jnp.zeros_like(grad_up)
	elif grad_up is None:
	grad_up = jnp.zeros_like(grad_gate)
	return (jnp.concatenate([grad_gate, grad_up], axis=-1),)