[DSV4][Kernel] Fuse shared experts into MegaMoE (#53040)

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
2026-08-20 12:07:02 -07:00
parent 3b829cf176
commit 4f6885fffc
5 changed files with 570 additions and 45 deletions
@@ -13,6 +13,7 @@ from vllm.models.deepseek_v4.nvidia.dspark import DSparkDeepseekV4ForCausalLM
 from vllm.models.deepseek_v4.nvidia.model import (
    DeepseekV4ForCausalLM,
    DeepseekV4MegaMoEExperts,
+    DeepseekV4MoE,
    make_deepseek_v4_expert_params_mapping,
 )
 from vllm.models.deepseek_v4.nvidia.mtp import DeepSeekV4MTP
@@ -168,6 +169,149 @@ def test_deepseek_v4_mega_moe_weight_loader_uses_ep_expert_ownership():
    assert torch.count_nonzero(experts.w13_weight[1]) == 0


+def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights(monkeypatch):
+    class FakeDeepGemm:
+        transformed_dims: list[tuple[int, int]] = []
+        scale_inputs: list[tuple[int, ...]] = []
+
+        @staticmethod
+        def get_symm_buffer_for_mega_moe(*args, num_shared_experts=0, **kwargs):
+            return None
+
+        @staticmethod
+        def get_block_m_for_mega_moe(*args, **kwargs):
+            return 128
+
+        @staticmethod
+        def fp8_fp4_mega_moe(
+            y,
+            l1_weights,
+            l2_weights,
+            sym_buffer,
+            shared_l1_weights=None,
+            shared_l2_weights=None,
+            **kwargs,
+        ):
+            return None
+
+        @classmethod
+        def transform_sf_into_required_layout(cls, sf, mn, k, *args, **kwargs):
+            cls.scale_inputs.append(tuple(sf.shape))
+            return torch.empty((sf.shape[0], mn, k // 128), dtype=torch.int32)
+
+        @classmethod
+        def transform_weights_for_mega_moe(cls, l1_weights, l2_weights):
+            cls.transformed_dims.append((l1_weights[0].dim(), l2_weights[0].dim()))
+            if l1_weights[0].dim() == 2:
+                return (l1_weights[0].clone(), l1_weights[1]), l2_weights
+            return l1_weights, l2_weights
+
+    vllm_config = SimpleNamespace(
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=4),
+        compilation_config=SimpleNamespace(static_forward_context={}),
+    )
+    experts = DeepseekV4MegaMoEExperts(
+        vllm_config,
+        num_experts=2,
+        num_local_experts=1,
+        experts_start_idx=0,
+        top_k=1,
+        hidden_size=128,
+        intermediate_size=128,
+        num_shared_experts=1,
+    )
+    experts._check_runtime_supported = lambda: None
+
+    def fp8_parameter(*shape):
+        return torch.nn.Parameter(
+            torch.empty(*shape, dtype=torch.float8_e4m3fn), requires_grad=False
+        )
+
+    def scale_parameter(*shape, dtype=torch.int32):
+        return torch.nn.Parameter(torch.ones(*shape, dtype=dtype), requires_grad=False)
+
+    shared_experts = SimpleNamespace(
+        gate_up_proj=SimpleNamespace(
+            weight=fp8_parameter(256, 128),
+            weight_block_size=(128, 128),
+            weight_scale_inv=scale_parameter(2, 1, dtype=torch.float8_e8m0fnu),
+        ),
+        down_proj=SimpleNamespace(
+            weight=fp8_parameter(128, 128),
+            weight_block_size=(128, 128),
+            weight_scale_inv=scale_parameter(1, 1, dtype=torch.float8_e8m0fnu),
+        ),
+    )
+    monkeypatch.setattr("vllm.utils.deep_gemm._import_deep_gemm", lambda: FakeDeepGemm)
+
+    original_gate_up_ptr = shared_experts.gate_up_proj.weight.data_ptr()
+    experts.finalize_weights(shared_experts)
+
+    assert FakeDeepGemm.transformed_dims == [(3, 3), (2, 2)]
+    assert FakeDeepGemm.scale_inputs[-2:] == [(1, 256, 4), (1, 128, 4)]
+    assert experts.has_fused_shared_experts
+    assert shared_experts.gate_up_proj.weight.data_ptr() != original_gate_up_ptr
+    assert (
+        experts._transformed_shared_l1_weights[0].data_ptr()
+        == shared_experts.gate_up_proj.weight.data_ptr()
+    )
+    assert (
+        experts._transformed_shared_l2_weights[0].data_ptr()
+        == shared_experts.down_proj.weight.data_ptr()
+    )
+
+
+@pytest.mark.parametrize("fused", [False, True])
+def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert(
+    monkeypatch, fused
+):
+    class FakeGate(torch.nn.Module):
+        tid2eid = None
+        e_score_correction_bias = None
+
+        def forward(self, hidden_states):
+            return torch.empty(hidden_states.shape[0], 2), None
+
+    class FakeExperts(torch.nn.Module):
+        has_fused_shared_experts = fused
+
+        def forward(self, hidden_states, *args, **kwargs):
+            return torch.ones_like(hidden_states)
+
+    class FakeSharedExperts(torch.nn.Module):
+        calls = 0
+
+        def forward(self, hidden_states):
+            self.calls += 1
+            return torch.full_like(hidden_states, 2)
+
+    moe = DeepseekV4MoE.__new__(DeepseekV4MoE)
+    torch.nn.Module.__init__(moe)
+    moe.use_mega_moe = True
+    moe.gate = FakeGate()
+    moe.experts = FakeExperts()
+    moe.shared_experts = FakeSharedExperts()
+    moe.scoring_func = "sqrtsoftplus"
+    moe.n_activated_experts = 1
+    moe.renormalize = True
+    moe.hash_indices_dtype = torch.int64
+    moe.routed_scaling_factor = 1.0
+    moe.swiglu_limit = 10.0
+    monkeypatch.setattr(
+        "vllm.models.deepseek_v4.nvidia.model.fused_topk_bias",
+        lambda **kwargs: (
+            torch.ones(kwargs["hidden_states"].shape[0], 1),
+            torch.zeros(kwargs["hidden_states"].shape[0], 1, dtype=torch.int64),
+        ),
+    )
+
+    output = moe(torch.zeros(2, 128))
+
+    expected = 1 if fused else 3
+    assert torch.all(output == expected)
+    assert moe.shared_experts.calls == (0 if fused else 1)
+
+
@pytest.mark.skipif(
    not torch.cuda.is_available(),
    reason="DeepSeek V4 MegaMoE fused input staging requires CUDA.",
@@ -246,6 +390,94 @@ def test_deepseek_v4_mega_moe_fused_input_staging_is_bitwise_exact():
    )


+@pytest.mark.parametrize("shared_block_m", [8, 32, 96, 128, 192])
+def test_deepseek_v4_mega_moe_stages_shared_scale_tma_layout(shared_block_m):
+    from vllm.third_party.deep_gemm.utils import per_token_cast_to_fp8
+
+    device = torch.device("cuda")
+    num_tokens = shared_block_m + 7
+    hidden_size = 256
+    top_k = 8
+    generator = torch.Generator(device=device)
+    generator.manual_seed(shared_block_m)
+    hidden_states = torch.randn(
+        num_tokens,
+        hidden_size,
+        device=device,
+        dtype=torch.bfloat16,
+        generator=generator,
+    )
+    topk_ids = torch.randint(
+        0,
+        256,
+        (num_tokens, top_k),
+        device=device,
+        dtype=torch.int32,
+        generator=generator,
+    )
+    topk_weights = torch.randn(
+        num_tokens,
+        top_k,
+        device=device,
+        dtype=torch.float32,
+        generator=generator,
+    )
+
+    ref_x, ref_x_sf = per_token_cast_to_fp8(
+        hidden_states,
+        use_ue8m0=True,
+        gran_k=32,
+        use_packed_ue8m0=True,
+    )
+    aligned_block_m = ((shared_block_m + 127) // 128) * 128
+    num_shared_rows = ((num_tokens + shared_block_m - 1) // shared_block_m) * (
+        aligned_block_m
+    )
+    ref_shared_x_sf = torch.zeros(
+        num_shared_rows,
+        hidden_size // 128,
+        dtype=torch.int32,
+        device=device,
+    )
+    for token_id in range(num_tokens):
+        m_in_block = token_id % shared_block_m
+        transposed_m = (
+            (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32
+        )
+        shared_row = token_id // shared_block_m * aligned_block_m + transposed_m
+        ref_shared_x_sf[shared_row].copy_(ref_x_sf[token_id])
+
+    fused_x = torch.empty_like(ref_x)
+    fused_x_sf = torch.empty_like(ref_x_sf)
+    fused_shared_storage = torch.full(
+        (hidden_size // 128, num_shared_rows),
+        -1,
+        dtype=torch.int32,
+        device=device,
+    )
+    fused_shared_x_sf = fused_shared_storage.t()
+    fused_topk_idx = torch.empty_like(topk_ids, dtype=torch.int64)
+    fused_topk_weights = torch.empty_like(topk_weights)
+
+    prepare_megamoe_inputs(
+        hidden_states,
+        topk_weights,
+        topk_ids,
+        fused_x,
+        fused_x_sf,
+        fused_topk_idx,
+        fused_topk_weights,
+        shared_x_sf=fused_shared_x_sf,
+        shared_block_m=shared_block_m,
+    )
+    torch.accelerator.synchronize()
+
+    populated = ref_shared_x_sf != 0
+    assert torch.equal(fused_x.view(torch.uint8), ref_x.view(torch.uint8))
+    assert torch.equal(fused_x_sf, ref_x_sf)
+    assert torch.equal(fused_shared_x_sf[populated], ref_shared_x_sf[populated])
+
+
 def test_deepseek_v4_pwal_hook_finalizes_mega_moe_and_mhc_broadcast():
    """The loader invokes the model-level PWAL hook for every load format,
    so it must finalize megamoe + mhc broadcast weights to cover dummy