Add chunked MoE support to MaxText.

It introduces the `num_moe_token_chunks` configuration to allow chunking of the MoE layer.
Also includes `equiv_chunk_test.py` to verify that chunked MoE produces equivalent results to the unchunked implementation.

- [x] I have tested these changes
- [x] I have updated the documentation

PiperOrigin-RevId: 950076344

Author: Google-ML-Automation

src/maxtext/configs/base.yml

@@ -213,6 +213,8 @@ num_experts_per_tok: 1
 megablox: true
 sparse_matmul: true
 capacity_factor: -1.0 # a factor to decide expert capacity for token dropping, and no dropping by default
+num_moe_token_chunks: 1 # split routed-MoE tokens into N chunks (ring-of-experts path) so each chunk's EP all-gather / reduce-scatter overlaps the previous chunk's GMM compute. 1 = disabled (identical to baseline).
+moe_chunk_barrier: false # wraps its iteration with optimization_barrier to chain the chunked ring-of-experts MoE loop so each chunk's input is fenced on the previous chunk's output, forcing sequential (no-interleave) chunk execution. Math unchanged (barrier is identity), loss bit-exact. Needs num_moe_token_chunks>1 + use_ring_of_experts=True to matter.
 ragged_buffer_factor: -1.0 # a factor to determine the size of the ragged buffer for routed MoE activations.
 # By default (-1), the routed buffer is worst case size to ensure no dropping.
 # When set to 1.0 this buffer if set to the size assuming perfectly balanced. If the routing dictates

src/maxtext/configs/types.py

@@ -793,6 +793,30 @@ class MoEGeneral(BaseModel):
       -1.0,
       description="Ragged buffer factor. If < 0, ragged buffer is worst case size.",
   )
+  num_moe_token_chunks: PositiveInt = Field(
+      1,
+      description=(
+          "Number of token chunks for the ring-of-experts MoE pipeline. 1"
+          " disables chunking (identical to baseline). >1 splits the per-shard"
+          " tokens along the sequence dimension so each chunk's EP all-gather /"
+          " reduce-scatter overlaps the previous chunk's GMM compute. Requires"
+          " use_ring_of_experts=True."
+      ),
+  )
+  moe_chunk_barrier: bool = Field(
+      False,
+      description=(
+          "Diagnostic (profiling, not production). When True, chain the chunked"
+          " ring-of-experts MoE loop so each chunk's input is fenced with"
+          " jax.lax.optimization_barrier on the previous chunk's output,"
+          " forcing XLA to run the chunks sequentially (no interleave/fusion)."
+          " Math is unchanged (barrier is identity), so loss stays bit-exact."
+          " Used to test whether the token-AG/RS chunks overlap at all today."
+          " Requires num_moe_token_chunks>1 and use_ring_of_experts=True to have any"
+          " effect."
+      ),
+  )
+
   moe_expert_input_dim: int = Field(
       -1,
       description="Dimension of tokens entering the MoE layer. If < 0, defaults to emb_dim.",
@@ -894,6 +918,12 @@ class MoEGeneral(BaseModel):
       "This can improve inference performance.",
   )
 
+  @model_validator(mode="after")
+  def validate_moe_chunks(self) -> "MoEGeneral":
+    if self.num_moe_token_chunks > 1 and not self.use_ring_of_experts:
+      raise ValueError("num_moe_token_chunks > 1 requires use_ring_of_experts=True.")
+    return self
+
 
 class MoEKernels(BaseModel):
   """Configuration for MoE-specific kernels like Megablox."""
@@ -3561,6 +3591,13 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       if self.share_kv_projections:
         raise ValueError("`num_kv_shared_layers > 0` is not compatible with `share_kv_projections`.")
 
