diff --git a/tests/models/test_deepseek_v4_aiter_mega_moe.py b/tests/models/test_deepseek_v4_aiter_mega_moe.py index a8f692e07bfb..19b3697811df 100644 --- a/tests/models/test_deepseek_v4_aiter_mega_moe.py +++ b/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]) diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index cff34139f968..3466700c3334 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -1012,7 +1012,11 @@ def __init__(self, cpu_group, all2all_backend: str): 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") diff --git a/vllm/envs.py b/vllm/envs.py index 02c282fd8ec2..4918b0b2714e 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -155,6 +155,7 @@ 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 @@ -1388,6 +1389,12 @@ def _resolve_rust_cli_path() -> str | None: 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. diff --git a/vllm/models/deepseek_v4/amd/mega_moe.py b/vllm/models/deepseek_v4/amd/mega_moe.py index 2fdff9335145..312c089e5aad 100644 --- a/vllm/models/deepseek_v4/amd/mega_moe.py +++ b/vllm/models/deepseek_v4/amd/mega_moe.py @@ -10,10 +10,51 @@ 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 @@ def __init__( 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 @@ def weight_loader( 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 @@ def weight_loader( 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 @@ def finalize_weights(self) -> None: 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 @@ def forward( ) -> 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 @@ def forward( 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() diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index 8b26fc0060e7..51d22b75e8c1 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -74,7 +74,10 @@ 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 @@ def __init__( 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 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: 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 @@ def process_weights_after_loading(self) -> None: 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",