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

Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>

tests/models/test_deepseek_v4_mega_moe.py

@@ -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

vllm/envs.py

@@ -288,6 +288,7 @@ if TYPE_CHECKING:
     VLLM_GC_DEBUG: str = ""
     VLLM_DEBUG_WORKSPACE: bool = False
     VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
+    VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION: bool = False
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 1024
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
@@ -1993,6 +1994,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_DISABLE_SHARED_EXPERTS_STREAM": lambda: bool(
         int(os.getenv("VLLM_DISABLE_SHARED_EXPERTS_STREAM", "0"))
     ),
+    # Emergency rollback for the DeepSeek-V4 NVIDIA MegaMoE path. By default,
+    # DeepGEMM computes replicated FP8 shared experts in the same persistent
+    # SM100 kernel as the routed FP4 experts.
+    "VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION": lambda: bool(
+        int(os.getenv("VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION", "0"))
+    ),
     # Limits when we run shared_experts in a separate stream.
     # We found out that for large batch sizes, the separate stream
     # execution is not beneficial (most likely because of the input clone)

vllm/models/deepseek_v4/nvidia/model.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import typing
 from collections.abc import Callable, Iterable
+from inspect import signature
 from itertools import islice
 
 import regex as re
@@ -18,6 +19,7 @@ from vllm.distributed import (
 )
 from vllm.distributed.eplb.eplb_state import EplbLayerState
 from vllm.forward_context import get_forward_context, is_forward_context_available
+from vllm.logger import init_logger
 from vllm.model_executor.kernels.mhc.tilelang import (
     hc_head_fused_kernel_tilelang,
     mhc_fused_post_pre_tilelang,
@@ -84,6 +86,8 @@ from vllm.utils.math_utils import cdiv
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 from vllm.v1.worker.ubatching import dbo_current_ubatch_id
 
+logger = init_logger(__name__)
+
 
 class DeepseekV4MLP(nn.Module):
     def __init__(
@@ -165,7 +169,7 @@ def make_deepseek_v4_expert_params_mapping(
 
 
 class DeepseekV4MegaMoEExperts(nn.Module):
-    _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {}
+    _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int, int], object] = {}
 
     def __init__(
         self,
@@ -177,6 +181,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         top_k: int,
         hidden_size: int,
         intermediate_size: int,
+        num_shared_experts: int = 0,
         prefix: str = "",
         num_logical_experts: int | None = None,
     ):
@@ -190,6 +195,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
         self.top_k = top_k
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
+        self.num_shared_experts = num_shared_experts
         self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
 
         self.num_logical_experts = (
@@ -247,6 +253,12 @@ class DeepseekV4MegaMoEExperts(nn.Module):
 
         self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None
         self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None
+        self._transformed_shared_l1_weights: (
+            tuple[torch.Tensor, torch.Tensor] | None
+        ) = None
+        self._transformed_shared_l2_weights: (
+            tuple[torch.Tensor, torch.Tensor] | None
+        ) = None
 
         # Register in the static forward context so the custom-op wrapper
         # can look up this module by name from within a torch.compile graph.
@@ -321,45 +333,220 @@ class DeepseekV4MegaMoEExperts(nn.Module):
                 "to be multiples of 128."
             )
 
-    def finalize_weights(self) -> None:
-        if self._transformed_l1_weights is not None:
+    @staticmethod
+    def _deep_gemm_supports_shared_experts(deep_gemm) -> bool:
+        """Check the Python API before touching a symmetric-memory group.
+
+        This also gives users of an older precompiled vLLM wheel a safe serial
+        fallback instead of failing halfway through multi-rank buffer setup.
+        """
+        try:
+            buffer_params = signature(deep_gemm.get_symm_buffer_for_mega_moe).parameters
+            kernel_params = signature(deep_gemm.fp8_fp4_mega_moe).parameters
+        except (TypeError, ValueError):
+            return False
+        return (
+            hasattr(deep_gemm, "get_block_m_for_mega_moe")
+            and hasattr(deep_gemm, "transform_weights_for_mega_moe")
+            and "num_shared_experts" in buffer_params
+            and "shared_l1_weights" in kernel_params
+            and "shared_l2_weights" in kernel_params
+        )
+
+    def _finalize_shared_expert_weights(
+        self, deep_gemm, shared_experts: DeepseekV4MLP
+    ) -> None:
+        gate_up = shared_experts.gate_up_proj
+        down = shared_experts.down_proj
+        gate_up_weight = gate_up.weight.data
+        gate_up_scale = gate_up.weight_scale_inv.data
+        down_weight = down.weight.data
+        down_scale = down.weight_scale_inv.data
+
+        # MegaMoE's shared FP8 MMA consumes a 1x32 scale for every weight row,
+        # while the checkpoint uses coarser block-FP8 scales (usually
+        # 128x128). Build a dedicated, numerically equivalent scale view before
+        # the generic linear post-load hook replaces the raw checkpoint scales
+        # with its 128x128 DeepGEMM layout.
+        checkpoint_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8)
+        if (
+            gate_up_scale.dtype in checkpoint_scale_dtypes
+            and down_scale.dtype in checkpoint_scale_dtypes
+        ):
+            gate_up_scale = self._prepare_shared_expert_scale(
+                deep_gemm,
+                gate_up,
+                gate_up_scale,
+                gate_up_weight.shape[0],
+                gate_up_weight.shape[1],
+            )
+            down_scale = self._prepare_shared_expert_scale(
+                deep_gemm,
+                down,
+                down_scale,
+                down_weight.shape[0],
+                down_weight.shape[1],
+            )
+
+        if gate_up_scale is None or down_scale is None:
+            self.num_shared_experts = 0
             return
 
-        self._check_runtime_supported()
+        shared_intermediate_size = self.intermediate_size * self.num_shared_experts
+        expected_gate_up_shape = (
+            2 * shared_intermediate_size,
+            self.hidden_size,
+        )
+        expected_down_shape = (self.hidden_size, shared_intermediate_size)
+        if (
+            gate_up_weight.dtype != torch.float8_e4m3fn
+            or down_weight.dtype != torch.float8_e4m3fn
+            or gate_up_scale.dtype != torch.int32
+            or down_scale.dtype != torch.int32
+            or tuple(gate_up_weight.shape) != expected_gate_up_shape
+            or tuple(down_weight.shape) != expected_down_shape
+        ):
+            logger.warning(
+                "Disabling native MegaMoE shared-expert fusion for %s: expected "
+                "replicated block-FP8 weights with gate_up=%s, down=%s, and "
+                "DeepGEMM int32 scales; got gate_up=%s/%s/%s and down=%s/%s/%s.",
+                self.prefix,
+                expected_gate_up_shape,
+                expected_down_shape,
+                tuple(gate_up_weight.shape),
+                gate_up_weight.dtype,
+                gate_up_scale.dtype,
+                tuple(down_weight.shape),
+                down_weight.dtype,
+                down_scale.dtype,
+            )
+            self.num_shared_experts = 0
+            return
+
+        transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(
+            (gate_up_weight, gate_up_scale),
+            (down_weight, down_scale),
+        )
+        # L1 interleaving allocates a full copy. Re-home the loader Parameter on
+        # that storage so the original 2*intermediate*hidden FP8 tensor can be
+        # released instead of adding roughly 0.7 GiB per rank on DSV4-Flash.
+        # The generic linear post-load hook may still repack the serial scales,
+        # but this shared MLP is never called after native fusion is enabled.
+        gate_up.weight.data = transformed_l1[0]
+        self._transformed_shared_l1_weights = (
+            gate_up.weight.data,
+            transformed_l1[1],
+        )
+        self._transformed_shared_l2_weights = transformed_l2
+
+    def _prepare_shared_expert_scale(
+        self,
+        deep_gemm,
+        linear: nn.Module,
+        scale: torch.Tensor,
+        mn: int,
+        k: int,
+    ) -> torch.Tensor | None:
+        block_size = getattr(linear, "weight_block_size", None)
+        if block_size is None or len(block_size) != 2:
+            logger.warning(
+                "Disabling native MegaMoE shared-expert fusion for %s: "
+                "shared FP8 weight block size is unavailable.",
+                self.prefix,
+            )
+            return None
+
+        block_m, block_k = block_size
+        expected_shape = (
+            (mn + block_m - 1) // block_m,
+            (k + block_k - 1) // block_k,
+        )
+        if block_k % 32 != 0 or tuple(scale.shape) != expected_shape:
+            logger.warning(
+                "Disabling native MegaMoE shared-expert fusion for %s: "
+                "cannot convert shared scale shape %s with block size %s "
+                "to MegaMoE's 1x32 layout for weight (%d, %d).",
+                self.prefix,
+                tuple(scale.shape),
+                tuple(block_size),
+                mn,
+                k,
+            )
+            return None
+
+        scale_fp32 = self._ue8m0_uint8_to_float(scale.view(torch.uint8))
+        scale_1x32 = (
+            scale_fp32.repeat_interleave(block_m, dim=0)
+            .repeat_interleave(block_k // 32, dim=1)[:mn, : k // 32]
+            .contiguous()
+        )
+        # The grouped API is used with a singleton dimension to request the
+        # MN-major, TMA-aligned packed-UE8M0 strides, then squeezed back to the
+        # 2D layout required for a shared expert.
+        return deep_gemm.transform_sf_into_required_layout(
+            scale_1x32.unsqueeze(0),
+            mn,
+            k,
+            (1, 32),
+            1,
+        ).squeeze(0)
+
+    def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
         from vllm.utils.deep_gemm import _import_deep_gemm
 
         deep_gemm = _import_deep_gemm()
 
-        w13_scale = deep_gemm.transform_sf_into_required_layout(
-            self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
-            2 * self.intermediate_size,
-            self.hidden_size,
-            (1, 32),
-            self.num_local_experts,
-        )
-        w2_scale = deep_gemm.transform_sf_into_required_layout(
-            self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
-            self.hidden_size,
-            self.intermediate_size,
-            (1, 32),
-            self.num_local_experts,
-        )
-        self._transformed_l1_weights, self._transformed_l2_weights = (
-            deep_gemm.transform_weights_for_mega_moe(
-                (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
-                (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
+        if self._transformed_l1_weights is None:
+            self._check_runtime_supported()
+            w13_scale = deep_gemm.transform_sf_into_required_layout(
+                self._ue8m0_uint8_to_float(self.w13_weight_scale.data).contiguous(),
+                2 * self.intermediate_size,
+                self.hidden_size,
+                (1, 32),
+                self.num_local_experts,
             )
-        )
-        # Drop the original loader-side parameters: the MegaMoE kernels only
-        # consume the transformed views above. transform_weights_for_mega_moe
-        # allocates a fresh tensor for the L1 weight (see _interleave_l1_weights)
-        # and fresh SF tensors for L1/L2; the L2 weight is the only tensor that
-        # aliases the original storage, and _transformed_l2_weights still holds
-        # it, so the storage stays live after we drop the Parameter.
-        self.w13_weight = None
-        self.w13_weight_scale = None
-        self.w2_weight = None
-        self.w2_weight_scale = None
+            w2_scale = deep_gemm.transform_sf_into_required_layout(
+                self._ue8m0_uint8_to_float(self.w2_weight_scale.data).contiguous(),
+                self.hidden_size,
+                self.intermediate_size,
+                (1, 32),
+                self.num_local_experts,
+            )
+            self._transformed_l1_weights, self._transformed_l2_weights = (
+                deep_gemm.transform_weights_for_mega_moe(
+                    (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale),
+                    (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale),
+                )
+            )
+            # Drop the original loader-side parameters: the MegaMoE kernels only
+            # consume the transformed views above. transform_weights_for_mega_moe
+            # allocates a fresh tensor for the L1 weight (see
+            # _interleave_l1_weights) and fresh SF tensors for L1/L2; the L2
+            # weight is the only tensor that aliases the original storage, and
+            # _transformed_l2_weights still holds it, so the storage stays live
+            # after we drop the Parameter.
+            self.w13_weight = None
+            self.w13_weight_scale = None
+            self.w2_weight = None
+            self.w2_weight_scale = None
+
+        if shared_experts is None or self.num_shared_experts == 0:
+            return
+        if self._transformed_shared_l1_weights is not None:
+            return
+        if not self._deep_gemm_supports_shared_experts(deep_gemm):
+            logger.warning_once(
+                "Disabling native MegaMoE shared-expert fusion because the "
+                "installed DeepGEMM Python API is older than the vLLM "
+                "source. Rebuild the vendored _deep_gemm_C extension to enable it.",
+            )
+            self.num_shared_experts = 0
+            return
+        self._finalize_shared_expert_weights(deep_gemm, shared_experts)
+
+    @property
+    def has_fused_shared_experts(self) -> bool:
+        return self._transformed_shared_l1_weights is not None
 
     def get_symm_buffer(self):
         from vllm.utils.deep_gemm import _import_deep_gemm
@@ -376,6 +563,7 @@ class DeepseekV4MegaMoEExperts(nn.Module):
             self.top_k,
             self.hidden_size,
             self.intermediate_size,
+            self.num_shared_experts if self.has_fused_shared_experts else 0,
         )
         symm_buffer = self._symm_buffer_cache.get(key)
         if symm_buffer is None:
@@ -386,6 +574,9 @@ class DeepseekV4MegaMoEExperts(nn.Module):
                 self.top_k,
                 self.hidden_size,
                 self.intermediate_size,
+                num_shared_experts=(
+                    self.num_shared_experts if self.has_fused_shared_experts else 0
+                ),
             )
             self._symm_buffer_cache[key] = symm_buffer
         return symm_buffer
@@ -492,6 +683,19 @@ class DeepseekV4MegaMoEExperts(nn.Module):
                 else None,
             )
 
+        shared_x_sf = None
+        shared_block_m = None
+        if self.has_fused_shared_experts:
+            shared_x_sf = symm_buffer.shared_l1_acts_sf
+            shared_block_m = deep_gemm.get_block_m_for_mega_moe(
+                get_ep_group().world_size,
+                self.num_experts,
+                symm_buffer.num_max_tokens_per_rank,
+                num_tokens,
+                self.top_k,
+                "fp8xfp4",
+            )
+
         prepare_megamoe_inputs(
             hidden_states,
             topk_weights,
@@ -501,18 +705,32 @@ class DeepseekV4MegaMoEExperts(nn.Module):
             symm_buffer.topk_idx[:num_tokens],
             symm_buffer.topk_weights[:num_tokens],
             is_padding=is_padding,
+            shared_x_sf=shared_x_sf,
+            shared_block_m=shared_block_m,
         )
 
         assert self._transformed_l1_weights is not None
         assert self._transformed_l2_weights is not None
-        deep_gemm.fp8_fp4_mega_moe(
-            y,
-            self._transformed_l1_weights,
-            self._transformed_l2_weights,
-            symm_buffer,
-            activation_clamp=activation_clamp,
-            fast_math=fast_math,
-        )
+        if self.has_fused_shared_experts:
+            deep_gemm.fp8_fp4_mega_moe(
+                y,
+                self._transformed_l1_weights,
+                self._transformed_l2_weights,
+                symm_buffer,
+                shared_l1_weights=self._transformed_shared_l1_weights,
+                shared_l2_weights=self._transformed_shared_l2_weights,
+                activation_clamp=activation_clamp,
+                fast_math=fast_math,
+            )
+        else:
+            deep_gemm.fp8_fp4_mega_moe(
+                y,
+                self._transformed_l1_weights,
+                self._transformed_l2_weights,
+                symm_buffer,
+                activation_clamp=activation_clamp,
+                fast_math=fast_math,
+            )
         return y
 
 
@@ -645,6 +863,16 @@ class DeepseekV4MoE(nn.Module):
         self.experts_start_idx = self.physical_expert_start
         self.experts_end_idx = self.physical_expert_end
 
+        # Native DeepGEMM fusion requires each EP rank to own the complete
+        # shared MLP. Sequence parallel replicates those weights while sharding
+        # tokens. TP=1 is also naturally replicated. With PP+TP the shared MLP
+        # remains tensor-sharded, so retain the serial path.
+        fuse_shared_experts = bool(
+            self.shared_experts is not None
+            and not envs.VLLM_DISABLE_DSV4_MEGAMOE_SHARED_EXPERT_FUSION
+            and (self.use_sequence_parallel or self.tp_size == 1)
+        )
+
         self.experts = DeepseekV4MegaMoEExperts(
             vllm_config,
             num_experts=self.n_physical_experts,
@@ -654,6 +882,7 @@ class DeepseekV4MoE(nn.Module):
             top_k=config.num_experts_per_tok,
             hidden_size=config.hidden_size,
             intermediate_size=config.moe_intermediate_size,
+            num_shared_experts=(self.n_shared_experts if fuse_shared_experts else 0),
             prefix=f"{prefix}.experts",
         )
 
@@ -739,7 +968,10 @@ class DeepseekV4MoE(nn.Module):
             activation_clamp=activation_clamp,
         )
 
