Overlap moe comms with collective matmul

Implements Overlapped All-Gather and GMM Computation (Collective Matmul) for Ring of Experts (RoE) in MoE layers to hide communication overhead and token sorting latency. The changes include:
- Chunks input activations along the contracting embedding dimension (embed_dim // num_moe_emb_chunks). While chunk i computes GMM up-projections, chunk i+1 concurrently executes Ring of Experts All-Gather and routing across expert parallel shards.
- Update kernel to use in-place Partial Sum Accumulation to avoid unnecessary copies using input_output_aliases.

Testing
- Layer Tests (tests/unit/moe_test.py): Verified chunked GMM v2 routing, partial sum accumulation, and MLP bias addition (test_moe_emb_chunking_gmm_v2 with mlp_bias=True).
- AOT Compilation Tests (tests/unit/train_compile_test.py): Verified ahead-of-time HLO compilation and sharding on multi-slice topologies for chunked MoE (test_moe_emb_chunking and test_moe_emb_chunking_with_mlp_bias).
- Low-Level Pallas Kernel Tests (tests/unit/pallas_mosaic_tpu_v2_kernel_test.py): Verified numerical correctness and in-place buffer aliasing for partial sum accumulation across GMM (test_gmm_partial_sum0/1) and TGMM with zero-initialized empty expert groups (test_tgmm_empty_group_with_partial_sum0-3).
- End-to-End Execution & Profiling: Verified E2E training execution on v6e-8 TPU VM (deepseek3-tiny with num_moe_emb_chunks=4) and generated XPlane / XProf traces confirming communication-compute overlap.

Author: Shuwen-Fang

src/maxtext/configs/base.yml

@@ -220,6 +220,7 @@ load_balance_loss_weight: 0.0 # weight for the load balance loss
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
 use_ring_of_experts: false # whether to use ring of experts for sparse matmul expert parallelism
+num_moe_emb_chunks: 0 # number of chunks for overlapping token all-gather and GMM computation along embedding dimension
 # If true, peel the 'expert' mesh axis off the MoE dispatch/MLP batch dim so the expert GEMM
 # stays expert-parallel (AllToAll); false keeps 'expert' on the batch dim (activation_batch_moe).
 # Only affects the dense (dense_matmul) MoE path; the sparse (shard_map) path is unaffected.

src/maxtext/configs/types.py

@@ -623,6 +623,9 @@ class Attention(BaseModel):
           " Requires use_tokamax_gmm: true. Currently incompatible with quantization."
       ),
   )
+  num_moe_emb_chunks: int = Field(
+      0, description="Number of chunks for overlapping token all-gather and GMM computation along embedding dimension."
+  )
   ragged_block_size: int = Field(256, description="Block size for ragged attention.")
   enable_padding_causal_mask: bool = Field(True, description="Temporary flag for TE padding.")
   use_tokamax_splash: bool = Field(False, description="Whether to use tokamax splash attention.")
@@ -2591,6 +2594,17 @@ def _validate_use_te_comm_gemm_overlap(self):
           "TE Collective GEMM operations are only supported for TE quantization recipes (i.e. starting with 'te_')."
       )
 
