From 326038ed7f641d814668251c2835e721c0430c21 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Olejniczak?= Date: Mon, 31 Aug 2026 10:35:35 +0300 Subject: [PATCH] [FIX_FOR_VLLM_CUSTOM=488b6da105222f4f8130b3aae4a0e67cfc61f522] Feature-detect MoE shared-experts overlap API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: vllm#52033 re-landed the dual-stream shared-experts work that vllm#52024 had reverted, removing MoERunner._maybe_sync_shared_experts_stream that our dp_size==1 MoE fast path calls, so every MoE forward raised AttributeError (19 of 59 hourly jobs red). Upstream: https://github.com/vllm-project/vllm/pull/52033 Fix: Route the aux-stream launch through _launch_shared_experts_overlap, which feature-detects _maybe_sync_shared_experts_stream vs SharedExperts.maybe_forward_async and only forwards shared_experts_overlapping when an async launch actually happened, so a future flip in either direction no longer breaks HPU. Signed-off-by: Paweł Olejniczak --- vllm_gaudi/ops/hpu_fused_moe.py | 59 ++++++++++++++++++++++++++------- 1 file changed, 47 insertions(+), 12 deletions(-) diff --git a/vllm_gaudi/ops/hpu_fused_moe.py b/vllm_gaudi/ops/hpu_fused_moe.py index 845d737b4d..111c82d438 100644 --- a/vllm_gaudi/ops/hpu_fused_moe.py +++ b/vllm_gaudi/ops/hpu_fused_moe.py @@ -627,6 +627,44 @@ def forward_oot( return output.view(*input_shape) +def _launch_shared_experts_overlap(runner: MoERunnerBase, shared_experts_input: torch.Tensor | None) -> bool: + """Start the shared-expert multi-stream overlap, tolerating either upstream API. + + Upstream keeps flip-flopping on how the aux-stream launch is spelled: vllm#51838 + added ``SharedExperts.maybe_forward_async`` plus a ``shared_experts_overlapping`` + kwarg on ``MoERunner._apply_quant_method``, vllm#52024 reverted to + ``MoERunner._maybe_sync_shared_experts_stream``, and vllm#52033 re-landed the + #51838 shape. Feature-detect so the next flip does not break the HPU fast path. + + On Gaudi every shape is a no-op — upstream gates the multi-stream path on + ``current_platform.is_cuda_alike()`` — so falling through to "no overlap" on an + unrecognised future shape stays functionally correct. + + Args: + runner: The upstream ``MoERunner`` (``self`` of the patched forward). + shared_experts_input: Shared-expert input, or None when there are none. + + Returns: + True iff the launch was asynchronous, meaning the caller must pass + ``shared_experts_overlapping=True`` down to ``_apply_quant_method``. + """ + sync_stream = getattr(runner, "_maybe_sync_shared_experts_stream", None) + if sync_stream is not None: + # vllm#52024 shape: runner-side sync, overlap decided internally. + sync_stream(shared_experts_input) + return False + + shared_experts = getattr(runner, "_shared_experts", None) + if shared_experts is None or shared_experts_input is None: + return False + + # vllm#51838 / vllm#52033 shape: launch on the aux stream, await it later. + forward_async = getattr(shared_experts, "maybe_forward_async", None) + if forward_async is None: + return False + return bool(forward_async(shared_experts_input)) + + def patched_fused_moe_forward( self, hidden_states: torch.Tensor, @@ -672,15 +710,10 @@ def patched_fused_moe_forward( # is initialized on routed_experts (which the runner holds directly), so # unlike the old FusedMoE-layer-based init we do NOT need the layer here. self.routed_experts._ensure_moe_quant_config_init() - # Sync the aux/main stream for shared-expert multi-stream overlap, - # mirroring upstream MoERunner._forward_impl. vllm PR #52024 reverted the - # dual-stream decode work: SharedExperts.maybe_forward_async was removed - # (folded back into maybe_sync_shared_experts_stream, re-exposed on the - # runner as _maybe_sync_shared_experts_stream), and _apply_quant_method no - # longer takes a shared_experts_overlapping flag — overlap is decided - # internally. On HPU (not CUDA-alike) this is a no-op and the shared - # experts run synchronously inside _apply_quant_method. - self._maybe_sync_shared_experts_stream(shared_experts_input) + # Start the shared-expert aux stream before routed dispatch, mirroring + # upstream MoERunner._forward_impl. Version-tolerant: see + # _launch_shared_experts_overlap. + shared_experts_overlapping = _launch_shared_experts_overlap(self, shared_experts_input) # Apply the gate if the runner holds it (mirrors _forward_impl). if self.gate is not None: if self._fse_fuse_gate: @@ -690,14 +723,16 @@ def patched_fused_moe_forward( router_logits, _ = self.gate(hidden_states) # Core MoERunner._apply_quant_method takes no layer argument — it reads # everything it needs off the runner (self.routed_experts / self.router). - # Call it exactly as upstream _forward_impl does. vllm PR #52024 dropped - # the shared_experts_overlapping argument: overlap is decided internally - # and the shared-expert output is stashed on self._shared_experts. + # Call it exactly as upstream _forward_impl does. Only pass + # shared_experts_overlapping when an async launch actually happened: that + # is the only case in which the kwarg is guaranteed to exist upstream. + extra_quant_kwargs = {"shared_experts_overlapping": True} if shared_experts_overlapping else {} shared_output, fused_hidden = self._apply_quant_method( hidden_states=hidden_states, router_logits=router_logits, shared_experts_input=shared_experts_input, input_ids=input_ids, + **extra_quant_kwargs, ) result = self._maybe_combine(shared_output, fused_hidden) else: