From 81198e97ba7eee2a22540caaa756b7fdddcb4d93 Mon Sep 17 00:00:00 2001 From: Hui Shi <70356247+huishi-hs@users.noreply.github.com> Date: Thu, 8 Oct 2026 13:05:39 +0800 Subject: [PATCH] [Refactor] Abstract DeepGemm MegaMOE backend (#58053) Signed-off-by: Hui Shi Co-authored-by: LuneExplorer Co-authored-by: ArmageddonKnight --- tests/models/test_deepseek_v4_mega_moe.py | 15 +- .../layers/fused_moe/deep_gemm_mega_moe.py | 466 ++++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 310 ++++-------- .../deepseek_v4/nvidia/ops/prepare_megamoe.py | 107 +++- vllm/models/kimi_k3/nvidia/model.py | 53 +- 5 files changed, 683 insertions(+), 268 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py diff --git a/tests/models/test_deepseek_v4_mega_moe.py b/tests/models/test_deepseek_v4_mega_moe.py index 104ad8e895..76915b94a3 100644 --- a/tests/models/test_deepseek_v4_mega_moe.py +++ b/tests/models/test_deepseek_v4_mega_moe.py @@ -6,6 +6,9 @@ from types import SimpleNamespace import pytest import torch +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmSm100MegaMoEBackend, +) from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( bind_routed_experts_capturer, ) @@ -451,7 +454,9 @@ def test_deepseek_v4_mega_moe_preserves_checkpoint_dimensions( ): param.data.random_(1, 128) originals.append(param.data.clone()) - monkeypatch.setattr(experts, "_check_runtime_supported", lambda: None) + monkeypatch.setattr( + experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + ) monkeypatch.setattr( "vllm.utils.deep_gemm._import_deep_gemm", lambda: SimpleNamespace( @@ -547,7 +552,9 @@ def test_deepseek_v4_mega_moe_finalizes_native_shared_expert_weights( intermediate_size=intermediate_size, num_shared_experts=1, ) - experts._check_runtime_supported = lambda: None + monkeypatch.setattr( + experts, "_ensure_backend", lambda: DeepGemmSm100MegaMoEBackend() + ) def fp8_parameter(*shape): return torch.nn.Parameter( @@ -664,6 +671,10 @@ def test_deepseek_v4_mega_moe_does_not_double_add_fused_shared_expert( class FakeExperts(torch.nn.Module): has_fused_shared_experts = fused + _backend = DeepGemmSm100MegaMoEBackend() + + def _ensure_backend(self): + return self._backend def forward(self, hidden_states, *args, **kwargs): return torch.ones_like(hidden_states) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py new file mode 100644 index 0000000000..633ced1082 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py @@ -0,0 +1,466 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DeepGEMM MegaMoE backends: fused EP dispatch, expert GEMMs and combine.""" + +from inspect import signature +from typing import Any + +import torch +import torch.nn as nn + +from vllm.logger import init_logger +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kMxfp8Dynamic, +) +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +def ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: + """Reinterpret a uint8 UE8M0 exponent tensor as float32 scale values.""" + return (sf.to(torch.int32) << 23).view(torch.float32) + + +class DeepGemmMegaMoEBackend: + """Abstract MegaMoE backend that hides deep_gemm-specific details. + + A backend advertises its expected hidden-state quantization layout via + ``hidden_quant`` and its expected MMA type via ``mma_type``. It is + responsible for the routed-expert weight transform run at load time, the + optional shared-expert weight transform, and the actual MegaMoE kernel + invocation at inference time. Runtime capability checks live in + :func:`get_deep_gemm_mega_moe_backend`: obtaining a backend already implies + that the current device and shapes are supported. + """ + + # Human-readable MMA type tag used by deep_gemm helpers. SM100 default. + mma_type: str = "fp8xfp4" + hidden_quant: QuantKey = kMxfp8Dynamic + + def transform_weights( + self, + *, + w13_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: str | None = None, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + raise NotImplementedError + + def supports_shared_experts(self, deep_gemm) -> bool: + return False + + def transform_shared_expert_weights( + self, + *, + shared_experts: nn.Module, + num_shared_experts: int, + hidden_size: int, + intermediate_size: int, + prefix: str, + ) -> ( + tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ] + | None + ): + """Return ``(l1, l2)`` transformed shared-expert weights, or ``None``. + + ``shared_experts`` must expose ``gate_up_proj`` and ``down_proj`` + linear layers. A ``None`` return means the backend cannot fuse shared + experts for this module; the caller should fall back to a separate + shared MLP. + """ + return None + + def run_mega_moe( + self, + *, + y: torch.Tensor, + l1_weights: tuple[torch.Tensor, torch.Tensor], + l2_weights: tuple[torch.Tensor, torch.Tensor], + symm_buffer, + activation_clamp: float | None, + fast_math: bool, + activation: str | None = None, + activation_alpha: float | None = None, + activation_beta: float | None = None, + shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + raise NotImplementedError + + def supports_bf16_mega_gate(self) -> bool: + """Whether this backend can run the DeepGEMM bf16 MegaMoE gate kernel. + + Defaults to ``False`` (the symbol is not guaranteed to exist in every + deep_gemm build). Backends that target a deep_gemm build exposing the + ``bf16_mega_gate`` symbol override this and :meth:`bf16_mega_gate`. + """ + return False + + def bf16_mega_gate( + self, + *, + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + + +class DeepGemmSm100MegaMoEBackend(DeepGemmMegaMoEBackend): + """Default MegaMoE backend targeting SM100 via deep_gemm fp8_fp4_mega_moe.""" + + mma_type = "fp8xfp4" + hidden_quant = kMxfp8Dynamic + + def transform_weights( + self, + *, + w13_weight: torch.Tensor, + w13_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + num_local_experts: int, + hidden_size: int, + intermediate_size: int, + activation: str | None = None, + ) -> tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ]: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + w13_scale = deep_gemm.transform_sf_into_required_layout( + ue8m0_uint8_to_float(w13_weight_scale).contiguous(), + 2 * intermediate_size, + hidden_size, + (1, 32), + num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + ue8m0_uint8_to_float(w2_weight_scale).contiguous(), + hidden_size, + intermediate_size, + (1, 32), + num_local_experts, + ) + kwargs = {} if activation is None else {"activation": activation} + return deep_gemm.transform_weights_for_mega_moe( + (w13_weight.view(torch.int8).contiguous(), w13_scale), + (w2_weight.view(torch.int8).contiguous(), w2_scale), + **kwargs, + ) + + def supports_shared_experts(self, 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 transform_shared_expert_weights( + self, + *, + shared_experts: nn.Module, + num_shared_experts: int, + hidden_size: int, + intermediate_size: int, + prefix: str, + ) -> ( + tuple[ + tuple[torch.Tensor, torch.Tensor], + tuple[torch.Tensor, torch.Tensor], + ] + | None + ): + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + + 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 + if hasattr(gate_up, "weight_scale") + else gate_up.weight_scale_inv + ).data + down_weight = down.weight.data + down_scale = ( + down.weight_scale + if hasattr(down, "weight_scale") + else 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. + ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) + if ( + gate_up_scale.dtype in ue8m0_scale_dtypes + and down_scale.dtype in ue8m0_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], + prefix, + ) + down_scale = self._prepare_shared_expert_scale( + deep_gemm, + down, + down_scale, + down_weight.shape[0], + down_weight.shape[1], + prefix, + ) + + if gate_up_scale is None or down_scale is None: + return None + + shared_intermediate_size = intermediate_size * num_shared_experts + expected_gate_up_shape = ( + 2 * shared_intermediate_size, + hidden_size, + ) + expected_down_shape = (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.", + 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, + ) + return None + + 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] + transformed_l1 = (gate_up.weight.data, transformed_l1[1]) + return transformed_l1, transformed_l2 + + @staticmethod + def _prepare_shared_expert_scale( + deep_gemm, + linear: nn.Module, + scale: torch.Tensor, + mn: int, + k: int, + prefix: str, + ) -> 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.", + prefix, + ) + return None + + block_m, block_k = block_size + num_k_blocks = (k + block_k - 1) // block_k + # The linear hook's DeepGEMM layout packs four k-block bytes per int32. + packed = scale.dtype == torch.int32 + expected_shape = ( + (mn, (num_k_blocks + 3) // 4) + if packed + else ((mn + block_m - 1) // block_m, num_k_blocks) + ) + 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).", + prefix, + tuple(scale.shape), + tuple(block_size), + mn, + k, + ) + return None + + if packed: + row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] + else: + row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] + scale_1x32 = ( + ue8m0_uint8_to_float(row_scale) + .repeat_interleave(block_k // 32, dim=1)[:, : 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 run_mega_moe( + self, + *, + y: torch.Tensor, + l1_weights: tuple[torch.Tensor, torch.Tensor], + l2_weights: tuple[torch.Tensor, torch.Tensor], + symm_buffer, + activation_clamp: float | None, + fast_math: bool, + activation: str | None = None, + activation_alpha: float | None = None, + activation_beta: float | None = None, + shared_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + shared_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> None: + from vllm.utils.deep_gemm import _import_deep_gemm + + deep_gemm = _import_deep_gemm() + kwargs: dict[str, Any] = { + name: value + for name, value in ( + ("activation", activation), + ("activation_alpha", activation_alpha), + ("activation_beta", activation_beta), + ) + if value is not None + } + if shared_l1_weights is not None and shared_l2_weights is not None: + kwargs["shared_l1_weights"] = shared_l1_weights + kwargs["shared_l2_weights"] = shared_l2_weights + deep_gemm.fp8_fp4_mega_moe( + y, + l1_weights, + l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + **kwargs, + ) + + def supports_bf16_mega_gate(self) -> bool: + """SM100 deep_gemm build is expected to provide the bf16 gate kernel.""" + return True + + def bf16_mega_gate( + self, + *, + x: torch.Tensor, + weight: torch.Tensor, + num_topk: int, + scoring_func: str, + routed_scaling_factor: float, + ep_rank: int, + bias: torch.Tensor | None = None, + image_bias: torch.Tensor | None = None, + image_token_mask: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + unmapped_topk_idx: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + from vllm.utils.deep_gemm import bf16_mega_gate + + return bf16_mega_gate( + x, + weight, + num_topk, + scoring_func=scoring_func, + routed_scaling_factor=routed_scaling_factor, + ep_rank=ep_rank, + bias=bias, + image_bias=image_bias, + image_token_mask=image_token_mask, + fix_routing_mask=fix_routing_mask, + unmapped_topk_idx=unmapped_topk_idx, + ) + + +def get_deep_gemm_mega_moe_backend( + device: torch.device, + hidden_size: int, + intermediate_size: int, +) -> DeepGemmMegaMoEBackend: + """Return the DeepGEMM MegaMoE backend for ``device``. + + A successful return implies the backend can run on this device with the + given ``hidden_size`` / ``intermediate_size``. All runtime capability + checks (arch + shape divisibility) live here; there is no separate check + on the returned backend. + + Currently only the SM100 deep_gemm backend is registered here. SM90 and + other backends are expected to be provided by new backends or plugins + that override this factory. + + Raises: + NotImplementedError: If the device architecture is unsupported. + ValueError: If the shapes are unsupported. + + """ + if not current_platform.is_device_capability_family( + 100, device_id=device.index or 0 + ): + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if hidden_size % 128 != 0 or intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + return DeepGemmSm100MegaMoEBackend() diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index b9996d007c..fc10cd07c1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -2,7 +2,6 @@ # 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 @@ -33,6 +32,11 @@ from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, fused_moe_make_expert_params_mapping, ) +from vllm.model_executor.layers.fused_moe.deep_gemm_mega_moe import ( + DeepGemmMegaMoEBackend, + get_deep_gemm_mega_moe_backend, + ue8m0_uint8_to_float, +) from vllm.model_executor.layers.fused_moe.router.base_router import ( eplb_map_to_physical_and_record, ) @@ -205,7 +209,9 @@ def make_deepseek_v4_expert_params_mapping( class DeepseekV4MegaMoEExperts(nn.Module): - _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int, int], object] = {} + _symm_buffer_cache: dict[ + tuple[int, int, int, int, int, int, int, int, str], object + ] = {} def __init__( self, @@ -296,6 +302,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): tuple[torch.Tensor, torch.Tensor] | None ) = None + self._backend: DeepGemmMegaMoEBackend | 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. compilation_config = vllm_config.compilation_config @@ -357,214 +365,71 @@ class DeepseekV4MegaMoEExperts(nn.Module): @staticmethod def _ue8m0_uint8_to_float(sf: torch.Tensor) -> torch.Tensor: - return (sf.to(torch.int32) << 23).view(torch.float32) + return ue8m0_uint8_to_float(sf) + + def _ensure_backend(self) -> DeepGemmMegaMoEBackend: + """Lazily resolve the MegaMoE backend. + + The backend is looked up via :func:`get_deep_gemm_mega_moe_backend`, which + performs all runtime capability checks. Failure propagates directly + to the caller. This must be called before :meth:`finalize_weights` + releases the loader-side parameters. + """ + if self._backend is not None: + return self._backend + assert self.w13_weight is not None, ( + "_ensure_backend() must be called before finalize_weights() " + "releases the loader-side parameters." + ) + self._backend = get_deep_gemm_mega_moe_backend( + self.w13_weight.device, + self.hidden_size, + self.intermediate_size, + ) + return self._backend def _check_runtime_supported(self) -> None: - device = self.w13_weight.device - if torch.cuda.get_device_capability(device)[0] != 10: - raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") - if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: - raise ValueError( - "DeepGEMM MegaMoE requires hidden and intermediate sizes " - "to be multiples of 128." - ) + """Preserved for API compatibility; delegates to the backend check.""" + self._ensure_backend() - @staticmethod - def _deep_gemm_supports_shared_experts(deep_gemm) -> bool: - """Check the Python API before touching a symmetric-memory group. + def _finalize_shared_expert_weights(self, shared_experts: DeepseekV4MLP) -> None: + """Delegate shared-expert weight transform to the backend and update state. - This also gives users of an older precompiled vLLM wheel a safe serial - fallback instead of failing halfway through multi-rank buffer setup. + On success, populates ``self._transformed_shared_l{1,2}_weights``. If + the backend cannot fuse shared experts (returns ``None``), disables + native fusion by setting ``self.num_shared_experts = 0``. """ - 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 + backend = self._ensure_backend() + result = backend.transform_shared_expert_weights( + shared_experts=shared_experts, + num_shared_experts=self.num_shared_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + prefix=self.prefix, ) - - 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 - if hasattr(gate_up, "weight_scale") - else gate_up.weight_scale_inv - ).data - down_weight = down.weight.data - down_scale = ( - down.weight_scale - if hasattr(down, "weight_scale") - else 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. - ue8m0_scale_dtypes = (torch.float8_e8m0fnu, torch.uint8, torch.int32) - if ( - gate_up_scale.dtype in ue8m0_scale_dtypes - and down_scale.dtype in ue8m0_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: + if result is None: self.num_shared_experts = 0 return - - 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 - num_k_blocks = (k + block_k - 1) // block_k - # The linear hook's DeepGEMM layout packs four k-block bytes per int32. - packed = scale.dtype == torch.int32 - expected_shape = ( - (mn, (num_k_blocks + 3) // 4) - if packed - else ((mn + block_m - 1) // block_m, num_k_blocks) - ) - 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 - - if packed: - row_scale = scale.flatten().view(torch.uint8).view(mn, -1)[:, :num_k_blocks] - else: - row_scale = scale.view(torch.uint8).repeat_interleave(block_m, dim=0)[:mn] - scale_1x32 = ( - self._ue8m0_uint8_to_float(row_scale) - .repeat_interleave(block_k // 32, dim=1)[:, : 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) + ( + self._transformed_shared_l1_weights, + self._transformed_shared_l2_weights, + ) = result def finalize_weights(self, shared_experts: DeepseekV4MLP | None = None) -> None: + backend = self._ensure_backend() from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() - 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, - ) - 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), + backend.transform_weights( + w13_weight=self.w13_weight.data, + w13_weight_scale=self.w13_weight_scale.data, + w2_weight=self.w2_weight.data, + w2_weight_scale=self.w2_weight_scale.data, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, ) ) # Drop the original loader-side parameters: the MegaMoE kernels only @@ -583,7 +448,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): return if self._transformed_shared_l1_weights is not None: return - if not self._deep_gemm_supports_shared_experts(deep_gemm): + if not backend.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 " @@ -591,7 +456,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): ) self.num_shared_experts = 0 return - self._finalize_shared_expert_weights(deep_gemm, shared_experts) + + self._finalize_shared_expert_weights(shared_experts) @property def has_fused_shared_experts(self) -> bool: @@ -602,6 +468,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): deep_gemm = _import_deep_gemm() + backend = self._ensure_backend() group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -613,6 +480,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): self.hidden_size, self.intermediate_size, self.num_shared_experts if self.has_fused_shared_experts else 0, + backend.mma_type, ) symm_buffer = self._symm_buffer_cache.get(key) if symm_buffer is None: @@ -626,6 +494,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): num_shared_experts=( self.num_shared_experts if self.has_fused_shared_experts else 0 ), + mma_type=backend.mma_type, ) self._symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -700,7 +569,6 @@ class DeepseekV4MegaMoEExperts(nn.Module): deep_gemm = _import_deep_gemm() - symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] is_padding = None if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -732,6 +600,8 @@ class DeepseekV4MegaMoEExperts(nn.Module): else None, ) + backend = self._ensure_backend() + symm_buffer = self.get_symm_buffer() shared_x_sf = None shared_block_m = None if self.has_fused_shared_experts: @@ -742,7 +612,7 @@ class DeepseekV4MegaMoEExperts(nn.Module): symm_buffer.num_max_tokens_per_rank, num_tokens, self.top_k, - "fp8xfp4", + backend.mma_type, ) prepare_megamoe_inputs( @@ -756,27 +626,28 @@ class DeepseekV4MegaMoEExperts(nn.Module): is_padding=is_padding, shared_x_sf=shared_x_sf, shared_block_m=shared_block_m, + hidden_quant=backend.hidden_quant, ) assert self._transformed_l1_weights is not None assert self._transformed_l2_weights is not None if self.has_fused_shared_experts: - deep_gemm.fp8_fp4_mega_moe( - y, - self._transformed_l1_weights, - self._transformed_l2_weights, - symm_buffer, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=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, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, activation_clamp=activation_clamp, fast_math=fast_math, ) @@ -1039,9 +910,18 @@ class DeepseekV4MoE(nn.Module): return self._forward_fused_moe(hidden_states, input_ids) org_shape = hidden_states.shape + # The MegaMoE experts module owns the backend; resolve and assert it. + self.experts._ensure_backend() + assert self.experts._backend is not None, ( + "MegaMoE backend must be resolved before routing." + ) + mega_backend = self.experts._backend # Small local padded batches favor GateLinear; 128-expert gates cross earlier. gate_threshold = 1 if self.gate.weight.shape[0] == 128 else 16 - if hidden_states.shape[0] <= gate_threshold: + if ( + not mega_backend.supports_bf16_mega_gate() + or hidden_states.shape[0] <= gate_threshold + ): router_logits, _ = self.gate(hidden_states) topk_weights, topk_ids = fused_topk_bias( hidden_states=hidden_states, @@ -1060,8 +940,6 @@ class DeepseekV4MoE(nn.Module): image_sentinel_lo=self.image_sentinel_lo if bias_vl is not None else 0, ) else: - from vllm.utils.deep_gemm import bf16_mega_gate - if mega_gate_metadata is None: raise RuntimeError( "MegaMoE routing metadata must be prepared once by the model." @@ -1079,10 +957,10 @@ class DeepseekV4MoE(nn.Module): ] fix_routing_mask = mega_gate_metadata.hash_token_mask - topk_weights, topk_ids = bf16_mega_gate( - hidden_states, - self.gate.weight, - self.n_activated_experts, + topk_weights, topk_ids = mega_backend.bf16_mega_gate( + x=hidden_states, + weight=self.gate.weight, + num_topk=self.n_activated_experts, scoring_func=self.scoring_func, routed_scaling_factor=self.routed_scaling_factor, ep_rank=self.ep_rank, diff --git a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py index aa9482bed6..2d93b0d430 100644 --- a/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py +++ b/vllm/models/deepseek_v4/nvidia/ops/prepare_megamoe.py @@ -9,8 +9,15 @@ MegaMoE kernels consume. import torch +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + FP8_DTYPE, + QuantKey, + kMxfp8Dynamic, +) from vllm.triton_utils import tl, triton +_E8M0_SCALE_DTYPES = (torch.uint8, torch.float8_e8m0fnu) + @triton.jit(do_not_specialize=["num_tokens"]) def _prepare_megamoe_inputs_kernel( @@ -48,6 +55,7 @@ def _prepare_megamoe_inputs_kernel( GROUP_K: tl.constexpr, BLOCK_TOPK: tl.constexpr, SHARED_BLOCK_M: tl.constexpr, + USE_UE8M0: tl.constexpr, ) -> None: token_id = tl.program_id(0) * BLOCK_M + tl.arange(0, BLOCK_M) token_mask = token_id < num_tokens @@ -69,15 +77,18 @@ def _prepare_megamoe_inputs_kernel( amax = tl.maximum(amax, 1.0e-4) scale = amax / 448.0 - scale_bits = scale.to(tl.uint32, bitcast=True) - scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( - tl.uint32 - ) - scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) - rounded_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + if USE_UE8M0: + scale_bits = scale.to(tl.uint32, bitcast=True) + scale_exp = ((scale_bits >> 23) & 0xFF) + ((scale_bits & 0x7FFFFF) != 0).to( + tl.uint32 + ) + scale_exp = tl.minimum(tl.maximum(scale_exp, 1), 254) + quant_scale = (scale_exp << 23).to(tl.float32, bitcast=True) + else: + quant_scale = scale hidden_groups = tl.reshape(hidden, [BLOCK_M, num_groups, GROUP_K]) - scaled = hidden_groups * (1.0 / rounded_scale)[:, :, None] + scaled = hidden_groups * (1.0 / quant_scale)[:, :, None] scaled = tl.reshape(scaled, [BLOCK_M, BLOCK_K]) fp8 = scaled.to(tl.float8e4nv) tl.store( @@ -86,21 +97,32 @@ def _prepare_megamoe_inputs_kernel( mask=k_mask, ) - scale_offsets = tl.arange(0, num_groups) - packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( - tl.int32 - ) - tl.store( - x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, - packed_scale, - mask=token_mask, - ) + if USE_UE8M0: + scale_offsets = tl.arange(0, num_groups) + packed_scale = tl.sum(scale_exp << (scale_offsets[None, :] * 8), axis=1).to( + tl.int32 + ) + tl.store( + x_sf + token_id * x_sf_stride_m + k_block_id * x_sf_stride_k, + packed_scale, + mask=token_mask, + ) + else: + sf_offset = tl.arange(0, num_groups) + tl.store( + x_sf + + token_id[:, None] * x_sf_stride_m + + k_block_id * x_sf_stride_k + + sf_offset[None, :], + quant_scale, + mask=token_mask[:, None], + ) # 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: + if USE_UE8M0 and 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 @@ -160,6 +182,27 @@ def _prepare_megamoe_inputs_kernel( ) +def _resolve_hidden_quant(hidden_quant: QuantKey, block_k: int) -> tuple[int, bool]: + """Map a hidden-state ``QuantKey`` to the kernel's ``(GROUP_K, USE_UE8M0)``.""" + scale = hidden_quant.scale + group_shape = scale.group_shape + if ( + hidden_quant.dtype != FP8_DTYPE + or not hidden_quant.symmetric + or hidden_quant.scale2 is not None + or scale.static + or not group_shape.is_per_group() + or block_k % group_shape.col != 0 + or scale.dtype not in (*_E8M0_SCALE_DTYPES, torch.float32) + ): + raise ValueError( + "DeepSeek V4 MegaMoE input staging requires symmetric dynamic fp8 " + f"per-group quantization with a group size dividing {block_k} and " + f"E8M0 or fp32 scales, got {hidden_quant}." + ) + return group_shape.col, scale.dtype in _E8M0_SCALE_DTYPES + + def prepare_megamoe_inputs( hidden_states: torch.Tensor, topk_weights: torch.Tensor, @@ -171,7 +214,26 @@ def prepare_megamoe_inputs( is_padding: torch.Tensor | None = None, shared_x_sf: torch.Tensor | None = None, shared_block_m: int | None = None, + hidden_quant: QuantKey = kMxfp8Dynamic, ) -> None: + """Quantize hidden states and repack top-k routing for DeepGEMM MegaMoE. + + Args: + hidden_states: Input activations of shape ``[num_tokens, hidden_dim]``. + topk_weights: Router top-k weights of shape ``[num_tokens, top_k]``. + topk_ids: Router top-k expert ids of shape ``[num_tokens, top_k]``. + x_fp8: Output buffer for the fp8-quantized hidden states. + x_sf: Output buffer for the hidden-state scale factors. + topk_idx_out: Output buffer for the repacked top-k expert ids. + topk_weights_out: Output buffer for the repacked top-k weights. + is_padding: Optional per-token mask; padded tokens are not routed. + shared_x_sf: Optional output buffer for shared-expert scale factors. + Must be given together with ``shared_block_m``. + shared_block_m: Block M used to lay out ``shared_x_sf``. + hidden_quant: Hidden-state quantization scheme. E8M0 scales are packed + four per int32 into ``x_sf``; fp32 scales are stored directly. + + """ num_tokens, hidden_size = hidden_states.shape if num_tokens == 0: return @@ -180,6 +242,8 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE input staging requires hidden_size to be " "a multiple of 128." ) + block_k = 128 + group_k, use_ue8m0 = _resolve_hidden_quant(hidden_quant, block_k) top_k = topk_ids.shape[1] if topk_weights.shape != topk_ids.shape: raise ValueError( @@ -191,6 +255,11 @@ def prepare_megamoe_inputs( "DeepSeek V4 MegaMoE shared input staging requires both " "shared_x_sf and shared_block_m." ) + if shared_x_sf is not None and not use_ue8m0: + raise ValueError( + "DeepSeek V4 MegaMoE shared input staging currently requires " + "UE8M0-packed hidden scales." + ) if shared_x_sf is not None: assert shared_block_m is not None if shared_block_m <= 0: @@ -209,7 +278,6 @@ def prepare_megamoe_inputs( f"{required_rows}, got {shared_x_sf.shape[0]}." ) - block_k = 128 # On GB200, eight-row tiles win from 64 tokens; keep smaller batches untiled. block_m = 8 if num_tokens >= 64 else 1 grid = (triton.cdiv(num_tokens, block_m), triton.cdiv(hidden_size, block_k)) @@ -247,8 +315,9 @@ def prepare_megamoe_inputs( top_k, BLOCK_M=block_m, BLOCK_K=block_k, - GROUP_K=32, + GROUP_K=group_k, BLOCK_TOPK=block_topk, SHARED_BLOCK_M=shared_block_m or 1, + USE_UE8M0=use_ue8m0, num_warps=4, ) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index cc1ca609a6..84e93b82fb 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -365,27 +365,16 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): return self._check_runtime_supported() - 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, - ) + backend = self._ensure_backend() 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), + backend.transform_weights( + w13_weight=self.w13_weight.data, + w13_weight_scale=self.w13_weight_scale.data, + w2_weight=self.w2_weight.data, + w2_weight_scale=self.w2_weight_scale.data, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, activation=self.activation, ) ) @@ -407,6 +396,7 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): from vllm.utils.deep_gemm import _import_deep_gemm deep_gemm = _import_deep_gemm() + backend = self._ensure_backend() group = get_ep_group().device_group device = torch.accelerator.current_device_index() key = ( @@ -418,6 +408,7 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): self.hidden_size, self.intermediate_size, self.activation, + backend.mma_type, ) symm_buffer = self._kimi_symm_buffer_cache.get(key) if symm_buffer is None: @@ -429,6 +420,7 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): self.hidden_size, self.intermediate_size, activation=self.activation, + mma_type=backend.mma_type, ) self._kimi_symm_buffer_cache[key] = symm_buffer return symm_buffer @@ -449,10 +441,6 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): f"but its symmetric buffer supports {self.max_num_tokens}." ) y = torch.empty_like(hidden_states, dtype=torch.bfloat16) - from vllm.utils.deep_gemm import _import_deep_gemm - - deep_gemm = _import_deep_gemm() - symm_buffer = self.get_symm_buffer() num_tokens = hidden_states.shape[0] is_padding = None if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available(): @@ -483,6 +471,8 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): else None, ) + backend = self._ensure_backend() + symm_buffer = self.get_symm_buffer() prepare_megamoe_inputs( hidden_states, topk_weights, @@ -492,22 +482,23 @@ class KimiK3MegaMoEExperts(DeepseekV4MegaMoEExperts): symm_buffer.topk_idx[:num_tokens], symm_buffer.topk_weights[:num_tokens], is_padding=is_padding, + hidden_quant=backend.hidden_quant, ) self.finalize_weights() 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, + backend.run_mega_moe( + y=y, + l1_weights=self._transformed_l1_weights, + l2_weights=self._transformed_l2_weights, + symm_buffer=symm_buffer, activation_clamp=activation_clamp, - activation=self.activation, + fast_math=fast_math, # DeepGEMM names the SiTU gate tanh scale `activation_alpha` and the # linear/up tanh scale `activation_beta`; beta=0 leaves up untouched. + activation=self.activation, activation_alpha=self.activation_beta or 1.0, activation_beta=self.activation_linear_beta or 0.0, - fast_math=fast_math, ) return y