From 84b2a8a7e7ec8c7dac98f9d3a300c21ca7d2d5e7 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Fri, 29 May 2026 13:11:17 -0400 Subject: [PATCH] [MoE Refactor] WNA16 MoE backend selection into oracle module (#42553) Signed-off-by: Bill Nell Co-authored-by: Claude --- .../model_executor/layers/fused_moe/config.py | 67 +--- .../layers/fused_moe/experts/marlin_moe.py | 7 +- .../fused_moe/experts/trtllm_mxint4_moe.py | 159 ++++++++ .../layers/fused_moe/oracle/int_wna16.py | 295 ++++++++++---- .../layers/quantization/auto_gptq.py | 3 +- .../layers/quantization/awq_marlin.py | 12 +- .../compressed_tensors_moe_wna16_marlin.py | 379 +++++++----------- .../layers/quantization/utils/quant_utils.py | 3 + 8 files changed, 547 insertions(+), 378 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index da9d5d3e23..b87b87f136 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -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, diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index a358796a7d..8bb9e5bdc0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -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 ( diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py new file mode 100644 index 0000000000..a65873aca4 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -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, + ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index 42b62f8e34..9de7a6ba11 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -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, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 85ee4061ad..1821fd5c7f 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -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( diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 14ebd920fd..81c0fcb331 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -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, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index c56962422b..2d629d73ed 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -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, ) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 29fa8bcbac..947e428ca8 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -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)