diff --git a/vllm/models/deepseek_v4/nvidia/dspark.py b/vllm/models/deepseek_v4/nvidia/dspark.py index 6ddf70e1c16a..69bb42c791cd 100644 --- a/vllm/models/deepseek_v4/nvidia/dspark.py +++ b/vllm/models/deepseek_v4/nvidia/dspark.py @@ -700,13 +700,17 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.max_batch = max(1, int(vllm_config.scheduler_config.max_num_seqs)) hidden_size = int(config.hidden_size) dtype = vllm_config.model_config.dtype - quant_config = get_draft_quant_config(vllm_config) + self.quant_config = get_draft_quant_config(vllm_config) + self.pad_shared_expert = ( + getattr(self.quant_config, "weight_block_size", None) is not None + and not vllm_config.parallel_config.use_sequence_parallel_moe + ) self.dspark_aux_hidden_size = hidden_size * len(self.target_layer_ids) self.embed_tokens = VocabParallelEmbedding( config.vocab_size, hidden_size, - quant_config=quant_config, + quant_config=self.quant_config, prefix=f"{prefix}embed_tokens", ) self.main_proj = ReplicatedLinear( @@ -714,7 +718,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): hidden_size, bias=False, params_dtype=dtype, - quant_config=quant_config, + quant_config=self.quant_config, prefix=f"{prefix}mtp.0.main_proj", return_bias=False, ) @@ -1044,6 +1048,12 @@ def map_name(name: str) -> str | None: else ".weight_scale_inv" ) name = name.removesuffix(".scale") + suffix + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if self.pad_shared_expert and ".shared_experts." in name: + loaded_weight = DeepseekV4Model._pad_shared_expert_weight( + self.quant_config, name, loaded_weight + ) for param_name, weight_name, shard_id in stacked_params_mapping: if not name.startswith("layers."): continue