From 7f3bcf4238ff2ea7b0d967a7e3d1d7438875318e Mon Sep 17 00:00:00 2001
From: Siddharth Singh <sidsingh@nvidia.com>
Date: Thu, 30 Apr 2026 21:14:34 -0700
Subject: [PATCH] overlap only for nvls dispatcher

---
 megatron/core/transformer/moe/moe_layer.py    | 14 +++++++-
 .../moe/token_dispatcher_inference.py         | 34 +++++++++++++++++--
 2 files changed, 44 insertions(+), 4 deletions(-)

diff --git a/megatron/core/transformer/moe/moe_layer.py b/megatron/core/transformer/moe/moe_layer.py
index 766dd049f45..1f830caf32c 100644
--- a/megatron/core/transformer/moe/moe_layer.py
+++ b/megatron/core/transformer/moe/moe_layer.py
@@ -371,6 +371,16 @@ def _setup_inference_mode(self, pg_collection):
             pg_collection=pg_collection,
         )
 
+        # Wire shared-expert overlap into the inference dispatcher (NVLS only).
+        # The dispatcher launches the shared-expert forward on SharedExpertMLP.stream
+        # concurrently with AGV+experts+RSV and adds it back in combine_postprocess.
+        if (
+            dispatcher_type == 'nvls'
+            and self.use_shared_expert
+            and self.config.moe_shared_expert_overlap
+        ):
+            self._inference_token_dispatcher.set_shared_experts(self.shared_experts)
+
     def train(self, mode: bool = True):
         """Swap token dispatcher when switching between train and eval modes."""
         super().train(mode)
@@ -380,7 +390,9 @@ def train(self, mode: bool = True):
                 self.shared_expert_overlap = self.config.moe_shared_expert_overlap
             else:
                 self.token_dispatcher = self._inference_token_dispatcher
-                self.shared_expert_overlap = False
+                self.shared_expert_overlap = (
+                    self._inference_token_dispatcher.shared_experts is not None
+                )
         return self
 
     def setup_delayed_wgrad_for_dispatch_backward_overlap(self):
diff --git a/megatron/core/transformer/moe/token_dispatcher_inference.py b/megatron/core/transformer/moe/token_dispatcher_inference.py
index da6f218a626..e31dbc877f3 100644
--- a/megatron/core/transformer/moe/token_dispatcher_inference.py
+++ b/megatron/core/transformer/moe/token_dispatcher_inference.py
@@ -38,8 +38,10 @@
     gather_from_sequence_parallel_region,
     reduce_scatter_to_sequence_parallel_region,
 )
+from megatron.core.transformer.moe.shared_experts import SharedExpertMLP
 from megatron.core.transformer.moe.token_dispatcher import MoEAllGatherTokenDispatcher
 from megatron.core.transformer.transformer_config import TransformerConfig
+from megatron.core.typed_torch import apply_module
 from megatron.core.utils import get_pg_rank, get_pg_size
 
 
@@ -420,12 +422,25 @@ def __init__(
         self.topk = config.moe_router_topk
         # Set in dispatch_preprocess; consumed by token_dispatch and token_combine.
         self._local_tokens: int = 0
+        # When shared_expert_overlap is enabled, the shared expert forward is launched
+        # on SharedExpertMLP.stream in dispatch_preprocess and joined in combine_postprocess.
+        self._shared_expert_output: Optional[torch.Tensor] = None
 
     # ── Dispatch path ─────────────────────────────────────────────────────────────
 
     def dispatch_preprocess(self, hidden_states, routing_map, probs):
-        """Store routing map and local token count; no communication."""
+        """Store routing map and local token count; no inter-rank communication.
+
+        If shared_expert_overlap is enabled (set_shared_experts has been called),
+        launch the entire shared-expert forward on SharedExpertMLP.stream so it
+        runs concurrently with AGV dispatch, expert GEMMs, and RSV combine.
+        """
         self.hidden_shape = hidden_states.shape
+        if self.shared_experts is not None:
+            stream = SharedExpertMLP.stream
+            stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(stream):
+                self._shared_expert_output = apply_module(self.shared_experts)(hidden_states)
         # [S/TP, B, H] -> [S*B/TP, H]
         hidden_states = hidden_states.view(-1, self.hidden_shape[-1])
         self._local_tokens = hidden_states.shape[0]
@@ -458,6 +473,9 @@ def token_dispatch(self, hidden_states, probs):
         rank_token_offset = self._rank_token_offset()
         ep_max_tokens = self._ep_max_tokens()
 
+        # Cap AGV CTAs when overlapping the shared expert so the AGV does not
+        # starve the shared-expert GEMMs running on the side stream.
+        agv_kwargs = {"max_num_blocks": 16} if self.shared_experts is not None else {}
         multimem_all_gatherv_3tensor(
             agv_h["tensor"],
             agv_r["tensor"],
@@ -471,6 +489,7 @@ def token_dispatch(self, hidden_states, probs):
             rank_token_offset=rank_token_offset,
             ep_max_tokens=ep_max_tokens,
             per_rank_max_tokens=per_rank_max,
+            **agv_kwargs,
         )
 
         topk = probs.shape[1]
@@ -523,5 +542,14 @@ def token_combine(self, hidden_states):
         return output.to(torch.bfloat16)
 
     def combine_postprocess(self, hidden_states):
-        """Restore original input shape (e.g. [S/TP, B, H] from [S*B/TP, H])."""
-        return hidden_states.view(self.hidden_shape)
+        """Restore original input shape (e.g. [S/TP, B, H] from [S*B/TP, H]).
+
+        If shared_expert_overlap is enabled, join SharedExpertMLP.stream and add
+        the shared-expert output produced concurrently during dispatch+combine.
+        """
+        output = hidden_states.view(self.hidden_shape)
+        if self._shared_expert_output is not None:
+            torch.cuda.current_stream().wait_stream(SharedExpertMLP.stream)
+            output = output + self._shared_expert_output
+            self._shared_expert_output = None
+        return output