[ROCm][DSv4] Fuse the shared expert into AITER MegaMoE (#61004)

Signed-off-by: Yanfei Wang <yanfei.wang@amd.com>
Co-authored-by: aoshen02 <aoshen@inferact.ai>

tests/models/test_deepseek_v4_aiter_mega_moe.py

@@ -29,6 +29,7 @@ def _make_experts(
     intermediate_size=INTER,
     max_num_batched_tokens=3000,
     swiglu_limit=None,
+    fuse_shared_expert=False,
 ):
     monkeypatch.setattr(
         mega_moe_module,
@@ -47,6 +48,7 @@ def _make_experts(
         hidden_size=HIDDEN,
         intermediate_size=intermediate_size,
         swiglu_limit=swiglu_limit,
+        fuse_shared_expert=fuse_shared_expert,
     )
 
 
@@ -108,6 +110,24 @@ def test_experts_load_into_local_slots(monkeypatch):
             assert torch.all(w2 == 10 * expert + 3)
 
 
+@pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm")
+def test_shared_expert_route_appended(monkeypatch):
+    experts = _make_experts(
+        monkeypatch, ep_rank=3, ep_size=8, num_experts=384, fuse_shared_expert=True
+    )
+    ids = torch.randint(0, 384, (37, 2), device="cuda", dtype=torch.int32)
+    ids[::5, 1] = -1
+    weights = torch.rand(37, 2, device="cuda")
+    out_weights, out_ids = experts.append_shared_expert(weights, ids)
+
+    shared = torch.full((37, 1), 3 * 49 + 48, device="cuda", dtype=torch.int32)
+    expected_ids = torch.cat([torch.where(ids >= 0, ids + ids // 48, ids), shared], 1)
+    torch.testing.assert_close(out_ids, expected_ids)
+    torch.testing.assert_close(
+        out_weights, torch.cat([weights, torch.ones(37, 1, device="cuda")], 1)
+    )
+
+
 @pytest.mark.skipif(not current_platform.is_rocm(), reason="Requires ROCm")
 @pytest.mark.parametrize("hash_routing", [False, True])
 @pytest.mark.parametrize("num_tokens", [1, 37])

vllm/distributed/device_communicators/all2all.py

@@ -1012,7 +1012,11 @@ class MoriAll2AllManager(All2AllManagerBase):
         if get_current_vllm_config().kernel_config.moe_backend == "aiter_mega_moe":
             # MegaMoEV2 places its dispatch/combine workspaces on the MoRI
             # symmetric heap, which defaults to 2 GB. MegaMoEV2 requires > 4GB
-            heap_size = os.environ.setdefault("MORI_SHMEM_HEAP_SIZE", "8G")
+            # (> 8GB with a fused shared expert).
+            heap_size = os.environ.setdefault(
+                "MORI_SHMEM_HEAP_SIZE",
+                "16G" if envs.VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS else "8G",
+            )
             logger.info_once("AITER MegaMoE: MORI_SHMEM_HEAP_SIZE=%s", heap_size)
         mori.shmem.shmem_torch_process_group_init("mori")
 

vllm/envs.py

@@ -155,6 +155,7 @@ if TYPE_CHECKING:
     VLLM_ROCM_USE_AITER_FP4BMM: bool = True
     VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION: bool = False
     VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: bool = False
+    VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS: bool = False
     VLLM_ROCM_MONO_DECODE: bool = False
     VLLM_ROCM_USE_AITER_TRITON_GEMM: bool = True
     VLLM_ROCM_USE_SKINNY_GEMM: bool = True
@@ -1391,6 +1392,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
         os.getenv("VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS", "False").lower()
         in ("true", "1")
     ),
+    # Fuse DeepSeek-V4's shared expert into AITER MegaMoE (requantized to
+    # MXFP4). By default is disabled.
+    "VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS": lambda: (
+        os.getenv("VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS", "False").lower()
+        in ("true", "1")
+    ),
     # Run the eligible layers of decode steps on the model's mono decode
     # kernels: persistent launches a layer, the TP all-reduces in-kernel
     # (DeepSeek-V4.1 on CDNA4). By default is disabled.

vllm/models/deepseek_v4/amd/mega_moe.py

@@ -10,10 +10,51 @@ import torch.nn as nn
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config import VllmConfig
 from vllm.distributed import get_ep_group
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    get_fp8_block_weight_scale,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import mxfp4_quantize
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    scaled_dequantize,
+)
 from vllm.model_executor.utils import set_weight_attrs
+from vllm.triton_utils import tl, triton
 from vllm.utils.math_utils import next_power_of_2
 
 
+@triton.jit
+def _append_shared_expert_kernel(
+    topk_ids_ptr,
+    topk_weights_ptr,
+    out_ids_ptr,
+    out_weights_ptr,
+    num_tokens,
+    shared_expert_id,
+    TOPK: tl.constexpr,
+    ROUTED_PER_RANK: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+):
+    rows = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M)
+    mask = rows < num_tokens
+    for k in tl.static_range(TOPK):
+        e = tl.load(topk_ids_ptr + rows * TOPK + k, mask=mask, other=0)
+        w = tl.load(topk_weights_ptr + rows * TOPK + k, mask=mask, other=0.0)
+        e = tl.where(e >= 0, e + e // ROUTED_PER_RANK, e)
+        tl.store(out_ids_ptr + rows * (TOPK + 1) + k, e, mask=mask)
+        tl.store(out_weights_ptr + rows * (TOPK + 1) + k, w, mask=mask)
+    tl.store(
+        out_ids_ptr + rows * (TOPK + 1) + TOPK,
+        tl.zeros_like(rows) + shared_expert_id,
+        mask=mask,
+    )
+    tl.store(
+        out_weights_ptr + rows * (TOPK + 1) + TOPK,
+        tl.full([BLOCK_M], 1.0, tl.float32),
+        mask=mask,
+    )
+
+
 class DeepseekV4AiterMegaMoEExperts(nn.Module):
     # MegaMoEV2 owns large symmetric-heap buffers, so every layer shares one
     # instance and passes its own weights per call.
@@ -28,16 +69,25 @@ class DeepseekV4AiterMegaMoEExperts(nn.Module):
         hidden_size: int,
         intermediate_size: int,
         swiglu_limit: float | None,
+        fuse_shared_expert: bool = False,
         prefix: str = "",
     ):
         super().__init__()
         ep_group = get_ep_group()
         self.ep_rank = ep_group.rank_in_group
         self.ep_size = ep_group.world_size
-        self.num_experts = num_experts
-        self.num_local_experts = num_experts // self.ep_size
-        self.experts_start_idx = self.ep_rank * self.num_local_experts
-        self.top_k = top_k
+        self.num_routed_local_experts = num_experts // self.ep_size
+        self.experts_start_idx = self.ep_rank * self.num_routed_local_experts
+        # A fused shared expert is one extra local expert per rank and one
+        # extra route per token, to the token's own rank.
+        self.fuse_shared_expert = fuse_shared_expert
+        self.num_local_experts = self.num_routed_local_experts + fuse_shared_expert
+        self.num_experts = self.num_local_experts * self.ep_size
+        self.routed_top_k = top_k
+        self.top_k = top_k + fuse_shared_expert
+        self.shared_expert_id = (
+            self.ep_rank * self.num_local_experts + self.num_routed_local_experts
+        )
         self.hidden_size = hidden_size
         self.intermediate_size = intermediate_size
         self.swiglu_limit = swiglu_limit or 0.0
@@ -70,7 +120,7 @@ class DeepseekV4AiterMegaMoEExperts(nn.Module):
         return_success: bool = False,
     ) -> bool | None:
         local_expert_id = expert_id - self.experts_start_idx
-        if not 0 <= local_expert_id < self.num_local_experts:
+        if not 0 <= local_expert_id < self.num_routed_local_experts:
             return False if return_success else None
         expert_data = param.data[local_expert_id]
         if shard_id in ("w1", "w3"):
@@ -81,6 +131,24 @@ class DeepseekV4AiterMegaMoEExperts(nn.Module):
         expert_data.copy_(loaded_weight.view(torch.uint8))
         return True if return_success else None
 
+    def load_shared_expert(self, shared_experts: nn.Module) -> None:
+        slot = self.num_routed_local_experts
+        for linear, weight, weight_scale in (
+            (shared_experts.gate_up_proj, self.w13_weight, self.w13_weight_scale),
+            (shared_experts.down_proj, self.w2_weight, self.w2_weight_scale),
+        ):
+            block_scale = get_fp8_block_weight_scale(linear)
+            assert block_scale is not None, "shared expert must be block-FP8"
+            group = GroupShape(
+                linear.weight.shape[0] // block_scale.shape[0],
+                linear.weight.shape[1] // block_scale.shape[1],
+            )
+            packed, scale = mxfp4_quantize(
+                scaled_dequantize(linear.weight, block_scale, group, torch.bfloat16)
+            )
+            weight.data[slot].copy_(packed.view(torch.uint8))
+            weight_scale.data[slot].copy_(scale.view(torch.uint8))
+
     def finalize_weights(self) -> None:
         if self._runtime is not None:
             return
@@ -129,6 +197,33 @@ class DeepseekV4AiterMegaMoEExperts(nn.Module):
             self._runtime_cache[key] = runtime
         self._runtime = runtime
 
+    def append_shared_expert(
+        self, topk_weights: torch.Tensor, topk_ids: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Remap routed ids to the per-rank (routed + shared) layout and append
+        a weight-1.0 route to this rank's shared slot."""
+        num_tokens = topk_ids.shape[0]
+        out_ids = torch.empty(
+            (num_tokens, self.top_k), dtype=topk_ids.dtype, device=topk_ids.device
+        )
+        out_weights = torch.empty(
+            (num_tokens, self.top_k), dtype=torch.float32, device=topk_ids.device
+        )
+        if num_tokens:
+            block_m = 64
+            _append_shared_expert_kernel[(triton.cdiv(num_tokens, block_m),)](
+                topk_ids.contiguous(),
+                topk_weights.to(torch.float32).contiguous(),
+                out_ids,
+                out_weights,
+                num_tokens,
+                self.shared_expert_id,
+                TOPK=self.routed_top_k,
+                ROUTED_PER_RANK=self.num_routed_local_experts,
+                BLOCK_M=block_m,
+            )
+        return out_weights, out_ids
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -137,6 +232,8 @@ class DeepseekV4AiterMegaMoEExperts(nn.Module):
     ) -> torch.Tensor:
         runtime = self._runtime
         assert runtime is not None
+        if self.fuse_shared_expert:
+            topk_weights, topk_ids = self.append_shared_expert(topk_weights, topk_ids)
         # MegaMoEV2 reads these per launch; rebind to this layer's experts.
         runtime._s1_w1 = self.w13_weight
         runtime._s1_w1_scale = self.w13_weight_scale
@@ -149,6 +246,10 @@ DeepseekV4AiterMegaMoEExperts.weight_loader.supports_moe_loading = True  # type:
 
 
 def finalize_mega_moe_weights(model: nn.Module) -> None:
-    for module in model.modules():
-        if isinstance(module, DeepseekV4AiterMegaMoEExperts):
-            module.finalize_weights()
+    for module in list(model.modules()):
+        experts = getattr(module, "experts", None)
+        if isinstance(experts, DeepseekV4AiterMegaMoEExperts):
+            if experts.fuse_shared_expert:
+                experts.load_shared_expert(module.shared_experts)
+                module.shared_experts = None
+            experts.finalize_weights()

vllm/models/deepseek_v4/amd/model.py

@@ -74,7 +74,10 @@ from vllm.model_executor.models.utils import (
     make_layers,
     maybe_prefix,
 )
-from vllm.models.deepseek_v4.amd.mega_moe import DeepseekV4AiterMegaMoEExperts
+from vllm.models.deepseek_v4.amd.mega_moe import (
+    DeepseekV4AiterMegaMoEExperts,
+    finalize_mega_moe_weights,
+)
 from vllm.models.deepseek_v4.amd.rocm import DeepseekV4ROCMAiterMLAAttention
 from vllm.platforms import current_platform
 from vllm.platforms.rocm import on_gfx950
@@ -690,6 +693,10 @@ class DeepseekV4MoE(nn.Module):
                 hidden_size=config.hidden_size,
                 intermediate_size=config.moe_intermediate_size,
                 swiglu_limit=self.swiglu_limit,
+                fuse_shared_expert=(
+                    config.n_shared_experts == 1
+                    and envs.VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS
+                ),
                 prefix=f"{prefix}.experts",
             )
             return
@@ -1524,6 +1531,8 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
         return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
 
     def process_weights_after_loading(self) -> None:
+        # Must precede the gate_up preshuffle of the fused shared experts.
+        finalize_mega_moe_weights(self)
         # After per-layer quant finalize, so we preshuffle the final fp8 weights.
         fused_compressor_layers = 0
         for module in self.modules():
@@ -1532,8 +1541,6 @@ class DeepseekV4ForCausalLM(nn.Module, SupportsPP, SupportsEagle3):
                 module.prepare_attn_preshuffle()
             elif isinstance(module, DeepseekV4MLP):
                 module.prepare_gateup_preshuffle()
-            elif isinstance(module, DeepseekV4AiterMegaMoEExperts):
-                module.finalize_weights()
         if fused_compressor_layers:
             logger.info(
                 "Fused the C4 compressor GEMMs in %d DeepSeek V4 layers",