Overlap moe comms with collective matmul by Shuwen-Fang · Pull Request #4295 · AI-Hypercomputer/maxtext

Shuwen-Fang · 2026-06-29T19:22:35Z

Description

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_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_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_chunking and test_moe_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 chunk_ag_gmm=True, num_chunks=4) and generated XPlane / XProf traces confirming communication-compute overlap.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
I have necessary comments in my code, particularly in hard-to-understand areas.
I have run end-to-end tests tests and provided workload links above if applicable.
I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

codecov · 2026-06-29T19:37:38Z

Codecov Report

❌ Patch coverage is 91.20879% with 8 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
...rnels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py	16.66%	4 Missing and 1 partial ⚠️
src/maxtext/layers/moe.py	97.22%	0 Missing and 2 partials ⚠️
...ernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py	88.88%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!

NuojCheng · 2026-07-09T23:21:10Z

 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
+chunk_ag_gmm: false # whether to chunk All-Gather and GMM computation to overlap communication
+num_chunks: 0 # number of chunks for overlapping All-Gather and GMM computation along embedding dimension


I suggest giving them more specific names to distinguish with the other chunking feature in cl/938756606

NuojCheng · 2026-07-09T23:22:35Z

  # 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
+  return dlhs, drhs, None, None, grad, grad if partial_sum_fwd is not None else None


maybe move the if-else logic out of the return for better readability

NuojCheng · 2026-07-09T23:23:38Z

-    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


for my understanding is this change specific due to chunking feature or a general fix?

This occurs when embedding dim is not a multiple of wi_tile_fwd_batch_seq, so m in padded but partial sum is not. This is specifically for chunking because partial sum is only used for chunking feature

NuojCheng · 2026-07-09T23:28:35Z

+              global_sorted_experts=selected_experts,
+              use_custom_sort_vjp=self.config.use_custom_sort_vjp,
+              use_ragged_sort=self.config.use_ragged_sort,
+          )


we could move these logics to a new function named ra2a_and_route, mirroring your changes earlier.

NuojCheng · 2026-07-09T23:29:48Z

+            and callable(getattr(rngs, "params"))
+            else [rngs] * self.config.num_chunks
+        )
+        # TODO: is this valid?


add a correctness test protecting this feature for convergence if not?

NuojCheng · 2026-07-09T23:31:02Z

+            weight_gather,
+            partial_accum0=ps0,
+            partial_accum1=ps1,
+        )


I think it might be a good idea to move all the logic in if-branch to a separate function, called moe_emb_chunking or better names.

NuojCheng · 2026-07-09T23:32:33Z

+        (last_x_chunk, ps0, ps1), _ = jax.lax.scan(
+            scan_fn,
+            (cur_x_chunk, partial_sum0, partial_sum1),
+            jnp.arange(self.config.num_chunks - 1),


I think you could use length=self.config.num_chunks-1 instead of arange. Also, have you tried performance when the scan is unrolled? e.g. add unroll=self.config.num_chunks-1, curious the performance diff.

That's a good idea! I'll give it a try. I'm still working on improving the performance of this feature and this PR is just to implement the feature. I can experiment & make changes if needed in a follow up PR

NuojCheng · 2026-07-09T23:35:00Z

+    mesh = Mesh(devices_array, cfg.mesh_axes)
+    variables, expected_output = self.get_expected_output(rng_model, hidden_states, cfg, mesh)
+    actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh)
+    self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False))


oh I see the correctness test, nice!

NuojCheng · 2026-07-09T23:36:28Z


    assert_arrays_all_close(actual, expected)

+  @pytest.mark.skip(reason="Test takes too long, can run locally to verify changes b/528087469")


emmm I honestly think these tests gonna never get run....

NuojCheng · 2026-07-09T23:36:43Z

    self.assertEqual(actual.shape, (num_local_groups, in_size, out_size))
    assert_arrays_all_close(actual, expected)

+  @pytest.mark.skip(reason="Test takes too long, can run locally to verify changes b/528087469")


emmm I honestly think these tests gonna never get run.... can we make it simple and fast?

We are going to use tokamax kernels directly in the future instead of forking it in maxtext, so this is only going to be here temporarily

NuojCheng

Overall LGTM. Leave some comments, thank you Shuwen and great work!

gobbleturk · 2026-07-13T19:00:22Z

