From 8c5d5f75bf8b62a9e5194db87a49d17b666d00a1 Mon Sep 17 00:00:00 2001 From: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Date: Tue, 11 Aug 2026 07:36:50 +0800 Subject: [PATCH] Fix DSV4 DSpark shared expert loading (#33312) --- .../sglang/srt/models/deepseek_v4_dspark.py | 19 +++- .../test_runai_model_streamer_loader.py | 1 + .../test_deepseek_v4_shared_expert_fusion.py | 107 +++++++++++++++++- 3 files changed, 122 insertions(+), 5 deletions(-) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 1af8df913..d7ca4e373 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -20,6 +20,7 @@ from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled from sglang.srt.layers.quantization.base_config import QuantizationConfig from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.vocab_parallel_embedding import VocabParallelEmbedding @@ -32,6 +33,7 @@ from sglang.srt.models.dbrx import ReplicatedLinear from sglang.srt.models.deepseek_v4 import ( DEEPSEEK_V4_STACKED_PARAMS_MAPPING, DeepseekV4DecoderLayer, + DeepseekV4ForCausalLM, MqaAttentionBase, _dequant_fp8_wo_a_streaming, hc_head_torch, @@ -571,6 +573,12 @@ class DSparkV4Stage(DeepseekV4DecoderLayer): class DeepseekV4ForCausalLMDSpark(nn.Module): + @classmethod + def shared_experts_fusion_disable_reason(cls, hf_config, quant_config): + return DeepseekV4ForCausalLM.shared_experts_fusion_disable_reason( + hf_config, quant_config + ) + def __init__( self, config: DeepSeekV4Config, @@ -580,6 +588,9 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): super().__init__() self.config = config self.quant_config = quant_config + self.num_fused_shared_experts = ( + 0 if is_shared_experts_fusion_disabled() else config.n_shared_experts + ) dspark_config = parse_dspark_draft_config(draft_hf_config=config) if not dspark_config.require_markov(): @@ -796,14 +807,18 @@ class DeepseekV4ForCausalLMDSpark(nn.Module): ckpt_gate_proj_name="gate_proj", ckpt_down_proj_name="down_proj", ckpt_up_proj_name="up_proj", - num_experts=self.config.n_routed_experts, + num_experts=(self.config.n_routed_experts + self.num_fused_shared_experts), ) for name, loaded_weight in weights: mapped = self._remap_dspark_weight_name(name) if mapped is None: continue - + if self.num_fused_shared_experts > 0 and ".mlp.shared_experts." in mapped: + mapped = mapped.replace( + ".mlp.shared_experts.", + f".mlp.experts.{self.config.n_routed_experts}.", + ) for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in mapped: continue diff --git a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py index 800fdc101..0793f8dc9 100644 --- a/test/registered/unit/model_loader/test_runai_model_streamer_loader.py +++ b/test/registered/unit/model_loader/test_runai_model_streamer_loader.py @@ -220,6 +220,7 @@ class TestRunaiModelStreamerLoader(CustomTestCase): remapper = SimpleNamespace(confidence_head=None) model = SimpleNamespace( config=SimpleNamespace(n_routed_experts=1), + num_fused_shared_experts=0, named_parameters=lambda: [ ("stages.0.self_attn.wo_a.weight", param), ], diff --git a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py index d7229cc98..b7486b7e2 100644 --- a/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py +++ b/test/registered/unit/models/test_deepseek_v4_shared_expert_fusion.py @@ -1,18 +1,25 @@ import unittest from types import SimpleNamespace +from unittest.mock import patch +import torch +from torch import nn + +from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.layers.moe.utils import ( install_shared_experts_fusion_decision, is_shared_experts_fusion_disabled, ) from sglang.srt.models.deepseek_v4 import DeepseekV4ForCausalLM +from sglang.srt.models.deepseek_v4_dspark import DeepseekV4ForCausalLMDSpark from sglang.srt.runtime_context import get_context, get_exec, get_flags from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=4, suite="base-a-test-cpu") -class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase): +class TestDeepseekV4SharedExpertFusionPolicy(CustomTestCase): """V4 fuses its shared expert only when explicitly asked to. The gate is a question the loader asks the model class before any layer @@ -31,13 +38,24 @@ class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase): lambda: setattr(get_flags().moe, "disable_shared_experts_fusion", None) ) - def _install(self, n_shared_experts=1): + def _install(self, model_class=DeepseekV4ForCausalLM, n_shared_experts=1): install_shared_experts_fusion_decision( - DeepseekV4ForCausalLM, + model_class, SimpleNamespace(n_shared_experts=n_shared_experts), None, ) + def _make_dspark_config(self): + return DeepSeekV4Config( + architectures=["DeepseekV4ForCausalLMDSpark"], + quantization_config={}, + rope_scaling={}, + compress_ratios=[], + n_shared_experts=1, + dspark_markov_rank=1, + num_nextn_predict_layers=1, + ) + def test_disables_shared_fusion_without_enforce(self): self._publish(enforce=False) self.assertEqual( @@ -68,6 +86,89 @@ class TestDeepseekV4SharedExpertFusionPolicy(unittest.TestCase): SimpleNamespace(n_shared_experts=2), None ) + def test_dspark_entry_class_uses_the_v4_gate(self): + """A DSV4 DSpark draft must inherit the target's default fusion policy.""" + self._publish(enforce=False) + + self._install(DeepseekV4ForCausalLMDSpark) + + self.assertTrue(is_shared_experts_fusion_disabled()) + + def test_dspark_records_explicitly_forced_fusion(self): + """A forced DSpark build must retain its fused shared-expert count.""" + self._publish(enforce=True) + self._install(DeepseekV4ForCausalLMDSpark) + + class Stage(nn.Module): + def __init__(self, **_kwargs): + super().__init__() + + class MarkovHead(nn.Module): + def __init__(self, **_kwargs): + super().__init__() + + with ( + patch( + "sglang.srt.models.deepseek_v4_dspark.DSparkV4Stage", + Stage, + ), + patch( + "sglang.srt.models.deepseek_v4_dspark.DSparkV4MarkovHead", + MarkovHead, + ), + patch( + "sglang.srt.models.deepseek_v4_dspark.build_dspark_v4_confidence_head", + return_value=None, + ), + ): + model = DeepseekV4ForCausalLMDSpark(self._make_dspark_config()) + + self.assertEqual(model.num_fused_shared_experts, 1) + + def test_dspark_loads_forced_shared_expert_into_fused_slot(self): + """Forced DSpark shared tensors must load instead of being skipped.""" + config = self._make_dspark_config() + loaded = [] + + class Param: + def weight_loader( + self, + _param, + loaded_weight, + candidate, + *, + shard_id, + expert_id, + ): + loaded.append((loaded_weight, candidate, shard_id, expert_id)) + + class DraftModel: + num_fused_shared_experts = 1 + confidence_head = None + + def __init__(self): + self.config = config + + def named_parameters(self): + return [("stages.0.mlp.experts.w13_weight", Param())] + + def _remap_dspark_weight_name(self, name): + return DeepseekV4ForCausalLMDSpark._remap_dspark_weight_name(self, name) + + def _assert_confidence_head_loaded(self, **_kwargs): + return None + + weight = torch.ones(1) + DeepseekV4ForCausalLMDSpark.load_weights( + DraftModel(), + [("mtp.0.ffn.shared_experts.w1.weight", weight)], + ) + + self.assertEqual(len(loaded), 1) + self.assertIs(loaded[0][0], weight) + self.assertEqual(loaded[0][1], "stages.0.mlp.experts.w13_weight") + self.assertEqual(loaded[0][2:], ("w1", config.n_routed_experts)) + if __name__ == "__main__": unittest.main()