A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix][MoE] Fix fused MoE expert mapping helper call sites (#44897)
Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com>
Mohammad Miadh Angkad committed
Jun 9, 2026 at 04:35 UTC
823a0ab75402ee9836ab228db0ad19c7102b97f3
6 files changed
+26
-11
vllm/lora/model_manager.py
+1
-1
@@ -1052,7 +1052,7 @@ class LoRAModelManager:
1052
"""Narrow a flat expert-major sub-module list to this rank's experts.
1053
1054
``new_module_names`` is produced by
1055
- ``FusedMoE.make_expert_params_mapping`` and is ordered
1055
+ ``fused_moe_make_expert_params_mapping`` and is ordered
1056
``[e=0,w1, e=0,w2, e=0,w3, e=1,w1, ...]`` (non-gated MoE has 2
1057
entries per expert instead of 3). When the module is a 2D
1058
``FusedMoEWithLoRA`` with EP enabled, we slice the list to the
vllm/model_executor/models/cohere2_moe.py
+5
-2
@@ -16,7 +16,10 @@ from vllm.distributed import (
16
)
17
from vllm.model_executor.layers.activation import SiluAndMul
18
from vllm.model_executor.layers.attention import Attention
19
-from vllm.model_executor.layers.fused_moe import FusedMoE
19
+from vllm.model_executor.layers.fused_moe import (
20
+ FusedMoE,
21
+ fused_moe_make_expert_params_mapping,
22
+)
23
from vllm.model_executor.layers.linear import (
24
MergedColumnParallelLinear,
25
QKVParallelLinear,
@@ -450,7 +453,7 @@ class Cohere2MoeModel(nn.Module):
453
("gate_up_proj", "up_proj", 1),
454
]
455
453
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
456
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
457
self,
458
ckpt_gate_proj_name="gate_proj",
459
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/hy_v3.py
+6
-2
@@ -43,7 +43,11 @@ from vllm.distributed import (
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.activation import SiluAndMul
45
from vllm.model_executor.layers.attention import Attention
46
-from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
46
+from vllm.model_executor.layers.fused_moe import (
47
+ FusedMoE,
48
+ GateLinear,
49
+ fused_moe_make_expert_params_mapping,
50
+)
51
from vllm.model_executor.layers.layernorm import RMSNorm
52
from vllm.model_executor.layers.linear import (
53
MergedColumnParallelLinear,
@@ -474,7 +478,7 @@ class HYV3Model(nn.Module):
478
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
479
# Params for weights, fp8 weight scales, fp8 activation scales
480
# (param_name, weight_name, expert_id, shard_id)
477
- return FusedMoE.make_expert_params_mapping(
481
+ return fused_moe_make_expert_params_mapping(
482
self,
483
ckpt_gate_proj_name="gate_proj",
484
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/hy_v3_mtp.py
+4
-2
@@ -32,7 +32,9 @@ from torch import nn
32
from transformers import PretrainedConfig
33
34
from vllm.config import CacheConfig, ModelConfig, VllmConfig
35
-from vllm.model_executor.layers.fused_moe import FusedMoE
35
+from vllm.model_executor.layers.fused_moe import (
36
+ fused_moe_make_expert_params_mapping,
37
+)
38
from vllm.model_executor.layers.layernorm import RMSNorm
39
from vllm.model_executor.layers.logits_processor import LogitsProcessor
40
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -294,7 +296,7 @@ class HYV3MTP(nn.Module):
296
]
297
298
if _is_moe(self.config):
297
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
299
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
300
self,
301
ckpt_gate_proj_name="gate_proj",
302
ckpt_down_proj_name="down_proj",
vllm/models/deepseek_v4/xpu/model.py
+6
-2
@@ -18,7 +18,11 @@ from vllm.distributed import (
18
)
19
from vllm.forward_context import get_forward_context
20
from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
21
-from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
21
+from vllm.model_executor.layers.fused_moe import (
22
+ FusedMoE,
23
+ GateLinear,
24
+ fused_moe_make_expert_params_mapping,
25
+)
26
from vllm.model_executor.layers.fused_moe.router.fused_topk_bias_router import (
27
fused_topk_bias,
28
)
@@ -1220,7 +1224,7 @@ class DeepseekV4Model(nn.Module):
1224
return make_deepseek_v4_expert_params_mapping(self.config.n_routed_experts)
1225
# Params for weights, fp8 weight scales, fp8 activation scales
1226
# (param_name, weight_name, expert_id, shard_id)
1223
- return FusedMoE.make_expert_params_mapping(
1227
+ return fused_moe_make_expert_params_mapping(
1228
self,
1229
ckpt_gate_proj_name="w1",
1230
ckpt_down_proj_name="w2",
vllm/models/deepseek_v4/xpu/mtp.py
+4
-2
@@ -25,7 +25,9 @@ from vllm.distributed import (
25
get_tensor_model_parallel_world_size,
26
)
27
from vllm.logger import init_logger
28
-from vllm.model_executor.layers.fused_moe import FusedMoE
28
+from vllm.model_executor.layers.fused_moe import (
29
+ fused_moe_make_expert_params_mapping,
30
+)
31
from vllm.model_executor.layers.layernorm import RMSNorm
32
from vllm.model_executor.layers.linear import ReplicatedLinear
33
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -332,7 +334,7 @@ class DeepSeekV4MTP(nn.Module):
334
self.config.n_routed_experts
335
)
336
else:
335
- expert_mapping = FusedMoE.make_expert_params_mapping(
337
+ expert_mapping = fused_moe_make_expert_params_mapping(
338
self,
339
ckpt_gate_proj_name="w1",
340
ckpt_down_proj_name="w2",