Upload folder using huggingface_hub

vllm-patches/A_wna16_marlin.py

@@ -0,0 +1,537 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from typing import Any
+
+import torch
+from compressed_tensors.quantization import (
+    QuantizationArgs,
+)
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe import (
+    RoutedExperts,
+    SharedExperts,
+)
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+)
+from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
+    WNA16MoEBackend,
+    convert_to_wna16_moe_kernel_format,
+    make_wna16_moe_kernel,
+    make_wna16_moe_quant_config,
+    select_wna16_moe_backend,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import (  # noqa E501
+    CompressedTensorsMoEMethod,
+)
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import (  # noqa
+    WNA16_SUPPORTED_TYPES_MAP,
+    WNA16_ZP_SUPPORTED_TYPES_MAP,
+)
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    get_marlin_input_dtype,
+    marlin_make_workspace_new,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kInt4Static32GroupScale,
+    kInt4StaticGroupScale,
+    kInt8StaticGroupScale,
+)
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+
+logger = init_logger(__name__)
+
+
+class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
+    def __init__(
+        self,
+        weight_quant: QuantizationArgs,
+        input_quant: QuantizationArgs | None,
+        moe: FusedMoEConfig,
+        layer_name: str | None = None,
+    ):
+        super().__init__(moe)
+        self.weight_quant = weight_quant
+        self.input_quant = input_quant
+        self.symmetric = weight_quant.symmetric
+        # Extract properties from weight_quant
+        self.num_bits = weight_quant.num_bits
+        self.packed_factor = 32 // weight_quant.num_bits
+        self.strategy = weight_quant.strategy
+        self.group_size = weight_quant.group_size
+        self.actorder = weight_quant.actorder
+
+        self.quant_type = (
+            WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
+            if self.symmetric
+            else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits]
+        )
+
+        self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
+
+        if self.num_bits == 4:
+            if self.group_size == 32:
+                scale = kInt4Static32GroupScale
+            else:
+                scale = kInt4StaticGroupScale
+        elif self.num_bits == 8:
+            assert self.group_size == -1
+            scale = kInt8StaticGroupScale
+        else:
+            raise ValueError(
+                "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now."
+            )
+
+        weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric)
+
+        # Select WNA16 MoE backend via oracle.
+        self.wna16_backend, self.experts_cls = select_wna16_moe_backend(
+            config=self.moe,
+            weight_key=weight_key,
+        )
+
+    def get_weight_shape(
+        self,
+        weight_name: str,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        num_groups_w2: int | None = None,
+        num_groups_w13: int | None = None,
+    ) -> tuple[int, int, int]:
+        """
+        Get the shape of the weight based on the weight name, number of experts
+        hidden size, intermediate size per partition, number of groups for w2,
+        and number of groups for w13. Pass in num_groups_w2 and num_groups_w13
+        for weight scales/zero_points.
+        """
+        if weight_name in ("w13_scale", "w13_zp"):
+            assert num_groups_w13 is not None, (
+                "num_groups_w13 must be provided for weight scales/zero_points"
+            )
+        if weight_name in ("w2_scale", "w2_zp"):
+            assert num_groups_w2 is not None, (
+                "num_groups_w2 must be provided for weight scales/zero_points"
+            )
+        w13_num_shards = 2 if self.moe.is_act_and_mul else 1
+        is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
+        shape_map = {
+            "w13_weight": {
+                "Flashinfer": (
+                    num_experts,
+                    w13_num_shards * intermediate_size_per_partition,
+                    hidden_size // self.packed_factor,
+                ),
+                "Marlin": (
+                    num_experts,
+                    hidden_size // self.packed_factor,
+                    w13_num_shards * intermediate_size_per_partition,
+                ),
+            },
+            "w13_scale": {
+                "Flashinfer": (
+                    num_experts,
+                    w13_num_shards * intermediate_size_per_partition,
+                    num_groups_w13,
+                ),
+                "Marlin": (
+                    num_experts,
+                    num_groups_w13,
+                    w13_num_shards * intermediate_size_per_partition,
+                ),
+            },
+            "w13_zp": {
+                "Marlin": (
+                    num_experts,
+                    num_groups_w13,
+                    w13_num_shards
+                    * intermediate_size_per_partition
+                    // self.packed_factor,
+                ),
+            },
+            "w2_weight": {
+                "Flashinfer": (
+                    num_experts,
+                    hidden_size,
+                    intermediate_size_per_partition // self.packed_factor,
+                ),
+                "Marlin": (
+                    num_experts,
+                    intermediate_size_per_partition // self.packed_factor,
+                    hidden_size,
+                ),
+            },
+            "w2_scale": {
+                "Flashinfer": (num_experts, hidden_size, num_groups_w2),
+                "Marlin": (num_experts, num_groups_w2, hidden_size),
+            },
+            "w2_zp": {
+                "Marlin": (
+                    num_experts,
+                    num_groups_w2,
+                    hidden_size // self.packed_factor,
+                ),
+            },
+        }
+        backend_key = "Flashinfer" if is_flashinfer else "Marlin"
+        return shape_map[weight_name][backend_key]
+
+    def create_weights(
+        self,
+        layer: torch.nn.Module,
+        num_experts: int,
+        hidden_size: int,
+        intermediate_size_per_partition: int,
+        params_dtype: torch.dtype,
+        **extra_weight_attrs,
+    ):
+        intermediate_size_full = extra_weight_attrs.pop("intermediate_size_full")
+
+        # Will transpose the loaded weight along the
+        # intermediate and hidden dim sizes. Will
+        # shard for TP along the transposed dims
+        is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
+        extra_weight_attrs.update(
+            {"is_transposed": is_transposed, "quant_method": self.strategy}
+        )
+
+        w13_weight = torch.nn.Parameter(
+            torch.empty(
+                *self.get_weight_shape(
+                    "w13_weight",
+                    num_experts,
+                    hidden_size,
+                    intermediate_size_per_partition,
+                ),
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_packed", w13_weight)
+        set_weight_attrs(w13_weight, extra_weight_attrs)
+
+        w2_weight = torch.nn.Parameter(
+            torch.empty(
+                *self.get_weight_shape(
+                    "w2_weight",
+                    num_experts,
+                    hidden_size,
+                    intermediate_size_per_partition,
+                ),
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_packed", w2_weight)
+        set_weight_attrs(w2_weight, extra_weight_attrs)
+
+        # In the case where we have actorder/g_idx,
+        # we do not partition the w2 scales
+        load_full_w2 = self.actorder and self.group_size != -1
+        w2_scales_size = (
+            intermediate_size_full if load_full_w2 else intermediate_size_per_partition
+        )
+
+        self.is_k_full = (not self.actorder) or (
+            intermediate_size_per_partition == intermediate_size_full
+        )
+
+        if self.strategy == "channel":
+            num_groups_w2 = num_groups_w13 = 1
+            self.group_size = -1
+        else:
+            num_groups_w2 = w2_scales_size // self.group_size
+            num_groups_w13 = hidden_size // self.group_size
+
+        layer.num_groups_w13 = num_groups_w13
+        layer.num_groups_w2 = num_groups_w2
+
+        w13_scale = torch.nn.Parameter(
+            torch.ones(
+                *self.get_weight_shape(
+                    "w13_scale",
+                    num_experts,
+                    hidden_size,
+                    intermediate_size_per_partition,
+                    num_groups_w13=num_groups_w13,
+                ),
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_scale", w13_scale)
+        set_weight_attrs(w13_scale, extra_weight_attrs)
+
+        w2_scale = torch.nn.Parameter(
+            torch.ones(
+                *self.get_weight_shape(
+                    "w2_scale",
+                    num_experts,
+                    hidden_size,
+                    intermediate_size_per_partition,
+                    num_groups_w2=num_groups_w2,
+                ),
+                dtype=params_dtype,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_scale", w2_scale)
+        set_weight_attrs(w2_scale, extra_weight_attrs)
+        set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2})
+
+        if not self.symmetric:
+            w13_zp = torch.nn.Parameter(
+                torch.zeros(
+                    *self.get_weight_shape(
+                        "w13_zp",
+                        num_experts,
+                        hidden_size,
+                        intermediate_size_per_partition,
+                        num_groups_w13=num_groups_w13,
+                    ),
+                    dtype=torch.int32,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w13_weight_zero_point", w13_zp)
+            set_weight_attrs(w13_zp, extra_weight_attrs)
+
+            w2_zp = torch.nn.Parameter(
+                torch.zeros(
+                    *self.get_weight_shape(
+                        "w2_zp",
+                        num_experts,
+                        hidden_size,
+                        intermediate_size_per_partition,
+                        num_groups_w2=num_groups_w2,
+                    ),
+                    dtype=torch.int32,
+                ),
+                requires_grad=False,
+            )
+            layer.register_parameter("w2_weight_zero_point", w2_zp)
+            set_weight_attrs(w2_zp, extra_weight_attrs)
+
+        w2_weight_shape = torch.nn.Parameter(
+            torch.empty(num_experts, 2), requires_grad=False
+        )
+        layer.register_parameter("w2_weight_shape", w2_weight_shape)
+        set_weight_attrs(w2_weight_shape, extra_weight_attrs)
+        w13_weight_shape = torch.nn.Parameter(
+            torch.empty(num_experts, 2), requires_grad=False
+        )
+
+        layer.register_parameter("w13_weight_shape", w13_weight_shape)
+        set_weight_attrs(w13_weight_shape, extra_weight_attrs)
+
+        w13_g_idx = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_weight_g_idx", w13_g_idx)
+        set_weight_attrs(w13_g_idx, extra_weight_attrs)
+
+        w2_g_idx = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                intermediate_size_per_partition,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_weight_g_idx", w2_g_idx)
+        set_weight_attrs(w2_g_idx, extra_weight_attrs)
+
+        w13_g_idx_sort_indices = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                hidden_size,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w13_g_idx_sort_indices", w13_g_idx_sort_indices)
+        set_weight_attrs(w13_g_idx_sort_indices, extra_weight_attrs)
+
+        w2_g_idx_sort_indices = torch.nn.Parameter(
+            torch.empty(
+                num_experts,
+                intermediate_size_per_partition,
+                dtype=torch.int32,
+            ),
+            requires_grad=False,
+        )
+        layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices)
+        set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs)
+
+        layer.a13_scale = None
+        layer.a2_scale = None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Process weights using the shared oracle infrastructure
+        is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM
+        (
+            w13_qweight,
+            w2_qweight,
+            w13_scales,
+            w2_scales,
+            w13_g_idx_processed,
+            w2_g_idx_processed,
+            w13_g_idx_sort_indices,
+            w2_g_idx_sort_indices,
+            w13_qzeros,
+            w2_qzeros,
+            w13_input_global_scale,
+            w2_input_global_scale,
+            _,  # w13_bias
+            _,  # w2_bias
+        ) = convert_to_wna16_moe_kernel_format(
+            backend=self.wna16_backend,
+            layer=layer,
+            quant_config=self.weight_quant,
+            input_dtype=self.marlin_input_dtype,
+            w13=layer.w13_weight_packed,
+            w2=layer.w2_weight_packed,
+            w13_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            w13_g_idx=layer.w13_weight_g_idx,
+            w2_g_idx=layer.w2_weight_g_idx,
+            w13_qzeros=getattr(layer, "w13_weight_zero_point", None),
+            w2_qzeros=getattr(layer, "w2_weight_zero_point", None),
+        )
+
+        # Replace common parameters
+        replace_parameter(layer, "w13_weight_packed", w13_qweight)
+        replace_parameter(layer, "w2_weight_packed", w2_qweight)
+        replace_parameter(layer, "w13_weight_scale", w13_scales)
+        replace_parameter(layer, "w2_weight_scale", w2_scales)
+
+        if not self.symmetric:
+            replace_parameter(layer, "w13_weight_zero_point", w13_qzeros)
+            replace_parameter(layer, "w2_weight_zero_point", w2_qzeros)
+
+        # Marlin-specific parameters (not needed for Flashinfer)
+        if not is_flashinfer:
+            replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed)
+            replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed)
+            replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices)
+            replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices)
+
+            # Register input global scales if present
+            if w13_input_global_scale is not None:
+                layer.register_parameter(
+                    "w13_input_global_scale",
+                    torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
+                )
+            if w2_input_global_scale is not None:
+                layer.register_parameter(
+                    "w2_input_global_scale",
+                    torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
+                )
+
+            layer.workspace = marlin_make_workspace_new(
+                layer.w13_weight_g_idx.device, 4
+            )
+
+        # Alias packed weights to w13_weight/w2_weight for the modular kernel interface
+        layer.w13_weight = layer.w13_weight_packed
+        layer.w2_weight = layer.w2_weight_packed
+
+        assert self.experts_cls is not None
+        self.moe_quant_config = self.get_fused_moe_quant_config(layer)
+        assert self.moe_quant_config is not None
+
+        # Add Marlin-specific arguments
+        marlin_args: dict[str, Any] = {}
+        if not is_flashinfer:
+            marlin_args = {
+                "w13_g_idx": layer.w13_weight_g_idx,
+                "w2_g_idx": layer.w2_weight_g_idx,
+                "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices,
+                "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices,
+                "is_k_full": self.is_k_full,
+            }
+
+        self.moe_kernel = make_wna16_moe_kernel(
+            moe_quant_config=self.moe_quant_config,
+            moe_config=self.moe,
+            experts_cls=self.experts_cls,
+            routing_tables=layer._expert_routing_tables(),
+            **marlin_args,
+        )
+
+    def get_fused_moe_quant_config(
+        self, layer: torch.nn.Module
+    ) -> FusedMoEQuantConfig | None:
+        return make_wna16_moe_quant_config(
+            w1_scale=layer.w13_weight_scale,
+            w2_scale=layer.w2_weight_scale,
+            group_size=self.group_size,
+            num_bits=self.num_bits,
+            w1_zp=getattr(layer, "w13_weight_zero_point", None),
+            w2_zp=getattr(layer, "w2_weight_zero_point", None),
+            # SwiGLU/swigluoai gate params live on the layer; plumb them into the
+            # quant config so the fused activation (swigluoai_uninterleave on
+            # MiniMax-M3) receives gemm1_clamp_limit/alpha/beta. Mirrors the
+            # fp8/nvfp4/mxfp8 CT MoE methods which already do this.
+            gemm1_clamp_limit=getattr(layer, "swiglu_limit", None),
+            gemm1_alpha=getattr(layer, "swiglu_alpha", None),
+            gemm1_beta=getattr(layer, "swiglu_beta", None),
+        )
+
+    def apply_monolithic(
+        self,
+        layer: RoutedExperts,
+        x: torch.Tensor,
+        router_logits: torch.Tensor,
+        input_ids: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        assert self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply_monolithic(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            router_logits,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            num_expert_group=layer.num_expert_group,
+            topk_group=layer.topk_group,
+            e_score_correction_bias=layer.e_score_correction_bias,
+            routed_scaling_factor=layer.routed_scaling_factor,
+        )
+
+    def apply(
+        self,
+        layer: RoutedExperts,
+        x: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+        shared_experts: SharedExperts | None,
+        shared_experts_input: torch.Tensor | None,
+    ) -> torch.Tensor:
+        assert not self.is_monolithic
+        assert self.moe_kernel is not None
+        return self.moe_kernel.apply(
+            x,
+            layer.w13_weight,
+            layer.w2_weight,
+            topk_weights,
+            topk_ids,
+            activation=layer.activation,
+            global_num_experts=layer.global_num_experts,
+            expert_map=layer.expert_map,
+            apply_router_weight_on_input=layer.apply_router_weight_on_input,
+            shared_experts=shared_experts,
+            shared_experts_input=shared_experts_input,
+        )

vllm-patches/B_int_wna16.py

@@ -0,0 +1,917 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import sys
+from enum import Enum
+from typing import Any
+
+import torch
+from compressed_tensors.quantization import (
+    QuantizationArgs,
+)
+
+import vllm._custom_ops as ops
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.config import (
+    FusedMoEConfig,
+    FusedMoEQuantConfig,
+    int4_w4a16_moe_quant_config,
+    int8_w8a16_moe_quant_config,
+)
+from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
+    BatchedMarlinExperts,
+    MarlinExperts,
+    MarlinExpertsBase,
+)
+from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import (
+    TrtLlmMxint4ExpertsMonolithic,
+)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    marlin_act_int8_process_scales,
+    marlin_moe_permute_scales,
+    marlin_permute_bias,
+    moe_awq_to_marlin_zero_points,
+    moe_packed_to_marlin_zero_points,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+)
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+class WNA16MoEBackend(Enum):
+    MARLIN = "MARLIN"
+    BATCHED_MARLIN = "BATCHED_MARLIN"
+    FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
+    XPU = "XPU"
+
+
+def backend_to_kernel_cls(
+    backend: WNA16MoEBackend,
+) -> list[type[mk.FusedMoEExperts]]:
+    """Return the experts class for the given backend, or None for NONE."""
+    if backend == WNA16MoEBackend.MARLIN:
+        return [MarlinExperts]
+    elif backend == WNA16MoEBackend.BATCHED_MARLIN:
+        return [BatchedMarlinExperts]
+    elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
+        return [TrtLlmMxint4ExpertsMonolithic]
+    elif backend == WNA16MoEBackend.XPU:
+        from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
+            XPUExpertsWNA16,
+        )
+
+        return [XPUExpertsWNA16]
+    else:
+        raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")
+
+
+def _get_priority_backends() -> list[WNA16MoEBackend]:
+    """
+    Get available backends in priority order based on platform and config.
+    """
+    if current_platform.is_xpu():
+        return [WNA16MoEBackend.XPU]
+
+    _AVAILABLE_BACKENDS = [
+        WNA16MoEBackend.FLASHINFER_TRTLLM,
+        WNA16MoEBackend.MARLIN,
+        WNA16MoEBackend.BATCHED_MARLIN,
+    ]
+    return _AVAILABLE_BACKENDS
+
+
+def select_wna16_moe_backend(
+    config: FusedMoEConfig,
+    weight_key: QuantKey,
+) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
+    """Select the WNA16 MoE backend.
+
+    Args:
+        config: the shared ``FusedMoEConfig`` for this layer.
+        weight_key: The QuantKey describing the weight quantization.
+                    Must have int4 or int8 type.
+
+    Returns:
+        A tuple of (``WNA16MoEBackend``, experts class or ``None``).
+    """
+
+    activation_format = (
+        mk.FusedMoEActivationFormat.BatchedExperts
+        if config.moe_parallel_config.use_batched_activation_format
+        else mk.FusedMoEActivationFormat.Standard
+    )
+
+    def _make_log_backend(backend: WNA16MoEBackend):
+        return f"Using '{backend.value}' WNA16 MoE backend."
+
+    def _make_log_unsupported(backend: WNA16MoEBackend, reason: str | None) -> str:
+        if reason:
+            return (
+                f"WNA16 MoE backend '{backend.value}' does not support the "
+                f"deployment configuration since {reason}."
+            )
+        return (
+            f"WNA16 MoE backend '{backend.value}' does not support the "
+            "deployment configuration."
+        )
+
+    def _return_or_raise(
+        backend: WNA16MoEBackend,
+        config: FusedMoEConfig,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+        activation_format: mk.FusedMoEActivationFormat,
+    ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
+        reason: str | None = None
+        for k_cls in backend_to_kernel_cls(backend):
+            supported, reason = k_cls.is_supported_config(
+                k_cls, config, weight_key, activation_key, activation_format
+            )
+            if supported:
+                logger.info_once(_make_log_backend(backend), scope="local")
+                return backend, k_cls
+        raise ValueError(_make_log_unsupported(backend, reason))
+
+    # Select kernels in order of backend.
+    AVAILABLE_BACKENDS = _get_priority_backends()
+
+    for backend in AVAILABLE_BACKENDS:
+        activation_key = None  # always BF16 activation for WNA16 MoE
+        for k_cls in backend_to_kernel_cls(backend):
+            supported, reason = k_cls.is_supported_config(
+                k_cls, config, weight_key, activation_key, activation_format
+            )
+            if supported:
+                logger.info_once(_make_log_backend(backend), scope="local")
+                return backend, k_cls
+            else:
+                logger.debug_once(_make_log_unsupported(backend, reason), scope="local")
+
+    raise NotImplementedError(
+        "No WNA16 MoE backend supports the deployment configuration."
+    )
+
+
+def make_wna16_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    group_size: int,
+    num_bits: int,
+    w1_zp: torch.Tensor | None = None,
+    w2_zp: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    a1_gscale: torch.Tensor | None = None,
+    a2_gscale: torch.Tensor | None = None,
+    gemm1_clamp_limit: float | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+) -> FusedMoEQuantConfig:
+    """Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE."""
+    if num_bits == 4:
+        return int4_w4a16_moe_quant_config(
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            w1_bias=w1_bias,
+            w2_bias=w2_bias,
+            block_shape=[0, group_size],
+            a1_gscale=a1_gscale,
+            a2_gscale=a2_gscale,
+            gemm1_clamp_limit=gemm1_clamp_limit,
+            gemm1_alpha=gemm1_alpha,
+            gemm1_beta=gemm1_beta,
+        )
+    else:
+        assert num_bits == 8
+        return int8_w8a16_moe_quant_config(
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            w1_zp=w1_zp,
+            w2_zp=w2_zp,
+            w1_bias=w1_bias,
+            w2_bias=w2_bias,
+            block_shape=[0, group_size],
+            a1_gscale=a1_gscale,
+            a2_gscale=a2_gscale,
+            gemm1_clamp_limit=gemm1_clamp_limit,
+            gemm1_alpha=gemm1_alpha,
+            gemm1_beta=gemm1_beta,
+        )
+
+
+def make_wna16_moe_kernel(
+    moe_quant_config: FusedMoEQuantConfig,
+    moe_config: FusedMoEConfig,
+    experts_cls: type[mk.FusedMoEExperts],
+    is_k_full: bool = False,
+    w13_g_idx: torch.Tensor | None = None,
+    w2_g_idx: torch.Tensor | None = None,
+    w13_g_idx_sort_indices: torch.Tensor | None = None,
+    w2_g_idx_sort_indices: torch.Tensor | None = None,
+    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
+) -> mk.FusedMoEKernel:
+    from vllm.model_executor.layers.fused_moe.all2all_utils import (
+        maybe_make_prepare_finalize,
+    )
+    from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
+        XPUExpertsWNA16,
+    )
+
+    # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts
+    # and BatchedMarlinExperts
+    assert experts_cls in (
+        MarlinExperts,
+        BatchedMarlinExperts,
+        TrtLlmMxint4ExpertsMonolithic,
+        XPUExpertsWNA16,
+    )
+
+    is_monolithic = experts_cls.is_monolithic()
+
+    prepare_finalize = maybe_make_prepare_finalize(
+        moe=moe_config,
+        quant_config=moe_quant_config,
+        routing_tables=routing_tables,
+        allow_new_interface=True,
+        use_monolithic=is_monolithic,
+    )
+    assert prepare_finalize is not None
+
+    logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local")
+
+    extra_args: dict[str, Any] = {}
+    if issubclass(experts_cls, MarlinExpertsBase):
+        extra_args = {
+            "w13_g_idx": w13_g_idx,
+            "w2_g_idx": w2_g_idx,
+            "w13_g_idx_sort_indices": w13_g_idx_sort_indices,
+            "w2_g_idx_sort_indices": w2_g_idx_sort_indices,
+            "is_k_full": is_k_full,
+        }
+
+    if experts_cls is XPUExpertsWNA16:
+        assert (
+            prepare_finalize.activation_format == mk.FusedMoEActivationFormat.Standard
+        ), (
+            "XPUExpertsWNA16 only supports the Standard activation format; "
+            "xpu_fused_moe(is_int4=True) does not implement BatchedExperts."
+        )
+        experts: mk.FusedMoEExperts = XPUExpertsWNA16(
+            moe_config=moe_config,
+            quant_config=moe_quant_config,
+        )
+    elif (
+        prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts
+    ):
+        max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
+        assert max_num_tokens is not None
+        experts = experts_cls(
+            max_num_tokens=max_num_tokens,
+            num_dispatchers=prepare_finalize.num_dispatchers(),
+            moe_config=moe_config,
+            quant_config=moe_quant_config,
+            **extra_args,
+        )
+    else:
+        experts = experts_cls(
+            moe_config=moe_config,
+            quant_config=moe_quant_config,
+            **extra_args,
+        )
+
+    return mk.FusedMoEKernel(
+        prepare_finalize,
+        experts,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Per-backend weight post-processing
+# ---------------------------------------------------------------------------
+
+
+def _process_weights_flashinfer(
+    w13_qweight: torch.Tensor,
+    w2_qweight: torch.Tensor,
+    w13_scales: torch.Tensor,
+    w2_scales: torch.Tensor,
+    w13_g_idx: torch.Tensor,
+    w2_g_idx: torch.Tensor,
+    w13_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,  # w13_qweight
+    torch.Tensor,  # w2_qweight
+    torch.Tensor,  # w13_scales
+    torch.Tensor,  # w2_scales
+    torch.Tensor,  # w13_g_idx
+    torch.Tensor,  # w2_g_idx
+    torch.Tensor | None,  # w13_g_idx_sort_indices
+    torch.Tensor | None,  # w2_g_idx_sort_indices
+    torch.Tensor | None,  # w13_qzeros
+    torch.Tensor | None,  # w2_qzeros
+    torch.Tensor | None,  # w13_input_global_scale
+    torch.Tensor | None,  # w2_input_global_scale
+    torch.Tensor | None,  # w13_bias
+    torch.Tensor | None,  # w2_bias
+]:
+    """Flashinfer (TRT-LLM MXINT4) weight post-processing.
+
+    Steps
+    -----
+    1. Transform weights/scales via ``prepare_static_weights_for_trtllm_mxint4_moe``.
+    2. Return transformed tensors, passing through g_idx/bias unchanged.
+    """
+    from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
+        prepare_static_weights_for_trtllm_mxint4_moe,
+    )
+
+    dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
+        w13_qweight,
+        w13_scales,
+        w2_qweight,
+        w2_scales,
+    )
+
+    return (
+        dict_weights_mxint4["gemm1_weights"],
+        dict_weights_mxint4["gemm2_weights"],
+        dict_weights_mxint4["gemm1_scales"],
+        dict_weights_mxint4["gemm2_scales"],
+        w13_g_idx,
+        w2_g_idx,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        w13_bias,
+        w2_bias,
+    )
+
+
+def _process_weights_marlin(
+    layer: torch.nn.Module,
+    input_dtype: torch.dtype | None,
+    num_bits: int,
+    pack_factor: int,
+    group_size: int,
+    actorder: str | None,
+    w13_qweight: torch.Tensor,
+    w2_qweight: torch.Tensor,
+    w13_scales: torch.Tensor,
+    w2_scales: torch.Tensor,
+    w13_g_idx: torch.Tensor,
+    w2_g_idx: torch.Tensor,
+    w13_qzeros: torch.Tensor | None = None,
+    w2_qzeros: torch.Tensor | None = None,
+    w13_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,  # w13_qweight
+    torch.Tensor,  # w2_qweight
+    torch.Tensor,  # w13_scales
+    torch.Tensor,  # w2_scales
+    torch.Tensor,  # w13_g_idx
+    torch.Tensor,  # w2_g_idx
+    torch.Tensor,  # w13_g_idx_sort_indices
+    torch.Tensor,  # w2_g_idx_sort_indices
+    torch.Tensor | None,  # w13_qzeros
+    torch.Tensor | None,  # w2_qzeros
+    torch.Tensor | None,  # w13_input_global_scale
+    torch.Tensor | None,  # w2_input_global_scale
+    torch.Tensor | None,  # w13_bias
+    torch.Tensor | None,  # w2_bias
+]:
+    """Standard Marlin weight post-processing shared by MARLIN and
+    BATCHED_MARLIN backends.
+
+    Steps
+    -----
+    1. Optional FP8 preprocessing of packed weights / scales.
+    2. Sort / reset g_idx tensors for act-order handling.
+    3. Repack weights via ``gptq_marlin_moe_repack``.
+    4. Permute scales (and optionally extract INT8 global scales).
+    5. Permute bias tensors.
+    """
+    is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
+
+    marlin_w13_qweight: torch.Tensor
+    marlin_w2_qweight: torch.Tensor
+    marlin_w13_scales: torch.Tensor
+    marlin_w2_scales: torch.Tensor
+    w13_g_idx_sort_indices: torch.Tensor | None = None
+    w2_g_idx_sort_indices: torch.Tensor | None = None
+    w13_input_global_scale: torch.Tensor | None = None
+    w2_input_global_scale: torch.Tensor | None = None
+    w13_bias_out: torch.Tensor | None = None
+    w2_bias_out: torch.Tensor | None = None
+
+    # --- FP8 weight / scale adjustment ---
+    if input_dtype == torch.float8_e4m3fn:
+        # NOTE: for non-zp quantization format only
+        marlin_w13_qweight = ops.marlin_int4_fp8_preprocess(w13_qweight, inplace=False)
+        marlin_w2_qweight = ops.marlin_int4_fp8_preprocess(w2_qweight, inplace=False)
+        marlin_w13_scales = w13_scales.data * 512
+        marlin_w2_scales = w2_scales.data * 512
+    else:
+        marlin_w13_qweight = w13_qweight
+        marlin_w2_qweight = w2_qweight
+        marlin_w13_scales = w13_scales
+        marlin_w2_scales = w2_scales
+
+    # --- Process act_order (g_idx) ---
+    if actorder == "group":
+        num_experts = w13_g_idx.shape[0]
+        w13_g_idx_sort_indices = torch.empty_like(w13_g_idx)
+        w2_g_idx_sort_indices = torch.empty_like(w2_g_idx)
+        w13_sorted_g_idx = torch.empty_like(w13_g_idx)
+        w2_sorted_g_idx = torch.empty_like(w2_g_idx)
+        for e in range(num_experts):
+            w13_g_idx_sort_indices[e] = torch.argsort(w13_g_idx[e]).to(torch.int32)
+            w2_g_idx_sort_indices[e] = torch.argsort(w2_g_idx[e]).to(torch.int32)
+            w13_sorted_g_idx[e] = w13_g_idx[e][w13_g_idx_sort_indices[e]]
+            w2_sorted_g_idx[e] = w2_g_idx[e][w2_g_idx_sort_indices[e]]
+        w13_g_idx = w13_sorted_g_idx
+        w2_g_idx = w2_sorted_g_idx
+    else:
+        num_experts = w13_g_idx.shape[0]
+        device = w13_g_idx.device
+        w13_g_idx = torch.nn.Parameter(
+            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+            requires_grad=False,
+        )
+        w2_g_idx = torch.nn.Parameter(
+            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+            requires_grad=False,
+        )
+        w13_g_idx_sort_indices = torch.nn.Parameter(
+            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+            requires_grad=False,
+        )
+        w2_g_idx_sort_indices = torch.nn.Parameter(
+            torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+            requires_grad=False,
+        )
+
+    # --- Repack weights ---
+    marlin_w13_qweight = ops.gptq_marlin_moe_repack(
+        marlin_w13_qweight,
+        w13_g_idx_sort_indices,
+        marlin_w13_qweight.shape[1] * pack_factor,
+        marlin_w13_qweight.shape[2],
+        num_bits,
+        is_a_8bit=is_a_8bit,
+    )
+    marlin_w2_qweight = ops.gptq_marlin_moe_repack(
+        marlin_w2_qweight,
+        w2_g_idx_sort_indices,
+        marlin_w2_qweight.shape[1] * pack_factor,
+        marlin_w2_qweight.shape[2],
+        num_bits,
+        is_a_8bit=is_a_8bit,
+    )
+
+    # --- Permute scales ---
+    marlin_w13_scales = marlin_moe_permute_scales(
+        s=marlin_w13_scales,
+        size_k=layer.intermediate_size_per_partition,
+        size_n=marlin_w13_scales.shape[2],
+        group_size=group_size,
+        is_a_8bit=is_a_8bit,
+    )
+    group_size_or_pack_factor = group_size if group_size != -1 else pack_factor
+    marlin_w2_scales = marlin_moe_permute_scales(
+        s=marlin_w2_scales,
+        size_k=marlin_w2_scales.shape[1] * group_size_or_pack_factor,
+        size_n=marlin_w2_scales.shape[2],
+        group_size=group_size,
+        is_a_8bit=is_a_8bit,
+    )
+
+    if input_dtype == torch.int8:
+        if layer.num_groups_w13 > 1:
+            marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
+                marlin_w13_scales
+            )
+        if layer.num_groups_w2 > 1:
+            marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
+                marlin_w2_scales
+            )
+
+    # --- Permute zero points ---
+    if w13_qzeros is not None and w2_qzeros is not None:
+        w13_qzeros = moe_packed_to_marlin_zero_points(
+            w13_qzeros,
+            size_k=w13_qzeros.shape[1],
+            size_n=w13_qzeros.shape[2] * pack_factor,
+            num_bits=num_bits,
+            is_a_8bit=is_a_8bit,
+        )
+        w2_qzeros = moe_packed_to_marlin_zero_points(
+            w2_qzeros,
+            size_k=w2_qzeros.shape[1],
+            size_n=w2_qzeros.shape[2] * pack_factor,
+            num_bits=num_bits,
+            is_a_8bit=is_a_8bit,
+        )
+
+    # --- Permute bias ---
+    if w13_bias is not None:
+        w13_bias_out = marlin_permute_bias(w13_bias)
+    if w2_bias is not None:
+        w2_bias_out = marlin_permute_bias(w2_bias)
+
+    return (
+        marlin_w13_qweight,
+        marlin_w2_qweight,
+        marlin_w13_scales,
+        marlin_w2_scales,
+        w13_g_idx,
+        w2_g_idx,
+        w13_g_idx_sort_indices,
+        w2_g_idx_sort_indices,
+        w13_qzeros,
+        w2_qzeros,
+        w13_input_global_scale,
+        w2_input_global_scale,
+        w13_bias_out,
+        w2_bias_out,
+    )
+
+
+def _process_awq_weights_marlin(
+    layer: torch.nn.Module,
+    weight_bits: int,
+    pack_factor: int,
+    group_size: int,
+    input_dtype: torch.dtype | None,
+    w13_qweight: torch.Tensor,
+    w2_qweight: torch.Tensor,
+    w13_scales: torch.Tensor,
+    w2_scales: torch.Tensor,
+    w13_qzeros: torch.Tensor,
+    w2_qzeros: torch.Tensor,
+    w13_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,  # w13_qweight
+    torch.Tensor,  # w2_qweight
+    torch.Tensor,  # w13_scales
+    torch.Tensor,  # w2_scales
+    torch.Tensor | None,  # w13_g_idx
+    torch.Tensor | None,  # w2_g_idx
+    torch.Tensor | None,  # w13_g_idx_sort_indices
+    torch.Tensor | None,  # w2_g_idx_sort_indices
+    torch.Tensor | None,  # w13_qzeros
+    torch.Tensor | None,  # w2_qzeros
+    torch.Tensor | None,  # w13_input_global_scale
+    torch.Tensor | None,  # w2_input_global_scale
+    torch.Tensor | None,  # w13_bias
+    torch.Tensor | None,  # w2_bias
+]:
+    """AWQ-specific Marlin weight post-processing.
+
+    AWQ checkpoints use a different packing order than GPTQ, so they need
+    AWQ-specific weight repacking and zero-point conversion before Marlin runs.
+    """
+    num_experts = w13_qweight.shape[0]
+    device = w13_qweight.device
+    is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1
+    w13_input_global_scale: torch.Tensor | None = None
+    w2_input_global_scale: torch.Tensor | None = None
+    w13_bias_out: torch.Tensor | None = None
+    w2_bias_out: torch.Tensor | None = None
+
+    if input_dtype == torch.float8_e4m3fn:
+        ops.marlin_int4_fp8_preprocess(
+            w13_qweight.view(-1, w13_qweight.size(2)),
+            w13_qzeros.view(-1, w13_qzeros.size(2)),
+            inplace=True,
+        )
+        ops.marlin_int4_fp8_preprocess(
+            w2_qweight.view(-1, w2_qweight.size(2)),
+            w2_qzeros.view(-1, w2_qzeros.size(2)),
+            inplace=True,
+        )
+        w13_scales = w13_scales.data * 512
+        w2_scales = w2_scales.data * 512
+
+    w13_g_idx_sort_indices = torch.nn.Parameter(
+        torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+        requires_grad=False,
+    )
+    w2_g_idx_sort_indices = torch.nn.Parameter(
+        torch.empty((num_experts, 0), dtype=torch.int32, device=device),
+        requires_grad=False,
+    )
+
+    marlin_w13_qweight = ops.awq_marlin_moe_repack(
+        w13_qweight,
+        w13_g_idx_sort_indices,
+        size_k=w13_qweight.shape[1],
+        size_n=w13_qweight.shape[2] * pack_factor,
+        num_bits=weight_bits,
+        is_a_8bit=is_a_8bit,
+    )
+    marlin_w2_qweight = ops.awq_marlin_moe_repack(
+        w2_qweight,
+        w2_g_idx_sort_indices,
+        size_k=w2_qweight.shape[1],
+        size_n=w2_qweight.shape[2] * pack_factor,
+        num_bits=weight_bits,
+        is_a_8bit=is_a_8bit,
+    )
+
+    marlin_w13_scales = marlin_moe_permute_scales(
+        s=w13_scales,
+        size_k=layer.intermediate_size_per_partition,
+        size_n=w13_scales.shape[2],
+        group_size=group_size,
+        is_a_8bit=is_a_8bit,
+    )
+    if input_dtype == torch.int8 and layer.num_groups_w13 > 1:
+        marlin_w13_scales, w13_input_global_scale = marlin_act_int8_process_scales(
+            marlin_w13_scales
+        )
+
+    marlin_w2_scales = marlin_moe_permute_scales(
+        s=w2_scales,
+        size_k=layer.intermediate_size_per_partition,
+        size_n=w2_scales.shape[2],
+        group_size=group_size,
+        is_a_8bit=is_a_8bit,
+    )
+    if input_dtype == torch.int8 and layer.num_groups_w2 > 1:
+        marlin_w2_scales, w2_input_global_scale = marlin_act_int8_process_scales(
+            marlin_w2_scales
+        )
+
+    marlin_w13_qzeros = moe_awq_to_marlin_zero_points(
+        w13_qzeros,
+        size_k=w13_qzeros.shape[1],
+        size_n=w13_qzeros.shape[2] * pack_factor,
+        num_bits=weight_bits,
+        is_a_8bit=is_a_8bit,
+    )
+    marlin_w2_qzeros = moe_awq_to_marlin_zero_points(
+        w2_qzeros,
+        size_k=w2_qzeros.shape[1],
+        size_n=w2_qzeros.shape[2] * pack_factor,
+        num_bits=weight_bits,
+        is_a_8bit=is_a_8bit,
+    )
+
+    if w13_bias is not None:
+        w13_bias_out = marlin_permute_bias(w13_bias)
+    if w2_bias is not None:
+        w2_bias_out = marlin_permute_bias(w2_bias)
+
+    return (
+        marlin_w13_qweight,
+        marlin_w2_qweight,
+        marlin_w13_scales,
+        marlin_w2_scales,
+        None,
+        None,
+        w13_g_idx_sort_indices,
+        w2_g_idx_sort_indices,
+        marlin_w13_qzeros,
+        marlin_w2_qzeros,
+        w13_input_global_scale,
+        w2_input_global_scale,
+        w13_bias_out,
+        w2_bias_out,
+    )
+
+
+def _process_weights_xpu(
+    layer: torch.nn.Module,
+    quant_config: QuantizationConfig,
+    w13_qweight: torch.Tensor,
+    w2_qweight: torch.Tensor,
+    w13_scales: torch.Tensor,
+    w2_scales: torch.Tensor,
+    w13_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,  # w13_qweight
+    torch.Tensor,  # w2_qweight
+    torch.Tensor,  # w13_scales
+    torch.Tensor,  # w2_scales
+    torch.Tensor | None,  # w13_bias
+    torch.Tensor | None,  # w2_bias
+]:
+    """Repack GPTQ-format INT4 MoE weights into the layout
+    `vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects:
+
+        w13: [E, 2*N, K] int4 (uint8 storage [E, 2*N, K // 2])
+        w13_scales: [E, 2*N, K // group_size] params_dtype
+        w2:  [E, K, N]   int4 (uint8 storage [E, K, N // 2])
+        w2_scales:  [E, K, N // group_size]   params_dtype
+
+    Input GPTQ layout from FusedMoE.weight_loader:
+        w13: [E, K // 8, 2*N] int32 (8 nibbles per int32 along the input dim)
+        w13_scales: [E, K // group_size, 2*N] params_dtype
+        w2:  [E, N // 8, K] int32
+        w2_scales:  [E, N // group_size, K] params_dtype
+
+    Transpose dim 1 ↔ dim 2 then view int32 → uint8 to recover sequential
+    int4-packed bytes along the input dim. Each packed int32 holds 8 nibbles
+    `(n7<<28)|(n6<<24)|...|(n1<<4)|n0` in ascending K order; on a
+    little-endian host the int32→uint8 view exposes them as bytes
+    `[n1<<4|n0, n3<<4|n2, n5<<4|n4, n7<<4|n6]`, i.e. two nibbles per byte
+    with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True)
+    expects this convention; on a big-endian host the byte order reverses
+    and the kernel would silently miscompute, so we hard-fail.
+    """
+    del layer, quant_config  # unused — kept for parity with the marlin helper
+
+    if sys.byteorder != "little":
+        raise NotImplementedError(
+            "_process_weights_xpu requires a little-endian host: the GPTQ "
+            "int32 → uint8 nibble repack relies on LE byte ordering."
+        )
+
+    w13_xpu = w13_qweight.transpose(1, 2).contiguous().view(torch.uint8)
+    w2_xpu = w2_qweight.transpose(1, 2).contiguous().view(torch.uint8)
+    w13_scales_xpu = w13_scales.transpose(1, 2).contiguous()
+    w2_scales_xpu = w2_scales.transpose(1, 2).contiguous()
+
+    return (
+        w13_xpu,
+        w2_xpu,
+        w13_scales_xpu,
+        w2_scales_xpu,
+        w13_bias,
+        w2_bias,
+    )
+
+
+def convert_to_wna16_moe_kernel_format(
+    backend: WNA16MoEBackend,
+    layer: torch.nn.Module,
+    quant_config: QuantizationConfig | QuantizationArgs | None,
+    input_dtype: torch.dtype | None,
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    w13_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    w13_g_idx: torch.Tensor | None = None,
+    w2_g_idx: torch.Tensor | None = None,
+    w13_qzeros: torch.Tensor | None = None,
+    w2_qzeros: torch.Tensor | None = None,
+    w13_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+) -> tuple[
+    torch.Tensor,  # w13_qweight
+    torch.Tensor,  # w2_qweight
+    torch.Tensor,  # w13_scales
+    torch.Tensor,  # w2_scales
+    torch.Tensor | None,  # w13_g_idx
+    torch.Tensor | None,  # w2_g_idx
+    torch.Tensor | None,  # w13_g_idx_sort_indices
+    torch.Tensor | None,  # w2_g_idx_sort_indices
+    torch.Tensor | None,  # w13_qzeros
+    torch.Tensor | None,  # w2_qzeros
+    torch.Tensor | None,  # w13_input_global_scale
+    torch.Tensor | None,  # w2_input_global_scale
+    torch.Tensor | None,  # w13_bias
+    torch.Tensor | None,  # w2_bias
+]:
+    """Dispatch weight post-processing to the appropriate per-backend handler.
+
+    To add a new backend, implement a ``_process_weights_<name>`` helper and
+    add a branch here.
+
+    Args:
+        backend: the selected ``WNA16MoEBackend``.
+        layer: the ``FusedMoE`` layer whose parameters are being prepared.
+        quant_config: the ``QuantizationConfig`` for this layer.
+        input_dtype: optional activation dtype, usually should be 16 bit.
+    """
+    if backend in (
+        WNA16MoEBackend.MARLIN,
+        WNA16MoEBackend.BATCHED_MARLIN,
+    ):
+        from vllm.model_executor.layers.quantization.auto_gptq import (
+            AutoGPTQConfig,
+        )
+        from vllm.model_executor.layers.quantization.awq_marlin import (
+            AWQMarlinConfig,
+        )
+
+        if isinstance(quant_config, AWQMarlinConfig):
+            if w13_qzeros is None or w2_qzeros is None:
+                raise ValueError("AWQ Marlin MoE requires zero-point tensors.")
+
+            weight_bits = quant_config.weight_bits
+            pack_factor = quant_config.pack_factor
+            group_size = quant_config.group_size
+
+            return _process_awq_weights_marlin(
+                layer,
+                weight_bits,
+                pack_factor,
+                group_size,
+                input_dtype,
+                w13,
+                w2,
+                w13_scale,
+                w2_scale,
+                w13_qzeros,
+                w2_qzeros,
+                w13_bias,
+                w2_bias,
+            )
+        elif isinstance(quant_config, AutoGPTQConfig):
+            num_bits = quant_config.quant_type.size_bits
+            pack_factor = quant_config.pack_factor
+            group_size = quant_config.group_size
+            actorder = "group" if quant_config.desc_act else None
+        elif isinstance(quant_config, QuantizationArgs):
+            num_bits = quant_config.num_bits
+            pack_factor = 32 // quant_config.num_bits
+            group_size = quant_config.group_size
+            actorder = quant_config.actorder
+        else:
+            raise TypeError(
+                "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or "
+                f"QuantizationArgs, got {type(quant_config).__name__}."
+            )
+        if w13_g_idx is None or w2_g_idx is None:
+            raise ValueError("GPTQ Marlin MoE requires g_idx tensors.")
+        return _process_weights_marlin(
+            layer,
+            input_dtype,
+            num_bits,
+            pack_factor,
+            group_size,
+            actorder,
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            w13_g_idx,
+            w2_g_idx,
+            w13_qzeros,
+            w2_qzeros,
+            w13_bias,
+            w2_bias,
+        )
+    elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
+        return _process_weights_flashinfer(
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            w13_g_idx,
+            w2_g_idx,
+            w13_bias,
+            w2_bias,
+        )
+    elif backend == WNA16MoEBackend.XPU:
+        assert quant_config is not None
+        (
+            w13_xpu,
+            w2_xpu,
+            w13_scale_xpu,
+            w2_scale_xpu,
+            w13_bias_out,
+            w2_bias_out,
+        ) = _process_weights_xpu(
+            layer,
+            quant_config,
+            w13,
+            w2,
+            w13_scale,
+            w2_scale,
+            w13_bias,
+            w2_bias,
+        )
+        empty = torch.empty((0,), dtype=torch.int32, device=w13.device)
+        return (
+            w13_xpu,
+            w2_xpu,
+            w13_scale_xpu,
+            w2_scale_xpu,
+            empty,  # w13_g_idx
+            empty,  # w2_g_idx
+            empty,  # w13_g_idx_sort_indices
+            empty,  # w2_g_idx_sort_indices
+            None,  # w13_qzeros — sym int4 on XPU has none; kernel does uint4b8→s4
+            None,  # w2_qzeros
+            None,  # w13_input_global_scale
+            None,  # w2_input_global_scale
+            w13_bias_out,
+            w2_bias_out,
+        )
+    else:
+        raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}")

