vllm.model_executor.layers.fused_moe.oracle.int_wna16

Classes:

• WNA16MoEKernelOracle –

  WNA16 MoE backend selection for one deployment configuration.

Functions:

• backend_to_kernel_cls –

  Return the experts class for the given backend, or None for NONE.

• convert_to_wna16_moe_kernel_format –

  Dispatch weight post-processing to the appropriate per-backend handler.

• make_wna16_moe_quant_config –

  Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE.

• map_wna16_backend –

  Map user's MoEBackend to WNA16MoEBackend.

• select_wna16_moe_backend –

  Select the WNA16 MoE backend.

WNA16MoEKernelOracle

Bases: MoEKernelOracle[WNA16MoEBackend]

WNA16 MoE backend selection for one deployment configuration.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

class WNA16MoEKernelOracle(MoEKernelOracle[WNA16MoEBackend]):
    """WNA16 MoE backend selection for one deployment configuration."""

    name = "WNA16"

    def backend_to_kernel_cls(
        self, backend: WNA16MoEBackend
    ) -> list[type[mk.FusedMoEExperts]]:
        return backend_to_kernel_cls(backend)

_MoeWNA16HummingWeightSchema

Adapter from MoeWNA16's generic packed layout to Humming's layout.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

class _MoeWNA16HummingWeightSchema:
    """Adapter from MoeWNA16's generic packed layout to Humming's layout."""

    def __init__(self, bits: int, group_size: int, has_zero_point: bool) -> None:
        self.bits = bits
        self.group_size = group_size
        self.has_zero_point = has_zero_point

    def convert_humming(
        self,
        tensors: dict[str, torch.Tensor],
        shape_n_stacks: list[int],
        shape_k_stacks: list[int],
        param_dtype: torch.dtype,
        num_experts: int | None = None,
    ) -> tuple[Any, dict[str, torch.Tensor]]:
        del shape_n_stacks, shape_k_stacks, num_experts
        from vllm.utils.humming import HummingWeightSchema, dtypes

        output = _convert_moe_wna16_humming_tensors(
            tensors, has_zero_point=self.has_zero_point
        )
        output["weight_scale"] = output["weight_scale"].to(param_dtype)
        schema = HummingWeightSchema(
            b_dtype=dtypes.DataType.from_str(f"uint{self.bits}"),
            weight_scale_group_size=self.group_size,
            has_zero_point=self.has_zero_point,
        )
        return schema, output

_convert_moe_wna16_humming_tensors(tensors, has_zero_point)

Convert MoeWNA16's N-first uint8 packing to Humming's int32 packing.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _convert_moe_wna16_humming_tensors(
    tensors: dict[str, torch.Tensor], has_zero_point: bool
) -> dict[str, torch.Tensor]:
    """Convert MoeWNA16's N-first uint8 packing to Humming's int32 packing."""
    if sys.byteorder != "little":
        raise NotImplementedError(
            "MoeWNA16 to Humming conversion requires a little-endian host."
        )

    output = {
        "weight": tensors["qweight"].contiguous().view(torch.int32),
        "weight_scale": tensors["scales"],
    }
    if has_zero_point:
        qzeros = tensors["qzeros"]
        output["zero_point"] = (
            qzeros.transpose(-1, -2)
            .contiguous()
            .view(torch.int32)
            .transpose(-1, -2)
            .contiguous()
        )
    return output

_get_priority_backends()

Get available backends in priority order based on platform and config.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _get_priority_backends() -> list[WNA16MoEBackend]:
    """Get available backends in priority order based on platform and config."""
    if current_platform.is_cpu():
        return [WNA16MoEBackend.ZEN_CPU, WNA16MoEBackend.CPU]
    if current_platform.is_xpu():
        return [WNA16MoEBackend.XPU]

    backends = [
        # Native HIP kernel, gated on gfx1100 by _supports_current_device().
        WNA16MoEBackend.RDNA3,
        WNA16MoEBackend.FLASHINFER_TRTLLM,
        WNA16MoEBackend.MARLIN,
        WNA16MoEBackend.BATCHED_MARLIN,
        WNA16MoEBackend.TRITON,
        WNA16MoEBackend.HUMMING,
        WNA16MoEBackend.EMULATION,
    ]
    return prioritize_humming(backends)

_humming_wna16_weight_schema(quant_config)

Humming weight schema for a WNA16 checkpoint, derived from the quant config rather than the running kernel.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _humming_wna16_weight_schema(
    quant_config: QuantizationConfig | QuantizationArgs | None,
) -> dict[str, Any]:
    """Humming weight schema for a WNA16 checkpoint, derived from the quant
    config rather than the running kernel."""
    from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
    from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig

    if isinstance(quant_config, AutoAWQConfig):
        return {
            "quant_method": "awq",
            "bits": quant_config.weight_bits,
            "group_size": quant_config.group_size,
            "zero_point": quant_config.zero_point,
        }
    if isinstance(quant_config, AutoGPTQConfig):
        return {
            "quant_method": "gptq",
            "bits": quant_config.weight_bits,
            "group_size": quant_config.group_size,
            "desc_act": quant_config.desc_act,
            "sym": quant_config.is_sym,
        }
    if isinstance(quant_config, QuantizationArgs):
        quant_type = getattr(quant_config.type, "value", quant_config.type)
        quant_strategy = getattr(quant_config.strategy, "value", quant_config.strategy)
        return {
            "quant_method": "compressed-tensors",
            "format": "pack-quantized",
            "type": str(quant_type),
            "num_bits": quant_config.num_bits,
            "strategy": str(quant_strategy),
            "group_size": quant_config.group_size,
            "symmetric": quant_config.symmetric,
        }
    raise TypeError(
        "Humming WNA16 checkpoint schema requires AutoAWQConfig, "
        "AutoGPTQConfig or QuantizationArgs, "
        f"got {type(quant_config).__name__}."
    )

_pad_rows(x, padded_rows)

Zero-pad a (E, rows, cols) tensor to padded_rows rows.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _pad_rows(x: torch.Tensor, padded_rows: int) -> torch.Tensor:
    """Zero-pad a ``(E, rows, cols)`` tensor to ``padded_rows`` rows."""
    if padded_rows == x.size(1):
        return x
    return torch.nn.functional.pad(x, (0, 0, 0, padded_rows - x.size(1)))

_pad_w13_bias(bias, n, padded_n)

Zero-pad each gate/up shard of a (E, 2 * n) bias to padded_n.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _pad_w13_bias(bias: torch.Tensor, n: int, padded_n: int) -> torch.Tensor:
    """Zero-pad each gate/up shard of a ``(E, 2 * n)`` bias to ``padded_n``."""
    if padded_n == n:
        return bias
    e = bias.size(0)
    bias = bias.view(e, 2, n)
    bias = torch.nn.functional.pad(bias, (0, padded_n - n))
    return bias.reshape(e, 2 * padded_n).contiguous()

