A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Fix MTP with deepep_low_latency (#25904)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 2, 2025 at 17:29 UTC
13cdc021732d69e7006cbe2ad564d987e749ce31
1 file changed
+16
-2
vllm/model_executor/layers/fused_moe/layer.py
+16
-2
@@ -1899,6 +1899,15 @@ class FusedMoE(CustomOp):
1899
staged_hidden_states.copy_(hidden_states, non_blocking=True)
1900
staged_router_logits.copy_(router_logits, non_blocking=True)
1901
1902
+ # If there are shared experts but we are not using a modular kernel,
1903
+ # the shared experts must be called here
1904
+ if (not isinstance(self.quant_method.fused_experts,
1905
+ FusedMoEModularKernel)
1906
+ and self.shared_experts is not None):
1907
+ shared_output = self.shared_experts(staged_hidden_states)
1908
+ else:
1909
+ shared_output = None
1910
+
1911
# Matrix multiply.
1912
final_hidden_states = self.quant_method.apply(
1913
layer=self,
@@ -1922,8 +1931,13 @@ class FusedMoE(CustomOp):
1931
logical_replica_count=self.logical_replica_count,
1932
)
1933
1925
- assert self.shared_experts is None or isinstance(
1926
- final_hidden_states, tuple)
1934
+ if shared_output is not None:
1935
+ assert not isinstance(final_hidden_states, tuple)
1936
+ assert self.shared_experts is not None
1937
+ final_hidden_states = (
1938
+ shared_output,
1939
+ final_hidden_states,
1940
+ )
1941
1942
if self.zero_expert_num is not None and self.zero_expert_num > 0:
1943
assert isinstance(final_hidden_states, tuple)