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 5, 13:30 UTC
[Bugfix] Fix DSV32 weight loading (#38870)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Apr 3, 2026 at 22:57 UTC
8617f8676b5ae936382ea00fa92693f59fbb9d69
2 files changed
+68
-27
vllm/model_executor/models/deepseek_mtp.py
+13
-3
@@ -184,11 +184,16 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
184
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
185
super().__init__()
186
self.config = vllm_config.model_config.hf_config
187
+ self.quant_config = vllm_config.quant_config
188
self.model = DeepSeekMultiTokenPredictor(
189
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
190
)
191
# Set MoE hyperparameters
192
self.set_moe_parameters()
193
+ self.is_fp4_ckpt = (
194
+ self.quant_config is not None
195
+ and self.quant_config.get_name() == "modelopt_fp4"
196
+ )
197
198
def set_moe_parameters(self):
199
self.expert_weights = []
@@ -241,11 +246,16 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
246
("gate_up_proj", "up_proj", 1),
247
("fused_qkv_a_proj", "q_a_proj", 0),
248
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
244
- # Fused indexer wk + weights_proj
245
- ("wk_weights_proj", "wk", 0),
246
- ("wk_weights_proj", "weights_proj", 1),
249
]
250
251
+ if self.is_fp4_ckpt:
252
+ # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
253
+ indexer_fused_mapping = [
254
+ ("wk_weights_proj", "wk", 0),
255
+ ("wk_weights_proj", "weights_proj", 1),
256
+ ]
257
+ stacked_params_mapping.extend(indexer_fused_mapping)
258
+
259
expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
260
self,
261
ckpt_gate_proj_name="gate_proj",
vllm/model_executor/models/deepseek_v2.py
+55
-24
@@ -625,6 +625,11 @@ class Indexer(nn.Module):
625
super().__init__()
626
self.vllm_config = vllm_config
627
self.config = config
628
+ self.quant_config = quant_config
629
+ self.is_fp4_ckpt = (
630
+ self.quant_config is not None
631
+ and self.quant_config.get_name() == "modelopt_fp4"
632
+ )
633
# self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
634
self.topk_tokens = config.index_topk
635
self.n_head = config.index_n_heads # 64
@@ -639,18 +644,36 @@ class Indexer(nn.Module):
644
quant_config=quant_config,
645
prefix=f"{prefix}.wq_b",
646
)
642
- # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
643
- # weights_proj does not get quantized, so we run both with quant_config=None
644
- # wk may be upcasted from the default quant; experiments show fusion is always
645
- # faster unless WK proj is in FP4, which is not the case for all known quants.
646
- self.wk_weights_proj = MergedColumnParallelLinear(
647
- hidden_size,
648
- [self.head_dim, self.n_head],
649
- bias=False,
650
- quant_config=None,
651
- disable_tp=True,
652
- prefix=f"{prefix}.wk_weights_proj",
653
- )
647
+ if self.is_fp4_ckpt:
648
+ # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
649
+ # weights_proj does not get quantized,
650
+ # so we run both with quant_config=None
651
+ # wk may be upcasted from the default quant;
652
+ # experiments show fusion is always faster unless WK proj is in FP4,
653
+ # which is not the case for all known quants.
654
+ self.wk_weights_proj = MergedColumnParallelLinear(
655
+ hidden_size,
656
+ [self.head_dim, self.n_head],
657
+ bias=False,
658
+ quant_config=None,
659
+ disable_tp=True,
660
+ prefix=f"{prefix}.wk_weights_proj",
661
+ )
662
+ else:
663
+ self.wk = ReplicatedLinear(
664
+ hidden_size,
665
+ self.head_dim,
666
+ bias=False,
667
+ quant_config=quant_config,
668
+ prefix=f"{prefix}.wk",
669
+ )
670
+ self.weights_proj = ReplicatedLinear(
671
+ hidden_size,
672
+ self.n_head,
673
+ bias=False,
674
+ quant_config=None,
675
+ prefix=f"{prefix}.weights_proj",
676
+ )
677
self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
678
self.softmax_scale = self.head_dim**-0.5
679
@@ -691,11 +714,14 @@ class Indexer(nn.Module):
714
q_pe, q_nope = torch.split(
715
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
716
)
694
-
695
- # Fused wk + weights_proj: one GEMM, then split
696
- kw, _ = self.wk_weights_proj(hidden_states)
697
- k = kw[:, : self.head_dim]
698
- weights_raw = kw[:, self.head_dim :]
717
+ if self.is_fp4_ckpt:
718
+ # Fused wk + weights_proj: one GEMM, then split
719
+ kw, _ = self.wk_weights_proj(hidden_states)
720
+ k = kw[:, : self.head_dim]
721
+ weights = kw[:, self.head_dim :]
722
+ else:
723
+ k, _ = self.wk(hidden_states)
724
+ weights, _ = self.weights_proj(hidden_states)
725
726
k = self.k_norm(k)
727
k_pe, k_nope = torch.split(
@@ -726,7 +752,7 @@ class Indexer(nn.Module):
752
q_scale = q_scale.view(-1, self.n_head, 1)
753
754
weights = (
729
- weights_raw.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
755
+ weights.unsqueeze(-1) * q_scale * self.softmax_scale * self.n_head**-0.5
756
)
757
weights = weights.squeeze(-1)
758
@@ -1314,6 +1340,10 @@ class DeepseekV2ForCausalLM(
1340
quant_config = vllm_config.quant_config
1341
self.config = config
1342
self.quant_config = quant_config
1343
+ self.is_fp4_ckpt = (
1344
+ self.quant_config is not None
1345
+ and self.quant_config.get_name() == "modelopt_fp4"
1346
+ )
1347
1348
qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
1349
qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
@@ -1439,12 +1469,13 @@ class DeepseekV2ForCausalLM(
1469
("qkv_proj", "k_proj", "k"),
1470
("qkv_proj", "v_proj", "v"),
1471
]
1442
- # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
1443
- indexer_fused_mapping = [
1444
- ("wk_weights_proj", "wk", 0),
1445
- ("wk_weights_proj", "weights_proj", 1),
1446
- ]
1447
- stacked_params_mapping.extend(indexer_fused_mapping)
1472
+ if self.is_fp4_ckpt:
1473
+ # Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
1474
+ indexer_fused_mapping = [
1475
+ ("wk_weights_proj", "wk", 0),
1476
+ ("wk_weights_proj", "weights_proj", 1),
1477
+ ]
1478
+ stacked_params_mapping.extend(indexer_fused_mapping)
1479
1480
if self.use_mha:
1481
stacked_params_mapping.extend(mha_params_mapping)