-        if self.shared_experts is not None:
+        if (
+            self.shared_experts is not None
+            and not self.experts.has_fused_shared_experts
+        ):
             shared_output = self.shared_experts(hidden_states)
             final_hidden_states += shared_output
 
@@ -759,7 +991,7 @@ class DeepseekV4MoE(nn.Module):
 
     def finalize_mega_moe_weights(self) -> None:
         if self.use_mega_moe:
-            self.experts.finalize_weights()
+            self.experts.finalize_weights(self.shared_experts)
 
 
 def _select_dsv4_attn_cls(vllm_config: VllmConfig) -> type[DeepseekV4Attention]:

vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py

@@ -17,6 +17,7 @@ def _prepare_megamoe_inputs_kernel(
     hidden_states,
     x_fp8,
     x_sf,
+    shared_x_sf,
     topk_ids,
     topk_weights,
     is_padding,
@@ -28,6 +29,8 @@ def _prepare_megamoe_inputs_kernel(
     x_stride_k: tl.constexpr,
     x_sf_stride_m: tl.constexpr,
     x_sf_stride_k: tl.constexpr,
+    shared_x_sf_stride_m: tl.constexpr,
+    shared_x_sf_stride_k: tl.constexpr,
     topk_ids_stride_m: tl.constexpr,
     topk_ids_stride_k: tl.constexpr,
     topk_weights_stride_m: tl.constexpr,
@@ -42,6 +45,7 @@ def _prepare_megamoe_inputs_kernel(
     BLOCK_K: tl.constexpr,
     GROUP_K: tl.constexpr,
     BLOCK_TOPK: tl.constexpr,
+    SHARED_BLOCK_M: tl.constexpr,
 ) -> None:
     token_id = tl.program_id(0)
     k_block_id = tl.program_id(1)
@@ -84,6 +88,25 @@ def _prepare_megamoe_inputs_kernel(
         packed_scale,
     )
 
+    # DeepGEMM's SM100 shared-expert TMA loads require the activation scales
+    # in an MN-major layout whose row permutation depends on the MegaMoE
+    # scheduler's runtime BLOCK_M. Write that view while the packed UE8M0 scale
+    # is already resident, avoiding another kernel and temporary tensor.
+    if shared_x_sf is not None:
+        m_block_id = token_id // SHARED_BLOCK_M
+        m_in_block = token_id % SHARED_BLOCK_M
+        aligned_block_m: tl.constexpr = triton.cdiv(SHARED_BLOCK_M, 128) * 128
+        transposed_m = (
+            (m_in_block // 128) * 128 + (m_in_block % 32) * 4 + (m_in_block % 128) // 32
+        )
+        shared_row = m_block_id * aligned_block_m + transposed_m
+        tl.store(
+            shared_x_sf
+            + shared_row * shared_x_sf_stride_m
+            + k_block_id * shared_x_sf_stride_k,
+            packed_scale,
+        )
+
     if k_block_id == 0:
         topk_offsets = tl.arange(0, BLOCK_TOPK)
         topk_mask = topk_offsets < top_k
@@ -131,6 +154,8 @@ def prepare_megamoe_inputs(
     topk_idx_out: torch.Tensor,
     topk_weights_out: torch.Tensor,
     is_padding: torch.Tensor | None = None,
+    shared_x_sf: torch.Tensor | None = None,
+    shared_block_m: int | None = None,
 ) -> None:
     num_tokens, hidden_size = hidden_states.shape
     if num_tokens == 0:
@@ -146,6 +171,28 @@ def prepare_megamoe_inputs(
             "DeepSeek V4 MegaMoE input staging requires topk_weights and "
             "topk_ids to have the same shape."
         )
+    if (shared_x_sf is None) != (shared_block_m is None):
+        raise ValueError(
+            "DeepSeek V4 MegaMoE shared input staging requires both "
+            "shared_x_sf and shared_block_m."
+        )
+    if shared_x_sf is not None:
+        assert shared_block_m is not None
+        if shared_block_m <= 0:
+            raise ValueError("MegaMoE shared_block_m must be positive.")
+        expected_sf_k = hidden_size // 128
+        if shared_x_sf.ndim != 2 or shared_x_sf.shape[1] != expected_sf_k:
+            raise ValueError(
+                "MegaMoE shared_x_sf must have shape "
+                f"(*, {expected_sf_k}), got {tuple(shared_x_sf.shape)}."
+            )
+        aligned_block_m = triton.cdiv(shared_block_m, 128) * 128
+        required_rows = triton.cdiv(num_tokens, shared_block_m) * aligned_block_m
+        if shared_x_sf.shape[0] < required_rows:
+            raise ValueError(
+                "MegaMoE shared_x_sf has insufficient rows: requires "
+                f"{required_rows}, got {shared_x_sf.shape[0]}."
+            )
 
     block_k = 128
     grid = (num_tokens, triton.cdiv(hidden_size, block_k))
@@ -155,6 +202,7 @@ def prepare_megamoe_inputs(
         hidden_states,
         x_fp8,
         x_sf,
+        shared_x_sf,
         topk_ids,
         topk_weights,
         is_padding,
@@ -166,6 +214,8 @@ def prepare_megamoe_inputs(
         x_fp8.stride(1),
         x_sf.stride(0),
         x_sf.stride(1),
+        shared_x_sf.stride(0) if shared_x_sf is not None else 0,
+        shared_x_sf.stride(1) if shared_x_sf is not None else 0,
         topk_ids.stride(0),
         topk_ids.stride(1),
         topk_weights.stride(0),
@@ -180,5 +230,6 @@ def prepare_megamoe_inputs(
         BLOCK_K=block_k,
         GROUP_K=32,
         BLOCK_TOPK=block_topk,
+        SHARED_BLOCK_M=shared_block_m or 1,
         num_warps=4,
     )

vllm/models/kimi_k3/nvidia/model.py

@@ -93,7 +93,10 @@ from vllm.models.common.ops.sequence_parallel import (
     sp_reduce_scatter,
     sp_shard,
 )
-from vllm.models.deepseek_v4.nvidia.model import DeepseekV4MegaMoEExperts
+from vllm.models.deepseek_v4.nvidia.model import (
+    DeepseekV4MegaMoEExperts,
+    DeepseekV4MLP,
+)
 from vllm.models.deepseek_v4.nvidia.ops.prepare_megamoe import prepare_megamoe_inputs
 from vllm.models.kimi_k3.nvidia.kda import KimiK3DeltaAttention
 from vllm.models.kimi_k3.nvidia.latent_moe_runner import (
@@ -354,7 +357,7 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts):
         torch.distributed.barrier(group=ep_group.cpu_group)
         self._synchronized_ep_groups.add(key)
 
-    def finalize_weights(self) -> None:
+    def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None:
         if self._transformed_l1_weights is not None:
             return