[Bugfix] Fix DeepSeek-V4 DSpark draft shared-expert padding for TP > 8 (#49415)
Signed-off-by: Mike G <180722391+mikekg@users.noreply.github.com>
Mike G committed
Jul 22, 2026 at 22:06 UTC
76bf55240cf8cb30aae3a3dedea89c4b8eda5830
3 files changed
+31
-11
vllm/models/deepseek_v4/nvidia/dspark.py
+12
-4
@@ -43,6 +43,7 @@ from vllm.model_executor.models.utils import maybe_prefix
43
44
from .model import (
45
DeepseekV4DecoderLayer,
46
+ DeepseekV4Model,
47
make_deepseek_v4_expert_params_mapping,
48
)
49
@@ -277,6 +278,11 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
278
assert vllm_config.speculative_config is not None
279
self.draft_model_config = vllm_config.speculative_config.draft_model_config
280
self.config = self.draft_model_config.hf_config
281
+ self.quant_config = vllm_config.quant_config
282
+ self.pad_shared_expert = (
283
+ getattr(self.quant_config, "weight_block_size", None) is not None
284
+ and not vllm_config.parallel_config.use_sequence_parallel_moe
285
+ )
286
self.model = DSparkDeepseekV4Model(
287
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
288
)
@@ -396,6 +402,12 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
402
else ".weight_scale_inv"
403
)
404
name = name.removesuffix(".scale") + suffix
405
+ if ".shared_experts.w2" in name:
406
+ name = name.replace(".shared_experts.w2", ".shared_experts.down_proj")
407
+ if self.pad_shared_expert and ".shared_experts." in name:
408
+ loaded_weight = DeepseekV4Model._pad_shared_expert_weight(
409
+ self.quant_config, name, loaded_weight
410
+ )
411
412
# E8M0 expert scales: keep raw exponent bytes.
413
if ".experts." in name:
@@ -440,10 +452,6 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
452
params_dict[name][: narrow.shape[0]].copy_(narrow)
453
loaded_params.add(name)
454
continue
443
- if ".shared_experts.w2" in name:
444
- name = name.replace(
445
- ".shared_experts.w2", ".shared_experts.down_proj"
446
- )
455
if name.endswith(".ffn.gate.bias"):
456
name = name.replace(
457
".ffn.gate.bias", ".ffn.gate.e_score_correction_bias"
vllm/models/deepseek_v4/nvidia/model.py
+8
-3
@@ -1181,7 +1181,9 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1181
1182
for name, loaded_weight in weights:
1183
if pad_shared_expert and ".shared_experts." in name:
1184
- loaded_weight = self._pad_shared_expert_weight(name, loaded_weight)
1184
+ loaded_weight = self._pad_shared_expert_weight(
1185
+ self.quant_config, name, loaded_weight
1186
+ )
1187
for param_name, weight_name, shard_id in stacked_params_mapping:
1188
# Skip non-stacked layers and experts (experts handled below).
1189
if ".experts." in name:
@@ -1256,15 +1258,18 @@ class DeepseekV4Model(nn.Module, EagleModelMixin):
1258
1259
return loaded_params
1260
1261
+ @staticmethod
1262
def _pad_shared_expert_weight(
1260
- self, name: str, loaded_weight: torch.Tensor
1263
+ quant_config: QuantizationConfig | None,
1264
+ name: str,
1265
+ loaded_weight: torch.Tensor,
1266
) -> torch.Tensor:
1267
"""Zero-pad a block-FP8 shared-expert weight/scale on its intermediate
1268
axis so the standard TP loaders split it into even, block-aligned shards
1269
(trailing ranks get the zero pad). gate (w1)/up (w3) [I, H] pad dim 0;
1270
down (w2 -> down_proj) [H, I] pads dim 1.
1271
"""
1267
- block_size = getattr(self.quant_config, "weight_block_size", None)
1272
+ block_size = getattr(quant_config, "weight_block_size", None)
1273
assert block_size is not None
1274
# Round the intermediate axis up to a whole number of TP shards. The axis
1275
# is in elements for weights (step = block) and in blocks for scales.
vllm/models/deepseek_v4/nvidia/mtp.py
+11
-4
@@ -52,6 +52,7 @@ from vllm.sequence import IntermediateTensors
52
53
from .model import (
54
DeepseekV4DecoderLayer,
55
+ DeepseekV4Model,
56
make_deepseek_v4_expert_params_mapping,
57
)
58
@@ -265,6 +266,10 @@ class DeepSeekV4MTP(nn.Module):
266
super().__init__()
267
self.config = vllm_config.model_config.hf_config
268
self.quant_config = vllm_config.quant_config
269
+ self.pad_shared_expert = (
270
+ getattr(self.quant_config, "weight_block_size", None) is not None
271
+ and not vllm_config.parallel_config.use_sequence_parallel_moe
272
+ )
273
self.model = DeepSeekV4MultiTokenPredictor(
274
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
275
)
@@ -387,6 +392,12 @@ class DeepSeekV4MTP(nn.Module):
392
else ".weight_scale_inv"
393
)
394
name = name.removesuffix(".scale") + suffix
395
+ if ".shared_experts.w2" in name:
396
+ name = name.replace(".shared_experts.w2", ".shared_experts.down_proj")
397
+ if self.pad_shared_expert and ".shared_experts." in name:
398
+ loaded_weight = DeepseekV4Model._pad_shared_expert_weight(
399
+ self.quant_config, name, loaded_weight
400
+ )
401
for param_name, weight_name, shard_id in stacked_params_mapping:
402
# Skip non-stacked layers and experts (experts handled below).
403
if ".experts." in name:
@@ -442,10 +453,6 @@ class DeepSeekV4MTP(nn.Module):
453
loaded_params.add(name)
454
continue
455
else:
445
- if ".shared_experts.w2" in name:
446
- name = name.replace(
447
- ".shared_experts.w2", ".shared_experts.down_proj"
448
- )
456
if name.endswith(".ffn.gate.bias"):
457
# ``e_score_correction_bias`` lives on the gate
458
# under a different attribute name.