obscura/vllm
2
0
Fork 0
mirror of https://github.com/vllm-project/vllm.git synced 2026-08-10 15:58:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

[MoE Refactor] WNA16 MoE backend selection into oracle module (#42553)

Browse Source 
Signed-off-by: Bill Nell <[email protected]>
Co-authored-by: Claude <[email protected]>
This commit is contained in:
bnellnm
2026-05-29 13:11:17 -04:00
committed by GitHub
co-authored by Claude
parent 4ff865c38e
commit 84b2a8a7e7
8 changed files with 547 additions and 378 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/fused_moe/config.py
+16 -51
View File
@@ -869,19 +869,23 @@ def nvfp4_w4a16_moe_quant_config(
def int4_w4a16_moe_quant_config(
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
w1_zp: torch.Tensor | None,
w2_zp: torch.Tensor | None,
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,
) -> 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),
_a2=FusedMoEQuantDesc(shape=group_shape),
_w1=FusedMoEQuantDesc("int4", group_shape, w1_scale, None, w1_zp),
_w2=FusedMoEQuantDesc("int4", group_shape, w2_scale, None, w2_zp),
_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),
)
@@ -922,19 +926,21 @@ def fp8_w8a16_moe_quant_config(
def int8_w8a16_moe_quant_config(
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
w1_zp: torch.Tensor | None,
w2_zp: torch.Tensor | None,
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,
) -> 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),
_a2=FusedMoEQuantDesc(shape=group_shape),
_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),
)
@@ -965,47 +971,6 @@ def int4_w4afp8_moe_quant_config(
)
def awq_marlin_moe_quant_config(
w1_scale: torch.Tensor,
w2_scale: torch.Tensor,
w1_zp: torch.Tensor | None,
w2_zp: torch.Tensor | None,
weight_bits: int,
group_size: int,
w1_bias: torch.Tensor | None = None,
w2_bias: torch.Tensor | None = None,
a1_gscale: torch.Tensor | None = None,
a2_gscale: torch.Tensor | None = None,
) -> FusedMoEQuantConfig:
"""
Construct a quant config for awq marlin quantization.
a1_gscale / a2_gscale are optional global scales applied to activation
quantization scales when Marlin runs with 8-bit activations.
"""
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 AWQ (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, alpha_or_gscale=a1_gscale),
_a2=FusedMoEQuantDesc(dtype=None, shape=a_shape, alpha_or_gscale=a2_gscale),
_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 biased_moe_quant_config(
w1_bias: torch.Tensor | None,
w2_bias: torch.Tensor | None,
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+6 -1
View File
@@ -44,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
kFp8StaticChannelSym,
kFp8StaticTensorSym,
kInt4Static,
kInt4Static32,
kInt8Static,
kMxfp4Static,
kMxfp8Static,
@@ -566,8 +567,9 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
quant_config.use_mxfp4_w4a16
or quant_config.use_nvfp4_w4a16
or quant_config.use_int4_w4a16
or quant_config.use_int8_w8a16
or quant_config.use_fp8_w8a16
), "Supports only {mxfp,nvfp,int}4_w4a16 or fp8_w8a16"
), "Supports only {mxfp,nvfp,int}4_w4a16, int8_w8a16 or fp8_w8a16"
self.w13_g_idx = w13_g_idx
self.w2_g_idx = w2_g_idx
self.w13_g_idx_sort_indices = w13_g_idx_sort_indices
@@ -608,6 +610,7 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
kNvfp4Static,
kInt4Static,
kInt8Static,
kInt4Static32,
]
return weight_key in SUPPORTED_W
@@ -640,6 +643,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular):
if self.w1_zp is not None or self.w2_zp is not None:
return scalar_types.uint4.id
return scalar_types.uint4b8.id
elif self.quant_config.use_int8_w8a16:
return scalar_types.uint8b128.id
elif self.quant_config.use_mxfp4_w4a16 or self.quant_config.use_nvfp4_w4a16:
return scalar_types.float4_e2m1f.id
elif (
vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py
+159
View File
@@ -0,0 +1,159 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEParallelConfig,
FusedMoEQuantConfig,
RoutingMethodType,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
kInt4Static32,
)
from vllm.platforms import current_platform
class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic):
"""
FlashInfer TRT-LLM MxInt4 MoE kernel. Monolithic interface
(fused router + experts).
Wraps flashinfer_trtllm_mxint4_moe().
"""
def __init__(
self,
moe_config: FusedMoEConfig,
quant_config: FusedMoEQuantConfig,
):
super().__init__(moe_config, quant_config)
self.topk = moe_config.experts_per_token
self.intermediate_size_per_partition = (
moe_config.intermediate_size_per_partition
)
self.local_num_experts = moe_config.num_local_experts
self.ep_rank = moe_config.ep_rank
self.routing_method = moe_config.routing_method
@staticmethod
def _supports_current_device() -> bool:
from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501
is_flashinfer_mxint4_moe_available,
)
p = current_platform
return (
p.is_cuda()
and p.is_device_capability_family(100)
and is_flashinfer_mxint4_moe_available()
)
@staticmethod
def _supports_no_act_and_mul() -> bool:
return False
@staticmethod
def _supports_quant_scheme(
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
return (weight_key, activation_key) == (kInt4Static32, None)
@staticmethod
def _supports_activation(activation: MoEActivation) -> bool:
# FlashInfer MxInt4 uses a fused SwiGLU activation.
return activation == MoEActivation.SWIGLUOAI
@staticmethod
def _supports_parallel_config(
moe_parallel_config: FusedMoEParallelConfig,
) -> bool:
return (
not moe_parallel_config.use_all2all_kernels
or moe_parallel_config.use_ag_rs_all2all_kernels
) and not moe_parallel_config.enable_eplb
@staticmethod
def _supports_routing_method(
routing_method: RoutingMethodType,
weight_key: QuantKey | None,
activation_key: QuantKey | None,
) -> bool:
return routing_method in [
RoutingMethodType.Renormalize,
RoutingMethodType.RenormalizeNaive,
RoutingMethodType.DeepSeekV3,
RoutingMethodType.Llama4,
RoutingMethodType.Simulated,
]
@staticmethod
def _supports_router_logits_dtype(
router_logits_dtype: torch.dtype | None,
routing_method: RoutingMethodType,
) -> bool:
if router_logits_dtype == torch.float32:
# DeepSeekV3 routing handles float32 logits internally.
# Simulated routing generates synthetic decisions.
return routing_method in (
RoutingMethodType.DeepSeekV3,
RoutingMethodType.Simulated,
)
return True
@staticmethod
def activation_format() -> mk.FusedMoEActivationFormat:
return mk.FusedMoEActivationFormat.Standard
def supports_expert_map(self) -> bool:
return False
@property
def expects_unquantized_inputs(self) -> bool:
# The kernel handles quantization internally.
return True
def apply(
self,
hidden_states: torch.Tensor,
w1: torch.Tensor,
w2: torch.Tensor,
router_logits: torch.Tensor,
activation: MoEActivation,
global_num_experts: int,
expert_map: torch.Tensor | None,
a1q_scale: torch.Tensor | None,
apply_router_weight_on_input: bool,
num_expert_group: int | None = None,
e_score_correction_bias: torch.Tensor | None = None,
routed_scaling_factor: float | None = None,
topk_group: int | None = None,
) -> torch.Tensor:
from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( # noqa: E501
flashinfer_trtllm_mxint4_moe,
)
assert self.w1_scale is not None
assert self.w2_scale is not None
return flashinfer_trtllm_mxint4_moe(
x=hidden_states,
router_logits=router_logits,
w13_weight_packed=w1,
w13_weight_scale=self.w1_scale,
w2_weight_packed=w2,
w2_weight_scale=self.w2_scale,
global_num_experts=global_num_experts,
top_k=self.topk,
intermediate_size_per_partition=self.intermediate_size_per_partition,
local_num_experts=self.local_num_experts,
ep_rank=self.ep_rank,
num_expert_group=num_expert_group,
topk_group=topk_group,
e_score_correction_bias=e_score_correction_bias,
routing_method_type=self.routing_method,
)
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
+220 -75
View File
@@ -2,9 +2,12 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import sys
from enum import Enum
from typing import TYPE_CHECKING
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
@@ -12,10 +15,16 @@ 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 (
@@ -29,16 +38,13 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
)
from vllm.platforms import current_platform
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
from vllm.model_executor.layers.quantization.awq_marlin import AWQMarlinConfig
logger = init_logger(__name__)
class WNA16MoEBackend(Enum):
MARLIN = "MARLIN"
BATCHED_MARLIN = "BATCHED_MARLIN"
FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM"
XPU = "XPU"
@@ -47,26 +53,17 @@ def backend_to_kernel_cls(
) -> list[type[mk.FusedMoEExperts]]:
"""Return the experts class for the given backend, or None for NONE."""
if backend == WNA16MoEBackend.MARLIN:
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
MarlinExperts,
)
return [MarlinExperts]
elif backend == WNA16MoEBackend.BATCHED_MARLIN:
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
BatchedMarlinExperts,
)
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}")
@@ -77,23 +74,25 @@ def _get_priority_backends() -> list[WNA16MoEBackend]:
"""
if current_platform.is_xpu():
return [WNA16MoEBackend.XPU]
return [
_AVAILABLE_BACKENDS = [
WNA16MoEBackend.FLASHINFER_TRTLLM,
WNA16MoEBackend.MARLIN,
WNA16MoEBackend.BATCHED_MARLIN,
]
return _AVAILABLE_BACKENDS
def select_wna16_moe_backend(
config: FusedMoEConfig,
weight_key: QuantKey,
weight_bits: int,
) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]:
"""Select the WNA16 MoE backend.
Args:
config: the shared ``FusedMoEConfig`` for this layer.
weight_bits: quantization bit-width (4 or 8). 8-bit weights are not
supported by the modular Marlin kernel, so ``NONE`` is returned.
weight_key: The QuantKey describing the weight quantization.
Must have int4 or int8 type.
Returns:
A tuple of (``WNA16MoEBackend``, experts class or ``None``).
@@ -156,15 +155,55 @@ def select_wna16_moe_backend(
)
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,
) -> 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,
)
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,
)
def make_wna16_moe_kernel(
moe_quant_config: FusedMoEQuantConfig,
moe_config: FusedMoEConfig,
experts_cls: type[mk.FusedMoEExperts] | None,
is_k_full: bool,
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,
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 (
@@ -174,16 +213,37 @@ def make_wna16_moe_kernel(
XPUExpertsWNA16,
)
assert experts_cls in (MarlinExperts, BatchedMarlinExperts, 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
assert isinstance(prepare_finalize, mk.FusedMoEPrepareAndFinalizeModular)
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 (
@@ -199,30 +259,20 @@ def make_wna16_moe_kernel(
elif (
prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts
):
assert experts_cls == BatchedMarlinExperts
max_num_tokens = prepare_finalize.max_num_tokens_per_rank()
assert max_num_tokens is not None
experts = BatchedMarlinExperts(
experts = experts_cls(
max_num_tokens=max_num_tokens,
num_dispatchers=prepare_finalize.num_dispatchers(),
moe_config=moe_config,
quant_config=moe_quant_config,
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,
**extra_args,
)
else:
assert experts_cls == MarlinExperts
experts = MarlinExperts(
experts = experts_cls(
moe_config=moe_config,
quant_config=moe_quant_config,
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,
**extra_args,
)
return mk.FusedMoEKernel(
@@ -236,10 +286,74 @@ def make_wna16_moe_kernel(
# ---------------------------------------------------------------------------
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,
quant_config: "AutoGPTQConfig",
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,
@@ -292,6 +406,7 @@ def _process_weights_marlin(
# --- 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
@@ -303,7 +418,7 @@ def _process_weights_marlin(
marlin_w2_scales = w2_scales
# --- Process act_order (g_idx) ---
if quant_config.desc_act:
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)
@@ -314,6 +429,8 @@ def _process_weights_marlin(
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
@@ -338,17 +455,17 @@ def _process_weights_marlin(
marlin_w13_qweight = ops.gptq_marlin_moe_repack(
marlin_w13_qweight,
w13_g_idx_sort_indices,
marlin_w13_qweight.shape[1] * quant_config.pack_factor,
marlin_w13_qweight.shape[1] * pack_factor,
marlin_w13_qweight.shape[2],
quant_config.quant_type.size_bits,
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] * quant_config.pack_factor,
marlin_w2_qweight.shape[1] * pack_factor,
marlin_w2_qweight.shape[2],
quant_config.quant_type.size_bits,
num_bits,
is_a_8bit=is_a_8bit,
)
@@ -357,19 +474,15 @@ def _process_weights_marlin(
s=marlin_w13_scales,
size_k=layer.intermediate_size_per_partition,
size_n=marlin_w13_scales.shape[2],
group_size=quant_config.group_size,
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]
* (
quant_config.group_size
if quant_config.group_size != -1
else quant_config.pack_factor
),
size_k=marlin_w2_scales.shape[1] * group_size_or_pack_factor,
size_n=marlin_w2_scales.shape[2],
group_size=quant_config.group_size,
group_size=group_size,
is_a_8bit=is_a_8bit,
)
@@ -409,7 +522,9 @@ def _process_weights_marlin(
def _process_awq_weights_marlin(
layer: torch.nn.Module,
quant_config: "AWQMarlinConfig",
weight_bits: int,
pack_factor: int,
group_size: int,
input_dtype: torch.dtype | None,
w13_qweight: torch.Tensor,
w2_qweight: torch.Tensor,
@@ -475,16 +590,16 @@ def _process_awq_weights_marlin(
w13_qweight,
w13_g_idx_sort_indices,
size_k=w13_qweight.shape[1],
size_n=w13_qweight.shape[2] * quant_config.pack_factor,
num_bits=quant_config.weight_bits,
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] * quant_config.pack_factor,
num_bits=quant_config.weight_bits,
size_n=w2_qweight.shape[2] * pack_factor,
num_bits=weight_bits,
is_a_8bit=is_a_8bit,
)
@@ -492,7 +607,7 @@ def _process_awq_weights_marlin(
s=w13_scales,
size_k=layer.intermediate_size_per_partition,
size_n=w13_scales.shape[2],
group_size=quant_config.group_size,
group_size=group_size,
is_a_8bit=is_a_8bit,
)
if input_dtype == torch.int8 and layer.num_groups_w13 > 1:
@@ -504,7 +619,7 @@ def _process_awq_weights_marlin(
s=w2_scales,
size_k=layer.intermediate_size_per_partition,
size_n=w2_scales.shape[2],
group_size=quant_config.group_size,
group_size=group_size,
is_a_8bit=is_a_8bit,
)
if input_dtype == torch.int8 and layer.num_groups_w2 > 1:
@@ -515,15 +630,15 @@ def _process_awq_weights_marlin(
marlin_w13_qzeros = moe_awq_to_marlin_zero_points(
w13_qzeros,
size_k=w13_qzeros.shape[1],
size_n=w13_qzeros.shape[2] * quant_config.pack_factor,
num_bits=quant_config.weight_bits,
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] * quant_config.pack_factor,
num_bits=quant_config.weight_bits,
size_n=w2_qzeros.shape[2] * pack_factor,
num_bits=weight_bits,
is_a_8bit=is_a_8bit,
)
@@ -616,7 +731,7 @@ def _process_weights_xpu(
def convert_to_wna16_moe_kernel_format(
backend: WNA16MoEBackend,
layer: torch.nn.Module,
quant_config: QuantizationConfig,
quant_config: QuantizationConfig | QuantizationArgs | None,
input_dtype: torch.dtype | None,
w13: torch.Tensor,
w2: torch.Tensor,
@@ -669,9 +784,16 @@ def convert_to_wna16_moe_kernel_format(
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,
quant_config,
weight_bits,
pack_factor,
group_size,
input_dtype,
w13,
w2,
@@ -682,19 +804,30 @@ def convert_to_wna16_moe_kernel_format(
w13_bias,
w2_bias,
)
if not isinstance(quant_config, AutoGPTQConfig):
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 or "
"AWQMarlinConfig, got "
f"{type(quant_config).__name__}."
"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,
quant_config,
input_dtype,
num_bits,
pack_factor,
group_size,
actorder,
w13,
w2,
w13_scale,
@@ -706,7 +839,19 @@ def convert_to_wna16_moe_kernel_format(
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,
vllm/model_executor/layers/quantization/auto_gptq.py
+2 -1
View File
@@ -483,7 +483,8 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):
weight_key = QuantKey(quant_type, scale)
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
moe, weight_key, quant_config.weight_bits
moe,
weight_key,
)
def create_weights(
vllm/model_executor/layers/quantization/awq_marlin.py
+5 -7
View File
@@ -29,6 +29,7 @@ from vllm.model_executor.layers.fused_moe.config import (
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
convert_to_wna16_moe_kernel_format,
make_wna16_moe_kernel,
make_wna16_moe_quant_config,
select_wna16_moe_backend,
)
from vllm.model_executor.layers.linear import (
@@ -521,7 +522,8 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
self.input_dtype = None
self.use_marlin = True
self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend(
moe, kInt4Static, quant_config.weight_bits
moe,
kInt4Static,
)
def create_weights(
@@ -706,15 +708,11 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):
)
def get_fused_moe_quant_config(self, layer: RoutedExperts) -> FusedMoEQuantConfig:
from vllm.model_executor.layers.fused_moe.config import (
awq_marlin_moe_quant_config,
)
return awq_marlin_moe_quant_config(
return make_wna16_moe_quant_config(
w1_scale=layer.w13_scales,
w2_scale=layer.w2_scales,
weight_bits=self.quant_config.weight_bits,
group_size=self.quant_config.group_size,
num_bits=self.quant_config.weight_bits,
w1_zp=getattr(layer, "w13_qzeros", None)
if self.quant_config.zero_point
else None,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
+136 -243
View File
@@ -1,16 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import enum
from enum import Enum
from typing import Any
import torch
from compressed_tensors.quantization import (
QuantizationArgs,
)
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
from vllm import _custom_ops as ops
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe import (
RoutedExperts,
@@ -19,12 +16,13 @@ from vllm.model_executor.layers.fused_moe import (
from vllm.model_executor.layers.fused_moe.config import (
FusedMoEConfig,
FusedMoEQuantConfig,
int4_w4a16_moe_quant_config,
)
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
BatchedMarlinExperts,
MarlinExperts,
fused_marlin_moe,
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,
@@ -32,27 +30,21 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa
WNA16_SUPPORTED_TYPES_MAP,
)
from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import (
flashinfer_trtllm_mxint4_moe,
is_flashinfer_mxint4_moe_available,
prepare_static_weights_for_trtllm_mxint4_moe,
)
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
get_marlin_input_dtype,
marlin_act_int8_process_scales,
marlin_make_workspace_new,
marlin_moe_permute_scales,
)
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 GPTQMarlinState(Enum):
REPACK = enum.auto()
READY = enum.auto()
class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
def __init__(
self,
@@ -75,19 +67,27 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
self.actorder = weight_quant.actorder
self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
self.use_flashinfer_mxint4_moe = (
is_flashinfer_mxint4_moe_available()
and self.group_size == 32
and weight_quant.num_bits == 4
)
self.kernel_backend = (
"Flashinfer" if self.use_flashinfer_mxint4_moe else "Marlin"
)
logger.info_once(
f"Using {self.kernel_backend} backend for WNA16 MoE "
f"(group_size={self.group_size}, num_bits={self.num_bits})",
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)
# 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(
@@ -114,6 +114,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
"num_groups_w2 must be provided for weight scales"
)
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": (
@@ -156,7 +157,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
"Marlin": (num_experts, num_groups_w2, hidden_size),
},
}
return shape_map[weight_name][self.kernel_backend]
backend_key = "Flashinfer" if is_flashinfer else "Marlin"
return shape_map[weight_name][backend_key]
def create_weights(
self,
@@ -172,7 +174,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
# Will transpose the loaded weight along the
# intermediate and hidden dim sizes. Will
# shard for TP along the transposed dims
is_transposed = self.kernel_backend != "Flashinfer"
is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM
extra_weight_attrs.update(
{"is_transposed": is_transposed, "quant_method": self.strategy}
)
@@ -319,200 +321,104 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
layer.a13_scale = None
layer.a2_scale = None
layer.marlin_state = GPTQMarlinState.REPACK
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
num_experts = layer.w13_weight_g_idx.shape[0]
device = layer.w13_weight_g_idx.device
if self.kernel_backend == "Flashinfer":
dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe(
layer.w13_weight_packed,
layer.w13_weight_scale,
layer.w2_weight_packed,
layer.w2_weight_scale,
)
replace_parameter(
layer, "w13_weight_packed", dict_weights_mxint4["gemm1_weights"]
)
replace_parameter(
layer, "w13_weight_scale", dict_weights_mxint4["gemm1_scales"]
)
replace_parameter(
layer, "w2_weight_packed", dict_weights_mxint4["gemm2_weights"]
)
replace_parameter(
layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"]
)
return None
is_a_8bit = (
self.marlin_input_dtype is not None
and self.marlin_input_dtype.itemsize == 1
# 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,
)
if self.marlin_input_dtype == torch.float8_e4m3fn:
# NOTE: for non-zp quantization format only
ops.marlin_int4_fp8_preprocess(layer.w13_weight_packed, inplace=True)
ops.marlin_int4_fp8_preprocess(layer.w2_weight_packed, inplace=True)
layer.w13_weight_scale.data = layer.w13_weight_scale.data * 512
layer.w2_weight_scale.data = layer.w2_weight_scale.data * 512
# 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)
# when running models with grouped act order,
# resort to g_idx values provided in checkpoint
if self.actorder == "group":
w13_g_idx_sort_indices = torch.empty_like(layer.w13_weight_g_idx)
w2_g_idx_sort_indices = torch.empty_like(layer.w2_weight_g_idx)
w13_sorted_g_idx = torch.empty_like(layer.w13_weight_g_idx)
w2_sorted_g_idx = torch.empty_like(layer.w2_weight_g_idx)
for e in range(num_experts):
w13_g_idx_sort_indices[e] = torch.argsort(layer.w13_weight_g_idx[e]).to(
torch.int32
)
w2_g_idx_sort_indices[e] = torch.argsort(layer.w2_weight_g_idx[e]).to(
torch.int32
)
w13_sorted_g_idx[e] = layer.w13_weight_g_idx[e][
w13_g_idx_sort_indices[e]
]
w2_sorted_g_idx[e] = layer.w2_weight_g_idx[e][w2_g_idx_sort_indices[e]]
replace_parameter(layer, "w13_weight_g_idx", w13_sorted_g_idx)
replace_parameter(layer, "w2_weight_g_idx", w2_sorted_g_idx)
# 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)
else:
layer.w13_weight_g_idx = torch.nn.Parameter(
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
requires_grad=False,
)
layer.w2_weight_g_idx = torch.nn.Parameter(
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
requires_grad=False,
)
layer.w13_g_idx_sort_indices = torch.nn.Parameter(
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
requires_grad=False,
)
layer.w2_g_idx_sort_indices = torch.nn.Parameter(
torch.empty((num_experts, 0), dtype=torch.int32, device=device),
requires_grad=False,
# 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
)
marlin_w13_qweight = ops.gptq_marlin_moe_repack(
layer.w13_weight_packed,
layer.w13_g_idx_sort_indices,
layer.w13_weight_packed.shape[1] * self.packed_factor,
layer.w13_weight_packed.shape[2],
self.num_bits,
is_a_8bit=is_a_8bit,
# 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,
)
replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight)
marlin_w2_qweight = ops.gptq_marlin_moe_repack(
layer.w2_weight_packed,
layer.w2_g_idx_sort_indices,
layer.w2_weight_packed.shape[1] * self.packed_factor,
layer.w2_weight_packed.shape[2],
self.num_bits,
is_a_8bit=is_a_8bit,
)
replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight)
# Repack scales
marlin_w13_scales = marlin_moe_permute_scales(
s=layer.w13_weight_scale,
size_k=layer.w13_weight_packed.shape[2],
size_n=layer.w13_weight_scale.shape[2],
group_size=self.group_size,
is_a_8bit=is_a_8bit,
)
if self.marlin_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
)
layer.register_parameter(
"w13_input_global_scale",
torch.nn.Parameter(w13_input_global_scale, requires_grad=False),
)
replace_parameter(layer, "w13_weight_scale", marlin_w13_scales)
marlin_w2_scales = marlin_moe_permute_scales(
s=layer.w2_weight_scale,
size_k=layer.w2_weight_scale.shape[1]
* (self.group_size if self.group_size != -1 else self.packed_factor),
size_n=layer.w2_weight_scale.shape[2],
group_size=self.group_size,
is_a_8bit=is_a_8bit,
)
if self.marlin_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
)
layer.register_parameter(
"w2_input_global_scale",
torch.nn.Parameter(w2_input_global_scale, requires_grad=False),
)
replace_parameter(layer, "w2_weight_scale", marlin_w2_scales)
layer.workspace = marlin_make_workspace_new(device, 4)
def get_fused_moe_quant_config(
self, layer: torch.nn.Module
) -> FusedMoEQuantConfig | None:
if self.num_bits != 4:
return None
return int4_w4a16_moe_quant_config(
return make_wna16_moe_quant_config(
w1_scale=layer.w13_weight_scale,
w2_scale=layer.w2_weight_scale,
w1_zp=None,
w2_zp=None,
block_shape=[0, self.group_size],
group_size=self.group_size,
num_bits=self.num_bits,
)
def select_gemm_impl(
self,
prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular,
layer: torch.nn.Module,
) -> mk.FusedMoEExpertsModular:
assert self.num_bits == 4, "only supporting w4"
layer.w13_weight = layer.w13_weight_packed
layer.w2_weight = layer.w2_weight_packed
assert all([w is not None for w in [layer.w13_weight, layer.w2_weight]])
assert self.moe_quant_config is not None
if (
prepare_finalize.activation_format
== mk.FusedMoEActivationFormat.BatchedExperts
):
max_num_tokens_per_rank = prepare_finalize.max_num_tokens_per_rank()
assert max_num_tokens_per_rank is not None
return BatchedMarlinExperts(
max_num_tokens=max_num_tokens_per_rank,
num_dispatchers=prepare_finalize.num_dispatchers(),
moe_config=self.moe,
quant_config=self.moe_quant_config,
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,
)
else:
return MarlinExperts(
moe_config=self.moe,
quant_config=self.moe_quant_config,
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,
)
@property
def is_monolithic(self) -> bool:
return self.kernel_backend == "Flashinfer"
def apply_monolithic(
self,
layer: RoutedExperts,
@@ -520,23 +426,21 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
router_logits: torch.Tensor,
input_ids: torch.Tensor | None = None,
) -> torch.Tensor:
assert self.kernel_backend == "Flashinfer"
return flashinfer_trtllm_mxint4_moe(
x=x,
router_logits=router_logits,
w13_weight_packed=layer.w13_weight_packed,
w13_weight_scale=layer.w13_weight_scale,
w2_weight_packed=layer.w2_weight_packed,
w2_weight_scale=layer.w2_weight_scale,
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,
top_k=layer.top_k,
intermediate_size_per_partition=layer.intermediate_size_per_partition,
local_num_experts=layer.local_num_experts,
ep_rank=layer.ep_rank,
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,
routing_method_type=layer.routing_method_type,
routed_scaling_factor=layer.routed_scaling_factor,
)
def apply(
@@ -548,29 +452,18 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
shared_experts: SharedExperts | None,
shared_experts_input: torch.Tensor | None,
) -> torch.Tensor:
assert self.kernel_backend == "Marlin"
return fused_marlin_moe(
assert not self.is_monolithic
assert self.moe_kernel is not None
return self.moe_kernel.apply(
x,
layer.w13_weight_packed,
layer.w2_weight_packed,
None,
None,
layer.w13_weight_scale,
layer.w2_weight_scale,
layer.w13_weight,
layer.w2_weight,
topk_weights,
topk_ids,
input_global_scale1=getattr(layer, "w13_input_global_scale", None),
input_global_scale2=getattr(layer, "w2_input_global_scale", None),
quant_type_id=self.quant_type.id,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
global_num_experts=layer.global_num_experts,
activation=layer.activation,
global_num_experts=layer.global_num_experts,
expert_map=layer.expert_map,
g_idx1=layer.w13_weight_g_idx,
g_idx2=layer.w2_weight_g_idx,
sort_indices1=layer.w13_g_idx_sort_indices,
sort_indices2=layer.w2_g_idx_sort_indices,
workspace=layer.workspace,
input_dtype=self.marlin_input_dtype,
is_k_full=self.is_k_full,
apply_router_weight_on_input=layer.apply_router_weight_on_input,
shared_experts=shared_experts,
shared_experts_input=shared_experts_input,
)
vllm/model_executor/layers/quantization/utils/quant_utils.py
+3
View File
@@ -178,6 +178,9 @@ kInt4Static = QuantKey(INT4_DTYPE, scale=kInt4StaticGroupScale, symmetric=True)
kInt8StaticGroupScale = ScaleDesc(torch.float16, True, GroupShape(1, -1))
kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True)
kInt4Static32GroupScale = ScaleDesc(torch.float16, True, GroupShape(1, 32))
kInt4Static32 = QuantKey(INT4_DTYPE, scale=kInt4Static32GroupScale, symmetric=True)
kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True)
kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)