[Refactor] Abstract DeepGemm MegaMOE backend (#58053)

Signed-off-by: Hui Shi <kalinshi@tencent.com>
Co-authored-by: LuneExplorer <liyihan23@mails.tsinghua.edu.cn>
Co-authored-by: ArmageddonKnight <bojian@cs.toronto.edu>

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)

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()

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,

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,
     )

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