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 4, 17:30 UTC
[MoE Refactor] Rename FusedMoE.make_expert_params_mapping to fused_moe_make_expert_params_mapping (#40671)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Apr 23, 2026 at 11:22 UTC
1c2c1eb8b9fdd2e67c45afb6123ccc07c0177555
53 files changed
+254
-98
vllm/model_executor/layers/fused_moe/__init__.py
+2
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe_method_base import (
19
from vllm.model_executor.layers.fused_moe.layer import (
20
FusedMoE,
21
FusedMoeWeightScaleSupported,
22
+ fused_moe_make_expert_params_mapping,
23
)
24
from vllm.model_executor.layers.fused_moe.modular_kernel import (
25
FusedMoEActivationFormat,
@@ -65,6 +66,7 @@ __all__ = [
66
"RoutingMethodType",
67
"activation_without_mul",
68
"apply_moe_activation",
69
+ "fused_moe_make_expert_params_mapping",
70
"override_config",
71
"get_config",
72
]
vllm/model_executor/layers/fused_moe/layer.py
+19
@@ -1618,6 +1618,25 @@ class FusedMoE(PluggableLayer):
1618
return s
1619
1620
1621
+# This is a temporary forwarding method which will be removed/modified layer.
1622
+def fused_moe_make_expert_params_mapping(
1623
+ model: torch.nn.Module,
1624
+ ckpt_gate_proj_name: str,
1625
+ ckpt_down_proj_name: str,
1626
+ ckpt_up_proj_name: str,
1627
+ num_experts: int,
1628
+ num_redundant_experts: int = 0,
1629
+) -> list[tuple[str, str, int, str]]:
1630
+ return FusedMoE.make_expert_params_mapping(
1631
+ model,
1632
+ ckpt_gate_proj_name,
1633
+ ckpt_down_proj_name,
1634
+ ckpt_up_proj_name,
1635
+ num_experts,
1636
+ num_redundant_experts,
1637
+ )
1638
+
1639
+
1640
# Mark the FusedMoE weight_loader as supporting MoE-specific parameters
1641
# to avoid expensive runtime reflection in model loading code
1642
FusedMoE.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
vllm/model_executor/models/AXK1.py
+6
-3
@@ -42,7 +42,10 @@ from vllm.distributed import (
42
)
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.attention import Attention
45
-from vllm.model_executor.layers.fused_moe import FusedMoE
45
+from vllm.model_executor.layers.fused_moe import (
46
+ FusedMoE,
47
+ fused_moe_make_expert_params_mapping,
48
+)
49
from vllm.model_executor.layers.layernorm import RMSNorm
50
from vllm.model_executor.layers.linear import (
51
ColumnParallelLinear,
@@ -916,7 +919,7 @@ class AXK1ForCausalLM(
919
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
920
# Params for weights, fp8 weight scales, fp8 activation scales
921
# (param_name, weight_name, expert_id, shard_id)
919
- return FusedMoE.make_expert_params_mapping(
922
+ return fused_moe_make_expert_params_mapping(
923
self,
924
ckpt_gate_proj_name="gate_proj",
925
ckpt_down_proj_name="down_proj",
@@ -950,7 +953,7 @@ class AXK1ForCausalLM(
953
954
# Params for weights, fp8 weight scales, fp8 activation scales
955
# (param_name, weight_name, expert_id, shard_id)
953
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
956
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
957
self,
958
ckpt_gate_proj_name="gate_proj",
959
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/afmoe.py
+5
-2
@@ -18,7 +18,10 @@ from vllm.distributed import (
18
)
19
from vllm.logger import init_logger
20
from vllm.model_executor.layers.attention import Attention
21
-from vllm.model_executor.layers.fused_moe import FusedMoE
21
+from vllm.model_executor.layers.fused_moe import (
22
+ FusedMoE,
23
+ fused_moe_make_expert_params_mapping,
24
+)
25
from vllm.model_executor.layers.layernorm import RMSNorm
26
from vllm.model_executor.layers.linear import (
27
ColumnParallelLinear,
@@ -479,7 +482,7 @@ class AfmoeModel(nn.Module, EagleModelMixin):
482
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
483
# Params for weights, fp8 weight scales, fp8 activation scales
484
# (param_name, weight_name, expert_id, shard_id)
482
- return FusedMoE.make_expert_params_mapping(
485
+ return fused_moe_make_expert_params_mapping(
486
self,
487
ckpt_gate_proj_name="gate_proj",
488
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/arctic.py
+4
-1
@@ -18,7 +18,10 @@ from vllm.distributed import (
18
)
19
from vllm.model_executor.layers.activation import SiluAndMul
20
from vllm.model_executor.layers.attention import Attention
21
-from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
21
+from vllm.model_executor.layers.fused_moe import (
22
+ fused_experts,
23
+ fused_topk,
24
+)
25
from vllm.model_executor.layers.layernorm import RMSNorm
26
from vllm.model_executor.layers.linear import (
27
MergedColumnParallelLinear,
vllm/model_executor/models/aria.py
+3
-1
@@ -14,7 +14,9 @@ from vllm.config.multimodal import BaseDummyOptions
14
from vllm.distributed import get_tensor_model_parallel_rank
15
from vllm.inputs import MultiModalDataDict
16
from vllm.model_executor.layers.activation import get_act_fn
17
-from vllm.model_executor.layers.fused_moe import FusedMoE
17
+from vllm.model_executor.layers.fused_moe import (
18
+ FusedMoE,
19
+)
20
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
21
from vllm.model_executor.layers.logits_processor import LogitsProcessor
22
from vllm.model_executor.layers.quantization import QuantizationConfig
vllm/model_executor/models/bailing_moe.py
+5
-2
@@ -41,7 +41,10 @@ from vllm.distributed import (
41
)
42
from vllm.model_executor.layers.activation import SiluAndMul
43
from vllm.model_executor.layers.attention import Attention
44
-from vllm.model_executor.layers.fused_moe import FusedMoE
44
+from vllm.model_executor.layers.fused_moe import (
45
+ FusedMoE,
46
+ fused_moe_make_expert_params_mapping,
47
+)
48
from vllm.model_executor.layers.layernorm import RMSNorm
49
from vllm.model_executor.layers.linear import (
50
MergedColumnParallelLinear,
@@ -461,7 +464,7 @@ class BailingMoeModel(nn.Module):
464
return hidden_states
465
466
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
464
- return FusedMoE.make_expert_params_mapping(
467
+ return fused_moe_make_expert_params_mapping(
468
self,
469
ckpt_gate_proj_name="gate_proj",
470
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/bailing_moe_linear.py
+5
-2
@@ -21,7 +21,10 @@ from vllm.model_executor.layers.fla.ops.layernorm_guard import (
21
RMSNormGated,
22
layernorm_fn,
23
)
24
-from vllm.model_executor.layers.fused_moe import FusedMoE
24
+from vllm.model_executor.layers.fused_moe import (
25
+ FusedMoE,
26
+ fused_moe_make_expert_params_mapping,
27
+)
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
from vllm.model_executor.layers.linear import (
30
ColumnParallelLinear,
@@ -990,7 +993,7 @@ class BailingMoeV25Model(nn.Module):
993
994
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
995
"""Get expert parameter mapping for MoE layers."""
993
- return FusedMoE.make_expert_params_mapping(
996
+ return fused_moe_make_expert_params_mapping(
997
self,
998
ckpt_gate_proj_name="gate_proj",
999
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/dbrx.py
+3
-1
@@ -15,7 +15,9 @@ from vllm.distributed import (
15
get_tensor_model_parallel_world_size,
16
)
17
from vllm.model_executor.layers.attention import Attention
18
-from vllm.model_executor.layers.fused_moe import FusedMoE
18
+from vllm.model_executor.layers.fused_moe import (
19
+ FusedMoE,
20
+)
21
from vllm.model_executor.layers.linear import (
22
QKVParallelLinear,
23
ReplicatedLinear,
vllm/model_executor/models/deepseek_eagle.py
+4
-2
@@ -8,7 +8,9 @@ import torch.nn as nn
8
9
from vllm.compilation.decorators import support_torch_compile
10
from vllm.config import VllmConfig
11
-from vllm.model_executor.layers.fused_moe import FusedMoE
11
+from vllm.model_executor.layers.fused_moe import (
12
+ fused_moe_make_expert_params_mapping,
13
+)
14
from vllm.model_executor.layers.layernorm import RMSNorm
15
from vllm.model_executor.layers.logits_processor import LogitsProcessor
16
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -105,7 +107,7 @@ class DeepseekV2Model(nn.Module):
107
108
# Params for weights, fp8 weight scales, fp8 activation scales
109
# (param_name, weight_name, expert_id, shard_id)
108
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
110
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
111
self,
112
ckpt_gate_proj_name="gate_proj",
113
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/deepseek_mtp.py
+4
-2
@@ -11,7 +11,9 @@ from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
-from vllm.model_executor.layers.fused_moe import FusedMoE
14
+from vllm.model_executor.layers.fused_moe import (
15
+ fused_moe_make_expert_params_mapping,
16
+)
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -252,7 +254,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
254
]
255
stacked_params_mapping.extend(indexer_fused_mapping)
256
255
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
257
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
258
self,
259
ckpt_gate_proj_name="gate_proj",
260
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/deepseek_v2.py
+3
-2
@@ -51,6 +51,7 @@ from vllm.model_executor.layers.fused_moe import (
51
FusedMoE,
52
GateLinear,
53
RoutingMethodType,
54
+ fused_moe_make_expert_params_mapping,
55
)
56
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
57
from vllm.model_executor.layers.linear import (
@@ -1432,7 +1433,7 @@ class DeepseekV2ForCausalLM(
1433
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1434
# Params for weights, fp8 weight scales, fp8 activation scales
1435
# (param_name, weight_name, expert_id, shard_id)
1435
- return FusedMoE.make_expert_params_mapping(
1436
+ return fused_moe_make_expert_params_mapping(
1437
self,
1438
ckpt_gate_proj_name="gate_proj",
1439
ckpt_down_proj_name="down_proj",
@@ -1474,7 +1475,7 @@ class DeepseekV2ForCausalLM(
1475
1476
# Params for weights, fp8 weight scales, fp8 activation scales
1477
# (param_name, weight_name, expert_id, shard_id)
1477
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
1478
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
1479
self,
1480
ckpt_gate_proj_name="gate_proj",
1481
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/dots1.py
+5
-2
@@ -40,7 +40,10 @@ from vllm.distributed import (
40
)
41
from vllm.model_executor.layers.activation import SiluAndMul
42
from vllm.model_executor.layers.attention import Attention
43
-from vllm.model_executor.layers.fused_moe import FusedMoE
43
+from vllm.model_executor.layers.fused_moe import (
44
+ FusedMoE,
45
+ fused_moe_make_expert_params_mapping,
46
+)
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
MergedColumnParallelLinear,
@@ -413,7 +416,7 @@ class Dots1Model(nn.Module):
416
return hidden_states
417
418
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
416
- return FusedMoE.make_expert_params_mapping(
419
+ return fused_moe_make_expert_params_mapping(
420
self,
421
ckpt_gate_proj_name="gate_proj",
422
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/ernie45_moe.py
+5
-2
@@ -42,7 +42,10 @@ from vllm.distributed import (
42
from vllm.logger import init_logger
43
from vllm.model_executor.layers.activation import SiluAndMul
44
from vllm.model_executor.layers.attention import Attention
45
-from vllm.model_executor.layers.fused_moe import FusedMoE
45
+from vllm.model_executor.layers.fused_moe import (
46
+ FusedMoE,
47
+ fused_moe_make_expert_params_mapping,
48
+)
49
from vllm.model_executor.layers.layernorm import RMSNorm
50
from vllm.model_executor.layers.linear import (
51
MergedColumnParallelLinear,
@@ -485,7 +488,7 @@ class Ernie4_5_MoeModel(nn.Module):
488
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
489
# Params for weights, fp8 weight scales, fp8 activation scales
490
# (param_name, weight_name, expert_id, shard_id)
488
- return FusedMoE.make_expert_params_mapping(
491
+ return fused_moe_make_expert_params_mapping(
492
self,
493
ckpt_gate_proj_name="gate_proj",
494
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/ernie45_vl_moe.py
+5
-2
@@ -36,7 +36,10 @@ from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.logger import init_logger
38
from vllm.model_executor.layers.attention import Attention
39
-from vllm.model_executor.layers.fused_moe import FusedMoE
39
+from vllm.model_executor.layers.fused_moe import (
40
+ FusedMoE,
41
+ fused_moe_make_expert_params_mapping,
42
+)
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
QKVParallelLinear,
@@ -649,7 +652,7 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):
652
653
# Params for weights, fp8 weight scales, fp8 activation scales
654
# (param_name, weight_name, expert_id, shard_id)
652
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
655
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
656
self,
657
ckpt_gate_proj_name="gate_proj",
658
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/exaone_moe.py
+5
-2
@@ -30,7 +30,10 @@ from vllm.distributed import (
30
get_pp_group,
31
get_tensor_model_parallel_world_size,
32
)
33
-from vllm.model_executor.layers.fused_moe import FusedMoE
33
+from vllm.model_executor.layers.fused_moe import (
34
+ FusedMoE,
35
+ fused_moe_make_expert_params_mapping,
36
+)
37
from vllm.model_executor.layers.layernorm import RMSNorm
38
from vllm.model_executor.layers.linear import ReplicatedLinear
39
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -326,7 +329,7 @@ class ExaoneMoeModel(nn.Module):
329
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
330
# Params for weights, fp8 weight scales, fp8 activation scales
331
# (param_name, weight_name, expert_id, shard_id)
329
- return FusedMoE.make_expert_params_mapping(
332
+ return fused_moe_make_expert_params_mapping(
333
self,
334
ckpt_gate_proj_name="gate_proj",
335
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/flex_olmo.py
+3
-1
@@ -20,7 +20,9 @@ from torch import nn
20
from vllm.config import VllmConfig
21
from vllm.distributed import get_tensor_model_parallel_world_size
22
from vllm.logger import init_logger
23
-from vllm.model_executor.layers.fused_moe import FusedMoE
23
+from vllm.model_executor.layers.fused_moe import (
24
+ FusedMoE,
25
+)
26
from vllm.model_executor.layers.layernorm import RMSNorm
27
from vllm.model_executor.layers.linear import ReplicatedLinear
28
from vllm.model_executor.models.olmoe import OlmoeAttention, OlmoeForCausalLM
vllm/model_executor/models/gemma4.py
+4
-1
@@ -37,7 +37,10 @@ from vllm.forward_context import get_forward_context
37
from vllm.logger import init_logger
38
from vllm.model_executor.layers.activation import GeluAndMul
39
from vllm.model_executor.layers.attention import Attention
40
-from vllm.model_executor.layers.fused_moe import FusedMoE, GateLinear
40
+from vllm.model_executor.layers.fused_moe import (
41
+ FusedMoE,
42
+ GateLinear,
43
+)
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (
46
ColumnParallelLinear,
vllm/model_executor/models/glm4_moe.py
+5
-2
@@ -42,7 +42,10 @@ from vllm.distributed import (
42
from vllm.logger import init_logger
43
from vllm.model_executor.layers.activation import SiluAndMul
44
from vllm.model_executor.layers.attention import Attention
45
-from vllm.model_executor.layers.fused_moe import FusedMoE
45
+from vllm.model_executor.layers.fused_moe import (
46
+ FusedMoE,
47
+ fused_moe_make_expert_params_mapping,
48
+)
49
from vllm.model_executor.layers.layernorm import RMSNorm
50
from vllm.model_executor.layers.linear import (
51
MergedColumnParallelLinear,
@@ -466,7 +469,7 @@ class Glm4MoeModel(nn.Module):
469
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
470
# Params for weights, fp8 weight scales, fp8 activation scales
471
# (param_name, weight_name, expert_id, shard_id)
469
- return FusedMoE.make_expert_params_mapping(
472
+ return fused_moe_make_expert_params_mapping(
473
self,
474
ckpt_gate_proj_name="gate_proj",
475
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/glm4_moe_lite.py
+6
-4
@@ -41,7 +41,9 @@ from vllm.distributed import (
41
get_pp_group,
42
)
43
from vllm.logger import init_logger
44
-from vllm.model_executor.layers.fused_moe import FusedMoE
44
+from vllm.model_executor.layers.fused_moe import (
45
+ fused_moe_make_expert_params_mapping,
46
+)
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.logits_processor import LogitsProcessor
49
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -308,7 +310,7 @@ class Glm4MoeLiteModel(nn.Module):
310
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
311
# Params for weights, fp8 weight scales, fp8 activation scales
312
# (param_name, weight_name, expert_id, shard_id)
311
- return FusedMoE.make_expert_params_mapping(
313
+ return fused_moe_make_expert_params_mapping(
314
self,
315
ckpt_gate_proj_name="gate_proj",
316
ckpt_down_proj_name="down_proj",
@@ -334,7 +336,7 @@ class Glm4MoeLiteModel(nn.Module):
336
337
# Params for weights, fp8 weight scales, fp8 activation scales
338
# (param_name, weight_name, expert_id, shard_id)
337
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
339
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
340
self,
341
ckpt_gate_proj_name="gate_proj",
342
ckpt_down_proj_name="down_proj",
@@ -616,7 +618,7 @@ class Glm4MoeLiteForCausalLM(
618
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
619
# Params for weights, fp8 weight scales, fp8 activation scales
620
# (param_name, weight_name, expert_id, shard_id)
619
- return FusedMoE.make_expert_params_mapping(
621
+ return fused_moe_make_expert_params_mapping(
622
self,
623
ckpt_gate_proj_name="gate_proj",
624
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/glm4_moe_lite_mtp.py
+5
-2
@@ -32,7 +32,10 @@ from transformers import PretrainedConfig
32
33
from vllm._aiter_ops import rocm_aiter_ops
34
from vllm.config import VllmConfig
35
-from vllm.model_executor.layers.fused_moe import FusedMoE
35
+from vllm.model_executor.layers.fused_moe import (
36
+ FusedMoE,
37
+ fused_moe_make_expert_params_mapping,
38
+)
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
41
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -260,7 +263,7 @@ class Glm4MoeLiteMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):
263
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
264
]
265
263
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
266
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
267
self,
268
ckpt_gate_proj_name="gate_proj",
269
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/glm4_moe_mtp.py
+5
-2
@@ -31,7 +31,10 @@ import torch.nn as nn
31
from transformers import PretrainedConfig
32
33
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
34
-from vllm.model_executor.layers.fused_moe import FusedMoE
34
+from vllm.model_executor.layers.fused_moe import (
35
+ FusedMoE,
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
@@ -247,7 +250,7 @@ class Glm4MoeMTP(nn.Module, Glm4MixtureOfExperts):
250
251
# Params for weights, fp8 weight scales, fp8 activation scales
252
# (param_name, weight_name, expert_id, shard_id)
250
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
253
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
254
self,
255
ckpt_gate_proj_name="gate_proj",
256
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/gpt_oss.py
+5
-2
@@ -20,7 +20,10 @@ from vllm.distributed import (
20
tensor_model_parallel_all_gather,
21
)
22
from vllm.model_executor.layers.attention import Attention
23
-from vllm.model_executor.layers.fused_moe import FusedMoE
23
+from vllm.model_executor.layers.fused_moe import (
24
+ FusedMoE,
25
+ fused_moe_make_expert_params_mapping,
26
+)
27
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
from vllm.model_executor.layers.linear import (
@@ -331,7 +334,7 @@ class GptOssModel(nn.Module, EagleModelMixin):
334
# Params for weights, weight scales, activation scales
335
# (param_name, weight_name, expert_id, shard_id)
336
# NOTE: this is only used for quark.
334
- return FusedMoE.make_expert_params_mapping(
337
+ return fused_moe_make_expert_params_mapping(
338
self,
339
ckpt_gate_proj_name="w1",
340
ckpt_down_proj_name="w2",
vllm/model_executor/models/granitemoe.py
+5
-2
@@ -39,7 +39,10 @@ from vllm.distributed import (
39
tensor_model_parallel_all_gather,
40
)
41
from vllm.model_executor.layers.attention import Attention
42
-from vllm.model_executor.layers.fused_moe import FusedMoE
42
+from vllm.model_executor.layers.fused_moe import (
43
+ FusedMoE,
44
+ fused_moe_make_expert_params_mapping,
45
+)
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
48
QKVParallelLinear,
@@ -351,7 +354,7 @@ class GraniteMoeModel(nn.Module):
354
355
# Params for weights, fp8 weight scales, fp8 activation scales
356
# (param_name, weight_name, expert_id, shard_id)
354
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
357
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
358
self,
359
ckpt_gate_proj_name="w1",
360
ckpt_down_proj_name="w2",
vllm/model_executor/models/grok1.py
+5
-2
@@ -38,7 +38,10 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
38
from vllm.logger import init_logger
39
from vllm.model_executor.layers.activation import GeluAndMul
40
from vllm.model_executor.layers.attention import Attention
41
-from vllm.model_executor.layers.fused_moe import FusedMoE
41
+from vllm.model_executor.layers.fused_moe import (
42
+ FusedMoE,
43
+ fused_moe_make_expert_params_mapping,
44
+)
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
47
MergedColumnParallelLinear,
@@ -519,7 +522,7 @@ class Grok1Model(nn.Module):
522
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
523
# Map expert parameter names to standard names
524
num_experts = _get_num_experts(self.config)
522
- return FusedMoE.make_expert_params_mapping(
525
+ return fused_moe_make_expert_params_mapping(
526
self,
527
ckpt_gate_proj_name=self.ckpt_gate_proj_name,
528
ckpt_down_proj_name=self.ckpt_down_proj_name,
vllm/model_executor/models/hunyuan_v1.py
+5
-2
@@ -42,7 +42,10 @@ from vllm.distributed import (
42
)
43
from vllm.model_executor.layers.activation import SiluAndMul
44
from vllm.model_executor.layers.attention import Attention
45
-from vllm.model_executor.layers.fused_moe import FusedMoE
45
+from vllm.model_executor.layers.fused_moe import (
46
+ FusedMoE,
47
+ fused_moe_make_expert_params_mapping,
48
+)
49
from vllm.model_executor.layers.layernorm import RMSNorm
50
from vllm.model_executor.layers.linear import (
51
ColumnParallelLinear,
@@ -712,7 +715,7 @@ class HunYuanModel(nn.Module, EagleModelMixin):
715
if _is_moe(self.config):
716
# Params for weights, fp8 weight scales, fp8 activation scales
717
# (param_name, weight_name, expert_id, shard_id)
715
- return FusedMoE.make_expert_params_mapping(
718
+ return fused_moe_make_expert_params_mapping(
719
self,
720
ckpt_gate_proj_name="gate_proj",
721
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/interns1_pro.py
+3
-1
@@ -41,7 +41,9 @@ from vllm.distributed import (
41
from vllm.logger import init_logger
42
from vllm.model_executor.layers.activation import SiluAndMul
43
from vllm.model_executor.layers.attention import Attention
44
-from vllm.model_executor.layers.fused_moe import FusedMoE
44
+from vllm.model_executor.layers.fused_moe import (
45
+ FusedMoE,
46
+)
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
MergedColumnParallelLinear,
vllm/model_executor/models/jamba.py
+5
-2
@@ -14,7 +14,10 @@ from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.model_executor.layers.attention import Attention
17
-from vllm.model_executor.layers.fused_moe import FusedMoE
17
+from vllm.model_executor.layers.fused_moe import (
18
+ FusedMoE,
19
+ fused_moe_make_expert_params_mapping,
20
+)
21
from vllm.model_executor.layers.layernorm import RMSNorm
22
from vllm.model_executor.layers.linear import (
23
QKVParallelLinear,
@@ -378,7 +381,7 @@ class JambaModel(nn.Module):
381
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
382
# Params for weights, fp8 weight scales, fp8 activation scales
383
# (param_name, weight_name, expert_id, shard_id)
381
- return FusedMoE.make_expert_params_mapping(
384
+ return fused_moe_make_expert_params_mapping(
385
self,
386
ckpt_gate_proj_name="gate_proj",
387
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/kimi_linear.py
+5
-2
@@ -14,7 +14,10 @@ from vllm.distributed import (
14
)
15
from vllm.logger import init_logger
16
from vllm.model_executor.layers.activation import SiluAndMul
17
-from vllm.model_executor.layers.fused_moe import FusedMoE
17
+from vllm.model_executor.layers.fused_moe import (
18
+ FusedMoE,
19
+ fused_moe_make_expert_params_mapping,
20
+)
21
from vllm.model_executor.layers.kda import KimiDeltaAttention
22
from vllm.model_executor.layers.layernorm import RMSNorm
23
from vllm.model_executor.layers.linear import (
@@ -476,7 +479,7 @@ class KimiLinearModel(nn.Module):
479
if self.config.is_moe:
480
# Params for weights, fp8 weight scales, fp8 activation scales
481
# (param_name, weight_name, expert_id, shard_id)
479
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
482
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
483
self,
484
ckpt_gate_proj_name="w1",
485
ckpt_down_proj_name="w2",
vllm/model_executor/models/lfm2_moe.py
+5
-2
@@ -15,7 +15,10 @@ from vllm.distributed import (
15
)
16
from vllm.model_executor.layers.activation import SiluAndMul
17
from vllm.model_executor.layers.attention import Attention
18
-from vllm.model_executor.layers.fused_moe import FusedMoE
18
+from vllm.model_executor.layers.fused_moe import (
19
+ FusedMoE,
20
+ fused_moe_make_expert_params_mapping,
21
+)
22
from vllm.model_executor.layers.layernorm import RMSNorm
23
from vllm.model_executor.layers.linear import (
24
MergedColumnParallelLinear,
@@ -482,7 +485,7 @@ class Lfm2MoeModel(nn.Module):
485
return hidden_states
486
487
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
485
- return FusedMoE.make_expert_params_mapping(
488
+ return fused_moe_make_expert_params_mapping(
489
self,
490
ckpt_gate_proj_name="w1",
491
ckpt_down_proj_name="w2",
vllm/model_executor/models/llama4.py
+7
-4
@@ -36,7 +36,10 @@ from vllm.model_executor.layers.attention import (
36
Attention,
37
ChunkedLocalAttention,
38
)
39
-from vllm.model_executor.layers.fused_moe import FusedMoE
39
+from vllm.model_executor.layers.fused_moe import (
40
+ FusedMoE,
41
+ fused_moe_make_expert_params_mapping,
42
+)
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
QKVParallelLinear,
@@ -414,7 +417,7 @@ class Llama4Model(LlamaModel):
417
params_dict: The dictionary of module parameters.
418
loaded_params: The set of already loaded parameters.
419
expert_params_mapping: The mapping of expert parameters. Must be
417
- generated by FusedMoE.make_expert_params_mapping().
420
+ generated by fused_moe_make_expert_params_mapping().
421
fused: Whether the expert weights are fused into a single weight
422
tensor or are separate weight tensors for each expert.
423
When fused is True, loaded_weight should have shape of:
@@ -554,7 +557,7 @@ class Llama4Model(LlamaModel):
557
fused_experts_params = False
558
# Expert parameter mapping for the case where the expert weights are
559
# not fused into a single weight tensor.
557
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
560
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
561
self,
562
ckpt_gate_proj_name="gate_proj",
563
ckpt_down_proj_name="down_proj",
@@ -564,7 +567,7 @@ class Llama4Model(LlamaModel):
567
)
568
# Expert parameter mapping for the case where the expert weights are
569
# fused into a single weight tensor.
567
- expert_params_mapping_fused = FusedMoE.make_expert_params_mapping(
570
+ expert_params_mapping_fused = fused_moe_make_expert_params_mapping(
571
self,
572
ckpt_gate_proj_name="gate_up_proj",
573
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/longcat_flash.py
+5
-2
@@ -46,7 +46,10 @@ from vllm.config import CacheConfig, VllmConfig
46
from vllm.distributed import get_pp_group
47
from vllm.logger import init_logger
48
from vllm.model_executor.layers.activation import SiluAndMul
49
-from vllm.model_executor.layers.fused_moe import FusedMoE
49
+from vllm.model_executor.layers.fused_moe import (
50
+ FusedMoE,
51
+ fused_moe_make_expert_params_mapping,
52
+)
53
from vllm.model_executor.layers.layernorm import RMSNorm
54
from vllm.model_executor.layers.linear import (
55
MergedColumnParallelLinear,
@@ -622,7 +625,7 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
625
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
626
# Params for weights, fp8 weight scales, fp8 activation scales
627
# (param_name, weight_name, expert_id, shard_id)
625
- return FusedMoE.make_expert_params_mapping(
628
+ return fused_moe_make_expert_params_mapping(
629
self,
630
ckpt_gate_proj_name="gate_proj",
631
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/mimo_v2_flash.py
+5
-2
@@ -22,7 +22,10 @@ from vllm.distributed import (
22
from vllm.logger import init_logger
23
from vllm.model_executor.layers.activation import SiluAndMul
24
from vllm.model_executor.layers.attention import Attention
25
-from vllm.model_executor.layers.fused_moe import FusedMoE
25
+from vllm.model_executor.layers.fused_moe import (
26
+ FusedMoE,
27
+ fused_moe_make_expert_params_mapping,
28
+)
29
from vllm.model_executor.layers.layernorm import RMSNorm
30
from vllm.model_executor.layers.linear import (
31
MergedColumnParallelLinear,
@@ -511,7 +514,7 @@ class MiMoV2Model(nn.Module):
514
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
515
# Params for weights, fp8 weight scales, fp8 activation scales
516
# (param_name, weight_name, expert_id, shard_id)
514
- return FusedMoE.make_expert_params_mapping(
517
+ return fused_moe_make_expert_params_mapping(
518
self,
519
ckpt_gate_proj_name="gate_proj",
520
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/minimax_m2.py
+5
-2
@@ -38,7 +38,10 @@ from vllm.distributed import (
38
get_tensor_model_parallel_world_size,
39
)
40
from vllm.model_executor.layers.attention import Attention
41
-from vllm.model_executor.layers.fused_moe import FusedMoE
41
+from vllm.model_executor.layers.fused_moe import (
42
+ FusedMoE,
43
+ fused_moe_make_expert_params_mapping,
44
+)
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
47
QKVParallelLinear,
@@ -393,7 +396,7 @@ class MiniMaxM2Model(nn.Module, EagleModelMixin):
396
return hidden_states
397
398
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
396
- return FusedMoE.make_expert_params_mapping(
399
+ return fused_moe_make_expert_params_mapping(
400
self,
401
ckpt_gate_proj_name="w1",
402
ckpt_down_proj_name="w2",
vllm/model_executor/models/minimax_text_01.py
+3
-1
@@ -24,7 +24,9 @@ from vllm.distributed.parallel_state import (
24
from vllm.forward_context import get_forward_context
25
from vllm.model_executor.layers.activation import SiluAndMul
26
from vllm.model_executor.layers.attention import Attention
27
-from vllm.model_executor.layers.fused_moe import FusedMoE
27
+from vllm.model_executor.layers.fused_moe import (
28
+ FusedMoE,
29
+)
30
from vllm.model_executor.layers.layernorm import RMSNorm
31
from vllm.model_executor.layers.linear import (
32
MergedColumnParallelLinear,
vllm/model_executor/models/mixtral.py
+5
-2
@@ -40,7 +40,10 @@ from vllm.distributed import (
40
get_tensor_model_parallel_world_size,
41
)
42
from vllm.model_executor.layers.attention import Attention
43
-from vllm.model_executor.layers.fused_moe import FusedMoE
43
+from vllm.model_executor.layers.fused_moe import (
44
+ FusedMoE,
45
+ fused_moe_make_expert_params_mapping,
46
+)
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
QKVParallelLinear,
@@ -364,7 +367,7 @@ class MixtralModel(nn.Module):
367
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
368
# Params for weights, fp8 weight scales, fp8 activation scales
369
# (param_name, weight_name, expert_id, shard_id)
367
- return FusedMoE.make_expert_params_mapping(
370
+ return fused_moe_make_expert_params_mapping(
371
self,
372
ckpt_gate_proj_name="w1",
373
ckpt_down_proj_name="w2",
vllm/model_executor/models/mllama4.py
+4
-2
@@ -40,7 +40,9 @@ from vllm.config.multimodal import BaseDummyOptions
40
from vllm.distributed import get_tensor_model_parallel_world_size
41
from vllm.inputs import MultiModalDataDict
42
from vllm.model_executor.layers.attention import MMEncoderAttention
43
-from vllm.model_executor.layers.fused_moe import FusedMoE
43
+from vllm.model_executor.layers.fused_moe import (
44
+ fused_moe_make_expert_params_mapping,
45
+)
46
from vllm.model_executor.layers.linear import (
47
ColumnParallelLinear,
48
QKVParallelLinear,
@@ -1072,7 +1074,7 @@ class Llama4ForConditionalGeneration(
1074
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
1075
# Params for weights, fp8 weight scales, fp8 activation scales
1076
# (param_name, weight_name, expert_id, shard_id)
1075
- return FusedMoE.make_expert_params_mapping(
1077
+ return fused_moe_make_expert_params_mapping(
1078
self,
1079
ckpt_gate_proj_name="gate_proj",
1080
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/nemotron_h.py
+2
-1
@@ -37,6 +37,7 @@ from vllm.model_executor.layers.fused_moe import (
37
FusedMoE,
38
GateLinear,
39
activation_without_mul,
40
+ fused_moe_make_expert_params_mapping,
41
)
42
from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.linear import (
@@ -652,7 +653,7 @@ class NemotronHModel(nn.Module):
653
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
654
if self.has_moe:
655
# (param_name, weight_name, expert_id, shard_id)
655
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
656
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
657
# - FusedMoe.w1 (aka gate_proj) should be up_proj since that's
658
# what the activation is applied to
659
# - FusedMoe.w3 (aka up_proj) should be ignored since we're
vllm/model_executor/models/nemotron_h_mtp.py
+4
-2
@@ -11,7 +11,9 @@ import torch.nn as nn
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, ModelConfig, VllmConfig
13
from vllm.config.parallel import ParallelConfig
14
-from vllm.model_executor.layers.fused_moe import FusedMoE
14
+from vllm.model_executor.layers.fused_moe import (
15
+ fused_moe_make_expert_params_mapping,
16
+)
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.linear import ColumnParallelLinear
19
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -399,7 +401,7 @@ class NemotronHMTP(nn.Module, SupportsPP):
401
if getattr(self.config, "model_type", None) == "nemotron_h_puzzle":
402
num_experts = self.config.mtp_n_routed_experts
403
if num_experts is not None:
402
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
404
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
405
self,
406
ckpt_gate_proj_name="up_proj",
407
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/olmoe.py
+5
-2
@@ -32,7 +32,10 @@ from vllm.distributed import (
32
from vllm.distributed.utils import split_tensor_along_last_dim
33
from vllm.logger import init_logger
34
from vllm.model_executor.layers.attention import Attention
35
-from vllm.model_executor.layers.fused_moe import FusedMoE
35
+from vllm.model_executor.layers.fused_moe import (
36
+ FusedMoE,
37
+ fused_moe_make_expert_params_mapping,
38
+)
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.linear import (
41
QKVParallelLinear,
@@ -336,7 +339,7 @@ class OlmoeModel(nn.Module):
339
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
340
# Params for weights, fp8 weight scales, fp8 activation scales
341
# (param_name, weight_name, expert_id, shard_id)
339
- return FusedMoE.make_expert_params_mapping(
342
+ return fused_moe_make_expert_params_mapping(
343
self,
344
ckpt_gate_proj_name="gate_proj",
345
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/openpangu.py
+5
-2
@@ -44,7 +44,10 @@ from vllm.model_executor.layers.attention import (
44
Attention,
45
StaticSinkAttention,
46
)
47
-from vllm.model_executor.layers.fused_moe import FusedMoE
47
+from vllm.model_executor.layers.fused_moe import (
48
+ FusedMoE,
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
ColumnParallelLinear,
@@ -1149,7 +1152,7 @@ class OpenPanguModel(nn.Module):
1152
]
1153
has_experts = hasattr(self.config, "n_routed_experts")
1154
if has_experts:
1152
- expert_merge_mapping = FusedMoE.make_expert_params_mapping(
1155
+ expert_merge_mapping = fused_moe_make_expert_params_mapping(
1156
self,
1157
ckpt_gate_proj_name="gate_proj",
1158
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/openpangu_mtp.py
+4
-2
@@ -28,7 +28,9 @@ from vllm.config import VllmConfig
28
29
# SPDX-License-Identifier: Apache-2.0
30
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
31
-from vllm.model_executor.layers.fused_moe import FusedMoE
31
+from vllm.model_executor.layers.fused_moe import (
32
+ fused_moe_make_expert_params_mapping,
33
+)
34
from vllm.model_executor.layers.layernorm import RMSNorm
35
from vllm.model_executor.layers.logits_processor import LogitsProcessor
36
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -147,7 +149,7 @@ class OpenPanguMTP(nn.Module):
149
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
150
]
151
150
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
152
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
153
self,
154
ckpt_gate_proj_name="gate_proj",
155
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/param2moe.py
+5
-2
@@ -32,7 +32,10 @@ from vllm.distributed import (
32
)
33
from vllm.model_executor.layers.activation import SiluAndMul
34
from vllm.model_executor.layers.attention import Attention
35
-from vllm.model_executor.layers.fused_moe import FusedMoE
35
+from vllm.model_executor.layers.fused_moe import (
36
+ FusedMoE,
37
+ fused_moe_make_expert_params_mapping,
38
+)
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.linear import (
41
MergedColumnParallelLinear,
@@ -690,7 +693,7 @@ class Param2MoEModel(nn.Module):
693
return loaded_params
694
695
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
693
- return FusedMoE.make_expert_params_mapping(
696
+ return fused_moe_make_expert_params_mapping(
697
self,
698
ckpt_gate_proj_name="gate_proj",
699
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/phimoe.py
+5
-2
@@ -35,7 +35,10 @@ from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.model_executor.layers.attention import Attention
38
-from vllm.model_executor.layers.fused_moe import FusedMoE
38
+from vllm.model_executor.layers.fused_moe import (
39
+ FusedMoE,
40
+ fused_moe_make_expert_params_mapping,
41
+)
42
from vllm.model_executor.layers.linear import (
43
QKVParallelLinear,
44
ReplicatedLinear,
@@ -514,7 +517,7 @@ class PhiMoEModel(nn.Module):
517
return hidden_states
518
519
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
517
- return FusedMoE.make_expert_params_mapping(
520
+ return fused_moe_make_expert_params_mapping(
521
self,
522
ckpt_gate_proj_name="w1",
523
ckpt_down_proj_name="w2",
vllm/model_executor/models/qwen2_moe.py
+5
-2
@@ -40,7 +40,10 @@ from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.activation import SiluAndMul
42
from vllm.model_executor.layers.attention import Attention
43
-from vllm.model_executor.layers.fused_moe import FusedMoE
43
+from vllm.model_executor.layers.fused_moe import (
44
+ FusedMoE,
45
+ fused_moe_make_expert_params_mapping,
46
+)
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
49
MergedColumnParallelLinear,
@@ -418,7 +421,7 @@ class Qwen2MoeModel(nn.Module):
421
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
422
# Params for weights, fp8 weight scales, fp8 activation scales
423
# (param_name, weight_name, expert_id, shard_id)
421
- return FusedMoE.make_expert_params_mapping(
424
+ return fused_moe_make_expert_params_mapping(
425
self,
426
ckpt_gate_proj_name="gate_proj",
427
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/qwen3_5_mtp.py
+4
-2
@@ -12,7 +12,9 @@ from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import VllmConfig
13
from vllm.distributed.parallel_state import get_pp_group
14
from vllm.logger import init_logger
15
-from vllm.model_executor.layers.fused_moe import FusedMoE
15
+from vllm.model_executor.layers.fused_moe import (
16
+ fused_moe_make_expert_params_mapping,
17
+)
18
from vllm.model_executor.layers.linear import ColumnParallelLinear
19
from vllm.model_executor.layers.logits_processor import LogitsProcessor
20
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -194,7 +196,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
196
197
# Params for weights, fp8 weight scales, fp8 activation scales
198
# (param_name, weight_name, expert_id, shard_id)
197
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
199
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
200
self,
201
ckpt_gate_proj_name="gate_proj",
202
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/qwen3_moe.py
+5
-2
@@ -43,7 +43,10 @@ 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
46
+from vllm.model_executor.layers.fused_moe import (
47
+ FusedMoE,
48
+ fused_moe_make_expert_params_mapping,
49
+)
50
from vllm.model_executor.layers.layernorm import RMSNorm
51
from vllm.model_executor.layers.linear import (
52
MergedColumnParallelLinear,
@@ -516,7 +519,7 @@ class Qwen3MoeModel(nn.Module, EagleModelMixin):
519
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
520
# Params for weights, fp8 weight scales, fp8 activation scales
521
# (param_name, weight_name, expert_id, shard_id)
519
- return FusedMoE.make_expert_params_mapping(
522
+ return fused_moe_make_expert_params_mapping(
523
self,
524
ckpt_gate_proj_name="gate_proj",
525
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/qwen3_next.py
+5
-2
@@ -23,7 +23,10 @@ from vllm.distributed import (
23
)
24
from vllm.logger import init_logger
25
from vllm.model_executor.layers.attention import Attention
26
-from vllm.model_executor.layers.fused_moe import FusedMoE
26
+from vllm.model_executor.layers.fused_moe import (
27
+ FusedMoE,
28
+ fused_moe_make_expert_params_mapping,
29
+)
30
from vllm.model_executor.layers.layernorm import (
31
GemmaRMSNorm as Qwen3NextRMSNorm,
32
)
@@ -533,7 +536,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
536
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
537
# Params for weights, fp8 weight scales, fp8 activation scales
538
# (param_name, weight_name, expert_id, shard_id)
536
- return FusedMoE.make_expert_params_mapping(
539
+ return fused_moe_make_expert_params_mapping(
540
self,
541
ckpt_gate_proj_name="gate_proj",
542
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/qwen3_next_mtp.py
+4
-2
@@ -11,7 +11,9 @@ from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import VllmConfig
12
from vllm.distributed.parallel_state import get_pp_group
13
from vllm.logger import init_logger
14
-from vllm.model_executor.layers.fused_moe import FusedMoE
14
+from vllm.model_executor.layers.fused_moe import (
15
+ fused_moe_make_expert_params_mapping,
16
+)
17
from vllm.model_executor.layers.linear import ColumnParallelLinear
18
from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.vocab_parallel_embedding import (
@@ -145,7 +147,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):
147
148
# Params for weights, fp8 weight scales, fp8 activation scales
149
# (param_name, weight_name, expert_id, shard_id)
148
- expert_params_mapping = FusedMoE.make_expert_params_mapping(
150
+ expert_params_mapping = fused_moe_make_expert_params_mapping(
151
self,
152
ckpt_gate_proj_name="gate_proj",
153
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/sarvam.py
+5
-2
@@ -35,7 +35,10 @@ from vllm.distributed import (
35
get_tensor_model_parallel_world_size,
36
)
37
from vllm.model_executor.layers.activation import SiluAndMul
38
-from vllm.model_executor.layers.fused_moe import FusedMoE
38
+from vllm.model_executor.layers.fused_moe import (
39
+ FusedMoE,
40
+ fused_moe_make_expert_params_mapping,
41
+)
42
from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.linear import (
44
ColumnParallelLinear,
@@ -529,7 +532,7 @@ class SarvamMLAModel(nn.Module):
532
return hidden_states
533
534
def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
532
- return FusedMoE.make_expert_params_mapping(
535
+ return fused_moe_make_expert_params_mapping(
536
self,
537
ckpt_gate_proj_name="gate_proj",
538
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/step3_text.py
+3
-1
@@ -18,7 +18,9 @@ from vllm.distributed import (
18
from vllm.logger import init_logger
19
from vllm.model_executor.layers.activation import SiluAndMul
20
from vllm.model_executor.layers.attention import Attention
21
-from vllm.model_executor.layers.fused_moe import FusedMoE
21
+from vllm.model_executor.layers.fused_moe import (
22
+ FusedMoE,
23
+)
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import (
26
ColumnParallelLinear,
vllm/model_executor/models/step3p5.py
+5
-2
@@ -23,7 +23,10 @@ from vllm.distributed import (
23
from vllm.logger import init_logger
24
from vllm.model_executor.layers.activation import SiluAndMul, SwigluStepAndMul
25
from vllm.model_executor.layers.attention import Attention
26
-from vllm.model_executor.layers.fused_moe import FusedMoE
26
+from vllm.model_executor.layers.fused_moe import (
27
+ FusedMoE,
28
+ fused_moe_make_expert_params_mapping,
29
+)
30
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
31
from vllm.model_executor.layers.linear import (
32
ColumnParallelLinear,
@@ -637,7 +640,7 @@ class Step3p5Model(nn.Module):
640
]
641
642
# New per-expert format: .moe.experts.E.gate_proj.weight_packed [out, in]
640
- per_expert_mapping = FusedMoE.make_expert_params_mapping(
643
+ per_expert_mapping = fused_moe_make_expert_params_mapping(
644
self,
645
ckpt_gate_proj_name="gate_proj",
646
ckpt_down_proj_name="down_proj",
vllm/model_executor/models/transformers/moe.py
+5
-2
@@ -25,7 +25,10 @@ from vllm.config.utils import getattr_iter
25
from vllm.distributed import get_dp_group, get_ep_group
26
from vllm.forward_context import ForwardContext, get_forward_context
27
from vllm.model_executor.custom_op import PluggableLayer
28
-from vllm.model_executor.layers.fused_moe import FusedMoE
28
+from vllm.model_executor.layers.fused_moe import (
29
+ FusedMoE,
30
+ fused_moe_make_expert_params_mapping,
31
+)
32
from vllm.model_executor.models.interfaces import MixtureOfExperts
33
from vllm.model_executor.models.utils import maybe_prefix
34
from vllm.platforms import current_platform
@@ -179,7 +182,7 @@ class MoEMixin(MixtureOfExperts):
182
num_redundant_experts = self.parallel_config.eplb_config.num_redundant_experts
183
for gate_proj, down_proj, up_proj in ckpt_names:
184
expert_mapping.extend(
182
- FusedMoE.make_expert_params_mapping(
185
+ fused_moe_make_expert_params_mapping(
186
self,
187
ckpt_gate_proj_name=gate_proj,
188
ckpt_down_proj_name=down_proj,