_pad_w13_shard_cols(x, unit, padded_unit)

Zero-pad each of the two gate/up shards of a (E, rows, 2 * unit) tensor along its last dim, from unit to padded_unit columns.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _pad_w13_shard_cols(x: torch.Tensor, unit: int, padded_unit: int) -> torch.Tensor:
    """Zero-pad each of the two gate/up shards of a ``(E, rows, 2 * unit)``
    tensor along its last dim, from ``unit`` to ``padded_unit`` columns."""
    if padded_unit == unit:
        return x
    e, rows, _ = x.shape
    x = x.view(e, rows, 2, unit)
    x = torch.nn.functional.pad(x, (0, padded_unit - unit))
    return x.reshape(e, rows, 2 * padded_unit).contiguous()

_process_awq_weights_marlin(layer, weight_bits, pack_factor, group_size, w13_qweight, w2_qweight, w13_scales, w2_scales, w13_qzeros, w2_qzeros, w13_bias=None, w2_bias=None)

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.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_awq_weights_marlin(
    layer: torch.nn.Module,
    weight_bits: int,
    pack_factor: int,
    group_size: int,
    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_qzeros
    torch.Tensor | None,  # w2_qzeros
    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.
    """
    w13_bias_out: torch.Tensor | None = None
    w2_bias_out: torch.Tensor | None = None

    # --- Pad the intermediate size to a valid Marlin thread tile ---
    # AWQ packs along N: w13's N is in the (shard) columns, w2's N in the rows.
    N = layer.intermediate_size_per_partition
    padded_N = marlin_moe_padded_intermediate(N, group_size)
    if padded_N != N:
        w13_qweight = _pad_w13_shard_cols(
            w13_qweight, N // pack_factor, padded_N // pack_factor
        )
        w2_qweight = _pad_rows(w2_qweight, padded_N)
        w13_scales = _pad_w13_shard_cols(w13_scales, N, padded_N)
        w13_qzeros = _pad_w13_shard_cols(
            w13_qzeros, N // pack_factor, padded_N // pack_factor
        )
        if group_size > 0:
            w2_scales = _pad_rows(w2_scales, padded_N // group_size)
            w2_qzeros = _pad_rows(w2_qzeros, padded_N // group_size)
        if w13_bias is not None:
            w13_bias = _pad_w13_bias(w13_bias, N, padded_N)

    marlin_w13_qweight = ops.awq_marlin_moe_repack(
        w13_qweight,
        size_k=w13_qweight.shape[1],
        size_n=w13_qweight.shape[2] * pack_factor,
        num_bits=weight_bits,
    )
    marlin_w2_qweight = ops.awq_marlin_moe_repack(
        w2_qweight,
        size_k=w2_qweight.shape[1],
        size_n=w2_qweight.shape[2] * pack_factor,
        num_bits=weight_bits,
    )

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

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

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

    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,
        marlin_w13_qzeros,
        marlin_w2_qzeros,
        w13_bias_out,
        w2_bias_out,
    )

_process_weights_cpu(quant_config, w13, w2, w13_scale, w2_scale, w13_qzeros=None, w2_qzeros=None, w13_bias=None, w2_bias=None)

CPU INT4 W4A16 weight post-processing.

Non-AWQ inputs arrive in canonical N-first layout and are transposed to K-first internally. AWQ uses a different packing axis and arrives in its native format.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_cpu(
    quant_config: QuantizationConfig | QuantizationArgs | None,
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: 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 | None,  # w13_qzeros
    torch.Tensor | None,  # w2_qzeros
    torch.Tensor | None,  # w13_bias
    torch.Tensor | None,  # w2_bias
]:
    """CPU INT4 W4A16 weight post-processing.

    Non-AWQ inputs arrive in canonical N-first layout and are transposed
    to K-first internally.  AWQ uses a different packing axis and arrives
    in its native format.
    """
    from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
        prepare_int4_moe_layer_for_cpu,
    )
    from vllm.model_executor.layers.quantization.auto_awq import (
        AutoAWQConfig,
    )
    from vllm.model_executor.layers.quantization.auto_gptq import (
        AutoGPTQConfig,
    )

    if isinstance(quant_config, AutoAWQConfig):
        cpu_quant_algo = ops.CPUQuantAlgo.AWQ
    elif isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)):
        if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act:
            raise NotImplementedError(
                "CPU WNA16 MoE backend does not support GPTQ with "
                "desc_act=True. The fused MoE kernel has no activation "
                "reordering support."
            )
        cpu_quant_algo = ops.CPUQuantAlgo.GPTQ
        # Canonical N-first → K-first for CPU kernel.
        w13 = w13.transpose(1, 2).contiguous()
        w2 = w2.transpose(1, 2).contiguous()
        w13_scale = w13_scale.transpose(1, 2).contiguous()
        w2_scale = w2_scale.transpose(1, 2).contiguous()
        if w13_qzeros is not None:
            w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()
        if w2_qzeros is not None:
            w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()
    else:
        raise TypeError(
            "CPU WNA16 MoE backend requires AutoAWQConfig, AutoGPTQConfig "
            f"or QuantizationArgs, got {type(quant_config).__name__}."
        )

    # Determine zero points for repacking.
    w13_zeros: torch.Tensor | None = None
    w2_zeros: torch.Tensor | None = None
    if w13_qzeros is not None:
        w13_zeros = (
            w13_qzeros.data.view(torch.int32)
            if w13_qzeros.dtype != torch.int32
            else w13_qzeros.data
        )
    if w2_qzeros is not None:
        w2_zeros = (
            w2_qzeros.data.view(torch.int32)
            if w2_qzeros.dtype != torch.int32
            else w2_qzeros.data
        )

    (
        blocked_w13,
        blocked_w2,
        blocked_s13,
        blocked_s2,
        blocked_z13,
        blocked_z2,
    ) = prepare_int4_moe_layer_for_cpu(
        w13,
        w2,
        w13_scale,
        w2_scale,
        quant_algo=cpu_quant_algo,
        w13_zeros=w13_zeros,
        w2_zeros=w2_zeros,
    )
    return (
        blocked_w13,
        blocked_w2,
        blocked_s13,
        blocked_s2,
        blocked_z13,
        blocked_z2,
        w13_bias.to(torch.float32) if w13_bias is not None else None,
        w2_bias.to(torch.float32) if w2_bias is not None else None,
    )

_process_weights_emulation_awq(w13, w2, w13_scale, w2_scale, w13_qzeros, w2_qzeros)

Dequantize AWQ int4 weights to BF16 for the emulation backend.

AWQ inputs

w13: [E, K, 2N//8] int32 (packed along N, gate+up on dim 2) w2: [E, N, K//8] int32 (packed along K) w13_scale: [E, K//gs, 2N] float16 w2_scale: [E, N//gs, K] float16

Outputs (what TritonExperts expects): w13_out: [E, 2*N, K] bfloat16 w2_out: [E, K, N] bfloat16

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_emulation_awq(
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    w13_qzeros: torch.Tensor | None,
    w2_qzeros: torch.Tensor | None,
) -> tuple:
    """Dequantize AWQ int4 weights to BF16 for the emulation backend.

    AWQ inputs:
        w13: [E, K, 2*N//8]       int32  (packed along N, gate+up on dim 2)
        w2:  [E, N, K//8]         int32  (packed along K)
        w13_scale: [E, K//gs, 2*N]  float16
        w2_scale:  [E, N//gs, K]    float16

    Outputs (what TritonExperts expects):
        w13_out: [E, 2*N, K]  bfloat16
        w2_out:  [E, K, N]    bfloat16
    """
    # w13: AWQ-packed along N (dim 2), K is unpacked in dim 1
    # _unpack_and_dequant_int4_awq with transpose_output=True yields [E, 2*N, K]
    w13_bf16 = _unpack_and_dequant_int4_awq(
        w13, w13_scale, w13_qzeros, transpose_output=True
    )

    # w2: AWQ packs along K (dim 2 is K//8), N is unpacked in dim 1.
    # AWQ w2 is [E, N, K//8] — same column-pack format applied to the K dim.
    # _unpack_and_dequant_int4_awq expects [E, rows, N_packed] where the
    # packed dim is columns. Treat dim 1 as rows and dim 2 as N_packed:
    # unpacking gives [E, N, K]. Then permute to [E, K, N].
    w2_unpacked = _unpack_and_dequant_int4_awq(
        w2, w2_scale, w2_qzeros, transpose_output=False
    )  # [E, N, K]
    w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous()  # [E, K, N]

    dummy = torch.ones(1, dtype=torch.float16, device=w13.device)
    return (
        w13_bf16,
        w2_bf16,
        dummy,
        dummy,
        None,
        None,
        None,
        None,
    )

_process_weights_emulation_gptq(w13, w2, w13_scale, w2_scale, w13_qzeros, w2_qzeros)

Dequantize int4 weights to BF16 for the emulation backend.

Inputs arrive in canonical N-first layout and are transposed to K-first internally for the GPTQ dequantization routines.

Outputs (what TritonExperts expects): w13_out: [E, 2*N, K] bfloat16 w2_out: [E, K, N] bfloat16

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_emulation_gptq(
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    w13_qzeros: torch.Tensor | None,
    w2_qzeros: torch.Tensor | None,
) -> tuple:
    """Dequantize int4 weights to BF16 for the emulation backend.

    Inputs arrive in canonical N-first layout and are transposed to
    K-first internally for the GPTQ dequantization routines.

    Outputs (what TritonExperts expects):
        w13_out: [E, 2*N, K]  bfloat16
        w2_out:  [E, K, N]    bfloat16
    """
    # Canonical N-first → K-first for dequantization.
    w13 = w13.transpose(1, 2).contiguous()
    w2 = w2.transpose(1, 2).contiguous()
    w13_scale = w13_scale.transpose(1, 2).contiguous()
    w2_scale = w2_scale.transpose(1, 2).contiguous()
    if w13_qzeros is not None:
        w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()
    if w2_qzeros is not None:
        w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()

    # w13: packed along K (dim 1), output cols are 2*N (dim 2)
    # transpose_output=True yields [E, 2*N, K]
    w13_bf16 = _unpack_and_dequant_int4_gptq(
        w13, w13_scale, w13_qzeros, transpose_output=True
    )

    # w2: packed along N (dim 1 is N//8), output cols are K (dim 2)
    # After unpacking we get [E, N, K]; we want [E, K, N] for TritonExperts
    # transpose_output=False gives [E, N, K], then we permute once more
    w2_unpacked = _unpack_and_dequant_int4_gptq(
        w2, w2_scale, w2_qzeros, transpose_output=False
    )  # [E, N, K]
    w2_bf16 = w2_unpacked.permute(0, 2, 1).contiguous()  # [E, K, N]

    dummy = torch.ones(1, dtype=torch.float16, device=w13.device)
    return (
        w13_bf16,  # w13_qweight  (now bf16, not int32)
        w2_bf16,  # w2_qweight   (now bf16, not int32)
        dummy,  # w13_scales   (unused; nulled out in Int4EmulationTritonExperts)
        dummy,  # w2_scales    (unused)
        None,  # w13_qzeros
        None,  # w2_qzeros
        None,  # w13_bias
        None,  # w2_bias
    )

_process_weights_flashinfer(w13_qweight, w2_qweight, w13_scales, w2_scales, w13_bias=None, w2_bias=None)

Flashinfer (TRT-LLM MXINT4) weight post-processing.

Steps

1. Transform weights/scales via prepare_static_weights_for_trtllm_mxint4_moe.
2. Return transformed tensors and biases.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_flashinfer(
    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_qzeros
    torch.Tensor | None,  # w2_qzeros
    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 and biases.
    """
    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"],
        None,
        None,
        w13_bias,
        w2_bias,
    )

_process_weights_marlin(layer, num_bits, pack_factor, group_size, w13_qweight, w2_qweight, w13_scales, w2_scales, w13_qzeros=None, w2_qzeros=None, w13_bias=None, w2_bias=None)

Standard Marlin weight post-processing shared by MARLIN and BATCHED_MARLIN backends.

Inputs arrive in canonical N-first layout [E, N, K_packed].

Steps

1. Transpose canonical N-first to K-first for Marlin.
2. Repack weights via gptq_marlin_moe_repack.
3. Permute scales.
4. Permute bias tensors.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_marlin(
    layer: torch.nn.Module,
    num_bits: int,
    pack_factor: int,
    group_size: int,
    w13_qweight: torch.Tensor,
    w2_qweight: torch.Tensor,
    w13_scales: torch.Tensor,
    w2_scales: 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 | None,  # w13_qzeros
    torch.Tensor | None,  # w2_qzeros
    torch.Tensor | None,  # w13_bias
    torch.Tensor | None,  # w2_bias
]:
    """Standard Marlin weight post-processing shared by MARLIN and
    BATCHED_MARLIN backends.

    Inputs arrive in canonical N-first layout ``[E, N, K_packed]``.

    Steps
    -----
    1. Transpose canonical N-first to K-first for Marlin.
    2. Repack weights via ``gptq_marlin_moe_repack``.
    3. Permute scales.
    4. Permute bias tensors.
    """
    # Canonical N-first → K-first for Marlin.
    w13_qweight = w13_qweight.transpose(1, 2).contiguous()
    w2_qweight = w2_qweight.transpose(1, 2).contiguous()
    w13_scales = w13_scales.transpose(1, 2).contiguous()
    w2_scales = w2_scales.transpose(1, 2).contiguous()
    if w13_qzeros is not None:
        w13_qzeros = w13_qzeros.transpose(1, 2).contiguous()
    if w2_qzeros is not None:
        w2_qzeros = w2_qzeros.transpose(1, 2).contiguous()

    marlin_w13_qweight = w13_qweight
    marlin_w2_qweight = w2_qweight
    marlin_w13_scales = w13_scales
    marlin_w2_scales = w2_scales
    w13_bias_out: torch.Tensor | None = None
    w2_bias_out: torch.Tensor | None = None

    # --- Pad the intermediate size to a valid Marlin thread tile ---
    # GPTQ packs along K: w13's N is in the (shard) columns, w2's N in the rows.
    N = layer.intermediate_size_per_partition
    padded_N = marlin_moe_padded_intermediate(N, group_size)
    if padded_N != N:
        marlin_w13_qweight = _pad_w13_shard_cols(marlin_w13_qweight, N, padded_N)
        marlin_w2_qweight = _pad_rows(marlin_w2_qweight, padded_N // pack_factor)
        marlin_w13_scales = _pad_w13_shard_cols(marlin_w13_scales, N, padded_N)
        if group_size > 0:
            marlin_w2_scales = _pad_rows(marlin_w2_scales, padded_N // group_size)
        if w13_qzeros is not None:
            w13_qzeros = _pad_w13_shard_cols(
                w13_qzeros, N // pack_factor, padded_N // pack_factor
            )
        if w2_qzeros is not None and group_size > 0:
            w2_qzeros = _pad_rows(w2_qzeros, padded_N // group_size)
        if w13_bias is not None:
            w13_bias = _pad_w13_bias(w13_bias, N, padded_N)

    # --- Repack weights ---
    marlin_w13_qweight = ops.gptq_marlin_moe_repack(
        marlin_w13_qweight,
        marlin_w13_qweight.shape[1] * pack_factor,
        marlin_w13_qweight.shape[2],
        num_bits,
    )
    marlin_w2_qweight = ops.gptq_marlin_moe_repack(
        marlin_w2_qweight,
        marlin_w2_qweight.shape[1] * pack_factor,
        marlin_w2_qweight.shape[2],
        num_bits,
    )

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

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

    # --- 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_qzeros,
        w2_qzeros,
        w13_bias_out,
        w2_bias_out,
    )

_process_weights_rdna3(w13, w2, w13_scale, w2_scale, group_size)

RDNA3 (gfx1100) W4A16 weight post-processing.

Transposes the canonical N-first inputs to the K-first layout of moe_gptq_gemm_rdna3 (weights [E, K // 8, N], scales [E, groups, N]), interleaves the packed nibbles per expert (the exllama shuffle the dense RDNA3 kernel also uses) and synthesizes the symmetric zero points the kernel dequantizes with.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_rdna3(
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    group_size: int,
) -> tuple[
    torch.Tensor,  # w13_qweight
    torch.Tensor,  # w2_qweight
    torch.Tensor,  # w13_scales
    torch.Tensor,  # w2_scales
    torch.Tensor | None,  # w13_qzeros
    torch.Tensor | None,  # w2_qzeros
    torch.Tensor | None,  # w13_bias
    torch.Tensor | None,  # w2_bias
]:
    """RDNA3 (gfx1100) W4A16 weight post-processing.

    Transposes the canonical N-first inputs to the K-first layout of
    ``moe_gptq_gemm_rdna3`` (weights ``[E, K // 8, N]``, scales
    ``[E, groups, N]``), interleaves the packed nibbles per expert (the exllama
    shuffle the dense RDNA3 kernel also uses) and synthesizes the symmetric
    zero points the kernel dequantizes with.
    """
    device = w13.device
    num_experts = w13.size(0)

    w13 = w13.transpose(1, 2).contiguous()
    w2 = w2.transpose(1, 2).contiguous()

    for e in range(num_experts):
        w13_e = w13[e].contiguous()
        ops.gptq_shuffle(w13_e, 4)
        w13[e] = w13_e
        w2_e = w2[e].contiguous()
        ops.gptq_shuffle(w2_e, 4)
        w2[e] = w2_e

    def _qzeros(w: torch.Tensor) -> torch.Tensor:
        qz = _synthesize_rdna3_qzeros((w.size(1) * 8) // group_size, w.size(2), device)
        return qz.unsqueeze(0).expand(num_experts, -1, -1).contiguous()

    return (
        w13,
        w2,
        w13_scale.transpose(1, 2).contiguous(),
        w2_scale.transpose(1, 2).contiguous(),
        _qzeros(w13),
        _qzeros(w2),
        None,  # w13_bias
        None,  # w2_bias
    )

_process_weights_xpu(layer, quant_config, w13_qweight, w2_qweight, w13_scales, w2_scales, w13_bias=None, w2_bias=None)

Repack INT4 MoE weights into the layout vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True) expects.

Inputs arrive in canonical N-first layout [E, N, K_packed] int32. The int32 → uint8 view recovers sequential int4-packed bytes along the input dim. Each packed int32 holds 8 nibbles in ascending K order; on a little-endian host the view exposes them as two nibbles per byte with the lower nibble = lower input-K index. xpu_fused_moe(is_int4=True) expects this convention; big-endian hosts are unsupported.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

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 INT4 MoE weights into the layout
    `vllm_xpu_kernels.fused_moe_interface.xpu_fused_moe(is_int4=True)` expects.

    Inputs arrive in canonical N-first layout ``[E, N, K_packed]`` int32.
    The int32 → uint8 view recovers sequential int4-packed bytes along
    the input dim.  Each packed int32 holds 8 nibbles in ascending K order;
    on a little-endian host the view exposes them as two nibbles per byte
    with the lower nibble = lower input-K index.  xpu_fused_moe(is_int4=True)
    expects this convention; big-endian hosts are unsupported.
    """
    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 "
            "int32 → uint8 nibble repack relies on LE byte ordering."
        )

    w13_xpu = w13_qweight.contiguous().view(torch.uint8)
    w2_xpu = w2_qweight.contiguous().view(torch.uint8)
    w13_scales_xpu = w13_scales.contiguous()
    w2_scales_xpu = w2_scales.contiguous()

    return (
        w13_xpu,
        w2_xpu,
        w13_scales_xpu,
        w2_scales_xpu,
        w13_bias,
        w2_bias,
    )

_process_weights_zen_cpu(w13, w2, w13_scale, w2_scale, w13_bias=None, w2_bias=None, num_bits=4)

Zen CPU DA8W4 (W4A8) / DA8W8 (W8A8) weight post-processing.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _process_weights_zen_cpu(
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    w13_bias: torch.Tensor | None = None,
    w2_bias: torch.Tensor | None = None,
    num_bits: int = 4,
) -> tuple[
    torch.Tensor,  # w13_qweight
    torch.Tensor,  # w2_qweight
    torch.Tensor,  # w13_scales
    torch.Tensor,  # w2_scales
    torch.Tensor | None,  # w13_qzeros
    torch.Tensor | None,  # w2_qzeros
    torch.Tensor | None,  # w13_bias
    torch.Tensor | None,  # w2_bias
]:
    """Zen CPU DA8W4 (W4A8) / DA8W8 (W8A8) weight post-processing."""
    # CT WNA16 buffers are N-first, so the scales are [E, N, G].
    hidden_size = w2_scale.shape[1]
    intermediate_size = w13_scale.shape[1] // 2

    # The kernel reads per-group scales as [E, G, N], so N-first needs a swap.
    return (
        _zen_decode_weight(w13.data, hidden_size, num_bits),
        _zen_decode_weight(w2.data, intermediate_size, num_bits),
        w13_scale.data.transpose(1, 2).to(torch.bfloat16).contiguous(),
        w2_scale.data.transpose(1, 2).to(torch.bfloat16).contiguous(),
        None,  # w13_qzeros (symmetric)
        None,  # w2_qzeros (symmetric)
        w13_bias,
        w2_bias,
    )

_synthesize_rdna3_qzeros(groups, out_features, device)

Create the packed zero-point tensor for symmetric quantization.

GPTQv1 +1 quirk: the kernel adds 1 to the stored zeros, so encode (bias - 1) = 7 for uint4b8 (bias=8).

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _synthesize_rdna3_qzeros(
    groups: int, out_features: int, device: torch.device
) -> torch.Tensor:
    """Create the packed zero-point tensor for symmetric quantization.

    GPTQv1 +1 quirk: the kernel adds 1 to the stored zeros, so encode
    (bias - 1) = 7 for uint4b8 (bias=8).
    """
    zeros = torch.full(
        (groups, out_features),
        scalar_types.uint4b8.bias - 1,
        dtype=torch.int32,
        device=device,
    )
    return pack_quantized_values_into_int32(zeros, scalar_types.uint4b8, packed_dim=1)

_unpack_and_dequant_int4_awq(w_int32, scale, qzeros, transpose_output, output_dtype=torch.bfloat16)

Unpack AWQ-packed int4 weights and dequantize to output_dtype.

AWQ packs along the N (column) dimension with an interleave permutation [0,2,4,6,1,3,5,7] applied before packing, so unpacking must undo that.

Parameters:

• w_int32

  (Tensor) –

  packed weights, shape [E, K, N_packed] where N_packed = N//8 (8 nibbles per int32, packed along N with AWQ interleaving).

• scale

  (Tensor) –

  per-group scales, shape [E, K//group_size, N], float16.

• qzeros

  (Tensor | None) –

  asymmetric zero-points, shape [E, K//gs, N_packed], int32. None for symmetric (uint4b8 with implicit bias 8).

• transpose_output

  (bool) –

  if True return [E, N, K]; if False return [E, K, N].

• output_dtype

  (dtype, default: bfloat16 ) –

  target floating-point dtype (bfloat16 or float16).

Returns:

• Tensor –

  Dequantized weight tensor in the requested layout.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _unpack_and_dequant_int4_awq(
    w_int32: torch.Tensor,
    scale: torch.Tensor,
    qzeros: torch.Tensor | None,
    transpose_output: bool,
    output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    """Unpack AWQ-packed int4 weights and dequantize to output_dtype.

    AWQ packs along the N (column) dimension with an interleave permutation
    [0,2,4,6,1,3,5,7] applied before packing, so unpacking must undo that.

    Args:
        w_int32: packed weights, shape [E, K, N_packed] where N_packed = N//8
                 (8 nibbles per int32, packed along N with AWQ interleaving).
        scale:   per-group scales, shape [E, K//group_size, N], float16.
        qzeros:  asymmetric zero-points, shape [E, K//gs, N_packed], int32.
                 None for symmetric (uint4b8 with implicit bias 8).
        transpose_output: if True return [E, N, K]; if False return [E, K, N].
        output_dtype: target floating-point dtype (bfloat16 or float16).

    Returns:
        Dequantized weight tensor in the requested layout.

    """
    E, K, N_packed = w_int32.shape
    N = N_packed * 8

    # Unpack 8 nibbles per int32 along the N dimension (LSB-first)
    shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4
    # [E, K, N_packed, 8] -> [E, K, N_packed*8] = [E, K, N_interleaved]
    nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF
    w_interleaved = nibbles.reshape(E, K, N)  # [E, K, N] but column-interleaved

    # Undo AWQ interleave: packed order is [0,2,4,6,1,3,5,7] within each group
    # of 8. Inverse: position i in packed -> original column interleave[i].
    # To reverse: we need the inverse permutation so that
    # w[:, :, inv_interleave] = w_interleaved gives the natural column order.
    interleave = torch.tensor([0, 2, 4, 6, 1, 3, 5, 7], device=w_int32.device)
    inv_interleave = torch.empty_like(interleave)
    inv_interleave[interleave] = torch.arange(8, device=w_int32.device)

    # Apply inverse interleave within each group of 8 columns
    w_reshaped = w_interleaved.reshape(E, K, N // 8, 8)  # [E, K, groups, 8]
    w_reordered = w_reshaped[:, :, :, inv_interleave]  # undo interleave
    w = w_reordered.reshape(E, K, N).to(torch.int16)  # [E, K, N]

    if qzeros is None:
        w = w - 8
    else:
        # qzeros: [E, K//gs, N_packed] int32, same AWQ column packing
        gs = K // scale.shape[1]
        n_gs = scale.shape[1]
        zp_nibbles = (qzeros.unsqueeze(-1) >> shifts) & 0xF  # [E, n_gs, N_packed, 8]
        zp_interleaved = zp_nibbles.reshape(E, n_gs, N)
        zp_reshaped = zp_interleaved.reshape(E, n_gs, N // 8, 8)
        zp_reordered = zp_reshaped[:, :, :, inv_interleave]
        zp = zp_reordered.reshape(E, n_gs, N).to(torch.int16)  # [E, n_gs, N]
        zp = zp.repeat_interleave(gs, dim=1)  # [E, K, N]
        w = w - zp

    gs = K // scale.shape[1]
    scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype)  # [E, K, N]

    w_dequant = w.to(output_dtype) * scale_broadcast  # [E, K, N]

    if transpose_output:
        return w_dequant.permute(0, 2, 1).contiguous()  # [E, N, K]
    return w_dequant.contiguous()  # [E, K, N]

_unpack_and_dequant_int4_gptq(w_int32, scale, qzeros, transpose_output, output_dtype=torch.bfloat16)

Unpack GPTQ-packed int4 weights and dequantize to output_dtype.

Parameters:

• w_int32

  (Tensor) –

  packed weights, shape [E, K_packed, N] where K_packed = K//8 (8 nibbles per int32, LSB-first in the K dimension).

• scale

  (Tensor) –

  per-group scales, shape [E, K//group_size, N], float16.

• qzeros

  (Tensor | None) –

  optional asymmetric zero-points, shape [E, K//gs, N//8], int32. None for symmetric (uint4b8 with implicit bias 8).

• transpose_output

  (bool) –

  if True return [E, N, K]; if False return [E, K, N].

• output_dtype

  (dtype, default: bfloat16 ) –

  target floating-point dtype (bfloat16 or float16).

Returns:

• Tensor –

  Dequantized weight tensor in the requested layout.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _unpack_and_dequant_int4_gptq(
    w_int32: torch.Tensor,
    scale: torch.Tensor,
    qzeros: torch.Tensor | None,
    transpose_output: bool,
    output_dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    """Unpack GPTQ-packed int4 weights and dequantize to output_dtype.

    Args:
        w_int32: packed weights, shape [E, K_packed, N] where K_packed = K//8
                 (8 nibbles per int32, LSB-first in the K dimension).
        scale:   per-group scales, shape [E, K//group_size, N], float16.
        qzeros:  optional asymmetric zero-points, shape [E, K//gs, N//8], int32.
                 None for symmetric (uint4b8 with implicit bias 8).
        transpose_output: if True return [E, N, K]; if False return [E, K, N].
        output_dtype: target floating-point dtype (bfloat16 or float16).

    Returns:
        Dequantized weight tensor in the requested layout.

    """
    E, K_packed, N = w_int32.shape
    K = K_packed * 8

    # Unpack: [E, K_packed, N] -> [E, K_packed, N, 8] via bit-shifts.
    # The nibble index (last dim) enumerates K rows within each packed column,
    # so we must fuse K_packed and the nibble dim, not N and the nibble dim.
    # Permute to [E, K_packed, 8, N] before reshaping to [E, K, N].
    shifts = torch.arange(8, device=w_int32.device, dtype=torch.int32) * 4
    nibbles = (w_int32.unsqueeze(-1) >> shifts) & 0xF  # [E, K_packed, N, 8]

    # Reshape to [E, K, N]: fuse K_packed and nibble index (dim 1 and 3)
    w = nibbles.permute(0, 1, 3, 2).reshape(E, K, N).to(torch.int16)

    if qzeros is None:
        # Symmetric uint4b8: subtract bias so the range is [-8, 7]
        w = w - 8
    else:
        # Asymmetric: unpack zero-points (same 8-nibble packing) and subtract
        # qzeros shape: [E, K//gs, N//8] int32
        gs = K // scale.shape[1]
        n_gs = scale.shape[1]
        zp_shifts = torch.arange(8, device=qzeros.device, dtype=torch.int32) * 4
        zp_nibbles = (qzeros.unsqueeze(-1) >> zp_shifts) & 0xF  # [E, n_gs, N//8, 8]
        zp = zp_nibbles.reshape(E, n_gs, N).to(torch.int16)  # [E, n_gs, N]
        zp = zp.repeat_interleave(gs, dim=1)  # [E, K, N]
        w = w - zp

    # Broadcast scale [E, K//gs, N] -> [E, K, N]
    gs = K // scale.shape[1]
    scale_broadcast = scale.repeat_interleave(gs, dim=1).to(output_dtype)

    w_dequant = w.to(output_dtype) * scale_broadcast  # [E, K, N]

    if transpose_output:
        return w_dequant.permute(0, 2, 1).contiguous()  # [E, N, K]
    return w_dequant.contiguous()  # [E, K, N]

_zen_cpu_num_bits(quant_config)

Weight bit width for the Zen CPU path, or None for a checkpoint format it cannot read one from; AWQ and GPTQ call it weight_bits.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _zen_cpu_num_bits(
    quant_config: QuantizationConfig | QuantizationArgs,
) -> int | None:
    """Weight bit width for the Zen CPU path, or None for a checkpoint format
    it cannot read one from; AWQ and GPTQ call it weight_bits."""
    from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
    from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig

    if isinstance(quant_config, AutoAWQConfig | AutoGPTQConfig):
        return quant_config.weight_bits
    if isinstance(quant_config, QuantizationArgs):
        return quant_config.num_bits
    return None

_zen_decode_weight(packed, in_features, num_bits)

Decode a CT [E, N, K//pack_factor] weight into the layout the op reads: packed s4 [E, N, K//8] for int4, plain s8 [E, N, K] for int8. The trailing dim is what makes the op pick DA8W4 over DA8W8.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def _zen_decode_weight(
    packed: torch.Tensor, in_features: int, num_bits: int
) -> torch.Tensor:
    """Decode a CT ``[E, N, K//pack_factor]`` weight into the layout the op
    reads: packed s4 ``[E, N, K//8]`` for int4, plain s8 ``[E, N, K]`` for
    int8. The trailing dim is what makes the op pick DA8W4 over DA8W8."""
    from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
        _import_unpack_from_int32,
    )

    num_experts, out_features, _ = packed.shape
    # packed_dim indexes the 2D expert slice, which is N-first, so K is dim 1.
    unpacked = _import_unpack_from_int32()(
        packed,
        num_bits,
        torch.Size([num_experts, out_features, in_features]),
        packed_dim=1,
    )
    # CT is already N-first, which is what the repack consumes.
    if num_bits == 8:
        return unpacked.to(torch.int8).contiguous()
    repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
    return torch.stack([repack_op(w) for w in unpacked])

backend_to_kernel_cls(backend)

Return the experts class for the given backend, or None for NONE.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

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.HUMMING:
        from vllm.model_executor.layers.fused_moe.experts.fused_humming_moe import (
            BatchedHummingGroupedExperts,
            HummingGroupedExperts,
            HummingIndexedExperts,
        )

        return [
            BatchedHummingGroupedExperts,
            HummingGroupedExperts,
            HummingIndexedExperts,
        ]
    elif backend == WNA16MoEBackend.MARLIN:
        return [MarlinExperts]
    elif backend == WNA16MoEBackend.BATCHED_MARLIN:
        return [BatchedMarlinExperts]
    elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
        return [TrtLlmMxint4ExpertsMonolithic]
    elif backend == WNA16MoEBackend.TRITON:
        return [TritonWNA16Experts]
    elif backend == WNA16MoEBackend.XPU:
        from vllm.model_executor.layers.fused_moe.experts.xpu_moe import (
            XPUExpertsWNA16,
        )

        return [XPUExpertsWNA16]
    elif backend == WNA16MoEBackend.CPU:
        from vllm.model_executor.layers.fused_moe.experts.cpu_moe import (
            CPUExpertsInt4,
        )

        return [CPUExpertsInt4]
    elif backend == WNA16MoEBackend.ZEN_CPU:
        from vllm.model_executor.layers.fused_moe.experts.zentorch_moe import (
            ZentorchExpertsWNA16,
        )

        return [ZentorchExpertsWNA16]
    elif backend == WNA16MoEBackend.EMULATION:
        from vllm.model_executor.layers.fused_moe.experts.int4_emulation_moe import (
            Int4EmulationTritonExperts,
        )

        return [Int4EmulationTritonExperts]
    elif backend == WNA16MoEBackend.RDNA3:
        from vllm.model_executor.layers.fused_moe.experts.rdna3_moe import (
            Rdna3WNA16Experts,
        )

        return [Rdna3WNA16Experts]
    else:
        raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}")

convert_to_wna16_moe_kernel_format(backend, layer, quant_config, w13, w2, w13_scale, w2_scale, w13_qzeros=None, w2_qzeros=None, w13_bias=None, w2_bias=None)

Dispatch weight post-processing to the appropriate per-backend handler.

All non-AWQ frontends must supply weights and scales in N-first canonical layout [E, N_out, K_packed]. AWQ uses a different packing axis and is handled by dedicated per-backend helpers.

To add a new backend, implement a _process_weights_<name> helper and add a branch here. Backends that rewrite the layer's parameters in place (e.g. Humming) return None; the caller then skips the param scatter.

Parameters:

• backend

  (WNA16MoEBackend) –

  the selected WNA16MoEBackend.

• layer

  (Module) –

  the MoERunner layer whose parameters are being prepared.

• quant_config

  (QuantizationConfig | QuantizationArgs | None) –

  the QuantizationConfig for this layer.

• w13

  (Tensor) –

  fused gate/up expert weights.

• w2

  (Tensor) –

  down-projection expert weights.

• w13_scale

  (Tensor) –

  quantization scales for w13.

• w2_scale

  (Tensor) –

  quantization scales for w2.

• w13_qzeros

  (Tensor | None, default: None ) –

  optional zero points for w13.

• w2_qzeros

  (Tensor | None, default: None ) –

  optional zero points for w2.

• w13_bias

  (Tensor | None, default: None ) –

  optional bias for w13.

• w2_bias

  (Tensor | None, default: None ) –

  optional bias for w2.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def convert_to_wna16_moe_kernel_format(
    backend: WNA16MoEBackend,
    layer: torch.nn.Module,
    quant_config: QuantizationConfig | QuantizationArgs | None,
    w13: torch.Tensor,
    w2: torch.Tensor,
    w13_scale: torch.Tensor,
    w2_scale: 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 | None,  # w13_qzeros
        torch.Tensor | None,  # w2_qzeros
        torch.Tensor | None,  # w13_bias
        torch.Tensor | None,  # w2_bias
    ]
    | None
):
    """Dispatch weight post-processing to the appropriate per-backend handler.

    All non-AWQ frontends must supply weights and scales in N-first
    canonical layout ``[E, N_out, K_packed]``.  AWQ uses a different
    packing axis and is handled by dedicated per-backend helpers.

    To add a new backend, implement a ``_process_weights_<name>`` helper and
    add a branch here. Backends that rewrite the layer's parameters in place
    (e.g. Humming) return ``None``; the caller then skips the param scatter.

    Args:
        backend: the selected ``WNA16MoEBackend``.
        layer: the ``MoERunner`` layer whose parameters are being prepared.
        quant_config: the ``QuantizationConfig`` for this layer.
        w13: fused gate/up expert weights.
        w2: down-projection expert weights.
        w13_scale: quantization scales for ``w13``.
        w2_scale: quantization scales for ``w2``.
        w13_qzeros: optional zero points for ``w13``.
        w2_qzeros: optional zero points for ``w2``.
        w13_bias: optional bias for ``w13``.
        w2_bias: optional bias for ``w2``.

    """
    if backend == WNA16MoEBackend.HUMMING:
        from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config
        from vllm.model_executor.layers.quantization.utils.humming import (
            convert_to_humming_moe_kernel_format,
        )

        if isinstance(quant_config, MoeWNA16Config):
            from vllm.utils.humming import HummingInputSchema

            convert_to_humming_moe_kernel_format(
                layer,
                weight_schema=_MoeWNA16HummingWeightSchema(
                    bits=quant_config.weight_bits,
                    group_size=layer.group_size,
                    has_zero_point=quant_config.has_zp,
                ),
                input_schema=HummingInputSchema(),
            )
        else:
            convert_to_humming_moe_kernel_format(
                layer, quant_config=_humming_wna16_weight_schema(quant_config)
            )
        return None

    if backend in (
        WNA16MoEBackend.MARLIN,
        WNA16MoEBackend.BATCHED_MARLIN,
    ):
        from vllm.model_executor.layers.quantization.auto_awq import (
            AutoAWQConfig,
        )
        from vllm.model_executor.layers.quantization.auto_gptq import (
            AutoGPTQConfig,
        )

        if isinstance(quant_config, AutoAWQConfig):
            num_bits = quant_config.weight_bits
            pack_factor = quant_config.pack_factor
            group_size = quant_config.group_size
        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
        elif isinstance(quant_config, QuantizationArgs):
            num_bits = quant_config.num_bits
            pack_factor = 32 // quant_config.num_bits
            group_size = (
                -1 if quant_config.group_size is None else quant_config.group_size
            )
        else:
            raise TypeError(
                "Marlin WNA16 MoE backend requires AutoGPTQConfig, AutoAWQConfig or "
                f"QuantizationArgs, got {type(quant_config).__name__}."
            )

        if isinstance(quant_config, AutoAWQConfig):
            if w13_qzeros is None or w2_qzeros is None:
                raise ValueError("AWQ Marlin MoE requires zero-point tensors.")

            return _process_awq_weights_marlin(
                layer,
                num_bits,
                pack_factor,
                group_size,
                w13,
                w2,
                w13_scale,
                w2_scale,
                w13_qzeros,
                w2_qzeros,
                w13_bias,
                w2_bias,
            )
        else:
            return _process_weights_marlin(
                layer,
                num_bits,
                pack_factor,
                group_size,
                w13,
                w2,
                w13_scale,
                w2_scale,
                w13_qzeros,
                w2_qzeros,
                w13_bias,
                w2_bias,
            )
    elif backend == WNA16MoEBackend.RDNA3:
        assert isinstance(quant_config, QuantizationArgs)
        return _process_weights_rdna3(
            w13,
            w2,
            w13_scale,
            w2_scale,
            quant_config.group_size,
        )
    elif backend == WNA16MoEBackend.CPU:
        return _process_weights_cpu(
            quant_config,
            w13,
            w2,
            w13_scale,
            w2_scale,
            w13_qzeros,
            w2_qzeros,
            w13_bias,
            w2_bias,
        )
    elif backend == WNA16MoEBackend.ZEN_CPU:
        # Backend selection already rejected widths outside _ZEN_CPU_NUM_BITS.
        zen_num_bits = _zen_cpu_num_bits(quant_config)
        assert zen_num_bits is not None
        # Names the width, since VLLM_CPU_INT4_W4A8 reads as int4-only.
        logger.info_once(
            "[zen_cpu] Using zentorch_fused_moe for DA8W%d MoE (W%dA16 checkpoint)",
            zen_num_bits,
            zen_num_bits,
        )
        return _process_weights_zen_cpu(
            w13,
            w2,
            w13_scale,
            w2_scale,
            w13_bias,
            w2_bias,
            num_bits=zen_num_bits,
        )
    elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM:
        return _process_weights_flashinfer(
            w13,
            w2,
            w13_scale,
            w2_scale,
            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,
        )
        return (
            w13_xpu,
            w2_xpu,
            w13_scale_xpu,
            w2_scale_xpu,
            None,  # w13_qzeros — sym int4 on XPU has none; kernel does uint4b8→s4
            None,  # w2_qzeros
            w13_bias_out,
            w2_bias_out,
        )
    elif backend == WNA16MoEBackend.EMULATION:
        from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig

        if isinstance(quant_config, AutoAWQConfig):
            return _process_weights_emulation_awq(
                w13,
                w2,
                w13_scale,
                w2_scale,
                w13_qzeros,
                w2_qzeros,
            )
        return _process_weights_emulation_gptq(
            w13,
            w2,
            w13_scale,
            w2_scale,
            w13_qzeros,
            w2_qzeros,
        )
    elif backend == WNA16MoEBackend.TRITON:
        # All inputs arrive in canonical N-first format; view as uint8.
        w13_uint8 = w13.contiguous().view(torch.uint8)
        w2_uint8 = w2.contiguous().view(torch.uint8)

        # Repack N-first int32 zero-points to uint8:
        # [E, N//pf, K//gs] int32 → [E, N*4//pf, K//gs] uint8
        if w13_qzeros is not None and w13_qzeros.dtype == torch.int32:
            E13, Np13, Kg13 = w13_qzeros.shape
            w13_qzeros = (
                w13_qzeros.contiguous()
                .view(torch.uint8)
                .reshape(E13, Np13, Kg13, 4)
                .permute(0, 1, 3, 2)
                .reshape(E13, Np13 * 4, Kg13)
                .contiguous()
            )
        if w2_qzeros is not None and w2_qzeros.dtype == torch.int32:
            E2, Np2, Kg2 = w2_qzeros.shape
            w2_qzeros = (
                w2_qzeros.contiguous()
                .view(torch.uint8)
                .reshape(E2, Np2, Kg2, 4)
                .permute(0, 1, 3, 2)
                .reshape(E2, Np2 * 4, Kg2)
                .contiguous()
            )

        return (
            w13_uint8,
            w2_uint8,
            w13_scale,
            w2_scale,
            w13_qzeros,
            w2_qzeros,
            w13_bias,
            w2_bias,
        )
    else:
        raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}")

make_wna16_moe_quant_config(w1_scale, w2_scale, group_size, num_bits, w1_zp=None, w2_zp=None, w1_bias=None, w2_bias=None, gemm1_clamp_limit=None, gemm1_alpha=None, gemm1_beta=None)

Create the FusedMoEQuantConfig for 4 or 8-bit WNA16 MoE.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

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,
    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],
            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],
            gemm1_clamp_limit=gemm1_clamp_limit,
            gemm1_alpha=gemm1_alpha,
            gemm1_beta=gemm1_beta,
        )

map_wna16_backend(runner_backend)

Map user's MoEBackend to WNA16MoEBackend.

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend:
    """Map user's MoEBackend to WNA16MoEBackend."""
    mapping = {
        "triton": WNA16MoEBackend.TRITON,
        "marlin": WNA16MoEBackend.MARLIN,
        "humming": WNA16MoEBackend.HUMMING,
        "flashinfer_trtllm": WNA16MoEBackend.FLASHINFER_TRTLLM,
        "emulation": WNA16MoEBackend.EMULATION,
        "rdna3": WNA16MoEBackend.RDNA3,
    }
    if backend := mapping.get(runner_backend):
        return backend
    raise ValueError(
        f"moe_backend='{runner_backend}' is not supported for WNA16 MoE. "
        f"Expected one of {list(mapping.keys())}."
    )

select_wna16_moe_backend(config, weight_key, quant_config, may_have_zp, may_have_bias, allow_tile_padding=False)

Select the WNA16 MoE backend.

Parameters:

• config

  (FusedMoEConfig) –

  the shared FusedMoEConfig for this layer.

• weight_key

  (QuantKey) –

  The QuantKey describing the weight quantization. Must have int4 or int8 type.

• quant_config

  (QuantizationConfig | QuantizationArgs) –

  Quantization structure and checkpoint format description.

• may_have_zp

  (bool) –

  Whether the integration can provide weight zero points.

• may_have_bias

  (bool) –

  Whether the integration can provide expert bias.

• allow_tile_padding

  (bool, default: False ) –

  Whether backends that require padding the weights up to a tile boundary may be selected.

Returns:

• tuple[WNA16MoEBackend, type[FusedMoEExperts]] –

  A tuple of (WNA16MoEBackend, experts class or None).

Source code in vllm/model_executor/layers/fused_moe/oracle/int_wna16.py

def select_wna16_moe_backend(
    config: FusedMoEConfig,
    weight_key: QuantKey,
    quant_config: QuantizationConfig | QuantizationArgs,
    may_have_zp: bool,
    may_have_bias: bool,
    allow_tile_padding: bool = False,
) -> 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.
        quant_config: Quantization structure and checkpoint format description.
        may_have_zp: Whether the integration can provide weight zero points.
        may_have_bias: Whether the integration can provide expert bias.
        allow_tile_padding: Whether backends that require padding the weights
            up to a tile boundary may be selected.

    Returns:
        A tuple of (``WNA16MoEBackend``, experts class or ``None``).

    """
    # WNA16 MoE always runs BF16 activations.
    oracle = WNA16MoEKernelOracle(config, weight_key, None)

    # Handle explicit moe_backend from user.
    runner_backend = config.moe_backend
    if runner_backend != "auto":
        requested_backend = map_wna16_backend(runner_backend)
        reason = _backend_incompatibility_reason(
            requested_backend,
            config,
            quant_config,
            may_have_zp,
            may_have_bias,
            allow_tile_padding,
        )
        if reason is not None:
            raise ValueError(oracle.make_log_unsupported(requested_backend, reason))
        return oracle.return_or_raise(requested_backend)

    # Select kernels in order of backend.
    AVAILABLE_BACKENDS = _get_priority_backends()

    for backend in AVAILABLE_BACKENDS:
        reason = _backend_incompatibility_reason(
            backend,
            config,
            quant_config,
            may_have_zp,
            may_have_bias,
            allow_tile_padding,
        )
        if reason is not None:
            logger.debug_once(oracle.make_log_unsupported(backend, reason))
            continue
        kernel_cls, _ = oracle.try_backend(backend)
        if kernel_cls is not None:
            logger.info_once(oracle.make_log_backend(backend))
            return backend, kernel_cls

    raise NotImplementedError(
        "No WNA16 MoE backend supports the deployment configuration."
    )