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 1d0cf91d42..64c68018f3 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,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, kInt4Static, kInt4Static32, + kInt4Static32Asym, + kInt4StaticAsym, kInt8Static, kMxfp4Static, kMxfp8Static, @@ -611,6 +613,8 @@ class MarlinExpertsBase(mk.FusedMoEExpertsModular): kInt4Static, kInt8Static, kInt4Static32, + kInt4StaticAsym, + kInt4Static32Asym, ] 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 9de7a6ba11..6ad60d62e9 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -32,6 +32,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import ( marlin_moe_permute_scales, marlin_permute_bias, moe_awq_to_marlin_zero_points, + moe_packed_to_marlin_zero_points, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -496,6 +497,23 @@ def _process_weights_marlin( marlin_w2_scales ) + # --- Permute zero points --- + if w13_qzeros is not None and w2_qzeros is not None: + w13_qzeros = moe_packed_to_marlin_zero_points( + w13_qzeros, + size_k=w13_qzeros.shape[1], + size_n=w13_qzeros.shape[2] * pack_factor, + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + w2_qzeros = moe_packed_to_marlin_zero_points( + w2_qzeros, + size_k=w2_qzeros.shape[1], + size_n=w2_qzeros.shape[2] * pack_factor, + num_bits=num_bits, + is_a_8bit=is_a_8bit, + ) + # --- Permute bias --- if w13_bias is not None: w13_bias_out = marlin_permute_bias(w13_bias) 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 2d629d73ed..2a98d444af 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 @@ -29,6 +29,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso ) from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa WNA16_SUPPORTED_TYPES_MAP, + WNA16_ZP_SUPPORTED_TYPES_MAP, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( get_marlin_input_dtype, @@ -56,9 +57,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): super().__init__(moe) self.weight_quant = weight_quant self.input_quant = input_quant - assert weight_quant.symmetric, ( - "Only symmetric quantization is supported for MoE" - ) + self.symmetric = weight_quant.symmetric # Extract properties from weight_quant self.num_bits = weight_quant.num_bits self.packed_factor = 32 // weight_quant.num_bits @@ -66,7 +65,12 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): self.group_size = weight_quant.group_size self.actorder = weight_quant.actorder - self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits] + self.quant_type = ( + WNA16_SUPPORTED_TYPES_MAP[self.num_bits] + if self.symmetric + else WNA16_ZP_SUPPORTED_TYPES_MAP[self.num_bits] + ) + self.marlin_input_dtype = get_marlin_input_dtype(layer_name) if self.num_bits == 4: @@ -82,7 +86,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): "CompressedTensorsWNA16MarlinMoEMethod only supports int4 and int8 now." ) - weight_key = QuantKey(self.quant_type, scale) + weight_key = QuantKey(self.quant_type, scale, symmetric=self.symmetric) # Select WNA16 MoE backend via oracle. self.wna16_backend, self.experts_cls = select_wna16_moe_backend( @@ -103,15 +107,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): Get the shape of the weight based on the weight name, number of experts hidden size, intermediate size per partition, number of groups for w2, and number of groups for w13. Pass in num_groups_w2 and num_groups_w13 - for weight scales. + for weight scales/zero_points. """ - if weight_name == "w13_scale": + if weight_name in ("w13_scale", "w13_zp"): assert num_groups_w13 is not None, ( - "num_groups_w13 must be provided for weight scales" + "num_groups_w13 must be provided for weight scales/zero_points" ) - if weight_name == "w2_scale": + if weight_name in ("w2_scale", "w2_zp"): assert num_groups_w2 is not None, ( - "num_groups_w2 must be provided for weight scales" + "num_groups_w2 must be provided for weight scales/zero_points" ) w13_num_shards = 2 if self.moe.is_act_and_mul else 1 is_flashinfer = self.wna16_backend == WNA16MoEBackend.FLASHINFER_TRTLLM @@ -140,6 +144,15 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w13_num_shards * intermediate_size_per_partition, ), }, + "w13_zp": { + "Marlin": ( + num_experts, + num_groups_w13, + w13_num_shards + * intermediate_size_per_partition + // self.packed_factor, + ), + }, "w2_weight": { "Flashinfer": ( num_experts, @@ -156,6 +169,13 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): "Flashinfer": (num_experts, hidden_size, num_groups_w2), "Marlin": (num_experts, num_groups_w2, hidden_size), }, + "w2_zp": { + "Marlin": ( + num_experts, + num_groups_w2, + hidden_size // self.packed_factor, + ), + }, } backend_key = "Flashinfer" if is_flashinfer else "Marlin" return shape_map[weight_name][backend_key] @@ -263,6 +283,39 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): set_weight_attrs(w2_scale, extra_weight_attrs) set_weight_attrs(w2_scale, {"load_full_w2": load_full_w2}) + if not self.symmetric: + w13_zp = torch.nn.Parameter( + torch.zeros( + *self.get_weight_shape( + "w13_zp", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w13=num_groups_w13, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_zero_point", w13_zp) + set_weight_attrs(w13_zp, extra_weight_attrs) + + w2_zp = torch.nn.Parameter( + torch.zeros( + *self.get_weight_shape( + "w2_zp", + num_experts, + hidden_size, + intermediate_size_per_partition, + num_groups_w2=num_groups_w2, + ), + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_zero_point", w2_zp) + set_weight_attrs(w2_zp, extra_weight_attrs) + w2_weight_shape = torch.nn.Parameter( torch.empty(num_experts, 2), requires_grad=False ) @@ -334,8 +387,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_g_idx_processed, 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 @@ -351,6 +404,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_scale=layer.w2_weight_scale, w13_g_idx=layer.w13_weight_g_idx, w2_g_idx=layer.w2_weight_g_idx, + w13_qzeros=getattr(layer, "w13_weight_zero_point", None), + w2_qzeros=getattr(layer, "w2_weight_zero_point", None), ) # Replace common parameters @@ -359,6 +414,10 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): replace_parameter(layer, "w13_weight_scale", w13_scales) replace_parameter(layer, "w2_weight_scale", w2_scales) + if not self.symmetric: + replace_parameter(layer, "w13_weight_zero_point", w13_qzeros) + replace_parameter(layer, "w2_weight_zero_point", w2_qzeros) + # Marlin-specific parameters (not needed for Flashinfer) if not is_flashinfer: replace_parameter(layer, "w13_weight_g_idx", w13_g_idx_processed) @@ -417,6 +476,8 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): w2_scale=layer.w2_weight_scale, group_size=self.group_size, num_bits=self.num_bits, + w1_zp=getattr(layer, "w13_weight_zero_point", None), + w2_zp=getattr(layer, "w2_weight_zero_point", None), ) def apply_monolithic( diff --git a/vllm/model_executor/layers/quantization/utils/marlin_utils.py b/vllm/model_executor/layers/quantization/utils/marlin_utils.py index eca04eed74..19f2605dc4 100644 --- a/vllm/model_executor/layers/quantization/utils/marlin_utils.py +++ b/vllm/model_executor/layers/quantization/utils/marlin_utils.py @@ -429,6 +429,30 @@ def maybe_warn_marlin_atomic_add(device, dtype): ) +def moe_packed_to_marlin_zero_points( + q_zp_packed: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, + is_a_8bit: 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. + """ + 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, + ) + for e in range(num_experts): + 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 + + def maybe_warn_marlin_atomic_add_env(): if torch.compiler.is_dynamo_compiling(): return diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index 947e428ca8..ba1016a4fb 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -181,6 +181,13 @@ kInt8Static = QuantKey(INT8_DTYPE, scale=kInt8StaticGroupScale, symmetric=True) kInt4Static32GroupScale = ScaleDesc(torch.float16, True, GroupShape(1, 32)) kInt4Static32 = QuantKey(INT4_DTYPE, scale=kInt4Static32GroupScale, symmetric=True) +kInt4StaticAsym = QuantKey( + scalar_types.uint4, scale=kInt4StaticGroupScale, symmetric=False +) +kInt4Static32Asym = QuantKey( + scalar_types.uint4, scale=kInt4Static32GroupScale, symmetric=False +) + kInt8StaticChannelSym = QuantKey(torch.int8, kStaticChannelScale, symmetric=True) kInt8DynamicTokenSym = QuantKey(torch.int8, kDynamicTokenScale, symmetric=True)