vllm-patches/C_config.py

@@ -0,0 +1,1436 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import Union
+
+import torch
+
+from vllm.config import ParallelConfig, SchedulerConfig
+from vllm.config.kernel import MoEBackend
+from vllm.distributed import get_dp_group, get_pcp_group, get_tensor_model_parallel_rank
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import (
+    OCP_MX_DTYPES,
+    OCP_MX_Scheme,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_triton_kernels
+from vllm.utils.math_utils import cdiv
+
+logger = init_logger(__name__)
+
+if has_triton_kernels():
+    try:
+        from triton_kernels.matmul_ogs import PrecisionConfig
+    except (ImportError, AttributeError) as e:
+        logger.error(
+            "Failed to import Triton kernels. Please make sure your triton "
+            "version is compatible. Error: %s",
+            e,
+        )
+
+
+def _get_config_dtype_str(
+    dtype: torch.dtype,
+    use_fp8_w8a8: bool = False,
+    use_fp8_w8a16: bool = False,
+    use_int8_w8a16: bool = False,
+    use_int4_w4a16: bool = False,
+    ocp_mx_scheme: str | None = None,
+) -> str | None:
+    """
+    Return a string used to construct the filename that contains the
+    tuning info for a particular quantization scheme.  See
+    try_get_optimal_moe_config in fused_moe.py.
+    """
+    if use_fp8_w8a8:
+        return "fp8_w8a8"
+    elif use_fp8_w8a16:
+        return "fp8_w8a16"
+    elif use_int8_w8a16:
+        return "int8_w8a16"
+    elif use_int4_w4a16:
+        return "int4_w4a16"
+    elif ocp_mx_scheme is not None:
+        # The output of this function is passed to `try_get_optimal_moe_config`,
+        # and as we only simulate OCP MX execution in fused_moe for now,
+        # we will NOT look for `*,dtype=w_mxfp4_a_mxfp4.json` for now.
+        return None
+    elif dtype == torch.float:
+        # avoiding cases where kernel fails when float32 MoE
+        # use fp16/bfloat16 configs
+        return "float32"
+    return None
+
+
+def _quant_flags_to_group_shape(
+    quant_dtype: torch.dtype | str | None,
+    per_act_token_quant: bool,
+    per_out_ch_quant: bool,
+    block_shape: list[int] | None,
+) -> tuple[GroupShape | None, GroupShape | None]:
+    """
+    Convert MoE quantization flags into more generic GroupShapes.
+    """
+    a_shape: GroupShape | None
+    w_shape: GroupShape | None
+    if block_shape is not None:
+        assert not per_act_token_quant
+        assert not per_out_ch_quant
+        # TODO(bnell): this is not quite right for activations since first
+        # dim should be 1.
+        a_shape = GroupShape(row=block_shape[0], col=block_shape[1])
+        w_shape = GroupShape(row=block_shape[0], col=block_shape[1])
+    else:
+        w_shape = None
+        a_shape = None if quant_dtype is None else GroupShape.PER_TENSOR
+
+        if per_act_token_quant:
+            a_shape = GroupShape.PER_TOKEN
+
+        if per_out_ch_quant:
+            w_shape = GroupShape.PER_TOKEN
+
+    return a_shape, w_shape
+
+
+# The type of method in top-K routing
+# Please keep this in sync with the counterpart defined in https://github.com/flashinfer-ai/flashinfer/blob/main/include/flashinfer/trtllm/fused_moe/runner.h
+class RoutingMethodType(IntEnum):
+    # Default: Softmax -> TopK
+    Default = (0,)
+    # Renormalize: TopK -> Softmax
+    Renormalize = (1,)
+    # DeepSeekV3: Sigmoid -> RoutingBiasAdd -> Top2 in group -> Top4 groups
+    # -> Top8 experts from the Top4 groups
+    DeepSeekV3 = (2,)
+    # Llama4: Top1 -> Sigmoid
+    Llama4 = (3,)
+    # RenormalizeNaive: Softmax -> TopK -> Renormalize
+    RenormalizeNaive = (4,)
+    # TopK: TopK (no softmax)
+    TopK = (5,)
+    # SigmoidRenorm: Sigmoid -> TopK -> Renormalize (divide by sum of top-K)
+    SigmoidRenorm = (6,)
+    # MiniMax2: Sigmoid + Bias -> TopK -> ScaledSumNormalize
+    # (routeScale=1.0, epsilon=1e-20)
+    MiniMax2 = (7,)
+    # Sigmoid: Sigmoid -> TopK (no renormalization)
+    Sigmoid = (8,)
+    # Unspecified
+    Unspecified = (9,)
+    # other routing types (not passed to FlashInfer kernels)
+    # Deepseek V4 -> sqrtsoftplus + Bias + Normalize
+    DeepseekV4 = (100,)
+    Custom = (101,)
+    Simulated = (102,)
+
+
+def get_routing_method_type(
+    scoring_func: str,
+    top_k: int,
+    renormalize: bool,
+    num_expert_group: int | None,
+    has_e_score_bias: bool,
+    routed_scaling_factor: float | None = 1.0,
+) -> RoutingMethodType:
+    if scoring_func == "sqrtsoftplus":
+        # DeepSeek V4 uses sqrtsoftplus routing with optional routing bias
+        # and top-k renormalization.
+        if renormalize:
+            return RoutingMethodType.DeepseekV4
+        else:
+            return RoutingMethodType.Unspecified
+
+    if has_e_score_bias:
+        if scoring_func == "sigmoid":
+            if not renormalize:
+                return RoutingMethodType.Unspecified
+            if (num_expert_group or 0) > 0:
+                return RoutingMethodType.DeepSeekV3
+            if routed_scaling_factor in (None, 1.0):
+                return RoutingMethodType.MiniMax2
+            return RoutingMethodType.Unspecified
+        else:
+            return RoutingMethodType.Unspecified
+
+    if scoring_func == "sigmoid":
+        if renormalize:
+            return RoutingMethodType.SigmoidRenorm
+        return RoutingMethodType.Sigmoid
+
+    if scoring_func == "softmax":
+        if renormalize:
+            return RoutingMethodType.RenormalizeNaive
+        else:
+            return RoutingMethodType.Default
+
+    return RoutingMethodType.Unspecified
+
+
+@dataclass
+class FusedMoEQuantDesc:
+    """
+    A quantization descriptor for fused MoE ops. This class can describe
+    either activations or weights.
+    """
+
+    # The quantized type of this parameters.  None means unquantized or
+    # already quantized.
+    # TODO (bnell): use scalar_type instead of Union.
+    dtype: torch.dtype | str | None = None
+
+    # A field that describes the quantization group shape, from quant_utils.py.
+    #  * (-1, -1)   for per-tensor quantization
+    #  * (1, -1)    for per-row quantization
+    #  * (-1, 1)    for per-column quantization
+    #  * (128, 128) for 128x128 deepseek style block quantization
+    #  * (1, 128)   for deepseek style activation quantization
+    #               (i.e. per-token-per-group)
+    shape: GroupShape | None = None
+
+    # Quantization scales.
+    # TODO(bnell): maybe put PrecisionConfigs in subclass of QuantDesc?
+    scale: Union[torch.Tensor, "PrecisionConfig", None] = None
+
+    # Quantization alphas or gscales, used for nvfp4 types.
+    # W4A8 FP8: used for per-channel scales
+    # TODO(bnell): put some of these in subclasses
+    alpha_or_gscale: torch.Tensor | None = None
+
+    # Zero points for int4/int8 types
+    zp: torch.Tensor | None = None
+
+    # Biases for GPT triton MoE
+    bias: torch.Tensor | None = None
+
+
+# TODO(bnell): have subclasses for specific moe methods?
+# e.g. for specific arguments bias, precision, etc.
+@dataclass
+class FusedMoEQuantConfig:
+    """
+    The FusedMoEQuantConfig contains all the quantization parameters for
+    a single FusedMoEMethodBase operation.  It consists of four
+    FusedMoEQuantDescs, one for each activation and set of weights.
+
+    Each FusedMoEMethodBase must implement a get_fused_moe_quant_config
+    method to construct a FusedMoEQuantConfig for use with that class.
+
+    FusedMoEQuant configs are only used for modular kernels, fused_experts
+    (from fused_moe.py), cutlass_moe_fp[48], rocm_aiter_fused_experts and
+    triton_kernel_moe_forward.  Other MoE methods can ignore the
+    FusedMoEQuantConfig (for now) and hardcode it to None.
+
+    There are currently some restrictions on what can be expressed:
+    - Most MoE ops only support similar quantization strategies for
+      each parameter, e.g. both weights must have the same GroupShape
+      and both activations must share the same GroupShape.  One exception to
+      this is the cutlass moe which allows per channel quantization on the
+      outputs.  Note: this restrictions are not always rigorously checked.
+    - Not all fused MoE functions support all the parameters, e.g. zero points,
+      global scales, alphas and biases are not universally supported.
+    - Fully general GroupShapes are not allowed.  Activations only support
+      per token, per tensor or K-blocked.
+    - Weights are not required to have a GroupShape since they have already
+      been quantized.
+
+    Other notes:
+    - PrecisionConfigs are specific to GPT OSS Triton.
+    - As a follow up it would probably make sense to subclass FusedMoEQuantDesc
+      or FusedMoEQuantConfig for particular FusedMoEMethodBase subclasses
+      so that only the required quantization parameters are used/stored.
+    """
+
+    # TODO(bnell) make sure a1_scales/a2_scales don't interfere with chunking
+    _a1: FusedMoEQuantDesc
+    _a2: FusedMoEQuantDesc
+    _w1: FusedMoEQuantDesc
+    _w2: FusedMoEQuantDesc
+    is_scale_swizzled: bool = True
+
+    # MXFP4-specific TRTLLM parameters for SwiGLU activation clamping.
+    # These correspond to gemm1_alpha, gemm1_beta, gemm1_clamp_limit
+    # in TrtLlmMxfp4ExpertsBase.
+    gemm1_alpha: float | None = None
+    gemm1_beta: float | None = None
+    gemm1_clamp_limit: float | None = None
+
+    mx_alignment: int = 0
+
+    def __post_init__(self):
+        assert not self.per_act_token_quant or self.block_shape is None, (
+            "illegal quantization"
+        )
+
+    #
+    # Convenience accessors for various properties.
+    #
+
+    @property
+    def quant_dtype(self) -> torch.dtype | str | None:
+        return self._a1.dtype
+
+    @property
+    def weight_quant_dtype(self) -> torch.dtype | str | None:
+        return self._w1.dtype
+
+    @property
+    def is_quantized(self) -> bool:
+        return self.quant_dtype is not None
+
+    @property
+    def is_per_act_token(self) -> bool:
+        return self._a1.shape == GroupShape.PER_TOKEN
+
+    @property
+    def per_act_token_quant(self) -> bool:
+        return self._a1.shape == GroupShape.PER_TOKEN
+
+    @property
+    def per_out_ch_quant(self) -> bool:
+        return self._w1.shape == GroupShape.PER_TOKEN
+
+    @property
+    def is_per_tensor(self) -> bool:
+        return self._a1.shape == GroupShape.PER_TENSOR
+
+    @property
+    def block_shape(self) -> list[int] | None:
+        if (
+            self._a1.shape is not None
+            and self._a1.shape != GroupShape.PER_TENSOR
+            and self._a1.shape != GroupShape.PER_TOKEN
+        ):
+            return [self._a1.shape.row, self._a1.shape.col]
+        else:
+            return None
+
+    @property
+    def is_block_quantized(self) -> bool:
+        return self.block_shape is not None
+
+    @property
+    def a1_scale(self) -> torch.Tensor | None:
+        assert self._a1.scale is None or isinstance(self._a1.scale, torch.Tensor)
+        return self._a1.scale
+
+    @property
+    def a1_gscale(self) -> torch.Tensor | None:
+        return self._a1.alpha_or_gscale
+
+    @property
+    def a2_scale(self) -> torch.Tensor | None:
+        assert self._a2.scale is None or isinstance(self._a2.scale, torch.Tensor)
+        return self._a2.scale
+
+    @property
+    def a2_gscale(self) -> torch.Tensor | None:
+        return self._a2.alpha_or_gscale
+
+    @property
+    def w1_scale(self) -> torch.Tensor | None:
+        assert self._w1.scale is None or isinstance(self._w1.scale, torch.Tensor)
+        return self._w1.scale
+
+    @property
+    def w1_zp(self) -> torch.Tensor | None:
+        return self._w1.zp
+
+    @property
+    def w1_bias(self) -> torch.Tensor | None:
+        return self._w1.bias
+
+    @property
+    def w1_precision(self) -> "PrecisionConfig | None":
+        assert self._w1.scale is None or isinstance(self._w1.scale, PrecisionConfig)
+        return self._w1.scale
+
+    @property
+    def g1_alphas(self) -> torch.Tensor | None:
+        return self._w1.alpha_or_gscale
+
+    @property
+    def w2_scale(self) -> torch.Tensor | None:
+        assert self._w2.scale is None or isinstance(self._w2.scale, torch.Tensor)
+        return self._w2.scale
+
+    @property
+    def w2_zp(self) -> torch.Tensor | None:
+        return self._w2.zp
+
+    @property
+    def w2_bias(self) -> torch.Tensor | None:
+        return self._w2.bias
+
+    @property
+    def w2_precision(self) -> "PrecisionConfig | None":
+        assert self._w2.scale is None or isinstance(self._w2.scale, PrecisionConfig)
+        return self._w2.scale
+
+    @property
+    def g2_alphas(self) -> torch.Tensor | None:
+        return self._w2.alpha_or_gscale
+
+    @property
+    def use_fp8_w8a8(self) -> bool:
+        return self.quant_dtype == current_platform.fp8_dtype()
+
+    @property
+    def use_int8_w8a8(self) -> bool:
+        return self.quant_dtype == torch.int8
+
+    @property
+    def use_int8_w8a16(self) -> bool:
+        return self._a1.dtype is None and self._w1.dtype == torch.int8
+
+    @property
+    def use_fp8_w8a16(self) -> bool:
+        return self._a1.dtype is None and self._w1.dtype == current_platform.fp8_dtype()
+
+    @property
+    def use_int4_w4a16(self) -> bool:
+        return self._a1.dtype is None and self._w1.dtype == "int4"
+
+    @property
+    def use_nvfp4_w4a16(self) -> bool:
+        return self._a1.dtype is None and self._w1.dtype == "nvfp4"
+
+    @property
+    def ocp_mx_scheme(self) -> str | None:
+        if not hasattr(self, "_ocp_mx_scheme"):
+            if (self._a1.dtype is not None and not isinstance(self._a1.dtype, str)) or (
+                self._w1.dtype is not None and not isinstance(self._w1.dtype, str)
+            ):
+                self._ocp_mx_scheme = None
+            else:
+                ocp_mx_scheme = OCP_MX_Scheme.from_quant_dtype(
+                    self._a1.dtype, self._w1.dtype
+                )
+
+                if ocp_mx_scheme is not None:
+                    ocp_mx_scheme = ocp_mx_scheme.value
+
+                self._ocp_mx_scheme = ocp_mx_scheme
+
+        return self._ocp_mx_scheme
+
+    @property
+    def use_mxfp4_w4a16(self) -> bool:
+        return self._a1.dtype is None and self._w1.dtype == "mxfp4"
+
+    @property
+    def use_mxfp4_w4a4(self) -> bool:
+        return self._a1.dtype == "mxfp4" and self._w1.dtype == "mxfp4"
+
+    @property
+    def use_nvfp4_w4a4(self) -> bool:
+        return self.quant_dtype == "nvfp4"
+
+    @property
+    def use_mxfp4_w4a8(self) -> bool:
+        return self._a1.dtype == "fp8" and self._w1.dtype == "mxfp4"
+
+    def config_name(self, dtype: torch.dtype) -> str | None:
+        """
+        Return a string used to construct the filename that contains the
+        tuning info for a particular quantization scheme.  See
+        try_get_optimal_moe_config in fused_moe.py.
+        """
+        return _get_config_dtype_str(
+            use_fp8_w8a8=self.use_fp8_w8a8,
+            use_fp8_w8a16=self.use_fp8_w8a16,
+            use_int8_w8a16=self.use_int8_w8a16,
+            use_int4_w4a16=self.use_int4_w4a16,
+            ocp_mx_scheme=self.ocp_mx_scheme,
+            dtype=dtype,
+        )
+
+    def scale_shape(
+        self,
+        max_tokens: int,
+        hidden_dim: int,
+    ) -> tuple[int, int] | None:
+        """
+        Construct the proper activation scale shape for this
+        config.
+        """
+        if self.is_quantized:
+            if self.is_block_quantized:
+                assert self.block_shape is not None
+                _, block_k = self.block_shape
+                k_tiles = cdiv(hidden_dim, block_k)
+                return (max_tokens, k_tiles)
+            elif self.is_per_act_token:
+                return (max_tokens, 1)
+            else:
+                return (1, 1)
+        else:
+            return None
+
+    def batched_scale_shape(
+        self,
+        num_experts: int,
+        max_tokens: int,
+        hidden_dim: int,
+    ) -> tuple[int, int, int] | None:
+        """
+        Construct the proper activation batched scale shape for this
+        config, e.g. (num experts, *scale_shape).
+        """
+        if self.is_quantized:
+            scale_shape = self.scale_shape(max_tokens, hidden_dim)
+            assert scale_shape is not None
+            return (num_experts, *scale_shape)
+        else:
+            return None
+
+    @staticmethod
+    def make(
+        quant_dtype: torch.dtype | str | None = None,
+        per_act_token_quant: bool = False,
+        per_out_ch_quant: bool = False,
+        block_shape: list[int] | None = None,
+        w1_scale: Union[torch.Tensor, "PrecisionConfig", None] = None,
+        w2_scale: Union[torch.Tensor, "PrecisionConfig", None] = None,
+        a1_scale: torch.Tensor | None = None,
+        a2_scale: torch.Tensor | None = None,
+        g1_alphas: torch.Tensor | None = None,
+        g2_alphas: torch.Tensor | None = None,
+        a1_gscale: torch.Tensor | None = None,
+        a2_gscale: torch.Tensor | None = None,
+        w1_bias: torch.Tensor | None = None,
+        w2_bias: torch.Tensor | None = None,
+        w1_zp: torch.Tensor | None = None,
+        w2_zp: torch.Tensor | None = None,
+        weight_dtype: torch.dtype | str | None = None,
+        is_scale_swizzled: bool = True,
+        gemm1_alpha: float | None = None,
+        gemm1_beta: float | None = None,
+        gemm1_clamp_limit: float | None = None,
+    ) -> "FusedMoEQuantConfig":
+        """
+        General builder function for a FusedMoEQuantConfig.
+        - quant_dtype: Optional quantization type. None if activations are
+          unquantized or quantized prior to calling.  Note: "nvfp4", "mxfp4",
+          "mxfp6_e3m2", "mxfp6_e2m3" are the only valid string values
+          for quant_dtype.
+        - per_act_token_quant: Activations have per token quantization.
+        - per_out_ch_quant: Outputs have per channel quantization. (only
+          for cutlass).
+        - block_shape: Optional block size for block-wise quantization.
+          Incompatible with per_act_token and per_out_ch quant.
+        - w1_scale: Optional scale to be used for w1.
+        - w2_scale: Optional scale to be used for w2.
+        - a1_scale: Optional scale to be used for a1.
+        - a2_scale: Optional scale to be used for a2.
+        - g1_alphas: Optional global quantization scales for w1 (for nvfp4).
+                     Optional per-channel scales for w1 (for W4A8 FP8).
+                     Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8).
+        - g2_alphas: Optional global quantization scales for w2 (for nvfp4).
+                     Optional per-channel scales for w2 (for W4A8 FP8).
+                     Optional dq scale i.e. w_scale * a_scale (for W8A8 fp8).
+        - a1_gscale: Optional global quantization scales for a1 (1.0 /a2_scale).
+        - a2_gscale: Optional global quantization scales for a2 (1.0 /a2_scale).
+
+        - w1_bias: Optional biases for w1 (GPT OSS Triton).
+        - w2_bias: Optional biases for w1 (GPT OSS Triton).
+        - w1_zp: Optional w1 zero points for int4/int8 quantization.
+        - w2_zp: Optional w2 zero points for int4/int8 quantization.
+        - is_scale_swizzled: Whether the activation scale-factor layout is
+          swizzled. Pass through to the underlying quantization kernel for
+          dtypes that distinguish layouts (nvfp4, mxfp8). Defaults to True.
+        - gemm1_alpha: Optional MXFP4 TRTLLM SwiGLU alpha parameter.
+        - gemm1_beta: Optional MXFP4 TRTLLM SwiGLU beta parameter.
+        - gemm1_clamp_limit: Optional MXFP4 TRTLLM SwiGLU clamp limit.
+        """
+        assert not isinstance(quant_dtype, str) or quant_dtype in {
+            "nvfp4",
+            "mxfp4",
+            "mxfp6_e3m2",
+            "mxfp6_e2m3",
+            "mxfp8",
+        }
+        assert not isinstance(weight_dtype, str) or weight_dtype in {
+            "nvfp4",
+            "mxfp4",
+            "mxfp6_e3m2",
+            "mxfp6_e2m3",
+            "int4",
+            "mxfp8",
+        }
+
+        if weight_dtype is None:
+            weight_dtype = quant_dtype
+
+        a_shape, w_shape = _quant_flags_to_group_shape(
+            quant_dtype, per_act_token_quant, per_out_ch_quant, block_shape
+        )
+        quant_config = FusedMoEQuantConfig(
+            _a1=FusedMoEQuantDesc(quant_dtype, a_shape, a1_scale, a1_gscale),
+            _a2=FusedMoEQuantDesc(quant_dtype, a_shape, a2_scale, a2_gscale),
+            _w1=FusedMoEQuantDesc(
+                weight_dtype, w_shape, w1_scale, g1_alphas, w1_zp, w1_bias
+            ),
+            _w2=FusedMoEQuantDesc(
+                weight_dtype, w_shape, w2_scale, g2_alphas, w2_zp, w2_bias
+            ),
+            is_scale_swizzled=is_scale_swizzled,
+            gemm1_alpha=gemm1_alpha,
+            gemm1_beta=gemm1_beta,
+            gemm1_clamp_limit=gemm1_clamp_limit,
+        )
+        assert quant_config.per_act_token_quant == per_act_token_quant
+        assert quant_config.per_out_ch_quant == per_out_ch_quant
+        assert quant_config.block_shape == block_shape
+        return quant_config
+
+
+def fp8_w8a8_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    per_act_token_quant: bool = False,
+    per_out_ch_quant: bool = False,
+    block_shape: list[int] | None = None,
+    a1_gscale: torch.Tensor | None = None,
+    a2_gscale: torch.Tensor | None = None,
+    g1_alphas: torch.Tensor | None = None,
+    g2_alphas: torch.Tensor | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for fp8 activations and fp8 weights.
+    """
+    return FusedMoEQuantConfig.make(
+        current_platform.fp8_dtype(),
+        w1_scale=w1_scale,
+        g1_alphas=g1_alphas,
+        w2_scale=w2_scale,
+        g2_alphas=g2_alphas,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        a1_scale=a1_scale,
+        a1_gscale=a1_gscale,
+        a2_scale=a2_scale,
+        a2_gscale=a2_gscale,
+        per_act_token_quant=per_act_token_quant,
+        per_out_ch_quant=per_out_ch_quant,
+        block_shape=block_shape,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def int8_w8a8_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    a1_scale: torch.Tensor | None,
+    a2_scale: torch.Tensor | None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    per_act_token_quant: bool = False,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for int8 activations and int8 weights.
+    """
+    return FusedMoEQuantConfig.make(
+        torch.int8,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        per_act_token_quant=per_act_token_quant,
+        per_out_ch_quant=False,
+        block_shape=None,
+    )
+
+
+def gptq_marlin_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    weight_bits: int,
+    group_size: int,
+    w1_zp: torch.Tensor | None = None,
+    w2_zp: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+):
+    """
+    Construct a quant config for gptq marlin quantization.
+    """
+    from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
+
+    w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size)
+
+    # Activations are NOT quantized for GPTQ (fp16/bf16)
+    a_shape = w_shape  # Same as weight shape for alignment
+
+    # Determine weight dtype
+    if weight_bits == 4:
+        weight_dtype = "int4"
+    elif weight_bits == 8:
+        weight_dtype = torch.int8
+    else:
+        raise ValueError(f"Unsupported weight_bits: {weight_bits}")
+
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc(dtype=None, shape=a_shape),
+        _a2=FusedMoEQuantDesc(dtype=None, shape=a_shape),
+        _w1=FusedMoEQuantDesc(weight_dtype, w_shape, w1_scale, None, w1_zp, w1_bias),
+        _w2=FusedMoEQuantDesc(weight_dtype, w_shape, w2_scale, None, w2_zp, w2_bias),
+    )
+
+
+def mxfp4_w4a16_moe_quant_config(
+    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for unquantized activations and mxfp4 weights.
+    """
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc(),
+        _a2=FusedMoEQuantDesc(),
+        _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
+        _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def mxfp4_mxfp8_moe_quant_config(
+    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+    mx_alignment: int = 0,
+    is_scale_swizzled: bool = True,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for mxfp4 activations and mxfp4 weights.
+    """
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc("mxfp8"),
+        _a2=FusedMoEQuantDesc("mxfp8"),
+        _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
+        _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+        mx_alignment=mx_alignment,
+        is_scale_swizzled=is_scale_swizzled,
+    )
+
+
+def mxfp4_w4a8_moe_quant_config(
+    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for fp8 activations and mxfp4 weights.
+    """
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc("fp8", None, a1_scale, None, None, None),
+        _a2=FusedMoEQuantDesc("fp8", None, a2_scale, None, None, None),
+        _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
+        _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def ocp_mx_moe_quant_config(
+    quant_dtype: str,
+    w1_scale: Union[torch.Tensor, "PrecisionConfig"],
+    w2_scale: Union[torch.Tensor, "PrecisionConfig"],
+    weight_dtype: str | None = None,
+    a1_scale: torch.Tensor | None = None,
+    a2_scale: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for mxfp4 activations and mxfp4 weights.
+    """
+    assert quant_dtype in OCP_MX_DTYPES
+    return FusedMoEQuantConfig.make(
+        quant_dtype=quant_dtype,
+        weight_dtype=weight_dtype,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        a1_scale=a1_scale,
+        a2_scale=a2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=block_shape,
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def nvfp4_moe_quant_config(
+    g1_alphas: torch.Tensor,
+    g2_alphas: torch.Tensor,
+    a1_gscale: torch.Tensor,
+    a2_gscale: torch.Tensor,
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    is_scale_swizzled: bool = True,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for mxfp4 activations and nvp4 weights.
+    """
+    return FusedMoEQuantConfig.make(
+        "nvfp4",
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        w1_bias=w1_bias,
+        w2_bias=w2_bias,
+        a1_gscale=a1_gscale,
+        a2_gscale=a2_gscale,
+        g1_alphas=g1_alphas,
+        g2_alphas=g2_alphas,
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=None,
+        is_scale_swizzled=is_scale_swizzled,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def mxfp4_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for MXFP4 x MXFP4 MoE.
+    MXFP4 uses block scaling only (E8M0 scales, 32-element groups), with no
+    separate alphas / global activation scales in this config.
+    """
+    return FusedMoEQuantConfig.make(
+        "mxfp4",
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        per_act_token_quant=False,
+        per_out_ch_quant=False,
+        block_shape=None,
+    )
+
+
+def nvfp4_w4a16_moe_quant_config(
+    g1_alphas: torch.Tensor,
+    g2_alphas: torch.Tensor,
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for 16-but activations and nvp4 weights.
+    """
+    return FusedMoEQuantConfig.make(
+        quant_dtype=None,
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        g1_alphas=g1_alphas,
+        g2_alphas=g2_alphas,
+        weight_dtype="nvfp4",
+    )
+
+
+def int4_w4a16_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    w1_zp: torch.Tensor | None = None,
+    w2_zp: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+    a1_gscale: torch.Tensor | None = None,
+    a2_gscale: torch.Tensor | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for 16-bit float activations and int4 weights.
+    """
+    group_shape = GroupShape(*block_shape) if block_shape is not None else None
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale),
+        _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale),
+        _w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp, w1_bias),
+        _w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp, w2_bias),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def fp8_w8a16_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for 16-bit float activations and fp8 weights.
+    """
+    group_shape = GroupShape(*block_shape) if block_shape is not None else None
+    fp8_dtype = current_platform.fp8_dtype()
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc(),
+        _a2=FusedMoEQuantDesc(),
+        _w1=FusedMoEQuantDesc(
+            fp8_dtype,
+            group_shape,
+            w1_scale,
+            None,
+            None,
+            w1_bias,
+        ),
+        _w2=FusedMoEQuantDesc(
+            fp8_dtype,
+            group_shape,
+            w2_scale,
+            None,
+            None,
+            w2_bias,
+        ),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def int8_w8a16_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    w1_zp: torch.Tensor | None = None,
+    w2_zp: torch.Tensor | None = None,
+    w1_bias: torch.Tensor | None = None,
+    w2_bias: torch.Tensor | None = None,
+    block_shape: list[int] | None = None,
+    a1_gscale: torch.Tensor | None = None,
+    a2_gscale: torch.Tensor | None = None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for 16-bit float activations and int8 weights.
+    """
+    group_shape = GroupShape(*block_shape) if block_shape is not None else None
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a1_gscale),
+        _a2=FusedMoEQuantDesc(shape=group_shape, alpha_or_gscale=a2_gscale),
+        _w1=FusedMoEQuantDesc(torch.int8, group_shape, w1_scale, None, w1_zp, w1_bias),
+        _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+def int4_w4afp8_moe_quant_config(
+    w1_scale: torch.Tensor,
+    w2_scale: torch.Tensor,
+    g1_alphas: torch.Tensor,
+    g2_alphas: torch.Tensor,
+    per_act_token_quant: bool = False,
+    per_out_ch_quant: bool = False,
+    block_shape: list[int] | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for fp8 activations and int4 weights.
+    """
+    return FusedMoEQuantConfig.make(
+        torch.float8_e4m3fn,  # quant dtype for activations
+        w1_scale=w1_scale,
+        w2_scale=w2_scale,
+        g1_alphas=g1_alphas,
+        g2_alphas=g2_alphas,
+        per_act_token_quant=per_act_token_quant,
+        per_out_ch_quant=per_out_ch_quant,
+        block_shape=block_shape,
+        weight_dtype="int4",  # weight dtype for weights
+    )
+
+
+def biased_moe_quant_config(
+    w1_bias: torch.Tensor | None,
+    w2_bias: torch.Tensor | None,
+    gemm1_alpha: float | None = None,
+    gemm1_beta: float | None = None,
+    gemm1_clamp_limit: float | None = None,
+) -> FusedMoEQuantConfig:
+    """
+    Construct a quant config for unquantized activations with biases.
+
+    gemm1_alpha/gemm1_beta/gemm1_clamp_limit carry the SwiGLU gate params
+    through to the fused activation kernel (e.g. swigluoai_uninterleave).
+    """
+    return FusedMoEQuantConfig(
+        _a1=FusedMoEQuantDesc(),
+        _a2=FusedMoEQuantDesc(),
+        _w1=FusedMoEQuantDesc(bias=w1_bias),
+        _w2=FusedMoEQuantDesc(bias=w2_bias),
+        gemm1_alpha=gemm1_alpha,
+        gemm1_beta=gemm1_beta,
+        gemm1_clamp_limit=gemm1_clamp_limit,
+    )
+
+
+# A FusedMoEQuantConfig constant for an unquantized MoE op.
+FUSED_MOE_UNQUANTIZED_CONFIG: FusedMoEQuantConfig = FusedMoEQuantConfig.make()
+
+
+@dataclass
+class FusedMoEParallelConfig:
+    tp_size: int
+    pcp_size: int
+    dp_size: int
+    ep_size: int
+    tp_rank: int
+    pcp_rank: int
+    dp_rank: int
+    ep_rank: int
+    sp_size: int
+
+    use_ep: bool  # whether to use EP or not
+    all2all_backend: str  # all2all backend for MoE communication
+    enable_eplb: bool  # whether to enable expert load balancing
+
+    @property
+    def is_sequence_parallel(self) -> bool:
+        return self.sp_size > 1
+
+    @property
+    def use_all2all_kernels(self):
+        return self.dp_size > 1 and self.use_ep
+
+    @property
+    def use_deepep_ht_kernels(self):
+        return (
+            self.use_all2all_kernels
+            and self.all2all_backend == "deepep_high_throughput"
+        )
+
+    @property
+    def use_deepep_ll_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend == "deepep_low_latency"
+
+    @property
+    def use_fi_nvl_two_sided_kernels(self):
+        return self.use_all2all_kernels and (
+            self.all2all_backend == "flashinfer_all2allv"
+            or self.all2all_backend == "flashinfer_nvlink_two_sided"
+        )
+
+    @property
+    def use_fi_nvl_one_sided_kernels(self):
+        return (
+            self.use_all2all_kernels
+            and self.all2all_backend == "flashinfer_nvlink_one_sided"
+        )
+
+    @property
+    def use_batched_activation_format(self):
+        return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels
+
+    @property
+    def needs_round_robin_routing_tables(self):
+        return self.use_deepep_ll_kernels or self.use_nixl_ep_kernels
+
+    @property
+    def use_ag_rs_all2all_kernels(self):
+        return (
+            self.use_all2all_kernels
+            and self.all2all_backend == "allgather_reducescatter"
+        )
+
+    @property
+    def use_mori_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend in (
+            "mori_high_throughput",
+            "mori_low_latency",
+        )
+
+    @property
+    def use_nixl_ep_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend == "nixl_ep"
+
+    @property
+    def use_deepep_v2_kernels(self):
+        return self.use_all2all_kernels and self.all2all_backend == "deepep_v2"
+
+    @staticmethod
+    def flatten_tp_across_dp_and_pcp(
+        tp_size: int, dp_size: int, dp_rank: int, pcp_size: int, pcp_rank: int
+    ) -> tuple[int, int]:
+        tp_rank = 0 if tp_size == 1 else get_tensor_model_parallel_rank()
+        # There are actually dp_size * pcp_size * tp_size devices.
+        # Update tp_size and tp_rank so we shard across all devices.
+        flatten_tp_size = dp_size * pcp_size * tp_size
+        flatten_tp_rank = dp_rank * pcp_size * tp_size + pcp_rank * tp_size + tp_rank
+        return flatten_tp_size, flatten_tp_rank
+
+    @staticmethod
+    def make(
+        tp_size_: int,
+        pcp_size_: int,
+        dp_size_: int,
+        sp_size_: int,
+        vllm_parallel_config: ParallelConfig,
+    ) -> "FusedMoEParallelConfig":
+        """
+        Determine MoE parallel configuration. Based on the input `tp_size_`,
+        `dp_size_` and vllm's parallel config, determine what
+        level's of parallelism to use in the fused moe layer.
+
+        Args:
+            tp_size_ (int): `tp_size` passed into the FusedMoE constructor.
+            pcp_size_ (int): `pcp_size` passed into the FusedMoE constructor.
+            dp_size_ (int): `dp_size` passed into the FusedMoE constructor.
+            vllm_parallel_config (ParallelConfig): vLLM's parallel config
+                object which contains the `enable_expert_parallel` flag.
+
+        Examples:
+            When there is no parallelism requested,
+            i.e. `tp_size_` = `pcp_size_` = `dp_size_` = 1, we simply return the sizes
+            unaltered and the ranks set to 0.
+
+            Expert Parallelism is considered only when either `dp_size_`, `pcp_size_` or
+            `tp_size_` is non trivial.
+
+            Note that PCP serves the same function as DP here.
+
+            When TP = 2, DP(PCP) = 1 and EP = False, the configuration on different
+            devices:
+
+            - device 0 : TP = {2, 0} DP = {1, 0} EP = {1, 0} //
+                legend : {size, rank}
+            - device 1 : TP = {2, 1} DP = {1, 0} EP = {1, 0}
+            - Comment : Tensors are sharded across 2 devices.
+
+            When TP = 1, DP(PCP) = 2 and EP = False, the configuration on different
+                devices:
+
+            - device 0 : TP = {2, 0} DP = {2, 0} EP = {1, 0}
+            - device 1 : TP = {2, 1} DP = {2, 1} EP = {1, 0}
+            - Comment: There are 2 engine instances and the tensors are sharded
+                across 2 decvices.
+
+            When TP = 2, DP(PCP) = 2 and EP = False, the configuration on different
+                devices:
+
+            - device 0: TP = {4, 0} DP = {2, 0} EP = {1, 0}
+            - device 1: TP = {4, 1} DP = {2, 0} EP = {1, 0}
+            - device 2: TP = {4, 2} DP = {2, 1} EP = {1, 0}
+            - device 3: TP = {4, 3} DP = {2, 1} EP = {1, 0}
+            - Comment: There are 2 engine instances and the tensors are sharded
+                across 4 devices.
+
+            When, TP = 2, DP(PCP) = 1 and EP = True, the configuration on different
+                devices:
+
+            - device 0: TP = {1, 0} DP = {1, 0} EP = {2, 0}
+            - device 1: TP = {1, 0} DP = {1, 0} EP = {2, 1}
+            - Comment: The experts are split between the 2 devices.
+
+            When, TP = 1, DP(PCP) = 2 and EP = True, the configuration on different
+                devices:
+
+            - device 0: TP = {1, 0} DP = {2, 0} EP = {2, 0}
+            - device 1: TP = {1, 0} DP = {2, 1} EP = {2, 1}
+            - Comment: There are 2 engine instances and the experts are split
+                between the 2 devices.
+
+            When TP = 2, DP(PCP) = 2 and EP = True, the configuration on different
+                devices:
+
+            - device 0: TP = {1, 0} DP = {2, 0} EP = {4, 0}
+            - device 1: TP = {1, 0} DP = {2, 0} EP = {4, 1}
+            - device 2: TP = {1, 0} DP = {2, 1} EP = {4, 2}
+            - device 3: TP = {1, 0} DP = {2, 1} EP = {4, 3}
+            - Comment: There are 2 engine instances and the experts are split
+                between the 4 devices.
+        """
+        use_ep = (
+            dp_size_ * pcp_size_ * tp_size_ > 1
+            and vllm_parallel_config.enable_expert_parallel
+        )
+
+        dp_size = dp_size_
+        dp_rank = get_dp_group().rank_in_group if dp_size > 1 else 0
+        pcp_size = pcp_size_
+        pcp_rank = get_pcp_group().rank_in_group if pcp_size > 1 else 0
+        tp_size, tp_rank = FusedMoEParallelConfig.flatten_tp_across_dp_and_pcp(
+            tp_size_, dp_size_, dp_rank, pcp_size_, pcp_rank
+        )
+
+        if not use_ep:
+            return FusedMoEParallelConfig(
+                tp_size=tp_size,
+                tp_rank=tp_rank,
+                pcp_size=pcp_size,
+                pcp_rank=pcp_rank,
+                dp_size=dp_size,
+                dp_rank=dp_rank,
+                ep_size=1,
+                ep_rank=0,
+                sp_size=sp_size_,
+                use_ep=False,
+                all2all_backend=vllm_parallel_config.all2all_backend,
+                enable_eplb=vllm_parallel_config.enable_eplb,
+            )
+        # DP + EP / TP + EP / DP + TP + EP
+        assert use_ep
+        # In EP, each device owns a set of experts fully. There is no tensor
+        # parallel update tp_size, tp_rank, ep_size and ep_rank to reflect that.
+        ep_size = tp_size
+        ep_rank = tp_rank
+        return FusedMoEParallelConfig(
+            tp_size=1,
+            tp_rank=0,
+            pcp_size=pcp_size,
+            pcp_rank=pcp_rank,
+            dp_size=dp_size,
+            dp_rank=dp_rank,
+            ep_size=ep_size,
+            ep_rank=ep_rank,
+            sp_size=sp_size_,
+            use_ep=True,
+            all2all_backend=vllm_parallel_config.all2all_backend,
+            enable_eplb=vllm_parallel_config.enable_eplb,
+        )
+
+    @classmethod
+    def make_no_parallel(cls) -> "FusedMoEParallelConfig":
+        """For usage in CI/CD and testing."""
+        return FusedMoEParallelConfig(
+            tp_size=1,
+            tp_rank=0,
+            pcp_size=1,
+            pcp_rank=0,
+            dp_size=1,
+            dp_rank=0,
+            ep_size=1,
+            ep_rank=0,
+            sp_size=1,
+            use_ep=False,
+            all2all_backend="allgather_reducescatter",
+            enable_eplb=False,
+        )
+
+
+# Adapted from pplx-kernels tests/all_to_all_utils.py
+@dataclass
+class FusedMoEConfig:
+    num_experts: int
+    experts_per_token: int
+    hidden_dim: int
+    intermediate_size: int
+    num_local_experts: int
+    num_logical_experts: int
+    activation: MoEActivation
+    device: torch.device | str
+    routing_method: RoutingMethodType
+    moe_parallel_config: FusedMoEParallelConfig
+
+    # The activation type.
+    in_dtype: torch.dtype
+
+    # Defaults to in_dtype if not specified.
+    router_logits_dtype: torch.dtype | None = None
+
+    # Defaults to hidden_dim if not specified.
+    hidden_dim_unpadded: int | None = None
+    # Defaults to intermediate_size_per_partition if not specified.
+    intermediate_size_per_partition_unpadded: int | None = None
+
+    moe_backend: MoEBackend = "auto"
+    max_num_tokens: int = SchedulerConfig.DEFAULT_MAX_NUM_BATCHED_TOKENS_FOR_BATCHED_DP
+    has_bias: bool = False
+    is_lora_enabled: bool = False
+
+    # SwiGLU clamp limit. When set, backends that do not implement the clamp
+    # are filtered out by `FusedMoEExperts.is_supported_config` so the oracle
+    # cannot silently select one and drop the clamp.
+    swiglu_limit: float | None = None
+    swiglu_alpha: float | None = None
+    swiglu_beta: float | None = None
+
+    max_capture_size: int = 0
+
+    # Set by __post_init__
+    intermediate_size_per_partition: int = -1
+    rocm_aiter_fmoe_enabled: bool = False
+    aiter_fmoe_shared_expert_enabled: bool = False
+
+    def __post_init__(self):
+        from vllm._aiter_ops import rocm_aiter_ops
+
+        tp_size = self.moe_parallel_config.tp_size
+        assert self.intermediate_size % tp_size == 0
+        self.intermediate_size_per_partition = self.intermediate_size // tp_size
+
+        if self.dp_size > 1:
+            logger.debug_once(
+                "Using FusedMoEConfig::max_num_tokens=%d", self.max_num_tokens
+            )
+
+        assert self.max_num_tokens > 0
+
+        if self.router_logits_dtype is None:
+            self.router_logits_dtype = self.in_dtype
+
+        if self.hidden_dim_unpadded is None:
+            self.hidden_dim_unpadded = self.hidden_dim
+        if self.intermediate_size_per_partition_unpadded is None:
+            self.intermediate_size_per_partition_unpadded = (
+                self.intermediate_size_per_partition
+            )
+
+        if self.is_act_and_mul:
+            self.rocm_aiter_fmoe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
+            self.aiter_fmoe_shared_expert_enabled = (
+                rocm_aiter_ops.is_fusion_moe_shared_experts_enabled()
+            )
+
+        if self.use_mori_kernels:
+            assert self.rocm_aiter_fmoe_enabled, (
+                "Mori needs to be used with aiter fused_moe for now."
+            )
+            assert not self.aiter_fmoe_shared_expert_enabled, (
+                "Mori does not support fusion shared expert now. "
+                "Turn it off by setting VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=0"
+            )
+
+        if not self.is_act_and_mul and not (
+            current_platform.is_cuda_alike() or current_platform.is_xpu()
+        ):
+            raise NotImplementedError(
+                "is_act_and_mul=False is supported only for CUDA, XPU and ROCm for now"
+            )
+
+    @property
+    def is_act_and_mul(self) -> bool:
+        return self.activation.is_gated
+
+    @property
+    def tp_size(self):
+        return self.moe_parallel_config.tp_size
+
+    @property
+    def dp_size(self):
+        return self.moe_parallel_config.dp_size
+
+    @property
+    def pcp_size(self):
+        return self.moe_parallel_config.pcp_size
+
+    @property
+    def ep_size(self):
+        return self.moe_parallel_config.ep_size
+
+    @property
+    def sp_size(self):
+        return self.moe_parallel_config.sp_size
+
+    @property
+    def is_sequence_parallel(self):
+        return self.moe_parallel_config.is_sequence_parallel
+
+    @property
+    def tp_rank(self):
+        return self.moe_parallel_config.tp_rank
+
+    @property
+    def dp_rank(self):
+        return self.moe_parallel_config.dp_rank
+
+    @property
+    def pcp_rank(self):
+        return self.moe_parallel_config.pcp_rank
+
+    @property
+    def ep_rank(self):
+        return self.moe_parallel_config.ep_rank
+
+    @property
+    def use_ep(self):
+        return self.moe_parallel_config.use_ep
+
+    @property
+    def use_deepep_ht_kernels(self):
+        return self.moe_parallel_config.use_deepep_ht_kernels
+
+    @property
+    def use_deepep_ll_kernels(self):
+        return self.moe_parallel_config.use_deepep_ll_kernels
+
+    @property
+    def use_mori_kernels(self):
+        return self.moe_parallel_config.use_mori_kernels
+
+    @property
+    def use_fi_nvl_two_sided_kernels(self):
+        return self.moe_parallel_config.use_fi_nvl_two_sided_kernels
+
+    @property
+    def use_fi_nvl_one_sided_kernels(self):
+        return self.moe_parallel_config.use_fi_nvl_one_sided_kernels
+
+    @property
+    def use_ag_rs_all2all_kernels(self):
+        return self.moe_parallel_config.use_ag_rs_all2all_kernels
+
+    @property
+    def use_nixl_ep_kernels(self):
+        return self.moe_parallel_config.use_nixl_ep_kernels
+
+    @property
+    def use_deepep_v2_kernels(self):
+        return self.moe_parallel_config.use_deepep_v2_kernels
+
+    @property
+    def needs_round_robin_routing_tables(self):
+        return self.moe_parallel_config.needs_round_robin_routing_tables

vllm-patches/PATCHES.md

@@ -0,0 +1,32 @@
+# Serving MiniMax-M3-AWQ-int4 on vLLM (A100 / Ampere)
+
+This is an **int4 AWQ** quant (group-size 128, routed-experts-only, GPTQ-mse refinement,
+Gemma-norm fold) of MiniMax-M3. It needs **three small vLLM patches** to serve coherently:
+they make the WNA16-Marlin fused-MoE kernel honor M3s clamped SwiGLU and dequantize the
+int4 packing correctly. **Stock vLLM will load the weights but emit incoherent output.**
+
+## Requirements
+- Base image: `vllm/vllm-openai:minimax-m3`
+- `--block-size 128` is **mandatory** (M3 MSA sparse-attention index-cache alignment).
+- 8x A100 80GB (TP8) for long context; smaller TP works at reduced context (e.g. 2x TP4 ~ up to 64K).
+
+## Patches (mount read-only over the image)
+Inside the container, `V=/usr/local/lib/python3.12/dist-packages/vllm`:
+
+| file in this folder | mount target |
+|---|---|
+| `A_wna16_marlin.py` | `$V/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py` |
+| `B_int_wna16.py`    | `$V/model_executor/layers/fused_moe/oracle/int_wna16.py` |
+| `C_config.py`       | `$V/model_executor/layers/fused_moe/config.py` |
+
+## Launch (TP8)
+See `launch_m3_awq.sh`. Core flags: `--tensor-parallel-size 8 --block-size 128
+--enable-expert-parallel --max-model-len 262144 --gpu-memory-utilization 0.95
+--reasoning-parser minimax_m3 --tool-call-parser minimax_m3 --enable-auto-tool-choice
+--trust-remote-code`.
+
+## Notes
+- KV dtype = bf16/auto. **fp8-KV is NOT available on A100** for this model: block-128 + fp8
+  routes to a Blackwell-only FlashInfer (trtllm-gen) kernel and fails at init.
+- Sampling: `temperature 1.0, top_p 0.95, top_k 40`.
+- The patches are modified Apache-2.0 vLLM source files.

vllm-patches/launch_m3_awq.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+# M3 AWQ-int4 at TP8 — identical to the proven RTN launch, only the model mount differs.
+set -u
+V=/usr/local/lib/python3.12/dist-packages/vllm
+P=/srv/models/m3-clamp-patch
+MODEL=${1:-/srv/quant/m3-awq-int4-v2}
+docker rm -f vllm-m3-awq 2>/dev/null
+docker run -d --name vllm-m3-awq \
+  --init \
+  --gpus all \
+  --ipc=host --shm-size=16g \
+  -p 127.0.0.1:8005:8000 \
+  -e HF_HUB_OFFLINE=1 \
+  -v $MODEL:/model:ro \
+  -v $P/A_wna16_marlin.py:$V/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py:ro \
+  -v $P/B_int_wna16.py:$V/model_executor/layers/fused_moe/oracle/int_wna16.py:ro \
+  -v $P/C_config.py:$V/model_executor/layers/fused_moe/config.py:ro \
+  vllm/vllm-openai:minimax-m3 \
+  --model /model --served-model-name m3 \
+  --tensor-parallel-size 8 \
+  --block-size 128 \
+  --enable-expert-parallel \
+  --enable-prefix-caching \
+  --max-model-len 262144 \
+  --max-num-seqs 64 \
+  --gpu-memory-utilization 0.95 \
+  --safetensors-load-strategy=prefetch \
+  --reasoning-parser minimax_m3 --tool-call-parser minimax_m3 --enable-auto-tool-choice \
+  --trust-remote-code
+echo "launched: $(docker ps --filter name=vllm-m3-awq --format '{{.Names}} {{.Status}}')"