[ROCm][DSv4] Overlap the unfused shared expert with AITER MegaMoE (#61819)

Signed-off-by: Taoyu Zhu <zhutaoyu97@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: vllm-agent <claw@inferact.ai>

tests/models/test_deepseek_v4_aiter_mega_moe.py

@@ -15,6 +15,7 @@ from vllm.models.deepseek_v4.amd.mega_moe import (
 )
 from vllm.models.deepseek_v4.amd.model import DeepseekV4MoE
 from vllm.platforms import current_platform
+from vllm.utils.torch_utils import aux_stream
 
 HIDDEN = 64
 INTER = 64
@@ -156,7 +157,10 @@ def test_shared_expert_route_appended(monkeypatch):
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm")
 @pytest.mark.parametrize("hash_routing", [False, True])
 @pytest.mark.parametrize("num_tokens", [1, 37])
-def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_tokens):
+@pytest.mark.parametrize("shared_stream", [False, True])
+def test_mega_moe_routing_matches_reference(
+    monkeypatch, hash_routing, num_tokens, shared_stream
+):
     num_experts, top_k, vocab = 16, 4, 64
     scaling = 2.5
     device = "cuda"
@@ -195,6 +199,10 @@ def test_mega_moe_routing_matches_reference(monkeypatch, hash_routing, num_token
     moe.gate = Gate()
     moe.experts = experts
     moe.shared_experts = lambda x: 3 * x
+    moe.shared_experts_stream = aux_stream() if shared_stream else None
+    moe.shared_experts_events = (
+        (torch.cuda.Event(), torch.cuda.Event()) if shared_stream else None
+    )
     moe.scoring_func = "sqrtsoftplus"
     moe.n_activated_experts = top_k
     moe.renormalize = True

vllm/models/deepseek_v4/amd/model.py

@@ -82,6 +82,8 @@ from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import on_gfx950, on_gfx1250
 from vllm.sequence import IntermediateTensors
+from vllm.utils.multi_stream_utils import maybe_execute_in_parallel
+from vllm.utils.torch_utils import aux_stream
 from vllm.v1.worker.ubatching import dbo_current_ubatch_token_offset
 
 from ..common.mm_preprocess import IMAGE_SENTINEL_BASE_ID
@@ -713,6 +715,18 @@ class DeepseekV4MoE(nn.Module):
                 ),
                 prefix=f"{prefix}.experts",
             )
+            self.shared_experts_stream = (
+                aux_stream()
+                if self.shared_experts is not None
+                and not self.experts.fuse_shared_expert
+                and not envs.VLLM_DISABLE_SHARED_EXPERTS_STREAM
+                else None
+            )
+            self.shared_experts_events = (
+                (torch.cuda.Event(), torch.cuda.Event())
+                if self.shared_experts_stream is not None
+                else None
+            )
             return
 
         fuse_shared_into_routed = (
@@ -777,6 +791,30 @@ class DeepseekV4MoE(nn.Module):
 
     def _forward_mega_moe(
         self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None
+    ) -> torch.Tensor:
+        shared_experts = self.shared_experts
+        if shared_experts is None:
+            return self._forward_mega_moe_routed(hidden_states, input_ids)
+        events = self.shared_experts_events
+        if (
+            events is None
+            or hidden_states.shape[0] > envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
+        ):
+            return self._forward_mega_moe_routed(
+                hidden_states, input_ids
+            ) + shared_experts(hidden_states)
+        # The unfused shared expert overlaps router + MegaMoE in small steps.
+        routed, shared = maybe_execute_in_parallel(
+            fn0=lambda: self._forward_mega_moe_routed(hidden_states, input_ids),
+            fn1=lambda: shared_experts(hidden_states),
+            event0=events[0],
+            event1=events[1],
+            aux_stream=self.shared_experts_stream,
+        )
+        return routed + shared
+
+    def _forward_mega_moe_routed(
+        self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None
     ) -> torch.Tensor:
         router_logits, _ = self.gate(hidden_states)
         topk_weights, topk_ids = fused_topk_bias(
@@ -792,12 +830,7 @@ class DeepseekV4MoE(nn.Module):
             bias_vl=self.gate.bias_vl,
             image_sentinel_lo=self.image_sentinel_lo,
         )
-        final_hidden_states = self.experts(hidden_states, topk_weights, topk_ids)
-        if self.shared_experts is not None:
-            final_hidden_states = final_hidden_states + self.shared_experts(
-                hidden_states
-            )
-        return final_hidden_states
+        return self.experts(hidden_states, topk_weights, topk_ids)
 
 
 # Hidden sizes supported by AITER mhc_pre_big_fuse_rmsnorm.