+  def validate_num_moe_emb_chunks(self):
+    """
+    Validates that num_moe_emb_chunks is used with supported settings.
+    """
+    if self.num_moe_emb_chunks > 0:
+      if not self.use_gmm_v2 or not self.use_ring_of_experts:
+        raise ValueError(
+            f"num_moe_emb_chunks > 0 requires use_gmm_v2=True and use_ring_of_experts=True. "
+            f"Got use_gmm_v2={self.use_gmm_v2}, use_ring_of_experts={self.use_ring_of_experts}."
+        )
+
   @model_validator(mode="after")
   def set_derived_and_validate_values(self) -> "MaxTextConfig":
     """
@@ -3200,6 +3214,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
       if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts:
         raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.")
       self.validate_ragged_buffer_factor()
+    self.validate_num_moe_emb_chunks()
 
     # Gemma 4 small (E2B / E4B) uses per-layer KV sharing, which is incompatible with nn.scan.
     if self.model_name in ("gemma4-e2b", "gemma4-e4b") and self.scan_layers:

src/maxtext/kernels/megablox/ops.py

@@ -74,6 +74,7 @@ def gmm(
     qwix_rule: qwix.QtRule | None = None,
     use_manual_quantization: bool = False,  # used in batchsplit
     use_gmm_v2: bool = False,
+    partial_sum: jnp.ndarray | None = None,
 ):
   """Grouped matrix multiplication operation."""
   quantization_rule = None
@@ -112,6 +113,7 @@ def gmm(
       lhs_vma_axes,
       rhs_vma_axes,
       use_gmm_v2,
+      partial_sum,
   )
 
 
@@ -142,6 +144,7 @@ def _gmm_fwd(
     lhs_vma_axes: tuple = tuple(),
     rhs_vma_axes: tuple = tuple(),
     use_gmm_v2: bool = False,
+    partial_sum: jnp.ndarray | None = None,
 ) -> tuple[
     jnp.ndarray,
     tuple[
@@ -200,6 +203,7 @@ def _gmm_fwd(
             tile_n=tiling[2],
         ),
         preferred_element_type=preferred_element_type,
+        partial_sum=partial_sum,
     )
   elif use_tokamax_backend:
     # manual_axis_type is for gmm with shard_map check_vma=True, needs tokamax > 0.0.12
@@ -234,7 +238,7 @@ def _gmm_fwd(
     for axis in lhs_vma_axes:
       out = jax.lax.pcast(out, axis_name=axis, to="varying")
 
-  return out, (lhs, rhs, group_sizes, group_offset)
+  return out, (lhs, rhs, group_sizes, group_offset, partial_sum)
 
 
 def _gmm_bwd(
@@ -256,12 +260,13 @@ def _gmm_bwd(
         jnp.ndarray | qpl.QArray,
         jnp.ndarray,
         jnp.ndarray | None,
+        jnp.ndarray | None,
     ],
     grad: jnp.ndarray,
-) -> tuple[jnp.ndarray, jnp.ndarray, None, None, jnp.ndarray]:
+) -> tuple[jnp.ndarray, jnp.ndarray, None, None, jnp.ndarray, jnp.ndarray | None]:
   """Backward function for throughput GMM VJP."""
   del preferred_element_type
-  lhs, rhs, group_sizes, group_offset = residual
+  lhs, rhs, group_sizes, group_offset, partial_sum_fwd = residual
   num_actual_groups = rhs.shape[0]
 
   # Jargon used here:
@@ -406,4 +411,5 @@ def _gmm_bwd(
   #
   # TODO(tgale, enriqueps, apaske): Fuse this transposition into the tgmm.
   drhs = drhs.swapaxes(1, 2) if transpose_rhs else drhs
-  return dlhs, drhs, None, None, grad
+  dpartial_sum = grad if partial_sum_fwd is not None else None
+  return dlhs, drhs, None, None, grad, dpartial_sum

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py

@@ -909,7 +909,9 @@ def calculate_tiling(
     tile_n_limit //= fuse_act_factor
 
   def _is_tile_k_quant_block_compatible(tk: int) -> bool:
-    if tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0:  # pyrefly: ignore[unsupported-operation]
+    if (
+        tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0
+    ):  # pyrefly: ignore[unsupported-operation]
       return False
     return True
 
@@ -999,7 +1001,9 @@ def validate_inputs(
   if rhs_bias is not None:
     assert rhs_bias.shape == (size_group, 1, size_n)
   if partial_sum is not None:
-    assert partial_sum.shape == (size_m, size_n)
+    assert partial_sum.shape[-1] == size_n
+    # lhs's m dimension can sometimes be padded to wi_tile_fwd_batch_seq
+    assert partial_sum.shape[0] <= size_m
   if rhs_scale is not None:
     num_quant_blocks = rhs_scale.shape[1]
     assert rhs_scale.shape == (size_group, num_quant_blocks, 1, size_n)
@@ -1341,6 +1345,13 @@ def gmm_v2(
       partial_sum_spec,  # partial_sum
   ]
 
+  input_output_aliases = {}
+  if partial_sum is not None:
+    flat_args_preceding = (group_sizes, group_offset, lhs, rhs_weights)
+    leaves = jax.tree_util.tree_leaves(flat_args_preceding)
+    partial_sum_idx = sum(1 for x in leaves if x is not None)
+    input_output_aliases = {partial_sum_idx: 0}
+
   return pl.pallas_call(
       functools.partial(kernel_main, cfgs=cfgs),
       out_shape=out_init,
@@ -1357,4 +1368,5 @@ def gmm_v2(
       name=get_scope_name(cfgs),
       cost_estimate=get_cost_estimate(cfgs),
       metadata=get_metadata(cfgs),
+      input_output_aliases=input_output_aliases,
   )(group_sizes, group_offset, lhs, rhs_weights, partial_sum)[:, : cfgs.out_size_n]

src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py

@@ -763,6 +763,13 @@ def tgmm_v2(
       partial_sum_spec,
   ]
 
+  input_output_aliases = {}
+  if partial_sum is not None:
+    flat_args_preceding = (group_sizes, group_offset, lhs, rhs)
+    leaves = jax.tree_util.tree_leaves(flat_args_preceding)
+    partial_sum_idx = sum(1 for x in leaves if x is not None)
+    input_output_aliases = {partial_sum_idx: 0}
+
   raw_out = pl.pallas_call(
       functools.partial(tgmm_kernel_main, cfgs=cfgs),
       out_shape=out_init,
@@ -781,6 +788,7 @@ def tgmm_v2(
       # the metadata here is for profiling, debugging, and cost modeling.
       # It does not affect the kernel's computation.
       metadata=gmm_v2.get_metadata(cfgs),
+      input_output_aliases=input_output_aliases,
   )(group_sizes, group_offset, lhs, rhs, partial_sum)[:, : dims.size_k, : dims.size_n]
 
   if partial_sum is not None: