A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 12, 03:30 UTC
[refactor] CTMoEMethods to use QuantizationArgs (#28871)
Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
HDCharles committed
Dec 3, 2025 at 06:00 UTC
b294e28db2c5dee61bc25157664edcada8b90b31
2 files changed
+86
-75
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+4
-2
@@ -767,8 +767,10 @@ class CompressedTensorsConfig(QuantizationConfig):
767
targets=self.target_scheme_map.keys(),
768
fused_mapping=self.packed_modules_mapping,
769
)
770
-
771
- return self.target_scheme_map[matched_target]
770
+ scheme_dict = self.target_scheme_map[matched_target]
771
+ if scheme_dict.get("format") is None:
772
+ scheme_dict["format"] = self.quant_format
773
+ return scheme_dict
774
775
return None
776
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+82
-73
@@ -7,7 +7,11 @@ from enum import Enum
7
8
import torch
9
from compressed_tensors import CompressionFormat
10
-from compressed_tensors.quantization import ActivationOrdering, QuantizationStrategy
10
+from compressed_tensors.quantization import (
11
+ ActivationOrdering,
12
+ QuantizationArgs,
13
+ QuantizationStrategy,
14
+)
15
from torch.nn.parameter import Parameter
16
17
import vllm.envs as envs
@@ -142,10 +146,26 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
146
# are supported + check if the layer is being ignored.
147
weight_quant = scheme_dict.get("weights")
148
input_quant = scheme_dict.get("input_activations")
149
+ format = scheme_dict.get("format")
150
151
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
152
# group_size=None means channelwise
153
group_size = weight_quant.group_size or -1
154
+
155
+ valid_format_and_bits = (
156
+ weight_quant.num_bits in WNA16_SUPPORTED_BITS
157
+ and format == CompressionFormat.pack_quantized.value
158
+ )
159
+
160
+ if not valid_format_and_bits:
161
+ raise ValueError(
162
+ "For Fused MoE layers, only format: ",
163
+ f"{CompressionFormat.pack_quantized.value} ",
164
+ f" and bits: {WNA16_SUPPORTED_BITS} is supported ",
165
+ f"but got format: {CompressionFormat.pack_quantized.value} "
166
+ f" and bits: {weight_quant.num_bits}",
167
+ )
168
+
169
# Prefer to use the MarlinMoE kernel when it is supported.
170
if (
171
not check_moe_marlin_supports_layer(layer, group_size)
@@ -161,12 +181,12 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
181
)
182
logger.info_once("Using CompressedTensorsWNA16MoEMethod")
183
return CompressedTensorsWNA16MoEMethod(
164
- quant_config, layer.moe_config, layer_name
184
+ weight_quant, input_quant, layer.moe_config
185
)
186
else:
187
logger.info_once("Using CompressedTensorsWNA16MarlinMoEMethod")
188
return CompressedTensorsWNA16MarlinMoEMethod(
169
- quant_config, layer.moe_config, layer_name
189
+ weight_quant, input_quant, layer.moe_config
190
)
191
elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
192
return CompressedTensorsW4A4Nvfp4MoEMethod(layer.moe_config, layer_name)
@@ -176,15 +196,15 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
196
or quant_config._is_fp8_w8a8(weight_quant, input_quant)
197
):
198
return CompressedTensorsW8A8Fp8MoEMethod(
179
- quant_config, layer.moe_config, layer_name
199
+ weight_quant, input_quant, layer.moe_config
200
)
201
elif quant_config._is_dynamic_token_w8a8(weight_quant, input_quant):
202
return CompressedTensorsW8A8Int8MoEMethod(
183
- quant_config, layer.moe_config, layer_name
203
+ weight_quant, input_quant, layer.moe_config
204
)
205
elif quant_config._is_dynamic_token_w4a8_int(weight_quant, input_quant):
206
return CompressedTensorsW4A8Int8MoEMethod(
187
- quant_config, layer.moe_config, layer_name
207
+ weight_quant, input_quant, layer.moe_config
208
)
209
else:
210
raise RuntimeError(
@@ -650,17 +670,19 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
670
class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
671
def __init__(
672
self,
653
- quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
673
+ weight_quant: QuantizationArgs,
674
+ input_quant: QuantizationArgs,
675
moe: FusedMoEConfig,
676
layer_name: str | None = None,
677
):
657
- super().__init__(moe)
658
- self.quant_config = quant_config
659
- self.weight_quant = self.quant_config.target_scheme_map["Linear"].get("weights")
660
- self.input_quant = self.quant_config.target_scheme_map["Linear"].get(
661
- "input_activations"
678
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
679
+ CompressedTensorsConfig,
680
)
681
682
+ super().__init__(moe)
683
+ self.weight_quant = weight_quant
684
+ self.input_quant = input_quant
685
+
686
per_tensor = (
687
self.weight_quant.strategy == QuantizationStrategy.TENSOR
688
and self.input_quant.strategy == QuantizationStrategy.TENSOR
@@ -698,11 +720,13 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
720
self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()
721
722
# cutlass path
701
- self.is_fp8_w8a8_sm100 = quant_config._is_fp8_w8a8_sm100(
723
+ self.is_fp8_w8a8_sm100 = CompressedTensorsConfig._is_fp8_w8a8_sm100(
724
self.weight_quant, self.input_quant
725
)
726
self.use_cutlass = not self.block_quant and (
705
- quant_config._is_fp8_w8a8_sm90(self.weight_quant, self.input_quant)
727
+ CompressedTensorsConfig._is_fp8_w8a8_sm90(
728
+ self.weight_quant, self.input_quant
729
+ )
730
or self.is_fp8_w8a8_sm100
731
)
732
self.disable_expert_map = False
@@ -1261,16 +1285,14 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
1285
class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1286
def __init__(
1287
self,
1264
- quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1288
+ weight_quant: QuantizationArgs,
1289
+ input_quant: QuantizationArgs,
1290
moe: FusedMoEConfig,
1291
layer_name: str | None = None,
1292
):
1293
super().__init__(moe)
1269
- self.quant_config = quant_config
1270
- self.weight_quant = self.quant_config.target_scheme_map["Linear"].get("weights")
1271
- self.input_quant = self.quant_config.target_scheme_map["Linear"].get(
1272
- "input_activations"
1273
- )
1294
+ self.weight_quant = weight_quant
1295
+ self.input_quant = input_quant
1296
1297
per_channel = (
1298
self.weight_quant.strategy == QuantizationStrategy.CHANNEL
@@ -1414,36 +1436,27 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):
1436
class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1437
def __init__(
1438
self,
1417
- quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1439
+ weight_quant: QuantizationArgs,
1440
+ input_quant: QuantizationArgs | None,
1441
moe: FusedMoEConfig,
1442
layer_name: str | None = None,
1443
):
1444
super().__init__(moe)
1422
- self.quant_config = quant_config
1423
- # TODO: @dsikka: refactor this to use schemes as other kernels
1424
- # are supported + check if the layer is being ignored.
1425
- config = self.quant_config.target_scheme_map["Linear"].get("weights")
1426
- self.num_bits = config.num_bits
1427
- self.packed_factor = 32 // config.num_bits
1428
- self.strategy = config.strategy
1429
- self.group_size = config.group_size
1430
- self.actorder = config.actorder
1431
- self.layer_name = layer_name
1432
- self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
1433
- assert config.symmetric, "Only symmetric quantization is supported for MoE"
1445
+ self.weight_quant = weight_quant
1446
+ self.input_quant = input_quant
1447
+ assert weight_quant.symmetric, (
1448
+ "Only symmetric quantization is supported for MoE"
1449
+ )
1450
+ # Extract properties from weight_quant
1451
+ self.num_bits = weight_quant.num_bits
1452
+ self.packed_factor = 32 // weight_quant.num_bits
1453
+ self.strategy = weight_quant.strategy
1454
+ self.group_size = weight_quant.group_size
1455
+ self.actorder = weight_quant.actorder
1456
1435
- if not (
1436
- self.quant_config.quant_format == CompressionFormat.pack_quantized.value
1437
- and self.num_bits in WNA16_SUPPORTED_BITS
1438
- ):
1439
- raise ValueError(
1440
- "For Fused MoE layers, only ",
1441
- f"{CompressionFormat.pack_quantized.value} ",
1442
- "is supported for the following bits: ",
1443
- f"{WNA16_SUPPORTED_BITS}",
1444
- )
1457
self.quant_type = WNA16_SUPPORTED_TYPES_MAP[self.num_bits]
1458
self.use_marlin = True
1459
+ self.marlin_input_dtype = get_marlin_input_dtype(layer_name)
1460
1461
def create_weights(
1462
self,
@@ -1812,35 +1825,26 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):
1825
class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
1826
def __init__(
1827
self,
1815
- quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
1828
+ weight_quant: QuantizationArgs,
1829
+ input_quant: QuantizationArgs | None,
1830
moe: FusedMoEConfig,
1831
layer_name: str | None = None,
1832
):
1833
super().__init__(moe)
1820
- self.quant_config = quant_config
1821
- # TODO: @dsikka: refactor this to use schemes as other kernels
1822
- # are supported + check if the layer is being ignored.
1823
- config = self.quant_config.target_scheme_map["Linear"].get("weights")
1824
- self.num_bits = config.num_bits
1825
- self.packed_factor = 32 // config.num_bits
1826
- self.strategy = config.strategy
1834
+ self.weight_quant = weight_quant
1835
+ self.input_quant = input_quant
1836
+ # Extract properties from weight_quant
1837
+ self.num_bits = weight_quant.num_bits
1838
+ self.packed_factor = 32 // weight_quant.num_bits
1839
+ self.strategy = weight_quant.strategy
1840
# channelwise is not supported by this kernel
1828
- assert config.strategy == "group"
1829
- self.group_size = config.group_size
1841
+ assert weight_quant.strategy == "group"
1842
+ self.group_size = weight_quant.group_size
1843
# grouped actorder isn't supported by this kernel
1831
- assert config.actorder != "group"
1832
- assert config.symmetric, "Only symmetric quantization is supported for MoE"
1833
-
1834
- if not (
1835
- self.quant_config.quant_format == CompressionFormat.pack_quantized.value
1836
- and self.num_bits in WNA16_SUPPORTED_BITS
1837
- ):
1838
- raise ValueError(
1839
- "For Fused MoE layers, only ",
1840
- f"{CompressionFormat.pack_quantized.value} ",
1841
- "is supported for the following bits: ",
1842
- f"{WNA16_SUPPORTED_BITS}",
1843
- )
1844
+ assert weight_quant.actorder != "group"
1845
+ assert weight_quant.symmetric, (
1846
+ "Only symmetric quantization is supported for MoE"
1847
+ )
1848
1849
def create_weights(
1850
self,
@@ -2065,28 +2069,33 @@ class CompressedTensorsW4A8Int8MoEMethod(CompressedTensorsMoEMethod):
2069
2070
def __init__(
2071
self,
2068
- quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
2072
+ weight_quant: QuantizationArgs,
2073
+ input_quant: QuantizationArgs,
2074
moe: FusedMoEConfig,
2075
layer_name: str | None = None,
2076
):
2077
super().__init__(moe)
2078
self.has_bias = self.moe.has_bias
2074
- self.quant_config = quant_config
2079
+ self.weight_quant = weight_quant
2080
+ self.input_quant = input_quant
2081
2082
# Validate scheme: weights=W4 (channel or group),
2083
# activations=dynamic TOKEN (A8)
2078
- wq = self.quant_config.target_scheme_map["Linear"].get("weights")
2079
- aq = self.quant_config.target_scheme_map["Linear"].get("input_activations")
2084
2085
# Must be dynamic per-token activations
2082
- if aq.strategy != QuantizationStrategy.TOKEN or not aq.dynamic:
2086
+ if (
2087
+ input_quant.strategy != QuantizationStrategy.TOKEN
2088
+ or not input_quant.dynamic
2089
+ ):
2090
raise ValueError(
2091
"W4A8-int MoE needs dynamic per-token activation quantization."
2092
)
2093
2094
# Weight can be channel-wise (group_size=None) or group-wise
2088
- self.group_size = wq.group_size if (wq.group_size is not None) else -1
2089
- if wq.num_bits != 4:
2095
+ self.group_size = (
2096
+ weight_quant.group_size if (weight_quant.group_size is not None) else -1
2097
+ )
2098
+ if weight_quant.num_bits != 4:
2099
raise ValueError("This method only supports 4-bit weights (num_bits=4).")
2100
2101
# CPU only