From 4f6885fffc931ce030d55a97ffb03bbeb0934d10 Mon Sep 17 00:00:00 2001 From: Canlin Guo Date: Fri, 21 Aug 2026 03:07:02 +0800 Subject: [PATCH] [DSV4][Kernel] Fuse shared experts into MegaMoE (#53040) Signed-off-by: Canlin Guo --- tests/models/test_deepseek_v4_mega_moe.py | 232 +++++++++++++ vllm/envs.py | 7 + vllm/models/deepseek_v4/nvidia/model.py | 318 +++++++++++++++--- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 51 +++ vllm/models/kimi_k3/nvidia/model.py | 7 +- 5 files changed, 570 insertions(+), 45 deletions(-) diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index ee1e765db6..f3c70dc88a 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/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 diff --git a/vllm/envs.py b/vllm/envs.py index a928461da2..90398ee4e0 100755 --- a/vllm/envs.py +++ b/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) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 27fded0407..be8214f21e 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/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]: diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index dac86be6ed..d254dad9d6 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/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, ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 472272a357..d2ceafd514 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/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