From e52950c776b9edeb5179cafcdb1628dcca167587 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 7 Apr 2026 13:01:13 +0000 Subject: [PATCH 01/53] Refactor CompressedTensorsWNA16MarlinMoEMethod to use kernel abstraction oracle Converts the WNA16 Marlin MoE method to the new oracle-based kernel abstraction pattern, matching the structure of CompressedTensorsW8A8Fp8MoEMethod. Key changes: - Add oracle/wna16.py with select_wna16_moe_backend(), make_wna16_moe_quant_config(), and make_wna16_moe_kernel() following the same pattern as oracle/fp8.py - Add experts/trtllm_mxint4_moe.py with TrtLlmMxint4ExpertsMonolithic wrapping flashinfer_trtllm_mxint4_moe() as a FusedMoEExpertsMonolithic subclass - Update CompressedTensorsWNA16MarlinMoEMethod.__init__ to call select_wna16_moe_backend() instead of inline backend selection - Update process_weights_after_loading to build moe_quant_config and moe_kernel for 4-bit paths (both Flashinfer and Marlin); non-4-bit Marlin retains the direct fused_marlin_moe() fallback in apply() - Remove GPTQMarlinState enum and select_gemm_impl() override (base class raises ValueError; moe_kernel is set so the old path is never reached for 4-bit) - apply() and apply_monolithic() delegate to moe_kernel when available Co-authored-by: Claude https://claude.ai/code/session_01FLcv14ytXvwBUCsKZqpze1 Signed-off-by: Bill Nell --- .../fused_moe/experts/trtllm_mxint4_moe.py | 153 +++++++++++++++ .../layers/fused_moe/oracle/wna16.py | 180 ++++++++++++++++++ .../compressed_tensors_moe_wna16_marlin.py | 174 ++++++++--------- 3 files changed, 420 insertions(+), 87 deletions(-) create mode 100644 vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py create mode 100644 vllm/model_executor/layers/fused_moe/oracle/wna16.py 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 000000000000..a370e582f9f6 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -0,0 +1,153 @@ +# 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 + + +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, + ) + + return 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: + # Selected directly by WNA16 oracle based on conditions, not via + # QuantKey matching. Always return True here. + return True + + @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 + and not moe_parallel_config.enable_eplb + and moe_parallel_config.dp_size <= 1 + ) + + @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/wna16.py b/vllm/model_executor/layers/fused_moe/oracle/wna16.py new file mode 100644 index 000000000000..d0867341eac9 --- /dev/null +++ b/vllm/model_executor/layers/fused_moe/oracle/wna16.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from enum import Enum + +import torch + +import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + maybe_make_prepare_finalize, +) +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEQuantConfig, + int4_w4a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( + BatchedMarlinExperts, + MarlinExperts, +) +from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( + SharedExperts, +) +from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( + is_flashinfer_mxint4_moe_available, +) + +logger = init_logger(__name__) + + +class WNA16MoeBackend(Enum): + MARLIN = "MARLIN" + FLASHINFER = "FLASHINFER" + + +def select_wna16_moe_backend( + config: FusedMoEConfig, + num_bits: int, + group_size: int, +) -> tuple[WNA16MoeBackend, type[mk.FusedMoEExperts] | None]: + """ + Select the WNA16 MoE backend and experts class. + + Returns (backend, experts_cls) where experts_cls may be None for + non-4-bit Marlin which falls back to the direct fused_marlin_moe() path. + """ + if ( + is_flashinfer_mxint4_moe_available() + and num_bits == 4 + and group_size == 32 + ): + from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( # noqa: E501 + TrtLlmMxint4ExpertsMonolithic, + ) + + logger.info_once( + "Using Flashinfer backend for WNA16 MoE " + f"(group_size={group_size}, num_bits={num_bits})", + scope="local", + ) + return WNA16MoeBackend.FLASHINFER, TrtLlmMxint4ExpertsMonolithic + + if num_bits != 4: + # Non-4-bit Marlin falls back to the direct fused_marlin_moe() path. + logger.info_once( + "Using Marlin backend for WNA16 MoE " + f"(group_size={group_size}, num_bits={num_bits})", + scope="local", + ) + return WNA16MoeBackend.MARLIN, None + + # 4-bit Marlin: select batched vs standard based on activation format. + activation_format = ( + mk.FusedMoEActivationFormat.BatchedExperts + if config.moe_parallel_config.use_batched_activation_format + else mk.FusedMoEActivationFormat.Standard + ) + experts_cls: type[mk.FusedMoEExperts] = ( + BatchedMarlinExperts + if activation_format == mk.FusedMoEActivationFormat.BatchedExperts + else MarlinExperts + ) + logger.info_once( + "Using Marlin backend for WNA16 MoE " + f"(group_size={group_size}, num_bits={num_bits})", + scope="local", + ) + return WNA16MoeBackend.MARLIN, experts_cls + + +def make_wna16_moe_quant_config( + layer: torch.nn.Module, + group_size: int, +) -> FusedMoEQuantConfig: + """Create the FusedMoEQuantConfig for 4-bit WNA16 MoE.""" + return int4_w4a16_moe_quant_config( + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) + + +def make_wna16_moe_kernel( + moe_quant_config: FusedMoEQuantConfig, + moe_config: FusedMoEConfig, + backend: WNA16MoeBackend, + experts_cls: type[mk.FusedMoEExperts], + 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, + is_k_full: bool = True, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] + | None = None, + shared_experts: SharedExperts | None = None, +) -> mk.FusedMoEKernel: + """ + Create the FusedMoEKernel for WNA16 MoE. + + For the Flashinfer (mxint4) backend, a monolithic kernel is created. + For the Marlin backend, a modular kernel is created using MarlinExperts + or BatchedMarlinExperts depending on the activation format. + """ + is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) + + prepare_finalize = maybe_make_prepare_finalize( + moe=moe_config, + quant_config=moe_quant_config, + routing_tables=routing_tables, + allow_new_interface=True, + use_monolithic=is_monolithic, + ) + assert prepare_finalize is not None + + logger.info_once( + "Using %s", prepare_finalize.__class__.__name__, scope="local" + ) + + if is_monolithic: + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + ) + elif ( + prepare_finalize.activation_format + == mk.FusedMoEActivationFormat.BatchedExperts + ): + max_num_tokens = prepare_finalize.max_num_tokens_per_rank() + assert max_num_tokens is not None + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + max_num_tokens=max_num_tokens, + num_dispatchers=prepare_finalize.num_dispatchers(), + 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, + ) + else: + 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, + ) + + return mk.FusedMoEKernel( + prepare_finalize, + experts, + shared_experts=shared_experts, + inplace=(not moe_config.disable_inplace and not is_monolithic), + ) 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 c60691df5835..18ab622d5d19 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,15 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import enum -from enum import Enum - 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 ( @@ -19,13 +15,16 @@ 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, +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( fused_marlin_moe, ) +from vllm.model_executor.layers.fused_moe.oracle.wna16 import ( + WNA16MoeBackend, + 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, ) @@ -33,8 +32,6 @@ 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 ( @@ -48,11 +45,6 @@ logger = init_logger(__name__) -class GPTQMarlinState(Enum): - REPACK = enum.auto() - READY = enum.auto() - - class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): def __init__( self, @@ -77,17 +69,12 @@ def __init__( 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})", + + # Select WNA16 MoE backend via oracle. + self.wna16_backend, self.experts_cls = select_wna16_moe_backend( + config=self.moe, + num_bits=self.num_bits, + group_size=self.group_size or -1, ) def get_weight_shape( @@ -114,6 +101,7 @@ def get_weight_shape( "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 shape_map = { "w13_weight": { "Flashinfer": ( @@ -156,7 +144,8 @@ def get_weight_shape( "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 +161,7 @@ def create_weights( # 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 extra_weight_attrs.update( {"is_transposed": is_transposed, "quant_method": self.strategy} ) @@ -319,12 +308,11 @@ def create_weights( 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": + if self.wna16_backend == WNA16MoeBackend.FLASHINFER: dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe( layer.w13_weight_packed, layer.w13_weight_scale, @@ -343,7 +331,24 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter( layer, "w2_weight_scale", dict_weights_mxint4["gemm2_scales"] ) - return None + # 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 + + # Setup moe_kernel for the Flashinfer monolithic path. + 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 + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + backend=self.wna16_backend, + experts_cls=self.experts_cls, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, + ) + return is_a_8bit = ( self.marlin_input_dtype is not None @@ -458,60 +463,42 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: 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( - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, self.group_size], - ) - - def select_gemm_impl( - self, - prepare_finalize: mk.FusedMoEPrepareAndFinalizeModular, - layer: torch.nn.Module, - ) -> mk.FusedMoEExpertsModular: - assert self.num_bits == 4, "only supporting w4" + # 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 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( + + # Setup moe_kernel for the 4-bit Marlin path. + if self.num_bits == 4 and 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 + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, moe_config=self.moe, - quant_config=self.moe_quant_config, + backend=self.wna16_backend, + experts_cls=self.experts_cls, 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, + routing_tables=layer._maybe_init_expert_routing_tables(), + shared_experts=layer.shared_experts, ) + def get_fused_moe_quant_config( + self, layer: torch.nn.Module + ) -> FusedMoEQuantConfig | None: + if self.num_bits != 4: + return None + return make_wna16_moe_quant_config( + layer=layer, + group_size=self.group_size, + ) + @property def is_monolithic(self) -> bool: - return self.kernel_backend == "Flashinfer" + return self.wna16_backend == WNA16MoeBackend.FLASHINFER def apply_monolithic( self, @@ -520,23 +507,21 @@ def apply_monolithic( 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.wna16_backend == WNA16MoeBackend.FLASHINFER + 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,7 +533,22 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - assert self.kernel_backend == "Marlin" + assert not self.is_monolithic + if self.moe_kernel is not None: + # 4-bit Marlin: use the modular kernel abstraction. + return self.moe_kernel.apply( + x, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, + ) + # Non-4-bit Marlin fallback: call fused_marlin_moe() directly. return fused_marlin_moe( x, layer.w13_weight_packed, From 5887f85bfc56dfa18c7f2c533437058c694a88e7 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 13 May 2026 16:58:52 +0000 Subject: [PATCH 02/53] wip Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/marlin_moe.py | 9 ++ .../layers/fused_moe/oracle/wna16.py | 72 +++++++------- .../compressed_tensors_moe_wna16_marlin.py | 93 ++++++------------- 3 files changed, 73 insertions(+), 101 deletions(-) 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 e006383e2a3f..b1b2fedc5f20 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -557,6 +557,8 @@ def __init__( 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, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, is_k_full: bool = True, ): # TODO (varun) : Enable activation quantization @@ -564,6 +566,7 @@ def __init__( 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" self.w13_g_idx = w13_g_idx @@ -636,6 +639,8 @@ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bo def quant_type_id(self) -> int: if self.quant_config.use_int4_w4a16: 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 ( @@ -752,6 +757,8 @@ def apply( w2_scale=self.w2_scale, topk_weights=topk_weights, topk_ids=topk_ids, + input_global_scale1=self.input_global_scale1, + input_global_scale2=self.input_global_scale2, global_scale1=self.g1_alphas, global_scale2=self.g2_alphas, quant_type_id=self.quant_type_id, @@ -851,6 +858,8 @@ def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None: w2_scale=self.w2_scale, topk_weights=topk_weights, topk_ids=topk_ids, + input_global_scale1=self.input_global_scale1, + input_global_scale2=self.input_global_scale2, global_scale1=self.g1_alphas, global_scale2=self.g2_alphas, quant_type_id=self.quant_type_id, diff --git a/vllm/model_executor/layers/fused_moe/oracle/wna16.py b/vllm/model_executor/layers/fused_moe/oracle/wna16.py index d0867341eac9..6141953ba398 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/wna16.py @@ -14,14 +14,12 @@ FusedMoEConfig, FusedMoEQuantConfig, int4_w4a16_moe_quant_config, + int8_w8a16_moe_quant_config, ) from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( BatchedMarlinExperts, MarlinExperts, ) -from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( - SharedExperts, -) from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( is_flashinfer_mxint4_moe_available, ) @@ -42,14 +40,9 @@ def select_wna16_moe_backend( """ Select the WNA16 MoE backend and experts class. - Returns (backend, experts_cls) where experts_cls may be None for - non-4-bit Marlin which falls back to the direct fused_marlin_moe() path. + Returns (backend, experts_cls). """ - if ( - is_flashinfer_mxint4_moe_available() - and num_bits == 4 - and group_size == 32 - ): + if is_flashinfer_mxint4_moe_available() and num_bits == 4 and group_size == 32: from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( # noqa: E501 TrtLlmMxint4ExpertsMonolithic, ) @@ -61,16 +54,7 @@ def select_wna16_moe_backend( ) return WNA16MoeBackend.FLASHINFER, TrtLlmMxint4ExpertsMonolithic - if num_bits != 4: - # Non-4-bit Marlin falls back to the direct fused_marlin_moe() path. - logger.info_once( - "Using Marlin backend for WNA16 MoE " - f"(group_size={group_size}, num_bits={num_bits})", - scope="local", - ) - return WNA16MoeBackend.MARLIN, None - - # 4-bit Marlin: select batched vs standard based on activation format. + # 4/8-bit Marlin: select batched vs standard based on activation format. activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if config.moe_parallel_config.use_batched_activation_format @@ -90,17 +74,29 @@ def select_wna16_moe_backend( def make_wna16_moe_quant_config( - layer: torch.nn.Module, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, group_size: int, + num_bits: int, ) -> FusedMoEQuantConfig: """Create the FusedMoEQuantConfig for 4-bit WNA16 MoE.""" - return int4_w4a16_moe_quant_config( - w1_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, group_size], - ) + if num_bits == 4: + return int4_w4a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) + else: + assert num_bits == 8 + return int8_w8a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) def make_wna16_moe_kernel( @@ -112,10 +108,10 @@ def make_wna16_moe_kernel( 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, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, is_k_full: bool = True, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] - | None = None, - shared_experts: SharedExperts | None = None, + routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: """ Create the FusedMoEKernel for WNA16 MoE. @@ -135,19 +131,20 @@ def make_wna16_moe_kernel( ) assert prepare_finalize is not None - logger.info_once( - "Using %s", prepare_finalize.__class__.__name__, scope="local" - ) + logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") if is_monolithic: + assert input_global_scale1 is None + assert input_global_scale2 is None experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, ) elif ( - prepare_finalize.activation_format - == mk.FusedMoEActivationFormat.BatchedExperts + prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts ): + assert input_global_scale1 is None + assert input_global_scale2 is None max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None experts = experts_cls( @@ -169,12 +166,13 @@ def make_wna16_moe_kernel( 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, + input_global_scale1=input_global_scale1, + input_global_scale2=input_global_scale2, is_k_full=is_k_full, ) return mk.FusedMoEKernel( prepare_finalize, experts, - shared_experts=shared_experts, inplace=(not moe_config.disable_inplace and not is_monolithic), ) 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 18ab622d5d19..e60cf94a2e46 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 @@ -16,9 +16,6 @@ FusedMoEConfig, FusedMoEQuantConfig, ) -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import ( - fused_marlin_moe, -) from vllm.model_executor.layers.fused_moe.oracle.wna16 import ( WNA16MoeBackend, make_wna16_moe_kernel, @@ -346,7 +343,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: backend=self.wna16_backend, experts_cls=self.experts_cls, routing_tables=layer._maybe_init_expert_routing_tables(), - shared_experts=layer.shared_experts, ) return @@ -468,38 +464,35 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.w13_weight = layer.w13_weight_packed layer.w2_weight = layer.w2_weight_packed - # Setup moe_kernel for the 4-bit Marlin path. - if self.num_bits == 4 and 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 - self.moe_kernel = make_wna16_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - backend=self.wna16_backend, - experts_cls=self.experts_cls, - 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, - routing_tables=layer._maybe_init_expert_routing_tables(), - shared_experts=layer.shared_experts, - ) + # Setup moe_kernel for the Marlin path. + 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 + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + backend=self.wna16_backend, + experts_cls=self.experts_cls, + 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, + input_global_scale1=getattr(layer, "w13_input_global_scale", None), + input_global_scale2=getattr(layer, "w2_input_global_scale", None), + is_k_full=self.is_k_full, + routing_tables=layer._maybe_init_expert_routing_tables(), + ) def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: - if self.num_bits != 4: - return None return make_wna16_moe_quant_config( - layer=layer, + w1_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, group_size=self.group_size, + num_bits=self.num_bits, ) - @property - def is_monolithic(self) -> bool: - return self.wna16_backend == WNA16MoeBackend.FLASHINFER - def apply_monolithic( self, layer: RoutedExperts, @@ -507,7 +500,7 @@ def apply_monolithic( router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, ) -> torch.Tensor: - assert self.wna16_backend == WNA16MoeBackend.FLASHINFER + assert self.is_monolithic assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( x, @@ -534,44 +527,16 @@ def apply( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: assert not self.is_monolithic - if self.moe_kernel is not None: - # 4-bit Marlin: use the modular kernel abstraction. - return self.moe_kernel.apply( - x, - layer.w13_weight, - layer.w2_weight, - topk_weights, - topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - shared_experts_input=shared_experts_input, - ) - # Non-4-bit Marlin fallback: call fused_marlin_moe() directly. - return fused_marlin_moe( + # 4-bit Marlin: use the modular kernel abstraction. + 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, - inplace=not self.moe.disable_inplace, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts_input=shared_experts_input, ) From f8d42acdd5ac91a53785a2064cb3b55d90d7ace1 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 13 May 2026 18:28:12 +0000 Subject: [PATCH 03/53] tweaks Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/marlin_moe.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) 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 b1b2fedc5f20..fbc1ec03fa21 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -403,6 +403,8 @@ def batched_fused_marlin_moe( expert_map: torch.Tensor | None = None, global_scale1: torch.Tensor | None = None, global_scale2: torch.Tensor | None = None, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, g_idx1: torch.Tensor | None = None, g_idx2: torch.Tensor | None = None, sort_indices1: torch.Tensor | None = None, @@ -527,6 +529,8 @@ def batched_fused_marlin_moe( num_tokens_post_padded=num_tokens_post_padded, global_scale1=global_scale1, global_scale2=global_scale2, + input_global_scale1=input_global_scale1, + input_global_scale2=input_global_scale2, g_idx1=g_idx1, g_idx2=g_idx2, sort_indices1=sort_indices1, @@ -568,7 +572,7 @@ def __init__( 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 @@ -898,6 +902,8 @@ def __init__( 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, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, is_k_full: bool = True, ): super().__init__( @@ -909,6 +915,8 @@ def __init__( 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, + input_global_scale1=input_global_scale1, + input_global_scale2=input_global_scale2, is_k_full=is_k_full, ) From 19cf0e7e0f2cb9c3c44478466a0ad91b15aaa2b0 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 13 May 2026 19:44:40 +0000 Subject: [PATCH 04/53] combine with int_wna16.py + cleanups Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/marlin_moe.py | 2 + .../fused_moe/experts/trtllm_mxint4_moe.py | 9 +- .../layers/fused_moe/oracle/int_wna16.py | 87 ++++++--- .../layers/fused_moe/oracle/wna16.py | 178 ------------------ .../compressed_tensors_moe_wna16_marlin.py | 40 ++-- .../layers/quantization/gptq_marlin.py | 4 +- .../layers/quantization/utils/quant_utils.py | 3 + 7 files changed, 106 insertions(+), 217 deletions(-) delete mode 100644 vllm/model_executor/layers/fused_moe/oracle/wna16.py 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 fbc1ec03fa21..8220a8052969 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -580,6 +580,8 @@ def __init__( self.is_k_full = is_k_full self.input_dtype = get_marlin_input_dtype() self.gemm1_clamp_limit = quant_config.gemm1_clamp_limit + self.input_global_scale1 = input_global_scale1 + self.input_global_scale2 = input_global_scale2 super().__init__( moe_config=moe_config, 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 index a370e582f9f6..b87b2607f7ce 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -11,7 +11,10 @@ FusedMoEQuantConfig, RoutingMethodType, ) -from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + QuantKey, + kInt4Static32, +) class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): @@ -53,9 +56,7 @@ def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - # Selected directly by WNA16 oracle based on conditions, not via - # QuantKey matching. Always return True here. - return True + return (weight_key, activation_key) == (kInt4Static32, None) @staticmethod def _supports_activation(activation: MoEActivation) -> bool: 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 9d25500ad078..b2c7dcd6fbc2 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch @@ -11,10 +11,13 @@ 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.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -35,6 +38,7 @@ class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" + FLASHINFER = "FLASHINFER" def backend_to_kernel_cls( @@ -55,6 +59,13 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] + elif backend == WNA16MoEBackend.FLASHINFER: + from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( # noqa: E501 + TrtLlmMxint4ExpertsMonolithic, + ) + + return [TrtLlmMxint4ExpertsMonolithic] + else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -64,6 +75,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: Get available backends in priority order based on platform and config. """ _AVAILABLE_BACKENDS = [ + WNA16MoEBackend.FLASHINFER, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] @@ -73,7 +85,6 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, - weight_bits: int, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. @@ -143,16 +154,43 @@ def _return_or_raise( ) +def make_wna16_moe_quant_config( + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + group_size: int, + num_bits: int, +) -> 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=None, + w2_zp=None, + block_shape=[0, group_size], + ) + else: + assert num_bits == 8 + return int8_w8a16_moe_quant_config( + w1_scale=w1_scale, + w2_scale=w2_scale, + w1_zp=None, + w2_zp=None, + block_shape=[0, group_size], + ) + + def make_wna16_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, experts_cls: type[mk.FusedMoEExperts] | None, - layer: torch.nn.Module, - 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, + 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, + input_global_scale1: torch.Tensor | None = None, + input_global_scale2: torch.Tensor | None = None, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: # Currently, we only support MarlinExperts and BatchedMarlinExperts @@ -169,10 +207,22 @@ def make_wna16_moe_kernel( allow_new_interface=True, ) 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, + "input_global_scale1": input_global_scale1, + "input_global_scale2": input_global_scale2, + } if 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: mk.FusedMoEExperts = BatchedMarlinExperts( @@ -180,28 +230,21 @@ def make_wna16_moe_kernel( 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( 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, ) + is_monolithic = experts_cls.is_monolithic() + return mk.FusedMoEKernel( prepare_finalize, experts, - inplace=not moe_config.disable_inplace, + inplace=not moe_config.disable_inplace and not is_monolithic, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/wna16.py b/vllm/model_executor/layers/fused_moe/oracle/wna16.py deleted file mode 100644 index 6141953ba398..000000000000 --- a/vllm/model_executor/layers/fused_moe/oracle/wna16.py +++ /dev/null @@ -1,178 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from enum import Enum - -import torch - -import vllm.model_executor.layers.fused_moe.modular_kernel as mk -from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe.all2all_utils import ( - maybe_make_prepare_finalize, -) -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.fused_marlin_moe import ( - BatchedMarlinExperts, - MarlinExperts, -) -from vllm.model_executor.layers.quantization.utils.flashinfer_mxint4_moe import ( - is_flashinfer_mxint4_moe_available, -) - -logger = init_logger(__name__) - - -class WNA16MoeBackend(Enum): - MARLIN = "MARLIN" - FLASHINFER = "FLASHINFER" - - -def select_wna16_moe_backend( - config: FusedMoEConfig, - num_bits: int, - group_size: int, -) -> tuple[WNA16MoeBackend, type[mk.FusedMoEExperts] | None]: - """ - Select the WNA16 MoE backend and experts class. - - Returns (backend, experts_cls). - """ - if is_flashinfer_mxint4_moe_available() and num_bits == 4 and group_size == 32: - from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( # noqa: E501 - TrtLlmMxint4ExpertsMonolithic, - ) - - logger.info_once( - "Using Flashinfer backend for WNA16 MoE " - f"(group_size={group_size}, num_bits={num_bits})", - scope="local", - ) - return WNA16MoeBackend.FLASHINFER, TrtLlmMxint4ExpertsMonolithic - - # 4/8-bit Marlin: select batched vs standard based on activation format. - activation_format = ( - mk.FusedMoEActivationFormat.BatchedExperts - if config.moe_parallel_config.use_batched_activation_format - else mk.FusedMoEActivationFormat.Standard - ) - experts_cls: type[mk.FusedMoEExperts] = ( - BatchedMarlinExperts - if activation_format == mk.FusedMoEActivationFormat.BatchedExperts - else MarlinExperts - ) - logger.info_once( - "Using Marlin backend for WNA16 MoE " - f"(group_size={group_size}, num_bits={num_bits})", - scope="local", - ) - return WNA16MoeBackend.MARLIN, experts_cls - - -def make_wna16_moe_quant_config( - w1_scale: torch.Tensor, - w2_scale: torch.Tensor, - group_size: int, - num_bits: int, -) -> FusedMoEQuantConfig: - """Create the FusedMoEQuantConfig for 4-bit WNA16 MoE.""" - if num_bits == 4: - return int4_w4a16_moe_quant_config( - w1_scale=w1_scale, - w2_scale=w2_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, group_size], - ) - else: - assert num_bits == 8 - return int8_w8a16_moe_quant_config( - w1_scale=w1_scale, - w2_scale=w2_scale, - w1_zp=None, - w2_zp=None, - block_shape=[0, group_size], - ) - - -def make_wna16_moe_kernel( - moe_quant_config: FusedMoEQuantConfig, - moe_config: FusedMoEConfig, - backend: WNA16MoeBackend, - experts_cls: type[mk.FusedMoEExperts], - 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, - input_global_scale1: torch.Tensor | None = None, - input_global_scale2: torch.Tensor | None = None, - is_k_full: bool = True, - routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, -) -> mk.FusedMoEKernel: - """ - Create the FusedMoEKernel for WNA16 MoE. - - For the Flashinfer (mxint4) backend, a monolithic kernel is created. - For the Marlin backend, a modular kernel is created using MarlinExperts - or BatchedMarlinExperts depending on the activation format. - """ - is_monolithic = issubclass(experts_cls, mk.FusedMoEExpertsMonolithic) - - prepare_finalize = maybe_make_prepare_finalize( - moe=moe_config, - quant_config=moe_quant_config, - routing_tables=routing_tables, - allow_new_interface=True, - use_monolithic=is_monolithic, - ) - assert prepare_finalize is not None - - logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") - - if is_monolithic: - assert input_global_scale1 is None - assert input_global_scale2 is None - experts = experts_cls( - moe_config=moe_config, - quant_config=moe_quant_config, - ) - elif ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts - ): - assert input_global_scale1 is None - assert input_global_scale2 is None - max_num_tokens = prepare_finalize.max_num_tokens_per_rank() - assert max_num_tokens is not None - experts = experts_cls( - moe_config=moe_config, - quant_config=moe_quant_config, - max_num_tokens=max_num_tokens, - num_dispatchers=prepare_finalize.num_dispatchers(), - 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, - ) - else: - 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, - input_global_scale1=input_global_scale1, - input_global_scale2=input_global_scale2, - is_k_full=is_k_full, - ) - - return mk.FusedMoEKernel( - prepare_finalize, - experts, - inplace=(not moe_config.disable_inplace and not is_monolithic), - ) 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 e60cf94a2e46..61997e035f54 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 @@ -16,8 +16,8 @@ FusedMoEConfig, FusedMoEQuantConfig, ) -from vllm.model_executor.layers.fused_moe.oracle.wna16 import ( - WNA16MoeBackend, +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, make_wna16_moe_kernel, make_wna16_moe_quant_config, select_wna16_moe_backend, @@ -37,6 +37,12 @@ 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__) @@ -64,14 +70,27 @@ def __init__( 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) + 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, - num_bits=self.num_bits, - group_size=self.group_size or -1, + weight_key=weight_key, ) def get_weight_shape( @@ -98,7 +117,7 @@ def get_weight_shape( "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 + is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER shape_map = { "w13_weight": { "Flashinfer": ( @@ -158,7 +177,7 @@ def create_weights( # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims - is_transposed = self.wna16_backend != WNA16MoeBackend.FLASHINFER + is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER extra_weight_attrs.update( {"is_transposed": is_transposed, "quant_method": self.strategy} ) @@ -309,7 +328,7 @@ def create_weights( 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.wna16_backend == WNA16MoeBackend.FLASHINFER: + if self.wna16_backend == WNA16MoEBackend.FLASHINFER: dict_weights_mxint4 = prepare_static_weights_for_trtllm_mxint4_moe( layer.w13_weight_packed, layer.w13_weight_scale, @@ -340,7 +359,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.moe_kernel = make_wna16_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, - backend=self.wna16_backend, experts_cls=self.experts_cls, routing_tables=layer._maybe_init_expert_routing_tables(), ) @@ -471,7 +489,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: self.moe_kernel = make_wna16_moe_kernel( moe_quant_config=self.moe_quant_config, moe_config=self.moe, - backend=self.wna16_backend, experts_cls=self.experts_cls, w13_g_idx=layer.w13_weight_g_idx, w2_g_idx=layer.w2_weight_g_idx, @@ -527,7 +544,7 @@ def apply( shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: assert not self.is_monolithic - # 4-bit Marlin: use the modular kernel abstraction. + assert self.moe_kernel is not None return self.moe_kernel.apply( x, layer.w13_weight, @@ -538,5 +555,6 @@ def apply( global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, shared_experts_input=shared_experts_input, ) diff --git a/vllm/model_executor/layers/quantization/gptq_marlin.py b/vllm/model_executor/layers/quantization/gptq_marlin.py index 26a41ef89e95..b3eee685d2d9 100644 --- a/vllm/model_executor/layers/quantization/gptq_marlin.py +++ b/vllm/model_executor/layers/quantization/gptq_marlin.py @@ -515,7 +515,8 @@ def __init__( 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( @@ -761,7 +762,6 @@ def _setup_kernel(self, layer: RoutedExperts) -> None: moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, - layer=layer, is_k_full=self.is_k_full, w13_g_idx=layer.w13_g_idx, w2_g_idx=layer.w2_g_idx, diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 0b1802522415..8db0a90ba60f 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 @@ def __str__(self): 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) From cbc26e3208bdb8f4bc1514aaa27c275a1c3fe5df Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 13 May 2026 23:52:15 +0000 Subject: [PATCH 05/53] review comments Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/marlin_moe.py | 2 ++ .../layers/fused_moe/oracle/int_wna16.py | 33 ++++++++----------- 2 files changed, 15 insertions(+), 20 deletions(-) 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 8220a8052969..9d1989580d34 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -45,6 +45,7 @@ kFp8StaticChannelSym, kFp8StaticTensorSym, kInt4Static, + kInt4Static32, kInt8Static, kMxfp4Static, kMxfp8Static, @@ -615,6 +616,7 @@ def _supports_quant_scheme( kNvfp4Static, kInt4Static, kInt8Static, + kInt4Static32, ] return weight_key in SUPPORTED_W 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 b2c7dcd6fbc2..65b00f31e0d4 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -19,6 +19,9 @@ MarlinExperts, MarlinExpertsBase, ) +from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( + TrtLlmMxint4ExpertsMonolithic, +) from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_act_int8_process_scales, @@ -46,26 +49,11 @@ 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: - from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( # noqa: E501 - TrtLlmMxint4ExpertsMonolithic, - ) - return [TrtLlmMxint4ExpertsMonolithic] - else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -183,7 +171,7 @@ def make_wna16_moe_quant_config( def make_wna16_moe_kernel( moe_quant_config: FusedMoEQuantConfig, moe_config: FusedMoEConfig, - experts_cls: type[mk.FusedMoEExperts] | 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, @@ -193,8 +181,13 @@ def make_wna16_moe_kernel( input_global_scale2: torch.Tensor | None = None, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: - # Currently, we only support MarlinExperts and BatchedMarlinExperts - assert experts_cls in (MarlinExperts, BatchedMarlinExperts) + # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts + # and BatchedMarlinExperts + assert experts_cls in ( + MarlinExperts, + BatchedMarlinExperts, + TrtLlmMxint4ExpertsMonolithic, + ) from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, @@ -225,7 +218,7 @@ def make_wna16_moe_kernel( if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None - experts: mk.FusedMoEExperts = BatchedMarlinExperts( + experts: mk.FusedMoEExperts = experts_cls( max_num_tokens=max_num_tokens, num_dispatchers=prepare_finalize.num_dispatchers(), moe_config=moe_config, @@ -233,7 +226,7 @@ def make_wna16_moe_kernel( **extra_args, ) else: - experts = MarlinExperts( + experts = experts_cls( moe_config=moe_config, quant_config=moe_quant_config, **extra_args, From c114fe864979bdb07fbf78ada5e8d6c5c461451b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 13:30:20 +0000 Subject: [PATCH 06/53] pass monolithic flag to maybe_make_pf Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) 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 65b00f31e0d4..d71f8b5cec66 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -193,11 +193,14 @@ def make_wna16_moe_kernel( maybe_make_prepare_finalize, ) + 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 @@ -232,8 +235,6 @@ def make_wna16_moe_kernel( **extra_args, ) - is_monolithic = experts_cls.is_monolithic() - return mk.FusedMoEKernel( prepare_finalize, experts, From eb74c1a130dbc4e5148d9b4776f7d89ce7990374 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 16:23:54 +0000 Subject: [PATCH 07/53] fix method name Signed-off-by: Bill Nell --- .../compressed_tensors_moe_wna16_marlin.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 61997e035f54..707fd2f1a452 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 @@ -360,7 +360,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, - routing_tables=layer._maybe_init_expert_routing_tables(), + routing_tables=layer._expert_routing_tables(), ) return @@ -497,7 +497,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: input_global_scale1=getattr(layer, "w13_input_global_scale", None), input_global_scale2=getattr(layer, "w2_input_global_scale", None), is_k_full=self.is_k_full, - routing_tables=layer._maybe_init_expert_routing_tables(), + routing_tables=layer._expert_routing_tables(), ) def get_fused_moe_quant_config( From 2cd4b0bd962fa02168209a20449d53d51363f5d9 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 19:34:43 +0000 Subject: [PATCH 08/53] wip Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) 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 d71f8b5cec66..7a5c5fde02e0 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -476,5 +476,21 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) + elif backend == WNA16MoEBackend.FLASHINFER: + # TODO(bnell): fill in + return ( + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + None, + None, + None, + None, + w13_bias, + w2_bias, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") From 79da17cdf036cf8d63afae624926398d242aa11c Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 20:44:04 +0000 Subject: [PATCH 09/53] wip Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 66 ++++++++++++++++++- .../compressed_tensors_moe_wna16_marlin.py | 50 ++++++++------ 2 files changed, 93 insertions(+), 23 deletions(-) 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 7a5c5fde02e0..893abc3e4057 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -247,6 +247,67 @@ 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_g_idx_sort_indices: torch.Tensor | None, + w2_g_idx_sort_indices: torch.Tensor | None, + w13_input_global_scale: torch.Tensor | None, + w2_input_global_scale: torch.Tensor | None, + w13_bias: torch.Tensor | None = None, + w2_bias: torch.Tensor | None = None, +) -> tuple[ + torch.Tensor, # w13_qweight + torch.Tensor, # w2_qweight + torch.Tensor, # w13_scales + torch.Tensor, # w2_scales + torch.Tensor, # w13_g_idx + torch.Tensor, # w2_g_idx + torch.Tensor, # w13_g_idx_sort_indices + torch.Tensor, # w2_g_idx_sort_indices + torch.Tensor | None, # w13_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, + w13_g_idx_sort_indices, + w2_g_idx_sort_indices, + w13_input_global_scale, + w2_input_global_scale, + w13_bias, + w2_bias, + ) + + def _process_weights_marlin( layer: torch.nn.Module, quant_config: "GPTQMarlinConfig", @@ -415,7 +476,7 @@ def _process_weights_marlin( def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, - quant_config: QuantizationConfig, + quant_config: QuantizationConfig | None, input_dtype: torch.dtype | None, w13: torch.Tensor, w2: torch.Tensor, @@ -477,8 +538,7 @@ def convert_to_wna16_moe_kernel_format( w2_bias, ) elif backend == WNA16MoEBackend.FLASHINFER: - # TODO(bnell): fill in - return ( + return _process_weights_flashinfer( w13, w2, w13_scale, 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 707fd2f1a452..61f52bbd24fe 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 @@ -18,6 +18,7 @@ ) 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, @@ -28,9 +29,6 @@ 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 ( - 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, @@ -329,24 +327,36 @@ 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.wna16_backend == WNA16MoEBackend.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"] + ( + marlin_w13_qweight, + marlin_w2_qweight, + marlin_w13_scales, + marlin_w2_scales, + _, # w13_g_idx + _, # w2_g_idx + _, # w13_g_idx_sort_indices + _, # w2_g_idx_sort_indices + _, # w13_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=None, # Not needed for flashinfer + input_dtype=None, # Not needed for flashinfer + 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, ) + replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight) + replace_parameter(layer, "w13_weight_scale", marlin_w13_scales) + replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight) + replace_parameter(layer, "w2_weight_scale", marlin_w2_scales) + # Alias packed weights to w13_weight/w2_weight for the modular # kernel interface. layer.w13_weight = layer.w13_weight_packed From fe3008daa8be3a24804f726fa941c2849d577591 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 21:10:28 +0000 Subject: [PATCH 10/53] refactor weight post processing Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 66 ++++-- .../compressed_tensors_moe_wna16_marlin.py | 216 ++++++------------ 2 files changed, 116 insertions(+), 166 deletions(-) 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 893abc3e4057..f74161269b38 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -4,6 +4,9 @@ from typing import TYPE_CHECKING, 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 @@ -33,7 +36,7 @@ ) if TYPE_CHECKING: - from vllm.model_executor.layers.quantization.gptq_marlin import GPTQMarlinConfig + pass logger = init_logger(__name__) @@ -310,8 +313,11 @@ def _process_weights_flashinfer( def _process_weights_marlin( layer: torch.nn.Module, - quant_config: "GPTQMarlinConfig", 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, @@ -360,6 +366,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 @@ -371,7 +378,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) @@ -382,6 +389,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 @@ -406,17 +415,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, ) @@ -425,19 +434,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, ) @@ -476,7 +481,7 @@ def _process_weights_marlin( def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, - quant_config: QuantizationConfig | None, + quant_config: QuantizationConfig | QuantizationArgs | None, input_dtype: torch.dtype | None, w13: torch.Tensor, w2: torch.Tensor, @@ -508,26 +513,41 @@ def convert_to_wna16_moe_kernel_format( Args: backend: the selected ``WNA16MoEBackend``. layer: the ``FusedMoE`` layer whose parameters are being prepared. - quant_config: the ``QuantizationConfig`` for this layer. + quant_config: the ``QuantizationConfig`` for this layer (GPTQ/AWQ). input_dtype: optional activation dtype, usually should be 16 bit. + num_bits: quantization bits (compressed tensors only). + pack_factor: packing factor (compressed tensors only). + group_size: group size for quantization (compressed tensors only). + actorder: activation ordering mode (compressed tensors only). """ if backend in ( WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ): + # GPTQ Marlin from vllm.model_executor.layers.quantization.gptq_marlin import ( GPTQMarlinConfig, ) - if not isinstance(quant_config, GPTQMarlinConfig): - raise TypeError( - "Marlin WNA16 MoE backend requires GPTQMarlinConfig, got " - f"{type(quant_config).__name__}." - ) + if isinstance(quant_config, GPTQMarlinConfig): + 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 + else: + assert 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 + return _process_weights_marlin( layer, - quant_config, input_dtype, + num_bits, + pack_factor, + group_size, + actorder, w13, w2, w13_scale, 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 61f52bbd24fe..10c175eb8041 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 @@ -6,7 +6,6 @@ QuantizationArgs, ) -from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( RoutedExperts, @@ -31,9 +30,7 @@ ) 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, @@ -324,14 +321,12 @@ def create_weights( layer.a2_scale = None 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.wna16_backend == WNA16MoEBackend.FLASHINFER: ( - marlin_w13_qweight, - marlin_w2_qweight, - marlin_w13_scales, - marlin_w2_scales, + w13_qweight, + w2_qweight, + w13_scales, + w2_scales, _, # w13_g_idx _, # w2_g_idx _, # w13_g_idx_sort_indices @@ -352,10 +347,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w13_g_idx=layer.w13_weight_g_idx, w2_g_idx=layer.w2_weight_g_idx, ) - replace_parameter(layer, "w13_weight_packed", marlin_w13_qweight) - replace_parameter(layer, "w13_weight_scale", marlin_w13_scales) - replace_parameter(layer, "w2_weight_packed", marlin_w2_qweight) - replace_parameter(layer, "w2_weight_scale", marlin_w2_scales) + replace_parameter(layer, "w13_weight_packed", w13_qweight) + replace_parameter(layer, "w13_weight_scale", w13_scales) + replace_parameter(layer, "w2_weight_packed", w2_qweight) + replace_parameter(layer, "w2_weight_scale", w2_scales) # Alias packed weights to w13_weight/w2_weight for the modular # kernel interface. @@ -372,143 +367,78 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: experts_cls=self.experts_cls, routing_tables=layer._expert_routing_tables(), ) - return - - is_a_8bit = ( - self.marlin_input_dtype is not None - and self.marlin_input_dtype.itemsize == 1 - ) - - 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 - - # 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) + else: + ( + 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_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, + ) + 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) + 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), + ) - 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, - ) - 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) + device = layer.w13_weight_g_idx.device + layer.workspace = marlin_make_workspace_new(device, 4) - # 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) + # 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 - 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), + # Setup moe_kernel for the Marlin path. + 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 + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + 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, + input_global_scale1=getattr(layer, "w13_input_global_scale", None), + input_global_scale2=getattr(layer, "w2_input_global_scale", None), + is_k_full=self.is_k_full, + routing_tables=layer._expert_routing_tables(), ) - replace_parameter(layer, "w2_weight_scale", marlin_w2_scales) - - layer.workspace = marlin_make_workspace_new(device, 4) - - # Alias packed weights to w13_weight/w2_weight for the modular - # kernel interface. - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - - # Setup moe_kernel for the Marlin path. - 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 - self.moe_kernel = make_wna16_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - 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, - input_global_scale1=getattr(layer, "w13_input_global_scale", None), - input_global_scale2=getattr(layer, "w2_input_global_scale", None), - is_k_full=self.is_k_full, - routing_tables=layer._expert_routing_tables(), - ) def get_fused_moe_quant_config( self, layer: torch.nn.Module From 177425479837eee8f7221b59255b8f6df8fd6cb7 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 21:28:39 +0000 Subject: [PATCH 11/53] cleanup Signed-off-by: Bill Nell --- .../compressed_tensors_moe_wna16_marlin.py | 169 +++++++----------- 1 file changed, 69 insertions(+), 100 deletions(-) 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 10c175eb8041..5fd1ef9c6e29 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,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import torch from compressed_tensors.quantization import ( QuantizationArgs, @@ -321,82 +323,42 @@ def create_weights( layer.a2_scale = None def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - if self.wna16_backend == WNA16MoEBackend.FLASHINFER: - ( - w13_qweight, - w2_qweight, - w13_scales, - w2_scales, - _, # w13_g_idx - _, # w2_g_idx - _, # w13_g_idx_sort_indices - _, # w2_g_idx_sort_indices - _, # 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=None, # Not needed for flashinfer - input_dtype=None, # Not needed for flashinfer - 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, - ) - replace_parameter(layer, "w13_weight_packed", w13_qweight) - replace_parameter(layer, "w13_weight_scale", w13_scales) - replace_parameter(layer, "w2_weight_packed", w2_qweight) - replace_parameter(layer, "w2_weight_scale", w2_scales) - - # 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 - - # Setup moe_kernel for the Flashinfer monolithic path. - 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 - 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(), - ) - else: - ( - 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_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, - ) - 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) + # Process weights using the shared oracle infrastructure + is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER + ( + 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_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, + ) + + # 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) + + # 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) @@ -414,32 +376,39 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: torch.nn.Parameter(w2_input_global_scale, requires_grad=False), ) - device = layer.w13_weight_g_idx.device - layer.workspace = marlin_make_workspace_new(device, 4) - - # Alias packed weights to w13_weight/w2_weight for the modular - # kernel interface. - layer.w13_weight = layer.w13_weight_packed - layer.w2_weight = layer.w2_weight_packed - - # Setup moe_kernel for the Marlin path. - 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 - self.moe_kernel = make_wna16_moe_kernel( - moe_quant_config=self.moe_quant_config, - moe_config=self.moe, - experts_cls=self.experts_cls, - 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, - input_global_scale1=getattr(layer, "w13_input_global_scale", None), - input_global_scale2=getattr(layer, "w2_input_global_scale", None), - is_k_full=self.is_k_full, - routing_tables=layer._expert_routing_tables(), + layer.workspace = marlin_make_workspace_new( + layer.w13_weight_g_idx.device, 4 ) + # Alias packed weights to w13_weight/w2_weight for the modular kernel interface + layer.w13_weight = layer.w13_weight_packed + layer.w2_weight = layer.w2_weight_packed + + assert self.experts_cls is not None + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + + # Add Marlin-specific arguments + marlin_args: dict[str, Any] = {} + if not is_flashinfer: + marlin_args = { + "w13_g_idx": layer.w13_weight_g_idx, + "w2_g_idx": layer.w2_weight_g_idx, + "w13_g_idx_sort_indices": layer.w13_g_idx_sort_indices, + "w2_g_idx_sort_indices": layer.w2_g_idx_sort_indices, + "input_global_scale1": getattr(layer, "w13_input_global_scale", None), + "input_global_scale2": getattr(layer, "w2_input_global_scale", None), + "is_k_full": self.is_k_full, + } + + self.moe_kernel = make_wna16_moe_kernel( + moe_quant_config=self.moe_quant_config, + moe_config=self.moe, + experts_cls=self.experts_cls, + routing_tables=layer._expert_routing_tables(), + **marlin_args, + ) + def get_fused_moe_quant_config( self, layer: torch.nn.Module ) -> FusedMoEQuantConfig | None: From 6224bed04867881b8420061c3cbecba3a1fd11ac Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 15 May 2026 19:04:29 +0000 Subject: [PATCH 12/53] try to fix doc Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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 5aabb775104e..31890761245b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -510,12 +510,8 @@ def convert_to_wna16_moe_kernel_format( Args: backend: the selected ``WNA16MoEBackend``. layer: the ``FusedMoE`` layer whose parameters are being prepared. - quant_config: the ``QuantizationConfig`` for this layer (GPTQ/AWQ). + quant_config: the ``QuantizationConfig`` for this layer. input_dtype: optional activation dtype, usually should be 16 bit. - num_bits: quantization bits (compressed tensors only). - pack_factor: packing factor (compressed tensors only). - group_size: group size for quantization (compressed tensors only). - actorder: activation ordering mode (compressed tensors only). """ if backend in ( WNA16MoEBackend.MARLIN, From e22637e3b4602ee2817f30b985d6d75d236e4036 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 15 May 2026 19:46:34 +0000 Subject: [PATCH 13/53] try to fix docs again Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 31890761245b..3bbb305fcaa2 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -78,8 +78,8 @@ def select_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``). From 765e31c235a9c5191fd57b0409994132532a9a40 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 20 May 2026 20:11:26 +0000 Subject: [PATCH 14/53] add missing args to config Signed-off-by: Bill Nell --- .../compressed_tensors_moe_wna16_marlin.py | 2 ++ 1 file changed, 2 insertions(+) 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 90f02108b221..46fc63ed40b6 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 @@ -417,6 +417,8 @@ def get_fused_moe_quant_config( w2_scale=layer.w2_weight_scale, group_size=self.group_size, num_bits=self.num_bits, + a1_gscale=getattr(layer, "w13_input_global_scale", None), + a2_gscale=getattr(layer, "w2_input_global_scale", None), ) def apply_monolithic( From 9da72db5f9c9d9b1aa8c4d9e99c7175c71909c01 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 21 May 2026 16:52:20 +0000 Subject: [PATCH 15/53] fix typo Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/config.py | 2 +- .../layers/fused_moe/experts/trtllm_mxint4_moe.py | 8 +++++++- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 8 ++++---- .../compressed_tensors_moe_wna16_marlin.py | 8 +++----- 4 files changed, 15 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 4b72fed577dd..33e99d394f1e 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -927,7 +927,7 @@ def int8_w8a16_moe_quant_config( _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, w1_bias), + _w2=FusedMoEQuantDesc(torch.int8, group_shape, w2_scale, None, w2_zp, w2_bias), ) 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 index b87b2607f7ce..d9ff72b8656c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -15,6 +15,7 @@ QuantKey, kInt4Static32, ) +from vllm.platforms import current_platform class TrtLlmMxint4ExpertsMonolithic(mk.FusedMoEExpertsMonolithic): @@ -45,7 +46,12 @@ def _supports_current_device() -> bool: is_flashinfer_mxint4_moe_available, ) - return 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: 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 ce42be11962c..14ae3126d526 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -42,7 +42,7 @@ class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" - FLASHINFER = "FLASHINFER" + FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" def backend_to_kernel_cls( @@ -53,7 +53,7 @@ def backend_to_kernel_cls( return [MarlinExperts] elif backend == WNA16MoEBackend.BATCHED_MARLIN: return [BatchedMarlinExperts] - elif backend == WNA16MoEBackend.FLASHINFER: + elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return [TrtLlmMxint4ExpertsMonolithic] else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -64,7 +64,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: Get available backends in priority order based on platform and config. """ _AVAILABLE_BACKENDS = [ - WNA16MoEBackend.FLASHINFER, + WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] @@ -748,7 +748,7 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) - elif backend == WNA16MoEBackend.FLASHINFER: + elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return _process_weights_flashinfer( w13, w2, 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 46fc63ed40b6..2d629d73edd6 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 @@ -114,7 +114,7 @@ def get_weight_shape( "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 + is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM shape_map = { "w13_weight": { "Flashinfer": ( @@ -174,7 +174,7 @@ def create_weights( # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims - is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER + is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM extra_weight_attrs.update( {"is_transposed": is_transposed, "quant_method": self.strategy} ) @@ -324,7 +324,7 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Process weights using the shared oracle infrastructure - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER + is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM ( w13_qweight, w2_qweight, @@ -417,8 +417,6 @@ def get_fused_moe_quant_config( w2_scale=layer.w2_weight_scale, group_size=self.group_size, num_bits=self.num_bits, - a1_gscale=getattr(layer, "w13_input_global_scale", None), - a2_gscale=getattr(layer, "w2_input_global_scale", None), ) def apply_monolithic( From a204a66d6da7fc4555011cee3d583b3884a6a9fb Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 13:27:41 +0000 Subject: [PATCH 16/53] wip Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/triton_moe.py | 45 +++++++++------- .../layers/fused_moe/oracle/int_wna16.py | 12 +++-- .../layers/quantization/moe_wna16.py | 51 +++++++++---------- 3 files changed, 60 insertions(+), 48 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 45b95b341020..a18502b84a4e 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -43,8 +43,11 @@ kFp8Static128BlockSym, kFp8StaticChannelSym, kFp8StaticTensorSym, + kInt4Static, + kInt4Static32, kInt8DynamicTokenSym, kInt8StaticChannelSym, + kInt8Static, ) from vllm.platforms import current_platform from vllm.triton_utils import tl @@ -459,40 +462,44 @@ def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: class TritonWNA16Experts(TritonExperts): @staticmethod def _supports_current_device() -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return current_platform.is_cuda_alike() or current_platform.is_xpu() @staticmethod def _supports_no_act_and_mul() -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return True @staticmethod def _supports_quant_scheme( weight_key: QuantKey | None, activation_key: QuantKey | None, ) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + SUPPORTED_W = [ + kInt4Static, + kInt8Static, + kInt4Static32, + # other group sizes? + ] + return weight_key in SUPPORTED_W @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." - ) + return activation in [ + MoEActivation.SILU, + MoEActivation.GELU, + MoEActivation.GELU_TANH, + MoEActivation.SWIGLUOAI, + MoEActivation.SWIGLUSTEP, + MoEActivation.SILU_NO_MUL, + MoEActivation.GELU_NO_MUL, + MoEActivation.GELU_TANH_NO_MUL, + MoEActivation.RELU2_NO_MUL, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: - raise NotImplementedError( - "TritonWNA16Experts is not yet used by an Oracle. " - "This method should not be called." + return not ( + moe_parallel_config.use_fi_nvl_two_sided_kernels + or moe_parallel_config.use_fi_nvl_one_sided_kernels ) def apply( 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 b300e8965e28..3ce7daffcb93 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -22,6 +22,9 @@ MarlinExperts, MarlinExpertsBase, ) +from vllm.model_executor.layers.fused_moe.experts.triton_moe import ( + TritonWNA16Experts, +) from vllm.model_executor.layers.fused_moe.experts.trtllm_mxint4_moe import ( TrtLlmMxint4ExpertsMonolithic, ) @@ -42,7 +45,8 @@ class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" - FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" + FLASHINFER = "FLASHINFER_TRTLLM" + TRITON = "TRITON" def backend_to_kernel_cls( @@ -55,6 +59,8 @@ def backend_to_kernel_cls( return [BatchedMarlinExperts] elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return [TrtLlmMxint4ExpertsMonolithic] + elif backend == WNA16MoEBackend.TRITON: + return [TritonWNA16Experts] else: raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") @@ -67,6 +73,7 @@ def _get_priority_backends() -> list[WNA16MoEBackend]: WNA16MoEBackend.FLASHINFER_TRTLLM, WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.TRITON, ] return _AVAILABLE_BACKENDS @@ -194,12 +201,11 @@ def make_wna16_moe_kernel( w2_g_idx_sort_indices: torch.Tensor | None = None, routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, ) -> mk.FusedMoEKernel: - # Currently, we only support TrtLlmMxint4ExpertsMonolithic, MarlinExperts - # and BatchedMarlinExperts assert experts_cls in ( MarlinExperts, BatchedMarlinExperts, TrtLlmMxint4ExpertsMonolithic, + TritonWNA16Experts, ) from vllm.model_executor.layers.fused_moe.all2all_utils import ( diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 471febab0443..87de30491ebe 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -13,11 +13,13 @@ RoutedExperts, SharedExperts, ) -from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, - int4_w4a16_moe_quant_config, - int8_w8a16_moe_quant_config, +) +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + make_wna16_moe_kernel, + make_wna16_moe_quant_config, + select_wna16_moe_backend, ) from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, @@ -216,6 +218,12 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: super().__init__(moe) self.quant_config = quant_config + # Select WNA16 MoE backend via oracle. + self.wna16_backend, self.experts_cls = select_wna16_moe_backend( + config=self.moe, + weight_key=weight_key, + ) + def create_weights( self, layer: RoutedExperts, @@ -339,21 +347,14 @@ def create_weights( def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: - weight_bits = self.quant_config.weight_bits has_zp = self.quant_config.has_zp - assert weight_bits == 4 or weight_bits == 8 - config_builder = ( - int4_w4a16_moe_quant_config - if weight_bits == 4 - else int8_w8a16_moe_quant_config - ) - - return config_builder( + return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, w2_scale=layer.w2_scales, w1_zp=layer.w13_qzeros if has_zp else None, w2_zp=layer.w2_qzeros if has_zp else None, - block_shape=[0, layer.group_size], + group_size=layer.group_size, + num_bits=self.quant_config.weight_bits, ) def apply( @@ -365,22 +366,20 @@ def apply( shared_experts: SharedExperts | None, shared_experts_input: torch.Tensor | None, ) -> torch.Tensor: - from vllm.model_executor.layers.fused_moe import fused_experts - - assert layer.activation == MoEActivation.SILU, ( - f"Only SiLU activation is supported, not {layer.activation}." - ) - - return fused_experts( + assert not self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply( x, - layer.w13_qweight, - layer.w2_qweight, - topk_weights=topk_weights, - topk_ids=topk_ids, - apply_router_weight_on_input=layer.apply_router_weight_on_input, + layer.w13_weight, + layer.w2_weight, + topk_weights, + topk_ids, + activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - quant_config=self.moe_quant_config, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + shared_experts=shared_experts, + shared_experts_input=shared_experts_input, ) @staticmethod From 6d275a98a63fc051214d4a8130ae1e697015b694 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 14:03:01 +0000 Subject: [PATCH 17/53] fix Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/triton_moe.py | 3 +- .../layers/quantization/moe_wna16.py | 37 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index a18502b84a4e..6f883b02c78d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -46,8 +46,8 @@ kInt4Static, kInt4Static32, kInt8DynamicTokenSym, - kInt8StaticChannelSym, kInt8Static, + kInt8StaticChannelSym, ) from vllm.platforms import current_platform from vllm.triton_utils import tl @@ -497,6 +497,7 @@ def _supports_activation(activation: MoEActivation) -> bool: @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: + # Why? return not ( moe_parallel_config.use_fi_nvl_two_sided_kernels or moe_parallel_config.use_fi_nvl_one_sided_kernels diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 87de30491ebe..01a05ccea44f 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -33,6 +33,14 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( check_marlin_supports_layer, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + INT4_DTYPE, + INT8_DTYPE, + QuantKey, + kInt4Static32GroupScale, + kInt4StaticGroupScale, + kInt8StaticGroupScale, +) from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform @@ -218,6 +226,24 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: super().__init__(moe) self.quant_config = quant_config + num_bits = (self.quant_config.weight_bits,) + group_size = self.quant_config.group_size + + if num_bits == 4: + quant_type = INT4_DTYPE + if group_size == 32: + scale = kInt4Static32GroupScale + else: + scale = kInt4StaticGroupScale + elif num_bits == 8: + assert group_size == -1 + quant_type = INT8_DTYPE + scale = kInt8StaticGroupScale + else: + raise ValueError("MoeWNA16Method only supports int4 and int8 now.") + + weight_key = QuantKey(quant_type, scale) + # Select WNA16 MoE backend via oracle. self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, @@ -344,6 +370,17 @@ def create_weights( layer.register_parameter(key, param) set_weight_attrs(param, extra_weight_attrs) + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + 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 + 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._maybe_init_expert_routing_tables(), + ) + def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: From db41077438f70b4cc85c0a10b7cca5d74331b20a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 16:23:18 +0000 Subject: [PATCH 18/53] fix Signed-off-by: Bill Nell --- vllm/model_executor/layers/quantization/moe_wna16.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 01a05ccea44f..0af91c9c4451 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -378,7 +378,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, - routing_tables=layer._maybe_init_expert_routing_tables(), + routing_tables=layer._expert_routing_tables(), ) def get_fused_moe_quant_config( From b3248d5bfc454c1a8fd0609d56a0c5514868793a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 19:32:17 +0000 Subject: [PATCH 19/53] wip Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 21 +++++++++++++++++++ .../layers/quantization/moe_wna16.py | 6 +++--- .../layers/quantization/utils/gptq_utils.py | 4 ++++ 3 files changed, 28 insertions(+), 3 deletions(-) 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 3ce7daffcb93..f4d678846ab5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -224,6 +224,7 @@ def make_wna16_moe_kernel( assert prepare_finalize is not None logger.info_once("Using %s", prepare_finalize.__class__.__name__, scope="local") + logger.info_once("Using %s", experts_cls.__name__, scope="local") extra_args: dict[str, Any] = {} if issubclass(experts_cls, MarlinExpertsBase): @@ -642,6 +643,7 @@ def _process_awq_weights_marlin( ) +# add zp? def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, @@ -764,5 +766,24 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) + elif backend in ( + WNA16MoEBackend.TRITON, + WNA16MoEBackend.FLASHINFER, + ): + # TODO(bnell): fill in + return ( + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + None, + None, + None, + None, + w13_bias, + w2_bias, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 0af91c9c4451..4736bee83966 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -226,7 +226,7 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: super().__init__(moe) self.quant_config = quant_config - num_bits = (self.quant_config.weight_bits,) + num_bits = self.quant_config.weight_bits group_size = self.quant_config.group_size if num_bits == 4: @@ -407,8 +407,8 @@ def apply( assert self.moe_kernel is not None return self.moe_kernel.apply( x, - layer.w13_weight, - layer.w2_weight, + layer.w13_qweight, + layer.w2_qweight, topk_weights, topk_ids, activation=layer.activation, diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index 691d80b0b747..c4f294fe09b3 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping from copy import deepcopy +from fractions import Fraction +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING @@ -83,6 +85,8 @@ def is_layer_gptq_quantized( proj_name = prefix.split(".")[-1] + quantized_layers = list(chain.from_iterable(quantized_layers)) + # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name # from the fused version to unfused + check to make sure that From 272bf3be5c78eca09f9e4394a7d7971b45dcfc04 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 23:15:15 +0000 Subject: [PATCH 20/53] zero point handling Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/marlin_moe.py | 6 ++ .../layers/fused_moe/oracle/int_wna16.py | 92 ++++++++++++++++--- .../layers/quantization/auto_gptq.py | 62 ++++++------- .../compressed_tensors_moe_wna16_marlin.py | 4 + .../layers/quantization/moe_wna16.py | 68 +++++++++++++- 5 files changed, 182 insertions(+), 50 deletions(-) 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 8bb9e5bdc06d..f2c8048097da 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -757,6 +757,8 @@ def apply( w2=w2, bias1=self.w1_bias, bias2=self.w2_bias, + w1_zeros=self.w1_zp, + w2_zeros=self.w2_zp, w1_scale=self.w1_scale, w2_scale=self.w2_scale, topk_weights=topk_weights, @@ -860,6 +862,8 @@ def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None: w2=w2, bias1=self.w1_bias, bias2=self.w2_bias, + w1_zeros=self.w1_zp, + w2_zeros=self.w2_zp, w1_scale=self.w1_scale, w2_scale=self.w2_scale, topk_weights=topk_weights, @@ -977,6 +981,8 @@ def apply( w2=w2, bias1=self.w1_bias, bias2=self.w2_bias, + w1_zeros=self.w1_zp, + w2_zeros=self.w2_zp, w1_scale=self.w1_scale, w2_scale=self.w2_scale, quant_type_id=self.quant_type_id, 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 f4d678846ab5..ffa9d1720385 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -34,6 +34,7 @@ marlin_moe_permute_scales, marlin_permute_bias, moe_awq_to_marlin_zero_points, + marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -65,22 +66,36 @@ def backend_to_kernel_cls( raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") -def _get_priority_backends() -> list[WNA16MoEBackend]: +def _get_priority_backends( + may_have_zp: bool, may_have_bias: bool +) -> list[WNA16MoEBackend]: """ Get available backends in priority order based on platform and config. """ - _AVAILABLE_BACKENDS = [ - WNA16MoEBackend.FLASHINFER_TRTLLM, + _AVAILABLE_BACKENDS = [] + + # Temp. first for testing + + if not may_have_bias: + _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON) + + if not may_have_zp and not may_have_bias: + _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER) + + # Marlin supports ZP and bias + _AVAILABLE_BACKENDS += [ WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, - WNA16MoEBackend.TRITON, ] + return _AVAILABLE_BACKENDS def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, + may_have_zp: bool, + may_have_bias: bool, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. @@ -93,6 +108,10 @@ def select_wna16_moe_backend( A tuple of (``WNA16MoEBackend``, experts class or ``None``). """ + # TODO: deal with ZP/bias - flashinfer does not support zp or bias + # triton wna16 does not support bias + # marlin appears to support both (ZP may need processing) + activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if config.moe_parallel_config.use_batched_activation_format @@ -131,7 +150,7 @@ def _return_or_raise( raise ValueError(_make_log_unsupported(backend, reason)) # Select kernels in order of backend. - AVAILABLE_BACKENDS = _get_priority_backends() + AVAILABLE_BACKENDS = _get_priority_backends(may_have_zp, may_have_bias) for backend in AVAILABLE_BACKENDS: activation_key = None # always BF16 activation for WNA16 MoE @@ -288,6 +307,8 @@ def _process_weights_flashinfer( torch.Tensor | None, # w2_input_global_scale torch.Tensor | None, # w13_bias torch.Tensor | None, # w2_bias + torch.Tensor | None, # w13_zp + torch.Tensor | None, # w2_zp ]: """Flashinfer (TRT-LLM MXINT4) weight post-processing. @@ -342,6 +363,8 @@ def _process_weights_marlin( w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, + w13_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, # w13_qweight torch.Tensor, # w2_qweight @@ -357,6 +380,8 @@ def _process_weights_marlin( torch.Tensor | None, # w2_input_global_scale torch.Tensor | None, # w13_bias torch.Tensor | None, # w2_bias + torch.Tensor | None, # w13_zp + torch.Tensor | None, # w2_zp ]: """Standard Marlin weight post-processing shared by MARLIN and BATCHED_MARLIN backends. @@ -381,6 +406,8 @@ def _process_weights_marlin( w2_input_global_scale: torch.Tensor | None = None w13_bias_out: torch.Tensor | None = None w2_bias_out: torch.Tensor | None = None + w13_zp_out: torch.Tensor | None = None + w2_zp_out: torch.Tensor | None = None # --- FP8 weight / scale adjustment --- if input_dtype == torch.float8_e4m3fn: @@ -480,6 +507,25 @@ def _process_weights_marlin( if w2_bias is not None: w2_bias_out = marlin_permute_bias(w2_bias) + # TODO: not sure about shapes here + if w13_zp is not None: + w13_zp_out = marlin_zero_points( + w13_zp, + size_k=layer.intermediate_size_per_partition, + size_n=w13_zp.shape[2], + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + + if w2_zp is not None: + w2_zp_out = marlin_zero_points( + w2_zp, + size_k=w2_zp.shape[1] * group_size_or_pack_factor, + size_n=w2_zp.shape[2], + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + return ( marlin_w13_qweight, marlin_w2_qweight, @@ -640,10 +686,11 @@ def _process_awq_weights_marlin( w2_input_global_scale, w13_bias_out, w2_bias_out, + w13_zp_out, + w2_zp_out, ) -# add zp? def convert_to_wna16_moe_kernel_format( backend: WNA16MoEBackend, layer: torch.nn.Module, @@ -659,6 +706,8 @@ def convert_to_wna16_moe_kernel_format( w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, + w13_zp: torch.Tensor | None = None, + w2_zp: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, # w13_qweight torch.Tensor, # w2_qweight @@ -674,6 +723,8 @@ def convert_to_wna16_moe_kernel_format( torch.Tensor | None, # w2_input_global_scale torch.Tensor | None, # w13_bias torch.Tensor | None, # w2_bias + torch.Tensor | None, # w13_zp + torch.Tensor | None, # w2_zp ]: """Dispatch weight post-processing to the appropriate per-backend handler. @@ -754,6 +805,8 @@ def convert_to_wna16_moe_kernel_format( w2_qzeros, w13_bias, w2_bias, + w13_zp, + w2_zp, ) elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return _process_weights_flashinfer( @@ -763,27 +816,40 @@ def convert_to_wna16_moe_kernel_format( w2_scale, w13_g_idx, w2_g_idx, +<<<<<<< variant A + w13_bias, + w2_bias, +>>>>>>> variant B + None, + None, + None, + None, +####### Ancestor + None, + None, + None, + None, w13_bias, w2_bias, +======= end ) - elif backend in ( - WNA16MoEBackend.TRITON, - WNA16MoEBackend.FLASHINFER, - ): - # TODO(bnell): fill in + elif backend == WNA16MoEBackend.TRITON: + # No processing needed for Triton. return ( w13, w2, w13_scale, w2_scale, - w13_g_idx, - w2_g_idx, + None, + None, None, None, None, None, w13_bias, w2_bias, + w13_zp, + w2_zp, ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 1821fd5c7f71..7903cbd03157 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -485,6 +485,8 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, weight_key, + may_have_zp=True, + may_have_bias=True, ) def create_weights( @@ -644,12 +646,22 @@ def create_weights( layer.workspace = marlin_make_workspace_new(device, 4) def process_weights_after_loading(self, layer: RoutedExperts) -> None: + def replace_or_register(name: str, val: torch.Tensor | None): + if val is None: + return + + if hasattr(layer, name): + replace_parameter(layer, name, val) + else: + layer.register_parameter( + name, torch.nn.Parameter(val, requires_grad=False) + ) + is_a_8bit = self.input_dtype is not None and self.input_dtype.itemsize == 1 - if is_a_8bit: - assert self.quant_config.quant_type.size_bits == 8, ( - "W8A8-INT8 is not supported by marlin kernel." - ) + assert not is_a_8bit or self.quant_config.quant_type.size_bits == 8, ( + "W8A8-INT8 is not supported by marlin kernel." + ) ( w13, @@ -666,6 +678,8 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2_input_global_scale, w13_bias, w2_bias, + w13_zp, + w2_zp, ) = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, layer=layer, @@ -679,6 +693,8 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2_g_idx=layer.w2_g_idx, w13_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), + w13_zp=getattr(layer, "w13_qzeros", None), + w2_zp=getattr(layer, "w2_qzeros", None), ) replace_parameter(layer, "w13_qweight", w13) @@ -689,38 +705,12 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w2_g_idx", w2_g_idx) 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) - if w13_input_global_scale is not None: - if hasattr(layer, "w13_input_global_scale"): - replace_parameter( - layer, "w13_input_global_scale", w13_input_global_scale - ) - else: - 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: - if hasattr(layer, "w2_input_global_scale"): - replace_parameter(layer, "w2_input_global_scale", w2_input_global_scale) - else: - layer.register_parameter( - "w2_input_global_scale", - torch.nn.Parameter(w2_input_global_scale, requires_grad=False), - ) - if w13_bias is not None: - if hasattr(layer, "w13_bias"): - replace_parameter(layer, "w13_bias", w13_bias) - else: - layer.register_parameter( - "w13_bias", torch.nn.Parameter(w13_bias, requires_grad=False) - ) - if w2_bias is not None: - if hasattr(layer, "w2_bias"): - replace_parameter(layer, "w2_bias", w2_bias) - else: - layer.register_parameter( - "w2_bias", torch.nn.Parameter(w2_bias, requires_grad=False) - ) + replace_or_register("w13_input_global_scale", w13_input_global_scale) + replace_or_register("w2_input_global_scale", w2_input_global_scale) + replace_or_register("w13_bias", w13_bias) + replace_or_register("w2_bias", w2_bias) + replace_or_register("w13_zp", w13_zp) + replace_or_register("w2_zp", w2_zp) self._setup_kernel(layer) 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 2d629d73edd6..d93d6a73b18f 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 @@ -88,6 +88,8 @@ def __init__( self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, weight_key=weight_key, + may_have_zp=False, + may_have_bias=False, ) def get_weight_shape( @@ -340,6 +342,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2_input_global_scale, _, # w13_bias _, # w2_bias + _, # w13_zp + _, # w2_zp ) = convert_to_wna16_moe_kernel_format( backend=self.wna16_backend, layer=layer, diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 4736bee83966..fd3806cf3d1c 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -17,6 +17,8 @@ FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + convert_to_wna16_moe_kernel_format, make_wna16_moe_kernel, make_wna16_moe_quant_config, select_wna16_moe_backend, @@ -41,7 +43,7 @@ kInt4StaticGroupScale, kInt8StaticGroupScale, ) -from vllm.model_executor.utils import set_weight_attrs +from vllm.model_executor.utils import replace_parameter, set_weight_attrs from vllm.platforms import current_platform @@ -245,9 +247,12 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: weight_key = QuantKey(quant_type, scale) # Select WNA16 MoE backend via oracle. + # handle ZP? self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, weight_key=weight_key, + may_have_zp=self.quant_config.has_zp, + may_have_bias=False, ) def create_weights( @@ -371,6 +376,67 @@ def create_weights( set_weight_attrs(param, extra_weight_attrs) def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + has_zp = self.quant_config.has_zp + ( + 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_input_global_scale, + w2_input_global_scale, + _, + _, + w13_zp, + w2_zp, + ) = convert_to_wna16_moe_kernel_format( + backend=self.wna16_backend, + layer=layer, + quant_config=self.quant_config, + input_dtype=None, # self.marlin_input_dtype, # XXXXXXXXXXXXX + w13=layer.w13_weight_packed, + w2=layer.w2_weight_packed, + w13_scale=layer.w13_weight_scale, + w2_scale=layer.w2_weight_scale, + w13_g_idx=layer.w13_weight_g_idx, + w2_g_idx=layer.w2_weight_g_idx, + w13_zp=layer.w13_qzeros if has_zp else None, + w2_zp=layer.w2_qzeros if has_zp else None, + ) + + # Replace common parameters + replace_parameter(layer, "w13_weight_packed", w13_qweight) + replace_parameter(layer, "w2_weight_packed", w2_qweight) + replace_parameter(layer, "w13_weight_scale", w13_scales) + replace_parameter(layer, "w2_weight_scale", w2_scales) + + if has_zp: + assert w13_zp is not None and w2_zp is not None + replace_parameter(layer, "w13_qzeros", w13_zp) + replace_parameter(layer, "w2_qzeros", w2_zp) + + # Marlin-specific parameters (not needed for Flashinfer) + if self.wna16_backend != WNA16MoEBackend.FLASHINFER: + replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) + replace_parameter(layer, "w2_weight_g_idx", w2_g_idx_processed) + replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) + replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) + + # Register input global scales if present + if w13_input_global_scale is not None: + layer.register_parameter( + "w13_input_global_scale", + torch.nn.Parameter(w13_input_global_scale, requires_grad=False), + ) + if w2_input_global_scale is not None: + layer.register_parameter( + "w2_input_global_scale", + torch.nn.Parameter(w2_input_global_scale, requires_grad=False), + ) + 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 From d1e1213cc2090193bb3e843b5c6b02ad23052e8b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 14 May 2026 23:46:42 +0000 Subject: [PATCH 21/53] fix Signed-off-by: Bill Nell --- vllm/model_executor/layers/quantization/moe_wna16.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index fd3806cf3d1c..599baf8abcc3 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -396,7 +396,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: backend=self.wna16_backend, layer=layer, quant_config=self.quant_config, - input_dtype=None, # self.marlin_input_dtype, # XXXXXXXXXXXXX + input_dtype=None, w13=layer.w13_weight_packed, w2=layer.w2_weight_packed, w13_scale=layer.w13_weight_scale, From e3a1503c8c7cd900163ed0834108fcb3ed839970 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 15 May 2026 17:20:54 +0000 Subject: [PATCH 22/53] fix Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 7 +------ .../model_executor/layers/quantization/utils/gptq_utils.py | 1 - 2 files changed, 1 insertion(+), 7 deletions(-) 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 ffa9d1720385..dbb613dd765e 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -74,8 +74,7 @@ def _get_priority_backends( """ _AVAILABLE_BACKENDS = [] - # Temp. first for testing - + # Temp. triton first for testing if not may_have_bias: _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON) @@ -108,10 +107,6 @@ def select_wna16_moe_backend( A tuple of (``WNA16MoEBackend``, experts class or ``None``). """ - # TODO: deal with ZP/bias - flashinfer does not support zp or bias - # triton wna16 does not support bias - # marlin appears to support both (ZP may need processing) - activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if config.moe_parallel_config.use_batched_activation_format diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index c4f294fe09b3..c20f8071729d 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping from copy import deepcopy -from fractions import Fraction from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING From f9fe9d83ddf295f1c860fbed3146a05b6212cc56 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 18 May 2026 19:30:43 +0000 Subject: [PATCH 23/53] fix merge Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 60 +++++-------------- .../layers/quantization/auto_gptq.py | 14 ++--- .../compressed_tensors_moe_wna16_marlin.py | 2 - .../layers/quantization/moe_wna16.py | 18 +++--- 4 files changed, 31 insertions(+), 63 deletions(-) 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 dbb613dd765e..1e3adf0cd2c2 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -33,8 +33,8 @@ marlin_act_int8_process_scales, marlin_moe_permute_scales, marlin_permute_bias, - moe_awq_to_marlin_zero_points, marlin_zero_points, + moe_awq_to_marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -302,8 +302,6 @@ def _process_weights_flashinfer( torch.Tensor | None, # w2_input_global_scale torch.Tensor | None, # w13_bias torch.Tensor | None, # w2_bias - torch.Tensor | None, # w13_zp - torch.Tensor | None, # w2_zp ]: """Flashinfer (TRT-LLM MXINT4) weight post-processing. @@ -358,8 +356,6 @@ def _process_weights_marlin( w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, - w13_zp: torch.Tensor | None = None, - w2_zp: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, # w13_qweight torch.Tensor, # w2_qweight @@ -375,8 +371,6 @@ def _process_weights_marlin( torch.Tensor | None, # w2_input_global_scale torch.Tensor | None, # w13_bias torch.Tensor | None, # w2_bias - torch.Tensor | None, # w13_zp - torch.Tensor | None, # w2_zp ]: """Standard Marlin weight post-processing shared by MARLIN and BATCHED_MARLIN backends. @@ -401,8 +395,8 @@ def _process_weights_marlin( w2_input_global_scale: torch.Tensor | None = None w13_bias_out: torch.Tensor | None = None w2_bias_out: torch.Tensor | None = None - w13_zp_out: torch.Tensor | None = None - w2_zp_out: torch.Tensor | None = None + w13_qzeros_out: torch.Tensor | None = None + w2_qzeros_out: torch.Tensor | None = None # --- FP8 weight / scale adjustment --- if input_dtype == torch.float8_e4m3fn: @@ -503,20 +497,20 @@ def _process_weights_marlin( w2_bias_out = marlin_permute_bias(w2_bias) # TODO: not sure about shapes here - if w13_zp is not None: - w13_zp_out = marlin_zero_points( - w13_zp, + if w13_qzeros is not None: + w13_qzeros_out = marlin_zero_points( + w13_qzeros, size_k=layer.intermediate_size_per_partition, - size_n=w13_zp.shape[2], + size_n=w13_qzeros.shape[2], num_bits=num_bits, is_a_8bit=is_a_8bit, ) - if w2_zp is not None: - w2_zp_out = marlin_zero_points( - w2_zp, - size_k=w2_zp.shape[1] * group_size_or_pack_factor, - size_n=w2_zp.shape[2], + if w2_qzeros is not None: + w2_qzeros_out = marlin_zero_points( + w2_qzeros, + size_k=w2_qzeros.shape[1] * group_size_or_pack_factor, + size_n=w2_qzeros.shape[2], num_bits=num_bits, is_a_8bit=is_a_8bit, ) @@ -530,8 +524,8 @@ def _process_weights_marlin( w2_g_idx, w13_g_idx_sort_indices, w2_g_idx_sort_indices, - w13_qzeros, - w2_qzeros, + w13_qzeros_out, + w2_qzeros_out, w13_input_global_scale, w2_input_global_scale, w13_bias_out, @@ -681,8 +675,6 @@ def _process_awq_weights_marlin( w2_input_global_scale, w13_bias_out, w2_bias_out, - w13_zp_out, - w2_zp_out, ) @@ -701,8 +693,6 @@ def convert_to_wna16_moe_kernel_format( w2_qzeros: torch.Tensor | None = None, w13_bias: torch.Tensor | None = None, w2_bias: torch.Tensor | None = None, - w13_zp: torch.Tensor | None = None, - w2_zp: torch.Tensor | None = None, ) -> tuple[ torch.Tensor, # w13_qweight torch.Tensor, # w2_qweight @@ -718,8 +708,6 @@ def convert_to_wna16_moe_kernel_format( torch.Tensor | None, # w2_input_global_scale torch.Tensor | None, # w13_bias torch.Tensor | None, # w2_bias - torch.Tensor | None, # w13_zp - torch.Tensor | None, # w2_zp ]: """Dispatch weight post-processing to the appropriate per-backend handler. @@ -800,8 +788,6 @@ def convert_to_wna16_moe_kernel_format( w2_qzeros, w13_bias, w2_bias, - w13_zp, - w2_zp, ) elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return _process_weights_flashinfer( @@ -811,22 +797,8 @@ def convert_to_wna16_moe_kernel_format( w2_scale, w13_g_idx, w2_g_idx, -<<<<<<< variant A - w13_bias, - w2_bias, ->>>>>>> variant B - None, - None, - None, - None, -####### Ancestor - None, - None, - None, - None, w13_bias, w2_bias, -======= end ) elif backend == WNA16MoEBackend.TRITON: # No processing needed for Triton. @@ -839,12 +811,12 @@ def convert_to_wna16_moe_kernel_format( None, None, None, + w13_qzeros, + w2_qzeros, None, None, w13_bias, w2_bias, - w13_zp, - w2_zp, ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 7903cbd03157..df9a413c9ebd 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -672,14 +672,12 @@ def replace_or_register(name: str, val: torch.Tensor | None): w2_g_idx, w13_g_idx_sort_indices, w2_g_idx_sort_indices, - _w13_qzeros, - _w2_qzeros, + w13_qzeros, + w2_qzeros, w13_input_global_scale, w2_input_global_scale, w13_bias, w2_bias, - w13_zp, - w2_zp, ) = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, layer=layer, @@ -693,8 +691,8 @@ def replace_or_register(name: str, val: torch.Tensor | None): w2_g_idx=layer.w2_g_idx, w13_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - w13_zp=getattr(layer, "w13_qzeros", None), - w2_zp=getattr(layer, "w2_qzeros", None), + w13_qzeros=getattr(layer, "w13_qzeros", None), + w2_qzeros=getattr(layer, "w2_qzeros", None), ) replace_parameter(layer, "w13_qweight", w13) @@ -709,8 +707,8 @@ def replace_or_register(name: str, val: torch.Tensor | None): replace_or_register("w2_input_global_scale", w2_input_global_scale) replace_or_register("w13_bias", w13_bias) replace_or_register("w2_bias", w2_bias) - replace_or_register("w13_zp", w13_zp) - replace_or_register("w2_zp", w2_zp) + replace_or_register("w13_zp", w13_qzeros) + replace_or_register("w2_zp", w2_qzeros) self._setup_kernel(layer) 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 d93d6a73b18f..d47f2e16af22 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 @@ -342,8 +342,6 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2_input_global_scale, _, # w13_bias _, # w2_bias - _, # w13_zp - _, # w2_zp ) = convert_to_wna16_moe_kernel_format( backend=self.wna16_backend, layer=layer, diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 599baf8abcc3..f604be93aa18 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -386,12 +386,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: 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_zp, - w2_zp, + _, # w13_bias + _, # w2_bias ) = convert_to_wna16_moe_kernel_format( backend=self.wna16_backend, layer=layer, @@ -403,8 +403,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2_scale=layer.w2_weight_scale, w13_g_idx=layer.w13_weight_g_idx, w2_g_idx=layer.w2_weight_g_idx, - w13_zp=layer.w13_qzeros if has_zp else None, - w2_zp=layer.w2_qzeros if has_zp else None, + w13_qzeros=layer.w13_qzeros if has_zp else None, + w2_qzeros=layer.w2_qzeros if has_zp else None, ) # Replace common parameters @@ -414,9 +414,9 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "w2_weight_scale", w2_scales) if has_zp: - assert w13_zp is not None and w2_zp is not None - replace_parameter(layer, "w13_qzeros", w13_zp) - replace_parameter(layer, "w2_qzeros", w2_zp) + assert w13_qzeros is not None and w2_qzeros is not None + replace_parameter(layer, "w13_qzeros", w13_qzeros) + replace_parameter(layer, "w2_qzeros", w2_qzeros) # Marlin-specific parameters (not needed for Flashinfer) if self.wna16_backend != WNA16MoEBackend.FLASHINFER: From 1b9356bf57de2b40532d400e016ca70edf3d6e3d Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Mon, 18 May 2026 20:19:03 +0000 Subject: [PATCH 24/53] fix lint Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/experts/marlin_moe.py | 6 ------ vllm/model_executor/layers/quantization/awq_marlin.py | 2 ++ 2 files changed, 2 insertions(+), 6 deletions(-) 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 f2c8048097da..8bb9e5bdc06d 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -757,8 +757,6 @@ def apply( w2=w2, bias1=self.w1_bias, bias2=self.w2_bias, - w1_zeros=self.w1_zp, - w2_zeros=self.w2_zp, w1_scale=self.w1_scale, w2_scale=self.w2_scale, topk_weights=topk_weights, @@ -862,8 +860,6 @@ def moe_sum_with_lora(moe_out: torch.Tensor, out: torch.Tensor) -> None: w2=w2, bias1=self.w1_bias, bias2=self.w2_bias, - w1_zeros=self.w1_zp, - w2_zeros=self.w2_zp, w1_scale=self.w1_scale, w2_scale=self.w2_scale, topk_weights=topk_weights, @@ -981,8 +977,6 @@ def apply( w2=w2, bias1=self.w1_bias, bias2=self.w2_bias, - w1_zeros=self.w1_zp, - w2_zeros=self.w2_zp, w1_scale=self.w1_scale, w2_scale=self.w2_scale, quant_type_id=self.quant_type_id, diff --git a/vllm/model_executor/layers/quantization/awq_marlin.py b/vllm/model_executor/layers/quantization/awq_marlin.py index 81c0fcb331e3..48fbe6335a34 100644 --- a/vllm/model_executor/layers/quantization/awq_marlin.py +++ b/vllm/model_executor/layers/quantization/awq_marlin.py @@ -524,6 +524,8 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, kInt4Static, + may_have_zp=self.quant_config.zero_point, + may_have_bias=True, ) def create_weights( From 0bde5a05170bfa815306161a71dbb33ba7d16d5e Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 21 May 2026 22:37:56 +0000 Subject: [PATCH 25/53] merge + review comments Signed-off-by: Bill Nell --- .../layers/quantization/auto_gptq.py | 4 ++-- .../layers/quantization/moe_wna16.py | 24 +++++++++---------- 2 files changed, 14 insertions(+), 14 deletions(-) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index df9a413c9ebd..cb2dad3d5249 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -707,8 +707,8 @@ def replace_or_register(name: str, val: torch.Tensor | None): replace_or_register("w2_input_global_scale", w2_input_global_scale) replace_or_register("w13_bias", w13_bias) replace_or_register("w2_bias", w2_bias) - replace_or_register("w13_zp", w13_qzeros) - replace_or_register("w2_zp", w2_qzeros) + replace_or_register("w13_qzeros", w13_qzeros) + replace_or_register("w2_qzeros", w2_qzeros) self._setup_kernel(layer) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index f604be93aa18..93d6401493d3 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -397,21 +397,21 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer=layer, quant_config=self.quant_config, input_dtype=None, - w13=layer.w13_weight_packed, - w2=layer.w2_weight_packed, - w13_scale=layer.w13_weight_scale, - w2_scale=layer.w2_weight_scale, - w13_g_idx=layer.w13_weight_g_idx, - w2_g_idx=layer.w2_weight_g_idx, + w13=layer.w13_qweight, + w2=layer.w2_qweight, + w13_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w13_g_idx=layer.w13_g_idx, + w2_g_idx=layer.w2_g_idx, w13_qzeros=layer.w13_qzeros if has_zp else None, w2_qzeros=layer.w2_qzeros if has_zp else None, ) # Replace common parameters - replace_parameter(layer, "w13_weight_packed", w13_qweight) - replace_parameter(layer, "w2_weight_packed", w2_qweight) - replace_parameter(layer, "w13_weight_scale", w13_scales) - replace_parameter(layer, "w2_weight_scale", w2_scales) + replace_parameter(layer, "w13_qweight", w13_qweight) + replace_parameter(layer, "w2_qweight", w2_qweight) + replace_parameter(layer, "w13_scales", w13_scales) + replace_parameter(layer, "w2_scales", w2_scales) if has_zp: assert w13_qzeros is not None and w2_qzeros is not None @@ -420,8 +420,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Marlin-specific parameters (not needed for Flashinfer) if self.wna16_backend != WNA16MoEBackend.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", w13_g_idx_processed) + replace_parameter(layer, "w2_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) From 9f01de36a2f6f5bcd466d96bd3347acef839cc38 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 21 May 2026 23:03:16 +0000 Subject: [PATCH 26/53] fix enums Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 4 ++-- vllm/model_executor/layers/quantization/moe_wna16.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) 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 1e3adf0cd2c2..b786b37ee254 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -46,7 +46,7 @@ class WNA16MoEBackend(Enum): MARLIN = "MARLIN" BATCHED_MARLIN = "BATCHED_MARLIN" - FLASHINFER = "FLASHINFER_TRTLLM" + FLASHINFER_TRTLLM = "FLASHINFER_TRTLLM" TRITON = "TRITON" @@ -79,7 +79,7 @@ def _get_priority_backends( _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON) if not may_have_zp and not may_have_bias: - _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER) + _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM) # Marlin supports ZP and bias _AVAILABLE_BACKENDS += [ diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 93d6401493d3..87232168304f 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -419,7 +419,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "w2_qzeros", w2_qzeros) # Marlin-specific parameters (not needed for Flashinfer) - if self.wna16_backend != WNA16MoEBackend.FLASHINFER: + if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM: replace_parameter(layer, "w13_g_idx", w13_g_idx_processed) replace_parameter(layer, "w2_g_idx", w2_g_idx_processed) replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) From cf40bbf4db01ea0f9386df17d25049777d0bec2e Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 22 May 2026 13:07:04 +0000 Subject: [PATCH 27/53] fix Signed-off-by: Bill Nell --- .../layers/quantization/utils/gptq_utils.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index c20f8071729d..b1c18f5fa4c9 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -2,9 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Mapping from copy import deepcopy -from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import regex as re import torch @@ -70,6 +69,20 @@ def get_dynamic_override( return default_value +def flatten_list(lst: list[Any]) -> list[Any]: + output = [] + + def _flatten(lst: list[Any]): + for i in lst: + if isinstance(type(i), list): + _flatten(i) + else: + output.append(i) + + _flatten(lst) + return output + + def is_layer_gptq_quantized( prefix: str, quantized_layers: list[str], @@ -84,7 +97,7 @@ def is_layer_gptq_quantized( proj_name = prefix.split(".")[-1] - quantized_layers = list(chain.from_iterable(quantized_layers)) + quantized_layers = flatten_list(quantized_layers) # Fused layers like gate_up_proj or qkv_proj will not be fused # in the safetensors checkpoint. So, we convert the name From 00cbe82968f6ea58c8ec8c48437e1ed1d254cb8b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 22 May 2026 13:08:41 +0000 Subject: [PATCH 28/53] fix Signed-off-by: Bill Nell --- vllm/model_executor/layers/quantization/utils/gptq_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/utils/gptq_utils.py b/vllm/model_executor/layers/quantization/utils/gptq_utils.py index b1c18f5fa4c9..c73c42f4ff02 100644 --- a/vllm/model_executor/layers/quantization/utils/gptq_utils.py +++ b/vllm/model_executor/layers/quantization/utils/gptq_utils.py @@ -74,7 +74,7 @@ def flatten_list(lst: list[Any]) -> list[Any]: def _flatten(lst: list[Any]): for i in lst: - if isinstance(type(i), list): + if isinstance(i, list): _flatten(i) else: output.append(i) From 0745aca086357514aec29a5143733e5d4e5af753 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Sat, 23 May 2026 07:48:15 +0000 Subject: [PATCH 29/53] fixes Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/triton_moe.py | 4 +++- .../layers/fused_moe/oracle/int_wna16.py | 14 +++++++++++--- .../compressed_tensors_moe_wna16_marlin.py | 15 ++++++++------- 3 files changed, 22 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index 6f883b02c78d..ddcef519da08 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -523,7 +523,9 @@ def apply( ): # Check constraints. if self.quant_config.use_int4_w4a16: - assert hidden_states.size(-1) // 2 == w1.size(2), "Hidden size mismatch" + assert hidden_states.size(-1) // 2 == w1.size(2), ( + f"Hidden size mismatch {hidden_states.size(-1) // 2} == {w1.size(2)}" + ) else: assert hidden_states.size(-1) == w1.size(2), ( f"Hidden size mismatch {hidden_states.size(-1)} != {w1.size(2)}" 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 b786b37ee254..d8cb0dfcdaa5 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -87,6 +87,9 @@ def _get_priority_backends( WNA16MoEBackend.BATCHED_MARLIN, ] + if not may_have_bias: + _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON) + return _AVAILABLE_BACKENDS @@ -769,6 +772,7 @@ def convert_to_wna16_moe_kernel_format( "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( @@ -801,10 +805,14 @@ def convert_to_wna16_moe_kernel_format( w2_bias, ) elif backend == WNA16MoEBackend.TRITON: - # No processing needed for Triton. + # Convert from int32 to uint8 format for Triton kernel. + # This changes the shape from (E, N, K // 8) to (E, N, K // 2) for int4, + # which matches what the Triton kernel expects. + w13_uint8 = w13.view(torch.uint8) + w2_uint8 = w2.view(torch.uint8) return ( - w13, - w2, + w13_uint8, + w2_uint8, w13_scale, w2_scale, 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 d47f2e16af22..4b591fabcc91 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 @@ -91,6 +91,10 @@ def __init__( may_have_zp=False, may_have_bias=False, ) + self.is_marlin = self.wna16_backend in [ + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + ] def get_weight_shape( self, @@ -116,7 +120,6 @@ def get_weight_shape( "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": ( @@ -159,7 +162,7 @@ def get_weight_shape( "Marlin": (num_experts, num_groups_w2, hidden_size), }, } - backend_key = "Flashinfer" if is_flashinfer else "Marlin" + backend_key = "Marlin" if self.is_marlin else "Flashinfer" return shape_map[weight_name][backend_key] def create_weights( @@ -176,9 +179,8 @@ def create_weights( # Will transpose the loaded weight along the # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims - is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM extra_weight_attrs.update( - {"is_transposed": is_transposed, "quant_method": self.strategy} + {"is_transposed": self.is_marlin, "quant_method": self.strategy} ) w13_weight = torch.nn.Parameter( @@ -326,7 +328,6 @@ def create_weights( def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Process weights using the shared oracle infrastructure - is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM ( w13_qweight, w2_qweight, @@ -362,7 +363,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: replace_parameter(layer, "w2_weight_scale", w2_scales) # Marlin-specific parameters (not needed for Flashinfer) - if not is_flashinfer: + if self.is_marlin: 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) @@ -394,7 +395,7 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: # Add Marlin-specific arguments marlin_args: dict[str, Any] = {} - if not is_flashinfer: + if self.is_marlin: marlin_args = { "w13_g_idx": layer.w13_weight_g_idx, "w2_g_idx": layer.w2_weight_g_idx, From 98e9c2e3d726dd7029de63ba69ced069561ac327 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Sat, 23 May 2026 08:06:26 +0000 Subject: [PATCH 30/53] more fixes Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/fused_moe.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index e13ef52c01df..a2d91757ae62 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -30,6 +30,7 @@ ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import next_power_of_2 from vllm.utils.torch_utils import direct_register_custom_op logger = init_logger(__name__) @@ -1243,7 +1244,11 @@ def get_default_config( bit = 4 if dtype == "int4_w4a16" else 8 use_moe_wna16_cuda = should_moe_wna16_use_cuda(M * topk, block_shape[1], E, bit) if use_moe_wna16_cuda: - config = {"BLOCK_SIZE_M": min(16, M), "SPLIT_K": 1} + config = { + "BLOCK_SIZE_M": min(16, next_power_of_2(M)), + "GROUP_SIZE_M": 1, + "SPLIT_K": 1, + } elif M <= 20: config = {"BLOCK_SIZE_M": 16, "GROUP_SIZE_M": 1, "SPLIT_K": 1} elif M <= 40: From 7946cd381b2a9677c8eca0c55ece3bee18b3084c Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 26 May 2026 16:43:41 +0000 Subject: [PATCH 31/53] remove debugging code Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 4 ---- 1 file changed, 4 deletions(-) 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 d8cb0dfcdaa5..abf165016036 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -74,10 +74,6 @@ def _get_priority_backends( """ _AVAILABLE_BACKENDS = [] - # Temp. triton first for testing - if not may_have_bias: - _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON) - if not may_have_zp and not may_have_bias: _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM) From bed43d67a13845372d632b27f6dd0a262d70d9f8 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 28 May 2026 15:19:59 +0000 Subject: [PATCH 32/53] fix Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 8 ++++++ .../layers/quantization/moe_wna16.py | 28 +++++++++++++++++-- 2 files changed, 34 insertions(+), 2 deletions(-) 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 abf165016036..ec56433eb017 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -729,6 +729,9 @@ def convert_to_wna16_moe_kernel_format( from vllm.model_executor.layers.quantization.awq_marlin import ( AWQMarlinConfig, ) + from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + ) if isinstance(quant_config, AWQMarlinConfig): if w13_qzeros is None or w2_qzeros is None: @@ -763,6 +766,11 @@ def convert_to_wna16_moe_kernel_format( pack_factor = 32 // quant_config.num_bits group_size = quant_config.group_size actorder = quant_config.actorder + elif isinstance(quant_config, MoeWNA16Config): + num_bits = quant_config.weight_bits + pack_factor = quant_config.bit8_pack_factor + group_size = quant_config.group_size + actorder = None else: raise TypeError( "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or " diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 87232168304f..fd99d520c640 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -401,8 +401,8 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: w2=layer.w2_qweight, w13_scale=layer.w13_scales, w2_scale=layer.w2_scales, - w13_g_idx=layer.w13_g_idx, - w2_g_idx=layer.w2_g_idx, + w13_g_idx=getattr(layer, "w13_g_idx", None), + w2_g_idx=getattr(layer, "w2_g_idx", None), w13_qzeros=layer.w13_qzeros if has_zp else None, w2_qzeros=layer.w2_qzeros if has_zp else None, ) @@ -485,6 +485,30 @@ def apply( shared_experts_input=shared_experts_input, ) + def apply_monolithic( + self, + layer: RoutedExperts, + x: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + assert self.is_monolithic + assert self.moe_kernel is not None + return self.moe_kernel.apply_monolithic( + x, + layer.w13_weight, + layer.w2_weight, + router_logits, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + expert_map=layer.expert_map, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + num_expert_group=layer.num_expert_group, + topk_group=layer.topk_group, + e_score_correction_bias=layer.e_score_correction_bias, + routed_scaling_factor=layer.routed_scaling_factor, + ) + @staticmethod def get_weight_loader(layer, weight_loader): def convert_awq_tensor(tensor, tensor_type): From 9c360d9a2b2a3ca748e9ab9835538e6cc2f96c2e Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 28 May 2026 15:40:20 +0000 Subject: [PATCH 33/53] copy parallel support from trtllm_fp8 Signed-off-by: Bill Nell --- .../layers/fused_moe/experts/trtllm_mxint4_moe.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) 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 index d9ff72b8656c..a65873aca49c 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_mxint4_moe.py @@ -75,9 +75,8 @@ def _supports_parallel_config( ) -> bool: return ( not moe_parallel_config.use_all2all_kernels - and not moe_parallel_config.enable_eplb - and moe_parallel_config.dp_size <= 1 - ) + or moe_parallel_config.use_ag_rs_all2all_kernels + ) and not moe_parallel_config.enable_eplb @staticmethod def _supports_routing_method( From 44bf4f2e50f00e880ec51c3365db791aff2fb83b Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 29 May 2026 17:20:57 +0000 Subject: [PATCH 34/53] remove cruft Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 1 - 1 file changed, 1 deletion(-) 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 b23c8040b9df..dd96921d7071 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -527,7 +527,6 @@ def _process_weights_marlin( if w2_bias is not None: w2_bias_out = marlin_permute_bias(w2_bias) - # TODO: not sure about shapes here if w13_qzeros is not None: w13_qzeros_out = marlin_zero_points( w13_qzeros, From bc4311442f52281b76eae9038d21734e1b9db305 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Sun, 31 May 2026 16:04:08 +0000 Subject: [PATCH 35/53] fix int_wna16.py Signed-off-by: Bill Nell --- .../model_executor/layers/fused_moe/config.py | 6 +- .../layers/fused_moe/oracle/int_wna16.py | 138 +++++++++++++++--- .../layers/quantization/humming.py | 3 +- .../layers/quantization/utils/marlin_utils.py | 41 ++++++ vllm/utils/import_utils.py | 5 + 5 files changed, 168 insertions(+), 25 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b87b87f136dd..7a9246113d79 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -665,8 +665,10 @@ def gptq_marlin_moe_quant_config( w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size) - # Activations are NOT quantized for GPTQ (fp16/bf16) - a_shape = w_shape # Same as weight shape for alignment + # Activations are NOT quantized for GPTQ (fp16/bf16). + # Use row=0 to signal "grouped along K only, no activation quantization", + # which makes block_shape[0] == 0 — required by the Triton WNA16 kernel. + a_shape = None if group_size == -1 else GroupShape(row=0, col=group_size) # Determine weight dtype if weight_bits == 4: 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 dd96921d7071..ab456eef4264 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -33,8 +33,8 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_act_int8_process_scales, marlin_moe_permute_scales, + marlin_moe_permute_zero_points, marlin_permute_bias, - marlin_zero_points, moe_awq_to_marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -91,13 +91,10 @@ def _get_priority_backends( # Marlin supports ZP and bias _AVAILABLE_BACKENDS += [ + WNA16MoEBackend.TRITON, # TEMP TESTING WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] - return _AVAILABLE_BACKENDS - - if not may_have_bias: - _AVAILABLE_BACKENDS.append(WNA16MoEBackend.TRITON) return _AVAILABLE_BACKENDS @@ -370,6 +367,87 @@ def _process_weights_flashinfer( ) +def _marlin_moe_process_zero_points( + qzeros: torch.Tensor, + pack_factor: int, + num_bits: int, + is_a_8bit: bool, +) -> torch.Tensor: + """Unpack and permute MoE zero points into Marlin format. + + Handles two input conventions: + + * **MoeWNA16** (``dtype == uint8``): shape ``(E, N_out // bit8_pack, K // gs)`` + where bit8_pack = 8 // num_bits. The packed dimension (axis 1) contains + output features; the group dimension (axis 2) contains input groups. The + nibbles are unpacked and transposed to produce ``(E, K // gs, N_out)`` + before applying the Marlin permutation. + + * **GPTQ / AutoGPTQ / QuantizationArgs** (other dtypes): shape + ``(E, K // gs, N_out // pack32)`` where pack32 = 32 // num_bits. The + group dimension (axis 1) and output features (axis 2) are already in the + right order; only int32 unpacking is needed. + + Returns Marlin-formatted zero points with shape + ``(E, K // gs, N_out // pack32)`` as int32. + """ + from vllm.model_executor.layers.quantization.utils.quant_utils import unpack_cols + + num_experts = qzeros.shape[0] + + if qzeros.dtype == torch.uint8: + # MoeWNA16 format: (E, N_out // bit8_pack, K // gs) + # axis 1 = packed output features, axis 2 = input groups + packed_out = qzeros.shape[1] # N_out // bit8_pack + num_groups = qzeros.shape[2] # K // group_size + size_n = packed_out * pack_factor # N_out (output features) + + unpacked = torch.empty( + (num_experts, num_groups, size_n), + dtype=torch.int32, + device=qzeros.device, + ) + for e in range(num_experts): + # Transpose so rows = groups, cols = packed outputs + zp_t = qzeros[e].T.to(torch.int32) # (num_groups, packed_out) + if num_bits == 4: + low = zp_t & 0xF # (num_groups, packed_out) + high = (zp_t >> 4) & 0xF # (num_groups, packed_out) + # Interleave: even channels from low, odd from high + unpacked[e] = torch.stack([low, high], dim=2).reshape( + num_groups, size_n + ) + else: # 8-bit: one zero point per byte, already transposed + unpacked[e] = zp_t + else: + # GPTQ/AutoGPTQ format: (E, K // gs, N_out // pack32) + # axis 1 = input groups, axis 2 = packed output features + # Note: the parameter may be stored as params_dtype (e.g. bfloat16) even + # though it holds packed-integer data. Use .to(int32) (value conversion) + # rather than .view(int32) (bitwise reinterpret) to avoid shape changes + # caused by the 2-byte vs 4-byte size difference. + num_groups = qzeros.shape[1] # K // group_size + size_n = qzeros.shape[2] * pack_factor # N_out (output features) + + unpacked = torch.empty( + (num_experts, num_groups, size_n), + dtype=torch.int32, + device=qzeros.device, + ) + for e in range(num_experts): + unpacked[e] = unpack_cols( + qzeros[e].to(torch.int32), num_bits, num_groups, size_n + ) + + return marlin_moe_permute_zero_points( + zp=unpacked, + size_k=num_groups, + size_n=size_n, + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + + def _process_weights_marlin( layer: torch.nn.Module, input_dtype: torch.dtype | None, @@ -528,21 +606,13 @@ def _process_weights_marlin( w2_bias_out = marlin_permute_bias(w2_bias) if w13_qzeros is not None: - w13_qzeros_out = marlin_zero_points( - w13_qzeros, - size_k=layer.intermediate_size_per_partition, - size_n=w13_qzeros.shape[2], - num_bits=num_bits, - is_a_8bit=is_a_8bit, + w13_qzeros_out = _marlin_moe_process_zero_points( + w13_qzeros, pack_factor, num_bits, is_a_8bit ) if w2_qzeros is not None: - w2_qzeros_out = marlin_zero_points( - w2_qzeros, - size_k=w2_qzeros.shape[1] * group_size_or_pack_factor, - size_n=w2_qzeros.shape[2], - num_bits=num_bits, - is_a_8bit=is_a_8bit, + w2_qzeros_out = _marlin_moe_process_zero_points( + w2_qzeros, pack_factor, num_bits, is_a_8bit ) return ( @@ -903,11 +973,35 @@ def convert_to_wna16_moe_kernel_format( w2_bias, ) elif backend == WNA16MoEBackend.TRITON: - # Convert from int32 to uint8 format for Triton kernel. - # This changes the shape from (E, N, K // 8) to (E, N, K // 2) for int4, - # which matches what the Triton kernel expects. - w13_uint8 = w13.view(torch.uint8) - w2_uint8 = w2.view(torch.uint8) + # Three possible input layouts depending on the quantization source: + # + # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first + # → just view as uint8 (no-op) + # + # AutoGPTQ (int32, K-first): (E, K // pack32, N_out) + # → transpose to N-first, then view as uint8 to get + # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s] + # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs) + # + # compressed-tensors (int32, N-first "Flashinfer" layout): + # (E, N_out, K // pack32) — already N-first + # → just view as uint8; scales are already (E, N_out, K // gs) + from vllm.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQConfig, + ) + + if isinstance(quant_config, AutoGPTQConfig): + # AutoGPTQ always builds in Marlin (K-first) format even when + # the Triton backend is selected. Transpose to N-first first. + w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8) + w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) + w13_scale = w13_scale.transpose(1, 2).contiguous() + w2_scale = w2_scale.transpose(1, 2).contiguous() + else: + # MoeWNA16 and compressed-tensors both use N-first layout when + # targeting the Triton backend. + w13_uint8 = w13.view(torch.uint8) + w2_uint8 = w2.view(torch.uint8) return ( w13_uint8, w2_uint8, diff --git a/vllm/model_executor/layers/quantization/humming.py b/vllm/model_executor/layers/quantization/humming.py index 12bb07a4022d..de2df2096963 100644 --- a/vllm/model_executor/layers/quantization/humming.py +++ b/vllm/model_executor/layers/quantization/humming.py @@ -44,8 +44,9 @@ ) from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +from vllm.utils.import_utils import has_humming -if current_platform.is_cuda(): +if has_humming() and current_platform.is_cuda(): from humming.dtypes import DataType from humming.layer import HummingMethod from humming.schema import ( diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index eca04eed74b6..3c57b62646fc 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -369,6 +369,47 @@ def marlin_zero_points( return zp +def marlin_moe_permute_zero_points( + zp: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: bool = False, +) -> torch.Tensor: + """Permute standard MoE zero points into Marlin format. + + Args: + zp: Unpacked zero points, shape (num_experts, num_groups, size_n). + Values should be integer zero point values (not packed). + num_groups = K // group_size. + size_k: Number of groups (K // group_size), NOT the original K. + size_n: Output feature dimension N (not packed). + num_bits: Quantization bit width (4 or 8). + is_a_8bit: Whether activations are 8-bit quantized. + + Returns: + Marlin-formatted zero points, shape + (num_experts, num_groups, size_n // pack_factor) as int32. + + The transformation per expert applies three steps: + 1. Permute columns with scale_perm (not scale_perm_single) to match + the Marlin dequantization MMA layout. + 2. Interleave columns for the dequantize code. + 3. Pack multiple values into int32. + """ + pack_factor = 32 // num_bits + num_experts = zp.shape[0] + # marlin_zero_points packs size_n values into size_n // pack_factor int32s + output = torch.empty( + (num_experts, size_k, size_n // pack_factor), + device=zp.device, + dtype=torch.int32, + ) + for e in range(num_experts): + output[e] = marlin_zero_points(zp[e], size_k, size_n, num_bits, is_a_8bit) + return output + + def awq_to_marlin_zero_points( q_zp_packed: torch.Tensor, size_k: int, diff --git a/vllm/utils/import_utils.py b/vllm/utils/import_utils.py index e97228bfa609..ed82326d507a 100644 --- a/vllm/utils/import_utils.py +++ b/vllm/utils/import_utils.py @@ -475,3 +475,8 @@ def has_fbgemm_gpu() -> bool: def has_cutedsl() -> bool: """Whether the optional `cutelass` package is available.""" return _has_module("cutlass") + + +def has_humming() -> bool: + """Whether the optional `humming` package is available.""" + return _has_module("humming") From ecce8d593d6d1e6a0a10f74278f90dd295822ffe Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 2 Jun 2026 23:39:34 +0000 Subject: [PATCH 36/53] cleanup merge Signed-off-by: Bill Nell --- vllm/model_executor/layers/quantization/moe_wna16.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index fd99d520c640..e973d866ab93 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -475,8 +475,8 @@ def apply( x, layer.w13_qweight, layer.w2_qweight, - topk_weights, - topk_ids, + topk_weights=topk_weights, + topk_ids=topk_ids, activation=layer.activation, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, From e0f7ad77e5d210733db9201fbedefdf75fdc5294 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 3 Jun 2026 01:29:22 +0000 Subject: [PATCH 37/53] fix merge Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 37 +++++-------------- 1 file changed, 9 insertions(+), 28 deletions(-) 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 88166c132b40..46eef044d1fb 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -36,7 +36,6 @@ marlin_moe_permute_zero_points, marlin_permute_bias, moe_awq_to_marlin_zero_points, - moe_packed_to_marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -505,8 +504,6 @@ def _process_weights_marlin( w2_input_global_scale: torch.Tensor | None = None w13_bias_out: torch.Tensor | None = None w2_bias_out: torch.Tensor | None = None - w13_qzeros_out: torch.Tensor | None = None - w2_qzeros_out: torch.Tensor | None = None # --- FP8 weight / scale adjustment --- if input_dtype == torch.float8_e4m3fn: @@ -601,20 +598,14 @@ def _process_weights_marlin( ) # --- Permute zero points --- - if w13_qzeros is not None and w2_qzeros is not None: - w13_qzeros = moe_packed_to_marlin_zero_points( - w13_qzeros, - size_k=w13_qzeros.shape[1], - size_n=w13_qzeros.shape[2] * pack_factor, - num_bits=num_bits, - is_a_8bit=is_a_8bit, + if w13_qzeros is not None: + w13_qzeros = _marlin_moe_process_zero_points( + w13_qzeros, pack_factor, num_bits, is_a_8bit ) - w2_qzeros = moe_packed_to_marlin_zero_points( - w2_qzeros, - size_k=w2_qzeros.shape[1], - size_n=w2_qzeros.shape[2] * pack_factor, - num_bits=num_bits, - is_a_8bit=is_a_8bit, + + if w2_qzeros is not None: + w2_qzeros = _marlin_moe_process_zero_points( + w2_qzeros, pack_factor, num_bits, is_a_8bit ) # --- Permute bias --- @@ -623,16 +614,6 @@ def _process_weights_marlin( if w2_bias is not None: w2_bias_out = marlin_permute_bias(w2_bias) - if w13_qzeros is not None: - w13_qzeros_out = _marlin_moe_process_zero_points( - w13_qzeros, pack_factor, num_bits, is_a_8bit - ) - - if w2_qzeros is not None: - w2_qzeros_out = _marlin_moe_process_zero_points( - w2_qzeros, pack_factor, num_bits, is_a_8bit - ) - return ( marlin_w13_qweight, marlin_w2_qweight, @@ -642,8 +623,8 @@ def _process_weights_marlin( w2_g_idx, w13_g_idx_sort_indices, w2_g_idx_sort_indices, - w13_qzeros_out, - w2_qzeros_out, + w13_qzeros, + w2_qzeros, w13_input_global_scale, w2_input_global_scale, w13_bias_out, From 57642d6ffdc90bc85e963e33f71f273a3434abf9 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 3 Jun 2026 04:20:18 +0000 Subject: [PATCH 38/53] move triton to the back Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/oracle/int_wna16.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 46eef044d1fb..25896c6b0ae8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -91,9 +91,9 @@ def _get_priority_backends( # Marlin supports ZP and bias _AVAILABLE_BACKENDS += [ - WNA16MoEBackend.TRITON, # TEMP TESTING WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.TRITON, ] return _AVAILABLE_BACKENDS From ee1d347467755a25ec90b4520567507400a1652e Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 5 Jun 2026 03:18:05 +0000 Subject: [PATCH 39/53] remove useless test Signed-off-by: Bill Nell --- tests/quantization/test_moe_wna16.py | 49 ------------------- .../layers/quantization/moe_wna16.py | 2 +- 2 files changed, 1 insertion(+), 50 deletions(-) delete mode 100644 tests/quantization/test_moe_wna16.py diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py deleted file mode 100644 index c4b0ab5a8464..000000000000 --- a/tests/quantization/test_moe_wna16.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from types import SimpleNamespace - -import pytest -import torch - -from vllm.model_executor.layers.fused_moe.activation import MoEActivation -from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method -from vllm.platforms import current_platform - - -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA") -def test_moe_wna16_apply_passes_layer_activation(monkeypatch): - captured_kwargs = {} - - def fake_fused_experts(*args, **kwargs): - captured_kwargs.update(kwargs) - return torch.empty(1, 2) - - monkeypatch.setattr( - "vllm.model_executor.layers.fused_moe.fused_experts", - fake_fused_experts, - ) - - method = object.__new__(MoeWNA16Method) - method.moe = SimpleNamespace(disable_inplace=False) - method.moe_quant_config = object() - layer = SimpleNamespace( - w13_qweight=torch.empty(1, 2), - w2_qweight=torch.empty(1, 2), - activation=MoEActivation.GELU_TANH, - apply_router_weight_on_input=False, - global_num_experts=1, - expert_map=None, - ) - - output = method.apply( - layer, - x=torch.empty(1, 2), - topk_weights=torch.empty(1, 1), - topk_ids=torch.empty(1, 1, dtype=torch.int32), - shared_experts=None, - shared_experts_input=None, - ) - - assert output.shape == (1, 2) - assert captured_kwargs["activation"] is MoEActivation.GELU_TANH diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index e973d866ab93..c521cc96ee24 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -375,7 +375,7 @@ def create_weights( layer.register_parameter(key, param) set_weight_attrs(param, extra_weight_attrs) - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + def process_weights_after_loading(self, layer: RoutedExperts) -> None: has_zp = self.quant_config.has_zp ( w13_qweight, From 3901126c6ade6c6ca5989e357d001b976137e741 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 5 Jun 2026 18:13:20 +0000 Subject: [PATCH 40/53] disallow marlin for moe_wna16 Signed-off-by: Bill Nell --- benchmarks/kernels/benchmark_moe.py | 3 - .../layers/fused_moe/oracle/int_wna16.py | 105 +++++++++--------- .../layers/quantization/moe_wna16.py | 22 ++-- 3 files changed, 63 insertions(+), 67 deletions(-) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 4463a23772ee..f7079af359cc 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -271,7 +271,6 @@ def run(): moe_config=moe_config, quant_config=quant_config, ), - inplace=not disable_inplace(), ) with override_config(config): @@ -279,7 +278,6 @@ def run(): x, input_gating, topk, renormalize=not use_deep_gemm ) - inplace = not disable_inplace() if use_deep_gemm: return deep_gemm_experts.apply( x, @@ -298,7 +296,6 @@ def run(): w2, topk_weights, topk_ids, - inplace=inplace, quant_config=quant_config, ) 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 25896c6b0ae8..8fcd7e3a165c 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -76,7 +76,9 @@ def backend_to_kernel_cls( def _get_priority_backends( - may_have_zp: bool, may_have_bias: bool + may_have_zp: bool, + may_have_bias: bool, + allow_marlin: bool, ) -> list[WNA16MoEBackend]: """ Get available backends in priority order based on platform and config. @@ -90,9 +92,13 @@ def _get_priority_backends( _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM) # Marlin supports ZP and bias + if allow_marlin: + _AVAILABLE_BACKENDS += [ + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + ] + _AVAILABLE_BACKENDS += [ - WNA16MoEBackend.MARLIN, - WNA16MoEBackend.BATCHED_MARLIN, WNA16MoEBackend.TRITON, ] @@ -104,6 +110,7 @@ def select_wna16_moe_backend( weight_key: QuantKey, may_have_zp: bool, may_have_bias: bool, + allow_marlin: bool = True, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. @@ -154,7 +161,9 @@ def _return_or_raise( raise ValueError(_make_log_unsupported(backend, reason)) # Select kernels in order of backend. - AVAILABLE_BACKENDS = _get_priority_backends(may_have_zp, may_have_bias) + AVAILABLE_BACKENDS = _get_priority_backends( + may_have_zp, may_have_bias, allow_marlin + ) for backend in AVAILABLE_BACKENDS: activation_key = None # always BF16 activation for WNA16 MoE @@ -892,21 +901,34 @@ def convert_to_wna16_moe_kernel_format( from vllm.model_executor.layers.quantization.awq_marlin import ( AWQMarlinConfig, ) - from vllm.model_executor.layers.quantization.moe_wna16 import ( - MoeWNA16Config, - ) 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 + 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 + actorder = "group" if quant_config.desc_act else None + elif isinstance(quant_config, QuantizationArgs): + num_bits = quant_config.num_bits + pack_factor = 32 // quant_config.num_bits + group_size = quant_config.group_size + actorder = quant_config.actorder + else: + raise TypeError( + "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or " + f"QuantizationArgs, got {type(quant_config).__name__}." + ) + + if isinstance(quant_config, AWQMarlinConfig): + 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, - weight_bits, + num_bits, pack_factor, group_size, input_dtype, @@ -919,47 +941,28 @@ def convert_to_wna16_moe_kernel_format( w13_bias, w2_bias, ) - elif isinstance(quant_config, AutoGPTQConfig): - num_bits = quant_config.quant_type.size_bits - pack_factor = quant_config.pack_factor - group_size = quant_config.group_size - actorder = "group" if quant_config.desc_act else None - elif isinstance(quant_config, QuantizationArgs): - num_bits = quant_config.num_bits - pack_factor = 32 // quant_config.num_bits - group_size = quant_config.group_size - actorder = quant_config.actorder - elif isinstance(quant_config, MoeWNA16Config): - num_bits = quant_config.weight_bits - pack_factor = quant_config.bit8_pack_factor - group_size = quant_config.group_size - actorder = None else: - raise TypeError( - "Marlin WNA16 MoE backend requires AutoGPTQConfig, AWQMarlinConfig or " - f"QuantizationArgs, got {type(quant_config).__name__}." - ) + if w13_g_idx is None or w2_g_idx is None: + raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") - if w13_g_idx is None or w2_g_idx is None: - raise ValueError("GPTQ Marlin MoE requires g_idx tensors.") - return _process_weights_marlin( - layer, - input_dtype, - num_bits, - pack_factor, - group_size, - actorder, - w13, - w2, - w13_scale, - w2_scale, - w13_g_idx, - w2_g_idx, - w13_qzeros, - w2_qzeros, - w13_bias, - w2_bias, - ) + return _process_weights_marlin( + layer, + input_dtype, + num_bits, + pack_factor, + group_size, + actorder, + w13, + w2, + w13_scale, + w2_scale, + w13_g_idx, + w2_g_idx, + w13_qzeros, + w2_qzeros, + w13_bias, + w2_bias, + ) elif backend == WNA16MoEBackend.FLASHINFER_TRTLLM: return _process_weights_flashinfer( w13, diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index c521cc96ee24..4cc1f557918f 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -253,6 +253,7 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: weight_key=weight_key, may_have_zp=self.quant_config.has_zp, may_have_bias=False, + allow_marlin=False, ) def create_weights( @@ -382,10 +383,10 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: 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, @@ -401,8 +402,8 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: w2=layer.w2_qweight, w13_scale=layer.w13_scales, w2_scale=layer.w2_scales, - w13_g_idx=getattr(layer, "w13_g_idx", None), - w2_g_idx=getattr(layer, "w2_g_idx", None), + w13_g_idx=None, + w2_g_idx=None, w13_qzeros=layer.w13_qzeros if has_zp else None, w2_qzeros=layer.w2_qzeros if has_zp else None, ) @@ -420,11 +421,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: # Marlin-specific parameters (not needed for Flashinfer) if self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM: - replace_parameter(layer, "w13_g_idx", w13_g_idx_processed) - replace_parameter(layer, "w2_g_idx", w2_g_idx_processed) - replace_parameter(layer, "w13_g_idx_sort_indices", w13_g_idx_sort_indices) - replace_parameter(layer, "w2_g_idx_sort_indices", w2_g_idx_sort_indices) - # Register input global scales if present if w13_input_global_scale is not None: layer.register_parameter( @@ -496,8 +492,8 @@ def apply_monolithic( assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( x, - layer.w13_weight, - layer.w2_weight, + layer.w13_qweight, + layer.w2_qweight, router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts, From 710eec1e453db56ca30a4d2dde1ac44b0d8da8ed Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 5 Jun 2026 19:09:25 +0000 Subject: [PATCH 41/53] fix lint Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 36 +++++-------------- 1 file changed, 9 insertions(+), 27 deletions(-) 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 8fcd7e3a165c..7d0b394eacef 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -274,35 +274,17 @@ def make_wna16_moe_kernel( "is_k_full": is_k_full, } - if experts_cls is XPUExpertsWNA16: - assert ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.Standard - ), ( - "XPUExpertsWNA16 only supports the Standard activation format; " - "xpu_fused_moe(is_int4=True) does not implement BatchedExperts." - ) - experts: mk.FusedMoEExperts = XPUExpertsWNA16( - moe_config=moe_config, - quant_config=moe_quant_config, - ) - elif ( - prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts - ): + if prepare_finalize.activation_format == mk.FusedMoEActivationFormat.BatchedExperts: max_num_tokens = prepare_finalize.max_num_tokens_per_rank() assert max_num_tokens is not None - experts = experts_cls( - max_num_tokens=max_num_tokens, - num_dispatchers=prepare_finalize.num_dispatchers(), - moe_config=moe_config, - quant_config=moe_quant_config, - **extra_args, - ) - else: - experts = experts_cls( - moe_config=moe_config, - quant_config=moe_quant_config, - **extra_args, - ) + extra_args["max_num_tokens"] = max_num_tokens + extra_args["num_dispatchers"] = prepare_finalize.num_dispatchers() + + experts = experts_cls( + moe_config=moe_config, + quant_config=moe_quant_config, + **extra_args, + ) return mk.FusedMoEKernel( prepare_finalize, From 8b998f705b3c675b79de379152de31a42de0a7c9 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Tue, 16 Jun 2026 20:47:52 +0000 Subject: [PATCH 42/53] cleanup Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 103 ++++++++---------- .../layers/quantization/moe_wna16.py | 26 ++--- 2 files changed, 59 insertions(+), 70 deletions(-) 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 9855c8164216..3d8c773c39f3 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -1069,63 +1069,6 @@ 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.TRITON: - # Three possible input layouts depending on the quantization source: - # - # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first - # → just view as uint8 (no-op) - # - # AutoGPTQ (int32, K-first): (E, K // pack32, N_out) - # → transpose to N-first, then view as uint8 to get - # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s] - # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs) - # - # compressed-tensors (int32, N-first "Flashinfer" layout): - # (E, N_out, K // pack32) — already N-first - # → just view as uint8; scales are already (E, N_out, K // gs) - from vllm.model_executor.layers.quantization.auto_gptq import ( - AutoGPTQConfig, - ) - - if isinstance(quant_config, AutoGPTQConfig): - # AutoGPTQ always builds in Marlin (K-first) format even when - # the Triton backend is selected. Transpose to N-first first. - w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8) - w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) - w13_scale = w13_scale.transpose(1, 2).contiguous() - w2_scale = w2_scale.transpose(1, 2).contiguous() - else: - # MoeWNA16 and compressed-tensors both use N-first layout when - # targeting the Triton backend. - w13_uint8 = w13.view(torch.uint8) - w2_uint8 = w2.view(torch.uint8) - return ( - w13_uint8, - w2_uint8, - w13_scale, - w2_scale, - None, - None, - None, - None, - w13_qzeros, - w2_qzeros, - None, - None, - w13_bias, - w2_bias, - ) elif backend == WNA16MoEBackend.CPU: return _process_weights_cpu( quant_config, @@ -1187,5 +1130,51 @@ def convert_to_wna16_moe_kernel_format( w13_bias_out, w2_bias_out, ) + elif backend == WNA16MoEBackend.TRITON: + # Three possible input layouts depending on the quantization source: + # + # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first + # → just view as uint8 (no-op) + # + # AutoGPTQ (int32, K-first): (E, K // pack32, N_out) + # → transpose to N-first, then view as uint8 to get + # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s] + # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs) + # + # compressed-tensors (int32, N-first "Flashinfer" layout): + # (E, N_out, K // pack32) — already N-first + # → just view as uint8; scales are already (E, N_out, K // gs) + from vllm.model_executor.layers.quantization.auto_gptq import ( + AutoGPTQConfig, + ) + + if isinstance(quant_config, AutoGPTQConfig): + # AutoGPTQ always builds in Marlin (K-first) format even when + # the Triton backend is selected. Transpose to N-first first. + w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8) + w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) + w13_scale = w13_scale.transpose(1, 2).contiguous() + w2_scale = w2_scale.transpose(1, 2).contiguous() + else: + # MoeWNA16 and compressed-tensors both use N-first layout when + # targeting the Triton backend. + w13_uint8 = w13.view(torch.uint8) + w2_uint8 = w2.view(torch.uint8) + return ( + w13_uint8, + w2_uint8, + w13_scale, + w2_scale, + None, + None, + None, + None, + w13_qzeros, + w2_qzeros, + None, + None, + w13_bias, + w2_bias, + ) else: raise ValueError(f"Unsupported wna16 MoE backend: {backend.value}") diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 4cc1f557918f..11470ddef4b4 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -376,6 +376,19 @@ def create_weights( layer.register_parameter(key, param) set_weight_attrs(param, extra_weight_attrs) + def get_fused_moe_quant_config( + self, layer: RoutedExperts + ) -> FusedMoEQuantConfig | None: + has_zp = self.quant_config.has_zp + return make_wna16_moe_quant_config( + w1_scale=layer.w13_scales, + w2_scale=layer.w2_scales, + w1_zp=layer.w13_qzeros if has_zp else None, + w2_zp=layer.w2_qzeros if has_zp else None, + group_size=layer.group_size, + num_bits=self.quant_config.weight_bits, + ) + def process_weights_after_loading(self, layer: RoutedExperts) -> None: has_zp = self.quant_config.has_zp ( @@ -443,19 +456,6 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: routing_tables=layer._expert_routing_tables(), ) - def get_fused_moe_quant_config( - self, layer: RoutedExperts - ) -> FusedMoEQuantConfig | None: - has_zp = self.quant_config.has_zp - return make_wna16_moe_quant_config( - w1_scale=layer.w13_scales, - w2_scale=layer.w2_scales, - w1_zp=layer.w13_qzeros if has_zp else None, - w2_zp=layer.w2_qzeros if has_zp else None, - group_size=layer.group_size, - num_bits=self.quant_config.weight_bits, - ) - def apply( self, layer: RoutedExperts, From 73a0c8ef133c6b8f931e7906b366e55b95386f06 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 17 Jun 2026 14:36:16 +0000 Subject: [PATCH 43/53] revert zero point code Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 103 +++--------------- .../layers/quantization/utils/marlin_utils.py | 41 ------- 2 files changed, 14 insertions(+), 130 deletions(-) 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 3d8c773c39f3..14ba13cf94a9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -33,9 +33,9 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_act_int8_process_scales, marlin_moe_permute_scales, - marlin_moe_permute_zero_points, marlin_permute_bias, moe_awq_to_marlin_zero_points, + moe_packed_to_marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -371,87 +371,6 @@ def _process_weights_flashinfer( ) -def _marlin_moe_process_zero_points( - qzeros: torch.Tensor, - pack_factor: int, - num_bits: int, - is_a_8bit: bool, -) -> torch.Tensor: - """Unpack and permute MoE zero points into Marlin format. - - Handles two input conventions: - - * **MoeWNA16** (``dtype == uint8``): shape ``(E, N_out // bit8_pack, K // gs)`` - where bit8_pack = 8 // num_bits. The packed dimension (axis 1) contains - output features; the group dimension (axis 2) contains input groups. The - nibbles are unpacked and transposed to produce ``(E, K // gs, N_out)`` - before applying the Marlin permutation. - - * **GPTQ / AutoGPTQ / QuantizationArgs** (other dtypes): shape - ``(E, K // gs, N_out // pack32)`` where pack32 = 32 // num_bits. The - group dimension (axis 1) and output features (axis 2) are already in the - right order; only int32 unpacking is needed. - - Returns Marlin-formatted zero points with shape - ``(E, K // gs, N_out // pack32)`` as int32. - """ - from vllm.model_executor.layers.quantization.utils.quant_utils import unpack_cols - - num_experts = qzeros.shape[0] - - if qzeros.dtype == torch.uint8: - # MoeWNA16 format: (E, N_out // bit8_pack, K // gs) - # axis 1 = packed output features, axis 2 = input groups - packed_out = qzeros.shape[1] # N_out // bit8_pack - num_groups = qzeros.shape[2] # K // group_size - size_n = packed_out * pack_factor # N_out (output features) - - unpacked = torch.empty( - (num_experts, num_groups, size_n), - dtype=torch.int32, - device=qzeros.device, - ) - for e in range(num_experts): - # Transpose so rows = groups, cols = packed outputs - zp_t = qzeros[e].T.to(torch.int32) # (num_groups, packed_out) - if num_bits == 4: - low = zp_t & 0xF # (num_groups, packed_out) - high = (zp_t >> 4) & 0xF # (num_groups, packed_out) - # Interleave: even channels from low, odd from high - unpacked[e] = torch.stack([low, high], dim=2).reshape( - num_groups, size_n - ) - else: # 8-bit: one zero point per byte, already transposed - unpacked[e] = zp_t - else: - # GPTQ/AutoGPTQ format: (E, K // gs, N_out // pack32) - # axis 1 = input groups, axis 2 = packed output features - # Note: the parameter may be stored as params_dtype (e.g. bfloat16) even - # though it holds packed-integer data. Use .to(int32) (value conversion) - # rather than .view(int32) (bitwise reinterpret) to avoid shape changes - # caused by the 2-byte vs 4-byte size difference. - num_groups = qzeros.shape[1] # K // group_size - size_n = qzeros.shape[2] * pack_factor # N_out (output features) - - unpacked = torch.empty( - (num_experts, num_groups, size_n), - dtype=torch.int32, - device=qzeros.device, - ) - for e in range(num_experts): - unpacked[e] = unpack_cols( - qzeros[e].to(torch.int32), num_bits, num_groups, size_n - ) - - return marlin_moe_permute_zero_points( - zp=unpacked, - size_k=num_groups, - size_n=size_n, - num_bits=num_bits, - is_a_8bit=is_a_8bit, - ) - - def _process_weights_marlin( layer: torch.nn.Module, input_dtype: torch.dtype | None, @@ -602,14 +521,20 @@ def _process_weights_marlin( ) # --- Permute zero points --- - if w13_qzeros is not None: - w13_qzeros = _marlin_moe_process_zero_points( - w13_qzeros, pack_factor, num_bits, is_a_8bit + if w13_qzeros is not None and w2_qzeros is not None: + w13_qzeros = moe_packed_to_marlin_zero_points( + w13_qzeros, + size_k=w13_qzeros.shape[1], + size_n=w13_qzeros.shape[2] * pack_factor, + num_bits=num_bits, + is_a_8bit=is_a_8bit, ) - - if w2_qzeros is not None: - w2_qzeros = _marlin_moe_process_zero_points( - w2_qzeros, pack_factor, num_bits, is_a_8bit + w2_qzeros = moe_packed_to_marlin_zero_points( + w2_qzeros, + size_k=w2_qzeros.shape[1], + size_n=w2_qzeros.shape[2] * pack_factor, + num_bits=num_bits, + is_a_8bit=is_a_8bit, ) # --- Permute bias --- diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index b4add39e01ec..1aba32621fc5 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -474,47 +474,6 @@ def marlin_zero_points( return zp -def marlin_moe_permute_zero_points( - zp: torch.Tensor, - size_k: int, - size_n: int, - num_bits: int, - is_a_8bit: bool = False, -) -> torch.Tensor: - """Permute standard MoE zero points into Marlin format. - - Args: - zp: Unpacked zero points, shape (num_experts, num_groups, size_n). - Values should be integer zero point values (not packed). - num_groups = K // group_size. - size_k: Number of groups (K // group_size), NOT the original K. - size_n: Output feature dimension N (not packed). - num_bits: Quantization bit width (4 or 8). - is_a_8bit: Whether activations are 8-bit quantized. - - Returns: - Marlin-formatted zero points, shape - (num_experts, num_groups, size_n // pack_factor) as int32. - - The transformation per expert applies three steps: - 1. Permute columns with scale_perm (not scale_perm_single) to match - the Marlin dequantization MMA layout. - 2. Interleave columns for the dequantize code. - 3. Pack multiple values into int32. - """ - pack_factor = 32 // num_bits - num_experts = zp.shape[0] - # marlin_zero_points packs size_n values into size_n // pack_factor int32s - output = torch.empty( - (num_experts, size_k, size_n // pack_factor), - device=zp.device, - dtype=torch.int32, - ) - for e in range(num_experts): - output[e] = marlin_zero_points(zp[e], size_k, size_n, num_bits, is_a_8bit) - return output - - def awq_to_marlin_zero_points( q_zp_packed: torch.Tensor, size_k: int, From 302c2383f873aebac2c05ebb68e9b0c4ff623616 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 17 Jun 2026 20:07:08 +0000 Subject: [PATCH 44/53] fix qwen1.5 crash Signed-off-by: Bill Nell --- .../layers/quantization/utils/marlin_utils.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 1aba32621fc5..17367d88f93c 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -540,20 +540,27 @@ def moe_packed_to_marlin_zero_points( size_n: int, num_bits: int, is_a_8bit: bool = False, + flag: bool = False, ): """Convert compressed-tensors packed zero points to Marlin format. Unlike AWQ, compressed-tensors uses standard bit packing without interleaving, so we just unpack and apply Marlin permutation directly. """ + use_int32 = (q_zp_packed.dtype != torch.uint8,) num_experts = q_zp_packed.shape[0] output = torch.empty( (num_experts, q_zp_packed.shape[1], q_zp_packed.shape[2]), device=q_zp_packed.device, - dtype=q_zp_packed.dtype, + dtype=torch.int32 if use_int32 else q_zp_packed.dtype, ) for e in range(num_experts): - q_zp = unpack_cols(q_zp_packed[e], num_bits, size_k, size_n) + q_zp = unpack_cols( + q_zp_packed[e] if not use_int32 else q_zp_packed[e].to(torch.int32), + num_bits, + size_k, + size_n, + ) output[e] = marlin_zero_points(q_zp, size_k, size_n, num_bits, is_a_8bit) return output From b723ecfc2d68fbefa916519c9ce1bd32dc4717b7 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Wed, 17 Jun 2026 22:49:54 +0000 Subject: [PATCH 45/53] fix tp_size Signed-off-by: Bill Nell --- vllm/model_executor/layers/quantization/moe_wna16.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 11470ddef4b4..7b61a601ff86 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -605,10 +605,10 @@ def moe_wna16_weight_loader( layer.group_size_div_factor, 1 ) + tp_size = layer.moe_config.tp_size + if "w13_qzeros" in weight_name: - tensor = loaded_weight.view(layer.tp_size, -1, loaded_weight.size(1))[ - tp_rank - ] + tensor = loaded_weight.view(tp_size, -1, loaded_weight.size(1))[tp_rank] if shard_id == "w1": param.data[expert_id, : shard_size // 2] = tensor else: @@ -616,7 +616,7 @@ def moe_wna16_weight_loader( return True if return_success else None elif "w2_qzeros" in weight_name: param.data[expert_id] = loaded_weight.view( - loaded_weight.size(0), layer.tp_size, -1 + loaded_weight.size(0), tp_size, -1 )[:, tp_rank] return True if return_success else None else: From 55134d41e03e4eee632c586e3ab8c76a6d42d797 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 18 Jun 2026 02:22:07 +0000 Subject: [PATCH 46/53] try to fix cpu zp Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 25 ++++++++++--------- .../layers/quantization/utils/marlin_utils.py | 6 ++--- 2 files changed, 16 insertions(+), 15 deletions(-) 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 14ba13cf94a9..544f3891921b 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -770,18 +770,19 @@ def _process_weights_cpu( # 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 - ) + + def convert_zp(t_zp: torch.Tensor | None) -> torch.Tensor | None: + if t_zp is None: + return t_zp + if t_zp.dtype.is_floating_point: + return t_zp.to(torch.int32) + elif t_zp.dtype == torch.int32: + return t_zp.data + else: + return t_zp.data.view(torch.int32) + + w13_zeros = convert_zp(w13_qzeros) + w2_zeros = convert_zp(w2_qzeros) ( blocked_w13, diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 17367d88f93c..88fa3e798736 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -547,16 +547,16 @@ def moe_packed_to_marlin_zero_points( Unlike AWQ, compressed-tensors uses standard bit packing without interleaving, so we just unpack and apply Marlin permutation directly. """ - use_int32 = (q_zp_packed.dtype != torch.uint8,) + cvt_int32 = q_zp_packed.dtype.is_floating_point num_experts = q_zp_packed.shape[0] output = torch.empty( (num_experts, q_zp_packed.shape[1], q_zp_packed.shape[2]), device=q_zp_packed.device, - dtype=torch.int32 if use_int32 else q_zp_packed.dtype, + dtype=torch.int32 if cvt_int32 else q_zp_packed.dtype, ) for e in range(num_experts): q_zp = unpack_cols( - q_zp_packed[e] if not use_int32 else q_zp_packed[e].to(torch.int32), + q_zp_packed[e] if not cvt_int32 else q_zp_packed[e].to(torch.int32), num_bits, size_k, size_n, From be21b02e10a98d9bcbdfe1d599d3a780a882b1a2 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 18 Jun 2026 19:43:42 +0000 Subject: [PATCH 47/53] fix Signed-off-by: Bill Nell --- .../layers/fused_moe/oracle/int_wna16.py | 25 +++++++++---------- .../layers/quantization/auto_gptq.py | 16 +++++++----- .../layers/quantization/utils/marlin_utils.py | 11 ++------ 3 files changed, 24 insertions(+), 28 deletions(-) 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 544f3891921b..14ba13cf94a9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -770,19 +770,18 @@ def _process_weights_cpu( # Determine zero points for repacking. w13_zeros: torch.Tensor | None = None w2_zeros: torch.Tensor | None = None - - def convert_zp(t_zp: torch.Tensor | None) -> torch.Tensor | None: - if t_zp is None: - return t_zp - if t_zp.dtype.is_floating_point: - return t_zp.to(torch.int32) - elif t_zp.dtype == torch.int32: - return t_zp.data - else: - return t_zp.data.view(torch.int32) - - w13_zeros = convert_zp(w13_qzeros) - w2_zeros = convert_zp(w2_qzeros) + 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, diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index 604bcfc4efa8..be23ffe5a955 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -60,6 +60,7 @@ PackedvLLMParameter, RowvLLMParameter, ) +from vllm.platforms import current_platform from vllm.scalar_type import scalar_types from vllm.transformers_utils.config import get_safetensors_params_metadata from vllm.utils.collection_utils import is_list_of @@ -576,25 +577,25 @@ def create_weights( set_weight_attrs(w2_scales, extra_weight_attrs) # don't shard the w2 scales when running act order set_weight_attrs(w2_scales, {"load_full_w2": self.quant_config.desc_act}) - # up_proj scales + # up_proj zero points w13_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size13, 2 * intermediate_size_per_partition // self.quant_config.pack_factor, - dtype=params_dtype, + dtype=torch.int32, ), requires_grad=False, ) layer.register_parameter("w13_qzeros", w13_qzeros) set_weight_attrs(w13_qzeros, extra_weight_attrs) - # down_proj scales + # down_proj zero points w2_qzeros = torch.nn.Parameter( torch.empty( num_experts, scales_size2, hidden_size // self.quant_config.pack_factor, - dtype=params_dtype, + dtype=torch.int32, ), requires_grad=False, ) @@ -667,6 +668,9 @@ def replace_or_register(name: str, val: torch.Tensor | None): "W8A8-INT8 is not supported by marlin kernel." ) + # CPU int_wna16 can't handle ZP correctly. + is_cpu = current_platform.is_cpu() + ( w13, w2, @@ -695,8 +699,8 @@ def replace_or_register(name: str, val: torch.Tensor | None): w2_g_idx=layer.w2_g_idx, w13_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - w13_qzeros=getattr(layer, "w13_qzeros", None), - w2_qzeros=getattr(layer, "w2_qzeros", None), + w13_qzeros=getattr(layer, "w13_qzeros", None) if not is_cpu else None, + w2_qzeros=getattr(layer, "w2_qzeros", None) if not is_cpu else None, ) replace_parameter(layer, "w13_qweight", w13) diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index 88fa3e798736..1aba32621fc5 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -540,27 +540,20 @@ def moe_packed_to_marlin_zero_points( size_n: int, num_bits: int, is_a_8bit: bool = False, - flag: bool = False, ): """Convert compressed-tensors packed zero points to Marlin format. Unlike AWQ, compressed-tensors uses standard bit packing without interleaving, so we just unpack and apply Marlin permutation directly. """ - cvt_int32 = q_zp_packed.dtype.is_floating_point num_experts = q_zp_packed.shape[0] output = torch.empty( (num_experts, q_zp_packed.shape[1], q_zp_packed.shape[2]), device=q_zp_packed.device, - dtype=torch.int32 if cvt_int32 else q_zp_packed.dtype, + dtype=q_zp_packed.dtype, ) for e in range(num_experts): - q_zp = unpack_cols( - q_zp_packed[e] if not cvt_int32 else q_zp_packed[e].to(torch.int32), - num_bits, - size_k, - size_n, - ) + q_zp = unpack_cols(q_zp_packed[e], num_bits, size_k, size_n) output[e] = marlin_zero_points(q_zp, size_k, size_n, num_bits, is_a_8bit) return output From 6f315c224c8035074f3342dabf17ad5905e8df20 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 19 Jun 2026 02:43:35 +0000 Subject: [PATCH 48/53] revert config changes Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/config.py | 6 ++---- vllm/model_executor/layers/quantization/auto_gptq.py | 8 ++------ 2 files changed, 4 insertions(+), 10 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 372ecf553066..905a9bea3c59 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -674,10 +674,8 @@ def gptq_marlin_moe_quant_config( w_shape = None if group_size == -1 else GroupShape(row=1, col=group_size) - # Activations are NOT quantized for GPTQ (fp16/bf16). - # Use row=0 to signal "grouped along K only, no activation quantization", - # which makes block_shape[0] == 0 — required by the Triton WNA16 kernel. - a_shape = None if group_size == -1 else GroupShape(row=0, col=group_size) + # Activations are NOT quantized for GPTQ (fp16/bf16) + a_shape = w_shape # Same as weight shape for alignment # Determine weight dtype if weight_bits == 4: diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index d8d1e968a7ae..aafb0e086bc0 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -61,7 +61,6 @@ PackedvLLMParameter, RowvLLMParameter, ) -from vllm.platforms import current_platform from vllm.scalar_type import scalar_types from vllm.transformers_utils.config import get_safetensors_params_metadata from vllm.utils.collection_utils import is_list_of @@ -669,9 +668,6 @@ def replace_or_register(name: str, val: torch.Tensor | None): "W8A8-INT8 is not supported by marlin kernel." ) - # CPU int_wna16 can't handle ZP correctly. - is_cpu = current_platform.is_cpu() - ( w13, w2, @@ -700,8 +696,8 @@ def replace_or_register(name: str, val: torch.Tensor | None): w2_g_idx=layer.w2_g_idx, w13_bias=getattr(layer, "w13_bias", None), w2_bias=getattr(layer, "w2_bias", None), - w13_qzeros=getattr(layer, "w13_qzeros", None) if not is_cpu else None, - w2_qzeros=getattr(layer, "w2_qzeros", None) if not is_cpu else None, + w13_qzeros=getattr(layer, "w13_qzeros", None), + w2_qzeros=getattr(layer, "w2_qzeros", None), ) replace_parameter(layer, "w13_qweight", w13) From cf8fc017d7c2bd3b9d3d4ea8196ca22699198749 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 19 Jun 2026 17:36:37 +0000 Subject: [PATCH 49/53] don't use marlin method for cpu Signed-off-by: Bill Nell --- .../compressed_tensors_moe/compressed_tensors_moe.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 2e45e0f298ba..4b98eedd3fc0 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -99,6 +99,7 @@ def get_moe_method( if ( not check_moe_marlin_supports_layer(layer, group_size) or current_platform.is_rocm() + or current_platform.is_cpu() ): if ( weight_quant.strategy == QuantizationStrategy.GROUP From e166d0fc5bda3b0fd4762c1c2a57b72d84e5ed7d Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Fri, 19 Jun 2026 18:41:36 +0000 Subject: [PATCH 50/53] try again Signed-off-by: Bill Nell --- .../compressed_tensors_moe/compressed_tensors_moe.py | 1 - .../compressed_tensors_moe_wna16_marlin.py | 5 +++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index 4b98eedd3fc0..2e45e0f298ba 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -99,7 +99,6 @@ def get_moe_method( if ( not check_moe_marlin_supports_layer(layer, group_size) or current_platform.is_rocm() - or current_platform.is_cpu() ): if ( weight_quant.strategy == QuantizationStrategy.GROUP 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 d443cd1ee078..2c24c135c95d 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 @@ -100,6 +100,7 @@ def __init__( WNA16MoEBackend.MARLIN, WNA16MoEBackend.BATCHED_MARLIN, ] + self.is_transposed = self.wna16_backend != WNA16MoEBackend.FLASHINFER_TRTLLM def get_weight_shape( self, @@ -183,7 +184,7 @@ def get_weight_shape( ), }, } - backend_key = "Marlin" if self.is_marlin else "Flashinfer" + backend_key = "Marlin" if self.is_transposed else "Flashinfer" return shape_map[weight_name][backend_key] def create_weights( @@ -201,7 +202,7 @@ def create_weights( # intermediate and hidden dim sizes. Will # shard for TP along the transposed dims extra_weight_attrs.update( - {"is_transposed": self.is_marlin, "quant_method": self.strategy} + {"is_transposed": self.is_transposed, "quant_method": self.strategy} ) w13_weight = torch.nn.Parameter( From 29610b521f194e47d16d68fd66a7bd678e94da2a Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 16 Jul 2026 19:46:32 +0000 Subject: [PATCH 51/53] don't try to load params that don't exist Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/routed_experts.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 80e4064ea6b6..e319f60a739d 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -884,7 +884,9 @@ def load_weights( matched = True weight_name = qual_name.replace(weight_name, param_name) param_name = weight_name.removeprefix(f"{self.layer_name}.") - param = getattr(self, param_name) + param = getattr(self, param_name, None) + if param is None: + continue if is_fused: # w1 and w3 share one fused tensor; use a local copy so the # transpose below doesn't mutate loaded_weight across From c7e71f4f79fb7e8aa08959be260a1c172e34e1c4 Mon Sep 17 00:00:00 2001 From: Bill Nell Date: Thu, 16 Jul 2026 19:59:36 +0000 Subject: [PATCH 52/53] revert change Signed-off-by: Bill Nell --- vllm/model_executor/layers/fused_moe/routed_experts.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index e319f60a739d..80e4064ea6b6 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -884,9 +884,7 @@ def load_weights( matched = True weight_name = qual_name.replace(weight_name, param_name) param_name = weight_name.removeprefix(f"{self.layer_name}.") - param = getattr(self, param_name, None) - if param is None: - continue + param = getattr(self, param_name) if is_fused: # w1 and w3 share one fused tensor; use a local copy so the # transpose below doesn't mutate loaded_weight across From b4efaa4fb8be190e839c438090dadecdffc06d90 Mon Sep 17 00:00:00 2001 From: mgoin Date: Tue, 21 Jul 2026 19:28:12 +0000 Subject: [PATCH 53/53] Fix WNA16 MoE backend compatibility --- tests/quantization/test_auto_gptq.py | 79 +++++++ tests/quantization/test_moe_wna16.py | 198 ++++++++++++++++++ .../layers/fused_moe/oracle/int_wna16.py | 177 ++++++++++++---- .../layers/fused_moe/routed_experts.py | 20 ++ .../layers/quantization/auto_awq.py | 1 + .../layers/quantization/auto_gptq.py | 34 ++- .../compressed_tensors_moe_wna16_marlin.py | 3 +- .../layers/quantization/moe_wna16.py | 21 +- 8 files changed, 487 insertions(+), 46 deletions(-) create mode 100644 tests/quantization/test_moe_wna16.py diff --git a/tests/quantization/test_auto_gptq.py b/tests/quantization/test_auto_gptq.py index b733ee486c18..9eefd88e915e 100644 --- a/tests/quantization/test_auto_gptq.py +++ b/tests/quantization/test_auto_gptq.py @@ -5,13 +5,17 @@ Run `pytest tests/quantization/test_auto_gptq.py -v -s`. """ +from types import SimpleNamespace + import pytest import torch from tests.quantization.utils import is_quant_method_supported +from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.quantization.auto_gptq import ( AutoGPTQConfig, AutoGPTQLinearMethod, + AutoGPTQMoEMethod, ) PROMPT = "On the surface of Mars, we found" @@ -54,3 +58,78 @@ def check_model(model): def test_auto_gptq_config_get_name(): """Test that AutoGPTQConfig.get_name() returns 'auto_gptq'.""" assert AutoGPTQConfig.get_name() == "auto_gptq" + + +def test_auto_gptq_moe_creates_zero_initialized_expert_biases(): + method = object.__new__(AutoGPTQMoEMethod) + method.quant_config = AutoGPTQConfig(4, 128, False, True, False, {}, {}) + method.input_dtype = None + method.experts_cls = None + layer = torch.nn.Module() + + method.create_weights( + layer=layer, + num_experts=2, + hidden_size=8, + intermediate_size_per_partition=4, + params_dtype=torch.float16, + intermediate_size_full=4, + weight_loader=lambda *args, **kwargs: None, + ) + + assert layer.w13_bias.shape == (2, 8) + assert layer.w2_bias.shape == (2, 8) + assert torch.count_nonzero(layer.w13_bias) == 0 + assert torch.count_nonzero(layer.w2_bias) == 0 + + +def test_routed_experts_loads_per_expert_biases(): + class Loader: + quant_config = None + quant_method = object() + moe_config = SimpleNamespace( + is_act_and_mul=True, + tp_rank=0, + moe_parallel_config=SimpleNamespace(tp_size=1), + ) + _get_hidden_dim = staticmethod(RoutedExperts._get_hidden_dim) + _narrow_expert_data_for_padding = staticmethod( + RoutedExperts._narrow_expert_data_for_padding + ) + _load_w13 = RoutedExperts._load_w13 + _loaded_expert_biases = set() + + @staticmethod + def _map_global_expert_id_to_local_expert_id(expert_id): + return expert_id + + loader = Loader() + w13_bias = torch.nn.Parameter(torch.zeros(1, 8), requires_grad=False) + w2_bias = torch.nn.Parameter(torch.zeros(1, 4), requires_grad=False) + + for shard_id, loaded in ( + ("w1", torch.tensor([1.0, 2.0, 3.0, 4.0])), + ("w3", torch.tensor([5.0, 6.0, 7.0, 8.0])), + ): + assert RoutedExperts.weight_loader( + loader, + w13_bias, + loaded, + weight_name="model.layers.0.mlp.experts.w13_bias", + shard_id=shard_id, + expert_id=0, + return_success=True, + ) + + assert RoutedExperts.weight_loader( + loader, + w2_bias, + torch.tensor([9.0, 10.0, 11.0, 12.0]), + weight_name="model.layers.0.mlp.experts.w2_bias", + shard_id="w2", + expert_id=0, + return_success=True, + ) + assert torch.equal(w13_bias, torch.arange(1, 9, dtype=torch.float32).reshape(1, 8)) + assert torch.equal(w2_bias, torch.arange(9, 13, dtype=torch.float32).reshape(1, 4)) + assert loader._loaded_expert_biases == {"w13_bias", "w2_bias"} diff --git a/tests/quantization/test_moe_wna16.py b/tests/quantization/test_moe_wna16.py new file mode 100644 index 000000000000..16c68fd1c79f --- /dev/null +++ b/tests/quantization/test_moe_wna16.py @@ -0,0 +1,198 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch +from compressed_tensors.quantization import ( + ActivationOrdering, + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) + +from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import ( + WNA16MoEBackend, + _backend_incompatibility_reason, + _convert_moe_wna16_humming_tensors, + convert_to_wna16_moe_kernel_format, + map_wna16_backend, +) +from vllm.model_executor.layers.quantization import moe_wna16 +from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig +from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig +from vllm.model_executor.layers.quantization.moe_wna16 import ( + MoeWNA16Config, + MoeWNA16Method, +) + + +def test_map_wna16_backend_supports_triton(): + assert map_wna16_backend("triton") == WNA16MoEBackend.TRITON + + +@pytest.mark.parametrize( + ("backend", "quant_config", "may_have_zp", "may_have_bias", "expected"), + [ + ( + WNA16MoEBackend.TRITON, + AutoAWQConfig(4, 128, True, False), + True, + False, + "AutoAWQ weight layout", + ), + ( + WNA16MoEBackend.TRITON, + AutoGPTQConfig(4, 128, True, True, False, {}, {}), + False, + False, + "activation ordering", + ), + ( + WNA16MoEBackend.TRITON, + QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=128, + actorder=ActivationOrdering.GROUP, + ), + False, + False, + "activation ordering", + ), + ( + WNA16MoEBackend.TRITON, + AutoGPTQConfig(4, 128, False, True, False, {}, {}), + False, + True, + "bias", + ), + ( + WNA16MoEBackend.MARLIN, + MoeWNA16Config( + linear_quant_method="gptq", + weight_bits=4, + group_size=128, + has_zp=False, + lm_head_quantized=False, + modules_to_not_convert=None, + full_config={}, + ), + False, + False, + "MoeWNA16 checkpoint layout", + ), + ], +) +def test_wna16_oracle_rejects_incompatible_quant_structures( + backend, quant_config, may_have_zp, may_have_bias, expected +): + reason = _backend_incompatibility_reason( + backend=backend, + quant_config=quant_config, + may_have_zp=may_have_zp, + may_have_bias=may_have_bias, + ) + + assert reason is not None + assert expected in reason + + +def test_compressed_tensors_weights_are_transposed_for_triton(): + quant_config = QuantizationArgs( + num_bits=4, + type=QuantizationType.INT, + strategy=QuantizationStrategy.GROUP, + symmetric=True, + dynamic=False, + group_size=32, + ) + w13 = torch.arange(16, dtype=torch.int32).reshape(1, 2, 8) + w2 = torch.arange(12, dtype=torch.int32).reshape(1, 2, 6) + w13_scale = torch.arange(32, dtype=torch.float16).reshape(1, 4, 8) + w2_scale = torch.arange(18, dtype=torch.float16).reshape(1, 3, 6) + + converted = convert_to_wna16_moe_kernel_format( + backend=WNA16MoEBackend.TRITON, + layer=torch.nn.Module(), + quant_config=quant_config, + input_dtype=None, + w13=w13, + w2=w2, + w13_scale=w13_scale, + w2_scale=w2_scale, + ) + + assert converted is not None + assert torch.equal(converted[0], w13.transpose(1, 2).contiguous().view(torch.uint8)) + assert torch.equal(converted[1], w2.transpose(1, 2).contiguous().view(torch.uint8)) + assert torch.equal(converted[2], w13_scale.transpose(1, 2).contiguous()) + assert torch.equal(converted[3], w2_scale.transpose(1, 2).contiguous()) + + +def test_moe_wna16_setup_forwards_selected_backend(monkeypatch): + method = object.__new__(MoeWNA16Method) + method.experts_cls = object + method.wna16_backend = WNA16MoEBackend.HUMMING + method.moe = object() + quant_config = object() + method.get_fused_moe_quant_config = lambda layer: quant_config + layer = SimpleNamespace(_expert_routing_tables=lambda: (None, None, None)) + captured = {} + kernel = object() + + def fake_make_wna16_moe_kernel(**kwargs): + captured.update(kwargs) + return kernel + + monkeypatch.setattr(moe_wna16, "make_wna16_moe_kernel", fake_make_wna16_moe_kernel) + + method._setup_kernel(layer) + + assert method.moe_kernel is kernel + assert captured["backend"] == WNA16MoEBackend.HUMMING + assert captured["layer"] is layer + + +def test_moe_wna16_humming_adapter_repacks_uint8_tensors(): + qweight = torch.arange(32, dtype=torch.uint8).reshape(1, 4, 8) + scales = torch.arange(16, dtype=torch.float16).reshape(1, 4, 4) + qzeros = torch.arange(16, dtype=torch.uint8).reshape(1, 8, 2) + + converted = _convert_moe_wna16_humming_tensors( + {"qweight": qweight, "scales": scales, "qzeros": qzeros}, + has_zero_point=True, + ) + + assert torch.equal(converted["weight"], qweight.view(torch.int32)) + assert converted["weight"].shape == (1, 4, 2) + assert torch.equal(converted["weight_scale"], scales) + expected_qzeros = ( + qzeros.transpose(-1, -2) + .contiguous() + .view(torch.int32) + .transpose(-1, -2) + .contiguous() + ) + assert torch.equal(converted["zero_point"], expected_qzeros) + assert converted["zero_point"].shape == (1, 2, 2) + + +def test_moe_wna16_uses_humming_quant_config(monkeypatch): + from vllm.model_executor.layers.quantization.utils import humming_utils + + method = object.__new__(MoeWNA16Method) + method.wna16_backend = WNA16MoEBackend.HUMMING + layer = object() + quant_config = object() + monkeypatch.setattr( + humming_utils, + "get_humming_moe_quant_config", + lambda actual_layer: quant_config if actual_layer is layer else None, + ) + + assert method.get_fused_moe_quant_config(layer) is quant_config 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 f0a2ad8c6dd4..9d9d3aa76569 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -104,11 +104,7 @@ def backend_to_kernel_cls( raise ValueError(f"Unknown WNA16 MoE backend: {backend.value}") -def _get_priority_backends( - may_have_zp: bool, - may_have_bias: bool, - allow_marlin: bool, -) -> list[WNA16MoEBackend]: +def _get_priority_backends() -> list[WNA16MoEBackend]: """ Get available backends in priority order based on platform and config. """ @@ -117,30 +113,56 @@ def _get_priority_backends( if current_platform.is_xpu(): return [WNA16MoEBackend.XPU] - _AVAILABLE_BACKENDS = [] - - if not may_have_zp and not may_have_bias: - _AVAILABLE_BACKENDS.append(WNA16MoEBackend.FLASHINFER_TRTLLM) - - # Marlin supports ZP and bias - if allow_marlin: - _AVAILABLE_BACKENDS += [ - WNA16MoEBackend.MARLIN, - WNA16MoEBackend.BATCHED_MARLIN, - ] - - _AVAILABLE_BACKENDS += [ + return [ + WNA16MoEBackend.FLASHINFER_TRTLLM, + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, WNA16MoEBackend.TRITON, WNA16MoEBackend.HUMMING, WNA16MoEBackend.EMULATION, ] - return _AVAILABLE_BACKENDS + +def _backend_incompatibility_reason( + backend: WNA16MoEBackend, + quant_config: QuantizationConfig | QuantizationArgs, + may_have_zp: bool, + may_have_bias: bool, +) -> str | None: + if backend == WNA16MoEBackend.FLASHINFER_TRTLLM and (may_have_zp or may_have_bias): + return "zero points and bias are not supported" + + from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig + from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config + + if backend == WNA16MoEBackend.TRITON: + if may_have_bias: + return "expert bias is not supported" + if isinstance(quant_config, AutoAWQConfig): + return "the AutoAWQ weight layout is not supported" + if isinstance(quant_config, AutoGPTQConfig) and quant_config.desc_act: + return "GPTQ activation ordering is not supported" + if ( + isinstance(quant_config, QuantizationArgs) + and quant_config.actorder == "group" + ): + return "group activation ordering is not supported" + + if isinstance(quant_config, MoeWNA16Config) and backend in ( + WNA16MoEBackend.MARLIN, + WNA16MoEBackend.BATCHED_MARLIN, + WNA16MoEBackend.EMULATION, + ): + return "the MoeWNA16 checkpoint layout is not supported" + + return None 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, @@ -157,9 +179,9 @@ def map_wna16_backend(runner_backend: MoEBackend) -> WNA16MoEBackend: def select_wna16_moe_backend( config: FusedMoEConfig, weight_key: QuantKey, + quant_config: QuantizationConfig | QuantizationArgs, may_have_zp: bool, may_have_bias: bool, - allow_marlin: bool = True, ) -> tuple[WNA16MoEBackend, type[mk.FusedMoEExperts]]: """Select the WNA16 MoE backend. @@ -167,6 +189,9 @@ def select_wna16_moe_backend( 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. Returns: A tuple of (``WNA16MoEBackend``, experts class or ``None``). @@ -213,16 +238,25 @@ def _return_or_raise( runner_backend = config.moe_backend if runner_backend != "auto": requested_backend = map_wna16_backend(runner_backend) + reason = _backend_incompatibility_reason( + requested_backend, quant_config, may_have_zp, may_have_bias + ) + if reason is not None: + raise ValueError(_make_log_unsupported(requested_backend, reason)) return _return_or_raise( requested_backend, config, weight_key, None, activation_format ) # Select kernels in order of backend. - AVAILABLE_BACKENDS = _get_priority_backends( - may_have_zp, may_have_bias, allow_marlin - ) + AVAILABLE_BACKENDS = _get_priority_backends() for backend in AVAILABLE_BACKENDS: + reason = _backend_incompatibility_reason( + backend, quant_config, may_have_zp, may_have_bias + ) + if reason is not None: + logger.debug_once(_make_log_unsupported(backend, reason), scope="local") + continue activation_key = None # always BF16 activation for WNA16 MoE for k_cls in backend_to_kernel_cls(backend): supported, reason = k_cls.is_supported_config( @@ -1044,11 +1078,68 @@ def _humming_wna16_weight_schema( "sym": quant_config.is_sym, } raise TypeError( - "Humming WNA16 MoE requires AutoAWQConfig or AutoGPTQConfig, " + "Humming WNA16 checkpoint schema requires AutoAWQConfig or " + "AutoGPTQConfig, " f"got {type(quant_config).__name__}." ) +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 + + +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 + + def _unpack_and_dequant_int4_gptq( w_int32: torch.Tensor, scale: torch.Tensor, @@ -1333,13 +1424,27 @@ def convert_to_wna16_moe_kernel_format( input_dtype: optional activation dtype, usually should be 16 bit. """ if backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Config from vllm.model_executor.layers.quantization.utils.humming_utils import ( convert_to_humming_moe_kernel_format, ) - convert_to_humming_moe_kernel_format( - layer, quant_config=_humming_wna16_weight_schema(quant_config) - ) + 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 ( @@ -1496,33 +1601,29 @@ def convert_to_wna16_moe_kernel_format( w2_qzeros, ) elif backend == WNA16MoEBackend.TRITON: - # Three possible input layouts depending on the quantization source: + # Two possible input layouts depending on the quantization source: # # MoeWNA16 (uint8): (E, N_out, K // bit8_pack) — N-first # → just view as uint8 (no-op) # - # AutoGPTQ (int32, K-first): (E, K // pack32, N_out) + # AutoGPTQ/compressed-tensors (int32, K-first): + # (E, K // pack32, N_out) # → transpose to N-first, then view as uint8 to get # (E, N_out, K // bit8_pack) [int32 = 4 bytes → 4 uint8s] # Scales: (E, K // gs, N_out) → transpose → (E, N_out, K // gs) - # - # compressed-tensors (int32, N-first "Flashinfer" layout): - # (E, N_out, K // pack32) — already N-first - # → just view as uint8; scales are already (E, N_out, K // gs) from vllm.model_executor.layers.quantization.auto_gptq import ( AutoGPTQConfig, ) - if isinstance(quant_config, AutoGPTQConfig): - # AutoGPTQ always builds in Marlin (K-first) format even when - # the Triton backend is selected. Transpose to N-first first. + if isinstance(quant_config, (AutoGPTQConfig, QuantizationArgs)): + # These integrations build in K-first format even when the Triton + # backend is selected. Transpose to N-first first. w13_uint8 = w13.transpose(1, 2).contiguous().view(torch.uint8) w2_uint8 = w2.transpose(1, 2).contiguous().view(torch.uint8) w13_scale = w13_scale.transpose(1, 2).contiguous() w2_scale = w2_scale.transpose(1, 2).contiguous() else: - # MoeWNA16 and compressed-tensors both use N-first layout when - # targeting the Triton backend. + # MoeWNA16 uses N-first uint8 weights and scales. w13_uint8 = w13.view(torch.uint8) w2_uint8 = w2.view(torch.uint8) return ( diff --git a/vllm/model_executor/layers/fused_moe/routed_experts.py b/vllm/model_executor/layers/fused_moe/routed_experts.py index 80e4064ea6b6..8bf0123e57fe 100644 --- a/vllm/model_executor/layers/fused_moe/routed_experts.py +++ b/vllm/model_executor/layers/fused_moe/routed_experts.py @@ -114,6 +114,7 @@ def __init__( self.e_score_correction_bias = e_score_correction_bias self.apply_router_weight_on_input = apply_router_weight_on_input # End random parameters + self._loaded_expert_biases: set[str] = set() self.quant_method = self._get_quant_method( self.layer_name, @@ -691,6 +692,25 @@ def weight_loader( expert_data = param.data if full_load else param.data[expert_id] + if "bias" in weight_name: + self._loaded_expert_biases.add(weight_name.rsplit(".", 1)[-1]) + if shard_id == "w2": + expert_data = self._narrow_expert_data_for_padding( + expert_data, + loaded_weight, + hidden_dim=0, + ) + expert_data.copy_(loaded_weight) + else: + self._load_w13( + shard_id=shard_id, + shard_dim=0, + loaded_weight=loaded_weight, + expert_data=expert_data, + tp_rank=self.moe_config.tp_rank, + ) + return True if return_success else None + # Case input scale: input_scale loading is only supported for fp8 if "input_scale" in weight_name: # this is needed for compressed-tensors only diff --git a/vllm/model_executor/layers/quantization/auto_awq.py b/vllm/model_executor/layers/quantization/auto_awq.py index 8b36313dd22a..a77fef05254e 100644 --- a/vllm/model_executor/layers/quantization/auto_awq.py +++ b/vllm/model_executor/layers/quantization/auto_awq.py @@ -560,6 +560,7 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, kInt4Static, + quant_config=self.quant_config, may_have_zp=self.quant_config.zero_point, may_have_bias=True, ) diff --git a/vllm/model_executor/layers/quantization/auto_gptq.py b/vllm/model_executor/layers/quantization/auto_gptq.py index c0a97e57ab99..f94bbb25b3ca 100644 --- a/vllm/model_executor/layers/quantization/auto_gptq.py +++ b/vllm/model_executor/layers/quantization/auto_gptq.py @@ -489,6 +489,7 @@ def __init__( self.wna16_moe_backend, self.experts_cls = select_wna16_moe_backend( moe, weight_key, + quant_config=self.quant_config, may_have_zp=True, may_have_bias=True, ) @@ -646,6 +647,26 @@ def create_weights( layer.register_parameter("w2_g_idx_sort_indices", w2_g_idx_sort_indices) set_weight_attrs(w2_g_idx_sort_indices, extra_weight_attrs) + # Some GPTQ checkpoints contain expert biases even when the model + # architecture does not declare them. Zero initialization keeps + # checkpoints without biases equivalent to the bias-free path. + w13_bias = torch.nn.Parameter( + torch.zeros( + num_experts, + 2 * intermediate_size_per_partition, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_bias", w13_bias) + set_weight_attrs(w13_bias, extra_weight_attrs) + w2_bias = torch.nn.Parameter( + torch.zeros(num_experts, hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w2_bias", w2_bias) + set_weight_attrs(w2_bias, extra_weight_attrs) + if self.experts_cls is not None and issubclass( self.experts_cls, FusedMoEExpertsModular ): @@ -670,6 +691,15 @@ def replace_or_register(name: str, val: torch.Tensor | None): "W8A8-INT8 is not supported by marlin kernel." ) + w13_bias = getattr(layer, "w13_bias", None) + if "w13_bias" not in layer._loaded_expert_biases: + layer.register_parameter("w13_bias", None) + w13_bias = None + w2_bias = getattr(layer, "w2_bias", None) + if "w2_bias" not in layer._loaded_expert_biases: + layer.register_parameter("w2_bias", None) + w2_bias = None + converted = convert_to_wna16_moe_kernel_format( backend=self.wna16_moe_backend, layer=layer, @@ -681,8 +711,8 @@ def replace_or_register(name: str, val: torch.Tensor | None): w2_scale=layer.w2_scales, w13_g_idx=layer.w13_g_idx, w2_g_idx=layer.w2_g_idx, - w13_bias=getattr(layer, "w13_bias", None), - w2_bias=getattr(layer, "w2_bias", None), + w13_bias=w13_bias, + w2_bias=w2_bias, w13_qzeros=getattr(layer, "w13_qzeros", None), w2_qzeros=getattr(layer, "w2_qzeros", 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 4643d45c1b19..516d94fae6e9 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 @@ -92,7 +92,8 @@ def __init__( self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, weight_key=weight_key, - may_have_zp=False, + quant_config=self.weight_quant, + may_have_zp=not self.symmetric, may_have_bias=False, ) self.is_marlin = self.wna16_backend in [ diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index 712e766398be..76f2d740bc73 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -235,9 +235,9 @@ def __init__(self, quant_config: MoeWNA16Config, moe: "FusedMoEConfig") -> None: self.wna16_backend, self.experts_cls = select_wna16_moe_backend( config=self.moe, weight_key=weight_key, + quant_config=self.quant_config, may_have_zp=self.quant_config.has_zp, may_have_bias=False, - allow_marlin=False, ) def create_weights( @@ -363,6 +363,13 @@ def create_weights( def get_fused_moe_quant_config( self, layer: RoutedExperts ) -> FusedMoEQuantConfig | None: + if self.wna16_backend == WNA16MoEBackend.HUMMING: + from vllm.model_executor.layers.quantization.utils.humming_utils import ( + get_humming_moe_quant_config, + ) + + return get_humming_moe_quant_config(layer) + has_zp = self.quant_config.has_zp return make_wna16_moe_quant_config( w1_scale=layer.w13_scales, @@ -381,6 +388,8 @@ def _setup_kernel(self, layer: RoutedExperts): moe_quant_config=self.moe_quant_config, moe_config=self.moe, experts_cls=self.experts_cls, + backend=self.wna16_backend, + layer=layer, routing_tables=layer._expert_routing_tables(), ) @@ -428,6 +437,8 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: replace_parameter(layer, "w2_qweight", w2_qweight) replace_parameter(layer, "w13_scales", w13_scales) replace_parameter(layer, "w2_scales", w2_scales) + layer.w13_weight = layer.w13_qweight + layer.w2_weight = layer.w2_qweight if has_zp: assert w13_qzeros is not None and w2_qzeros is not None @@ -463,8 +474,8 @@ def apply( assert self.moe_kernel is not None return self.moe_kernel.apply( x, - layer.w13_qweight, - layer.w2_qweight, + layer.w13_weight, + layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, activation=layer.activation, @@ -486,8 +497,8 @@ def apply_monolithic( assert self.moe_kernel is not None return self.moe_kernel.apply_monolithic( x, - layer.w13_qweight, - layer.w2_qweight, + layer.w13_weight, + layer.w2_weight, router_logits, activation=layer.activation, global_num_experts=layer.global_num_experts,