482914849c · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/model_executor/layers/fused_moe/layer.py

+10 -3

@@ -2007,6 +2007,7 @@ class FusedMoE(CustomOp):

2007 @classmethod

2008 def make_expert_params_mapping(

2009 cls,

2010 + model: torch.nn.Module,

2011 ckpt_gate_proj_name: str,

2012 ckpt_down_proj_name: str,

2013 ckpt_up_proj_name: str,

@@ -2025,13 +2026,19 @@ class FusedMoE(CustomOp):

2026 )

2027 )

2028

2029 + base_layer = (

2030 + "base_layer."

2031 + if any(".base_layer." in name for name, _ in model.named_parameters())

2032 + else ""

2033 + )

2034 +

2035 return [

2036 # (param_name, weight_name, expert_id, shard_id)

2037 (

2031 - "experts.w13_"

2038 + f"experts.{base_layer}w13_"

2039 if weight_name in [ckpt_gate_proj_name, ckpt_up_proj_name]

2033 - else "experts.w2_",

2034 - f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.",

2040 + else f"experts.{base_layer}w2_",

2041 + f"experts.{physical_to_logical_map[expert_id]}.{weight_name}.{base_layer}",

2042 expert_id,

2043 shard_id,

2044 )

vllm/model_executor/models/afmoe.py

+1

@@ -475,6 +475,7 @@ class AfmoeModel(nn.Module):

475 # Params for weights, fp8 weight scales, fp8 activation scales

476 # (param_name, weight_name, expert_id, shard_id)

477 return SharedFusedMoE.make_expert_params_mapping(

478 + self,

479 ckpt_gate_proj_name="gate_proj",

480 ckpt_down_proj_name="down_proj",

481 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/bailing_moe.py

+1

@@ -476,6 +476,7 @@ class BailingMoeModel(nn.Module):

476

477 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

478 return SharedFusedMoE.make_expert_params_mapping(

479 + self,

480 ckpt_gate_proj_name="gate_proj",

481 ckpt_down_proj_name="down_proj",

482 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/deepseek_eagle.py

+1

@@ -106,6 +106,7 @@ class DeepseekV2Model(nn.Module):

106 # Params for weights, fp8 weight scales, fp8 activation scales

107 # (param_name, weight_name, expert_id, shard_id)

108 expert_params_mapping = FusedMoE.make_expert_params_mapping(

109 + self,

110 ckpt_gate_proj_name="gate_proj",

111 ckpt_down_proj_name="down_proj",

112 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/deepseek_mtp.py

+1

@@ -245,6 +245,7 @@ class DeepSeekMTP(nn.Module, SupportsPP, DeepseekV2MixtureOfExperts):

245 ]

246

247 expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

248 + self,

249 ckpt_gate_proj_name="gate_proj",

250 ckpt_down_proj_name="down_proj",

251 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/deepseek_v2.py

+2

@@ -1486,6 +1486,7 @@ class DeepseekV2ForCausalLM(

1486 # Params for weights, fp8 weight scales, fp8 activation scales

1487 # (param_name, weight_name, expert_id, shard_id)

1488 return SharedFusedMoE.make_expert_params_mapping(

1489 + self,

1490 ckpt_gate_proj_name="gate_proj",

1491 ckpt_down_proj_name="down_proj",

1492 ckpt_up_proj_name="up_proj",

@@ -1519,6 +1520,7 @@ class DeepseekV2ForCausalLM(

1520 # Params for weights, fp8 weight scales, fp8 activation scales

1521 # (param_name, weight_name, expert_id, shard_id)

1522 expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

1523 + self,

1524 ckpt_gate_proj_name="gate_proj",

1525 ckpt_down_proj_name="down_proj",

1526 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/dots1.py

+1

@@ -424,6 +424,7 @@ class Dots1Model(nn.Module):

424

425 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

426 return SharedFusedMoE.make_expert_params_mapping(

427 + self,

428 ckpt_gate_proj_name="gate_proj",

429 ckpt_down_proj_name="down_proj",

430 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/ernie45_moe.py

+1

@@ -497,6 +497,7 @@ class Ernie4_5_MoeModel(nn.Module):

497 # Params for weights, fp8 weight scales, fp8 activation scales

498 # (param_name, weight_name, expert_id, shard_id)

499 return SharedFusedMoE.make_expert_params_mapping(

500 + self,

501 ckpt_gate_proj_name="gate_proj",

502 ckpt_down_proj_name="down_proj",

503 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/ernie45_vl_moe.py

+1

@@ -675,6 +675,7 @@ class Ernie4_5_VLMoeForCausalLM(nn.Module, SupportsPP):

675 # Params for weights, fp8 weight scales, fp8 activation scales

676 # (param_name, weight_name, expert_id, shard_id)

677 expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

678 + self,

679 ckpt_gate_proj_name="gate_proj",

680 ckpt_down_proj_name="down_proj",

681 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/glm4_moe.py

+1

@@ -496,6 +496,7 @@ class Glm4MoeModel(nn.Module):

496 # Params for weights, fp8 weight scales, fp8 activation scales

497 # (param_name, weight_name, expert_id, shard_id)

498 return SharedFusedMoE.make_expert_params_mapping(

499 + self,

500 ckpt_gate_proj_name="gate_proj",

501 ckpt_down_proj_name="down_proj",

502 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/glm4_moe_mtp.py

+1

@@ -248,6 +248,7 @@ class Glm4MoeMTP(nn.Module, SupportsPP, Glm4MixtureOfExperts):

248 # Params for weights, fp8 weight scales, fp8 activation scales

249 # (param_name, weight_name, expert_id, shard_id)

250 expert_params_mapping = FusedMoE.make_expert_params_mapping(

251 + self,

252 ckpt_gate_proj_name="gate_proj",

253 ckpt_down_proj_name="down_proj",

254 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/gpt_oss.py

+1

@@ -729,6 +729,7 @@ class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):

729 # Params for weights, weight scales, activation scales

730 # (param_name, weight_name, expert_id, shard_id)

731 return FusedMoE.make_expert_params_mapping(

732 + self,

733 ckpt_gate_proj_name="gate_proj",

734 ckpt_down_proj_name="down_proj",

735 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/granitemoe.py

+1

@@ -353,6 +353,7 @@ class GraniteMoeModel(nn.Module):

353 # Params for weights, fp8 weight scales, fp8 activation scales

354 # (param_name, weight_name, expert_id, shard_id)

355 expert_params_mapping = FusedMoE.make_expert_params_mapping(

356 + self,

357 ckpt_gate_proj_name="w1",

358 ckpt_down_proj_name="w2",

359 ckpt_up_proj_name="w3",

vllm/model_executor/models/grok1.py

+1

@@ -369,6 +369,7 @@ class Grok1Model(nn.Module):

369 # Grok1 uses "num_experts" in its config

370 num_experts = getattr(self.config, "num_experts", 8)

371 return FusedMoE.make_expert_params_mapping(

372 + self,

373 ckpt_gate_proj_name="linear", # Grok1 specific

374 ckpt_down_proj_name="linear_1", # Grok1 specific

375 ckpt_up_proj_name="linear_v", # Grok1 specific

vllm/model_executor/models/hunyuan_v1.py

+1

@@ -706,6 +706,7 @@ class HunYuanModel(nn.Module):

706 # Params for weights, fp8 weight scales, fp8 activation scales

707 # (param_name, weight_name, expert_id, shard_id)

708 return SharedFusedMoE.make_expert_params_mapping(

709 + self,

710 ckpt_gate_proj_name="gate_proj",

711 ckpt_down_proj_name="down_proj",

712 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/jamba.py

+1

@@ -378,6 +378,7 @@ class JambaModel(nn.Module):

378 # Params for weights, fp8 weight scales, fp8 activation scales

379 # (param_name, weight_name, expert_id, shard_id)

380 return FusedMoE.make_expert_params_mapping(

381 + self,

382 ckpt_gate_proj_name="gate_proj",

383 ckpt_down_proj_name="down_proj",

384 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/kimi_linear.py

+1

@@ -560,6 +560,7 @@ class KimiLinearForCausalLM(

560 # Params for weights, fp8 weight scales, fp8 activation scales

561 # (param_name, weight_name, expert_id, shard_id)

562 expert_params_mapping = FusedMoE.make_expert_params_mapping(

563 + self,

564 ckpt_gate_proj_name="w1",

565 ckpt_down_proj_name="w2",

566 ckpt_up_proj_name="w3",

vllm/model_executor/models/kimi_vl.py

+1

@@ -462,6 +462,7 @@ class KimiVLForConditionalGeneration(nn.Module, SupportsMultiModal, SupportsPP):

462 # Params for weights, fp8 weight scales, fp8 activation scales

463 # (param_name, weight_name, expert_id, shard_id)

464 expert_params_mapping = FusedMoE.make_expert_params_mapping(

465 + self,

466 ckpt_gate_proj_name="gate_proj",

467 ckpt_down_proj_name="down_proj",

468 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/lfm2_moe.py

+1

@@ -486,6 +486,7 @@ class Lfm2MoeModel(nn.Module):

486

487 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

488 return FusedMoE.make_expert_params_mapping(

489 + self,

490 ckpt_gate_proj_name="w1",

491 ckpt_down_proj_name="w2",

492 ckpt_up_proj_name="w3",

vllm/model_executor/models/llama4.py

+2

@@ -539,6 +539,7 @@ class Llama4Model(LlamaModel):

539 # Expert parameter mapping for the case where the expert weights are

540 # not fused into a single weight tensor.

541 expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(

542 + self,

543 ckpt_gate_proj_name="gate_proj",

544 ckpt_down_proj_name="down_proj",

545 ckpt_up_proj_name="up_proj",

@@ -548,6 +549,7 @@ class Llama4Model(LlamaModel):

549 # Expert parameter mapping for the case where the expert weights are

550 # fused into a single weight tensor.

551 expert_params_mapping_fused = SharedFusedMoE.make_expert_params_mapping(

552 + self,

553 ckpt_gate_proj_name="gate_up_proj",

554 ckpt_down_proj_name="down_proj",

555 ckpt_up_proj_name="gate_up_proj",

vllm/model_executor/models/longcat_flash.py

+1

@@ -626,6 +626,7 @@ class LongcatFlashForCausalLM(nn.Module, SupportsLoRA, SupportsPP):

626 # Params for weights, fp8 weight scales, fp8 activation scales

627 # (param_name, weight_name, expert_id, shard_id)

628 return FusedMoE.make_expert_params_mapping(

629 + self,

630 ckpt_gate_proj_name="gate_proj",

631 ckpt_down_proj_name="down_proj",

632 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/mimo_v2_flash.py

+1

@@ -512,6 +512,7 @@ class MiMoV2Model(nn.Module):

512 # Params for weights, fp8 weight scales, fp8 activation scales

513 # (param_name, weight_name, expert_id, shard_id)

514 return FusedMoE.make_expert_params_mapping(

515 + self,

516 ckpt_gate_proj_name="gate_proj",

517 ckpt_down_proj_name="down_proj",

518 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/minimax_m2.py

+1

@@ -392,6 +392,7 @@ class MiniMaxM2Model(nn.Module):

392

393 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

394 return FusedMoE.make_expert_params_mapping(

395 + self,

396 ckpt_gate_proj_name="w1",

397 ckpt_down_proj_name="w2",

398 ckpt_up_proj_name="w3",

vllm/model_executor/models/mixtral.py

+1

@@ -366,6 +366,7 @@ class MixtralModel(nn.Module):

366 # Params for weights, fp8 weight scales, fp8 activation scales

367 # (param_name, weight_name, expert_id, shard_id)

368 return FusedMoE.make_expert_params_mapping(

369 + self,

370 ckpt_gate_proj_name="w1",

371 ckpt_down_proj_name="w2",

372 ckpt_up_proj_name="w3",

vllm/model_executor/models/mllama4.py

+1

@@ -1084,6 +1084,7 @@ class Llama4ForConditionalGeneration(

1084 # Params for weights, fp8 weight scales, fp8 activation scales

1085 # (param_name, weight_name, expert_id, shard_id)

1086 return FusedMoE.make_expert_params_mapping(

1087 + self,

1088 ckpt_gate_proj_name="gate_proj",

1089 ckpt_down_proj_name="down_proj",

1090 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/nemotron_h.py

+1

@@ -636,6 +636,7 @@ class NemotronHModel(nn.Module):

636 # what the activation is applied to

637 # - FusedMoe.w3 (aka up_proj) should be ignored since we're

638 # using non-gated MoE

639 + self,

640 ckpt_gate_proj_name="up_proj",

641 ckpt_down_proj_name="down_proj",

642 ckpt_up_proj_name="",

vllm/model_executor/models/olmoe.py

+1

@@ -338,6 +338,7 @@ class OlmoeModel(nn.Module):

338 # Params for weights, fp8 weight scales, fp8 activation scales

339 # (param_name, weight_name, expert_id, shard_id)

340 return FusedMoE.make_expert_params_mapping(

341 + self,

342 ckpt_gate_proj_name="gate_proj",

343 ckpt_down_proj_name="down_proj",

344 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/openpangu.py

+1

@@ -1161,6 +1161,7 @@ class OpenPanguModel(nn.Module):

1161 has_experts = hasattr(self.config, "n_routed_experts")

1162 if has_experts:

1163 expert_merge_mapping = SharedFusedMoE.make_expert_params_mapping(

1164 + self,

1165 ckpt_gate_proj_name="gate_proj",

1166 ckpt_down_proj_name="down_proj",

1167 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/openpangu_mtp.py

+1

@@ -149,6 +149,7 @@ class OpenPanguMTP(nn.Module, SupportsPP):

149 ]

150

151 expert_params_mapping = FusedMoE.make_expert_params_mapping(

152 + self,

153 ckpt_gate_proj_name="gate_proj",

154 ckpt_down_proj_name="down_proj",

155 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/phimoe.py

+1

@@ -516,6 +516,7 @@ class PhiMoEModel(nn.Module):

516

517 def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:

518 return FusedMoE.make_expert_params_mapping(

519 + self,

520 ckpt_gate_proj_name="w1",

521 ckpt_down_proj_name="w2",

522 ckpt_up_proj_name="w3",

vllm/model_executor/models/qwen2_moe.py

+1

@@ -423,6 +423,7 @@ class Qwen2MoeModel(nn.Module):

423 # Params for weights, fp8 weight scales, fp8 activation scales

424 # (param_name, weight_name, expert_id, shard_id)

425 return SharedFusedMoE.make_expert_params_mapping(

426 + self,

427 ckpt_gate_proj_name="gate_proj",

428 ckpt_down_proj_name="down_proj",

429 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/qwen3_moe.py

+1

@@ -470,6 +470,7 @@ class Qwen3MoeModel(nn.Module):

470 # Params for weights, fp8 weight scales, fp8 activation scales

471 # (param_name, weight_name, expert_id, shard_id)

472 return FusedMoE.make_expert_params_mapping(

473 + self,

474 ckpt_gate_proj_name="gate_proj",

475 ckpt_down_proj_name="down_proj",

476 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/qwen3_next.py

+1

@@ -1031,6 +1031,7 @@ class Qwen3NextModel(nn.Module):

1031 # Params for weights, fp8 weight scales, fp8 activation scales

1032 # (param_name, weight_name, expert_id, shard_id)

1033 return SharedFusedMoE.make_expert_params_mapping(

1034 + self,

1035 ckpt_gate_proj_name="gate_proj",

1036 ckpt_down_proj_name="down_proj",

1037 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/qwen3_next_mtp.py

+1

@@ -147,6 +147,7 @@ class Qwen3NextMultiTokenPredictor(nn.Module):

147 # Params for weights, fp8 weight scales, fp8 activation scales

148 # (param_name, weight_name, expert_id, shard_id)

149 expert_params_mapping = FusedMoE.make_expert_params_mapping(

150 + self,

151 ckpt_gate_proj_name="gate_proj",

152 ckpt_down_proj_name="down_proj",

153 ckpt_up_proj_name="up_proj",

vllm/model_executor/models/transformers/moe.py

+1

@@ -165,6 +165,7 @@ class MoEMixin(MixtureOfExperts):

165 for gate_proj, down_proj, up_proj in ckpt_names:

166 expert_mapping.extend(

167 FusedMoE.make_expert_params_mapping(

168 + self,

169 ckpt_gate_proj_name=gate_proj,

170 ckpt_down_proj_name=down_proj,

171 ckpt_up_proj_name=up_proj,

[BugFix] LoRA: Support loading base_layer of experts (#31104)