+    if self.num_moe_token_chunks > 1:
+      if self.max_target_length % self.num_moe_token_chunks != 0:
+        raise ValueError(
+            f"num_moe_token_chunks={self.num_moe_token_chunks} must evenly divide "
+            f"max_target_length={self.max_target_length}."
+        )
+
     # I. FINAL TYPE CONVERSIONS AND DERIVED LISTS
     ici_map = {
         "diloco": self.ici_diloco_parallelism,

src/maxtext/kernels/megablox/ops.py

@@ -62,7 +62,7 @@ def gmm(
     group_offset: jnp.ndarray | None = None,
     existing_out: jnp.ndarray | None = None,
     transpose_rhs: bool = False,
-    interpret: bool = False,
+    interpret: bool | None = None,
     lhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None,  # pyrefly: ignore[invalid-literal]
     rhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None,  # pyrefly: ignore[invalid-literal]
     use_qwix_quantization: bool = False,
@@ -77,6 +77,13 @@ def gmm(
     partial_sum: jnp.ndarray | None = None,
 ):
   """Grouped matrix multiplication operation."""
+  if interpret is None:
+    # Default to native (TPU) lowering. `jax.devices()[0]` is NOT the compile TARGET:
+    # during train_compile the local backend is CPU (JAX_PLATFORMS=cpu) while the mesh
+    # targets tpu7x, and interpret-mode there breaks check_vma (exposes the kernel's
+    # internal dynamic_slice VMA) and balloons HBM temporaries. Callers that genuinely
+    # run off-TPU (e.g. equiv_chunk_test) pass interpret based on their target mesh.
+    interpret = False
   quantization_rule = None
   if use_qwix_quantization:
     # get_current_rule has to be called outside of the _gmm_fwd function.

src/maxtext/kernels/ragged/ragged_gather.py

@@ -396,9 +396,14 @@ def ragged_gather(
 
   dtype = x.dtype
 
+  # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests).
+  # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU.
+  if enforce_fallback or jax.devices()[0].platform != "tpu":
+    return _fallback_implementation(x, indices, weights, has_weights)
+
   sc_info = pltpu.get_tpu_info().sparse_core
-  if sc_info is None or enforce_fallback:
-    # Sparse core is not available or fallback is enforced. Use JAX reference.
+  if sc_info is None:
+    # Sparse core is not available. Use JAX reference.
     return _fallback_implementation(x, indices, weights, has_weights)
 
   hidden_size = x.shape[-1]

src/maxtext/kernels/ragged/ragged_gather_reduce.py

@@ -457,9 +457,14 @@ def ragged_gather_reduce(
   assert topk_weights.ndim == 1, "ragged_gather_reduce only supports 1d topk_weights."
   assert valid_rows_mask.ndim == 1, "ragged_gather_reduce only supports 1d valid_rows_mask."
 
+  # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests).
+  # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU.
+  if enforce_fallback or jax.devices()[0].platform != "tpu":
+    return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size)
+
   sc_info = pltpu.get_tpu_info().sparse_core
-  if sc_info is None or enforce_fallback:
-    # Sparse core is not available or fallback is enforced. Use JAX reference.
+  if sc_info is None:
+    # Sparse core is not available. Use JAX reference.
     return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size)
 
   # Heuristic threshold on whether to fallback for small inputs.

src/maxtext/kernels/ragged/ragged_gather_reduce_v2.py

@@ -650,8 +650,13 @@ def ragged_gather_reduce(
     Reduced output, ``(input_size // reduce_group_size, hidden_size)``.
   """
   # Step 1: Choose the implementation (TensorCore fallback or SparseCore).
+  # Guard against eager initialization on non-TPU hardware (e.g. during CPU tests).
+  # pltpu.get_tpu_info() expects TPU hardware and will crash if executed on CPU.
+  if enforce_fallback or jax.devices()[0].platform != "tpu":
+    return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size)
+
   sc_info = pltpu.get_tpu_info().sparse_core
-  if sc_info is None or enforce_fallback:
+  if sc_info is None:
     return _fallback_implementation(x, indices, topk_weights, valid_rows_mask, reduce_group_size)
 
   # For a small {input + output} both likely fit in TensorCore VMEM, where a