+    Validates that use_moe_emb_chunking is used with supported settings.
+    """
+    if self.use_moe_emb_chunking:
+      logger.warning(


I would prefer to not have this warning, but up to you

gobbleturk · 2026-07-13T19:01:02Z

 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
+use_moe_emb_chunking: false # whether to chunk All-Gather and GMM computation to overlap communication
+num_moe_emb_chunks: 0 # number of chunks for overlapping All-Gather and GMM computation along embedding dimension


I would only have oneuflag - num_moe_emb_chunks - when its set to 0 its equivalent to false

good idea, updated

I would specify "overlapping all-gather" with "overlapping token all-gather" to distinguish it from FSDP all gathers

gobbleturk · 2026-07-14T00:34:51Z

-      """Performs both across device and within device token routing/sorting"""
-      num_ep = self.get_expert_parallelism_size()
-      expert_shard_id = jax.lax.axis_index(self._expert_parallelism_name) if num_ep > 1 else 0
+    def roe_ag_and_route(x, logits, pre_bias_logits, num_ep, expert_shard_id, rngs, input_ids=None):


Nice breakdown!

gobbleturk · 2026-07-14T00:35:47Z

          ),
      )

+    def route(x, logits, pre_bias_logits, rngs, input_ids=None):


This looks so good!

gobbleturk · 2026-07-14T00:36:40Z


+    def moe_emb_chunking(
+        x, logits, pre_bias_logits, w0, w1, w0_bias, w1_bias, wo_bias, sharded_input_ids, rngs, embed_dim
+    ):


can you include a comment for this method?

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.

Shuwen-Fang added the pull ready label Jun 29, 2026

Shuwen-Fang self-assigned this Jun 29, 2026

Shuwen-Fang force-pushed the chunk_moe branch from 63ce931 to 4ab6553 Compare June 29, 2026 19:32

Shuwen-Fang added pull ready and removed pull ready labels Jun 29, 2026

Shuwen-Fang force-pushed the chunk_moe branch from 4ab6553 to 6a5354d Compare July 8, 2026 22:09

Shuwen-Fang requested a review from xibinliu as a code owner July 8, 2026 22:09

Shuwen-Fang force-pushed the chunk_moe branch from 6a5354d to eed3d2f Compare July 8, 2026 23:10

Shuwen-Fang changed the title ~~DO NOT MERGE Chunk moe~~ Overlap moe comms with collective matmul Jul 9, 2026

Shuwen-Fang force-pushed the chunk_moe branch from 1f2e152 to 58bb8cd Compare July 9, 2026 02:40

NuojCheng reviewed Jul 9, 2026

View reviewed changes

Shuwen-Fang requested a review from NuojCheng July 10, 2026 00:36

Shuwen-Fang force-pushed the chunk_moe branch from 58bb8cd to 6da8e0a Compare July 10, 2026 00:41

NuojCheng approved these changes Jul 10, 2026

View reviewed changes

gobbleturk reviewed Jul 13, 2026

View reviewed changes

Shuwen-Fang force-pushed the chunk_moe branch from 6da8e0a to 345d252 Compare July 13, 2026 21:11

Shuwen-Fang requested a review from huytransformer as a code owner July 13, 2026 21:11

Shuwen-Fang requested a review from gobbleturk July 13, 2026 21:12

Shuwen-Fang force-pushed the chunk_moe branch from 345d252 to 33293a4 Compare July 13, 2026 21:14

gobbleturk reviewed Jul 14, 2026

View reviewed changes

gobbleturk approved these changes Jul 14, 2026

View reviewed changes

Shuwen-Fang force-pushed the chunk_moe branch from 33293a4 to afcb5ff Compare July 14, 2026 18:12

copybara-service Bot merged commit fa79905 into main Jul 14, 2026
12 checks passed

copybara-service Bot deleted the chunk_moe branch July 14, 2026 19:47


		assert_arrays_all_close(actual, expected)

		@pytest.mark.skip(reason="Test takes too long, can run locally to verify changes b/528087469")

Uh oh!

Conversation

Shuwen-Fang commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Testing

Checklist

Uh oh!

codecov Bot commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng Jul 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

NuojCheng left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Shuwen-Fang commented Jun 29, 2026 •

edited

Loading

codecov Bot commented Jun 29, 2026 •

edited

Loading

NuojCheng Jul 9, 2026 •

edited

Loading