diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py index 1a53476694..c3ffded512 100644 --- a/tests/models/test_deepseek_v4_aiter_mega_moe.py +++ b/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 diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index d1a8ca1d9d..569cc0c9f5 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/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.