Fix DSV4 DSpark shared expert loading (#33312)

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

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),
             ],

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()