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 8, 08:45 UTC
Fix MoE for the Transformers modelling backend (#34436)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Feb 12, 2026 at 18:29 UTC
679ca5d8d346ede84c9cbba5d6a8789723c295c0
1 file changed
+3
-3
vllm/model_executor/models/transformers/moe.py
+3
-3
@@ -45,7 +45,6 @@ class TransformersFusedMoE(FusedMoE):
45
# --8<-- [end:transformers_fused_moe]
46
47
def __init__(self, *args, **kwargs):
48
- super().__init__(*args, **kwargs)
48
self._topk_ids: torch.Tensor = None
49
50
def custom_routing_function(hidden_states, gating_output, topk, renormalize):
@@ -63,7 +62,8 @@ class TransformersFusedMoE(FusedMoE):
62
(topk_ids,) = dist_group.all_gatherv([topk_ids], 0, sizes)
63
return topk_weights, topk_ids
64
66
- self.custom_routing_function = custom_routing_function
65
+ kwargs["custom_routing_function"] = custom_routing_function
66
+ super().__init__(*args, **kwargs)
67
68
def forward(
69
self,
@@ -94,7 +94,7 @@ def transformers_moe_forward(
94
self = forward_context.no_compile_layers[layer_name]
95
self._topk_ids = topk_ids
96
# Clone hidden_states because it will be mutated in-place in FusedMoE
97
- return self.forward_impl(hidden_states.clone(), topk_weights)
97
+ return self.runner.forward(hidden_states.clone(), topk_weights)
98
99
100
def transformers_moe_forward_fake(