diff --git a/vllm_gaudi/models/qwen3_next.py b/vllm_gaudi/models/qwen3_next.py index 16181c4866..baeae5a208 100644 --- a/vllm_gaudi/models/qwen3_next.py +++ b/vllm_gaudi/models/qwen3_next.py @@ -136,7 +136,10 @@ def _hpu_qwen3next_sparse_moe_forward( if self.is_sequence_parallel: hidden_states = sequence_parallel_chunk(hidden_states) - if self.experts.is_internal_router: + # Upstream removed the MoERunner.is_internal_router property (vllm PR + # #51838); its body was simply `self.gate is not None`. Inline that check + # so the runner-holds-the-gate path still routes through the internal gate. + if getattr(self.experts, "gate", None) is not None: final_hidden_states = self.experts(hidden_states=hidden_states, router_logits=hidden_states) else: router_logits, _ = self.gate(hidden_states) diff --git a/vllm_gaudi/ops/hpu_fused_moe.py b/vllm_gaudi/ops/hpu_fused_moe.py index c5caf04778..332b5c7230 100644 --- a/vllm_gaudi/ops/hpu_fused_moe.py +++ b/vllm_gaudi/ops/hpu_fused_moe.py @@ -658,7 +658,15 @@ 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() - self._maybe_sync_shared_experts_stream(shared_experts_input) + # Upstream removed MoERunner._maybe_sync_shared_experts_stream (vllm + # PR #51838); the multi-stream shared-expert launch now lives on + # SharedExperts.maybe_forward_async, mirrored by _forward_impl. On HPU + # (not CUDA-alike) maybe_forward_async returns False, so shared experts + # run synchronously inside _apply_quant_method — behaviourally identical + # to the old no-op stream sync, but kept in sync with upstream. + shared_experts_overlapping = False + if self._shared_experts is not None: + shared_experts_overlapping = self._shared_experts.maybe_forward_async(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: @@ -668,12 +676,14 @@ 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. + # Call it exactly as upstream _forward_impl does, threading the overlap + # flag so an async-launched shared expert is awaited (not recomputed). 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, + shared_experts_overlapping=shared_experts_overlapping, ) result = self._maybe_combine(shared_output, fused_hidden) else: @@ -696,13 +706,21 @@ def patched_fused_moe_forward( if og_hidden_dim_pre_xform is not None: fused_output = fused_output[..., :og_hidden_dim_pre_xform] - shared_output = self._maybe_reduce_shared_expert_output(shared_output) + # Mirror upstream MoERunner.forward's two mutually-exclusive all-reduce + # points, threaded by _fused_output_is_reduced: reduce a latent routed + # output before its (possibly non-linear) transform, then reduce shared + # output to match, else all-reduce the combined sum in the final step. + fused_output_is_reduced = self._fused_output_is_reduced + fused_output, fused_output_is_reduced = self._maybe_reduce_routed_output_before_transform( + fused_output, fused_output_is_reduced) + + shared_output = self._maybe_reduce_shared_expert_output(shared_output, fused_output_is_reduced) shared_output, fused_output = self._maybe_apply_routed_scale_to_output(shared_output, fused_output) fused_output = self.apply_routed_output_transform(fused_output) combined = (shared_output + fused_output) if shared_output is not None else fused_output - combined = self._maybe_reduce_final_output(combined, og_hidden_dim_post_xform) + combined = self._maybe_reduce_final_output(combined, og_hidden_dim_post_xform, fused_output_is_reduced) return self._maybe_add_zero_expert_output(combined)