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 10, 11:30 UTC
[Compressed-Tensors] Simplify NVFP4 Conditions, enable marlin support for NVFP4A16 MoEs (#30881)
Dipika Sikka committed
Jan 8, 2026 at 17:45 UTC
5d3b6097ad6f8ddb156b92525cce19e08e5c7b89
2 files changed
+42
-42
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+11
-38
@@ -331,25 +331,17 @@ class CompressedTensorsConfig(QuantizationConfig):
331
return False
332
333
@staticmethod
334
- def _is_fp4a4_nvfp4(weight_quant: QuantizationArgs, input_quant: QuantizationArgs):
335
- if weight_quant is None or input_quant is None:
334
+ def _is_nvfp4_format(quant_args: QuantizationArgs):
335
+ if quant_args is None:
336
return False
337
-
337
is_tensor_group_quant = (
339
- weight_quant.strategy == QuantizationStrategy.TENSOR_GROUP.value
340
- and input_quant.strategy == QuantizationStrategy.TENSOR_GROUP.value
338
+ quant_args.strategy == QuantizationStrategy.TENSOR_GROUP.value
339
)
342
- is_symmetric = weight_quant.symmetric and input_quant.symmetric
343
-
344
- is_group_size_16 = (
345
- weight_quant.group_size == 16 and input_quant.group_size == 16
346
- )
347
- is_float_type = (
348
- weight_quant.type == QuantizationType.FLOAT
349
- and input_quant.type == QuantizationType.FLOAT
350
- )
351
- is_4_bits = weight_quant.num_bits == 4 and input_quant.num_bits == 4
340
+ is_symmetric = quant_args.symmetric
341
342
+ is_group_size_16 = quant_args.group_size == 16
343
+ is_float_type = quant_args.type == QuantizationType.FLOAT
344
+ is_4_bits = quant_args.num_bits == 4
345
return (
346
is_tensor_group_quant
347
and is_float_type
@@ -358,27 +350,6 @@ class CompressedTensorsConfig(QuantizationConfig):
350
and is_symmetric
351
)
352
361
- @staticmethod
362
- def _is_fp4a16_nvfp4(weight_quant: QuantizationArgs, input_quant: QuantizationArgs):
363
- is_weight_only = weight_quant is not None and input_quant is None
364
- is_tensor_group_quant = (
365
- weight_quant.strategy == QuantizationStrategy.TENSOR_GROUP.value
366
- )
367
- is_symmetric = weight_quant.symmetric
368
-
369
- is_group_size_16 = weight_quant.group_size == 16
370
- is_float_type = weight_quant.type == QuantizationType.FLOAT
371
- is_4_bits = weight_quant.num_bits == 4
372
-
373
- return (
374
- is_weight_only
375
- and is_tensor_group_quant
376
- and is_float_type
377
- and is_4_bits
378
- and is_group_size_16
379
- and is_symmetric
380
- )
381
-
353
@staticmethod
354
def _is_static_tensor_w8a8(
355
weight_quant: QuantizationArgs, input_quant: QuantizationArgs
@@ -576,7 +547,7 @@ class CompressedTensorsConfig(QuantizationConfig):
547
format = format if format is not None else self.quant_format
548
549
# Detect If Mixed Precision
579
- if self._is_fp4a16_nvfp4(weight_quant, input_quant):
550
+ if self._is_nvfp4_format(weight_quant) and input_quant is None:
551
return CompressedTensorsW4A16Fp4()
552
553
if self._is_fp8_w4a8_sm90(weight_quant, input_quant):
@@ -614,7 +585,9 @@ class CompressedTensorsConfig(QuantizationConfig):
585
586
act_quant_format = is_activation_quantization_format(format)
587
if act_quant_format:
617
- if self._is_fp4a4_nvfp4(weight_quant, input_quant):
588
+ if self._is_nvfp4_format(weight_quant) and self._is_nvfp4_format(
589
+ input_quant
590
+ ):
591
if cutlass_fp4_supported() or envs.VLLM_USE_NVFP4_CT_EMULATIONS:
592
return CompressedTensorsW4A4Fp4()
593
else:
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+31
-4
@@ -77,6 +77,9 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (
77
marlin_make_workspace_new,
78
marlin_moe_permute_scales,
79
)
80
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
81
+ is_fp4_marlin_supported,
82
+)
83
from vllm.model_executor.layers.quantization.utils.quant_utils import (
84
convert_bf16_scales_to_fp8,
85
convert_packed_uint4b8_to_signed_int4_inplace,
@@ -182,8 +185,18 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
185
return CompressedTensorsWNA16MarlinMoEMethod(
186
weight_quant, input_quant, layer.moe_config
187
)
185
- elif quant_config._is_fp4a4_nvfp4(weight_quant, input_quant):
186
- return CompressedTensorsW4A4Nvfp4MoEMethod(layer.moe_config, layer_name)
188
+ elif quant_config._is_nvfp4_format(weight_quant):
189
+ _is_valid_nvfp4_activations = (
190
+ quant_config._is_nvfp4_format(input_quant) or input_quant is None
191
+ )
192
+ if not _is_valid_nvfp4_activations:
193
+ raise ValueError(
194
+ "For NVFP4 weights, input quantization must also be NVFP4 format ",
195
+ f"or None for NVFP4A16, found {input_quant}",
196
+ )
197
+ return CompressedTensorsW4A4Nvfp4MoEMethod(
198
+ layer.moe_config, layer_name, use_marlin=input_quant is None
199
+ )
200
elif (
201
quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
202
or quant_config._is_fp8_w8a8_sm100(weight_quant, input_quant)
@@ -212,7 +225,12 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
225
226
227
class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
215
- def __init__(self, moe: FusedMoEConfig, layer_name: str | None = None):
228
+ def __init__(
229
+ self,
230
+ moe: FusedMoEConfig,
231
+ layer_name: str | None = None,
232
+ use_marlin: bool = False,
233
+ ):
234
if not moe.is_act_and_mul:
235
raise ValueError(
236
"CompressedTensorsW4A4Nvfp4MoEMethod does not yet "
@@ -221,7 +239,16 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):
239
240
super().__init__(moe)
241
self.group_size = 16
224
- self.nvfp4_backend = select_nvfp4_moe_backend()
242
+ if use_marlin:
243
+ if is_fp4_marlin_supported():
244
+ self.nvfp4_backend = NvFp4MoeBackend.MARLIN
245
+ else:
246
+ raise ValueError(
247
+ "Marlin FP4 MoE kernel requested but not ",
248
+ "supported on current platform.",
249
+ )
250
+ else:
251
+ self.nvfp4_backend = select_nvfp4_moe_backend()
252
self.use_global_sf = is_global_sf_supported_for_nvfp4_backend(
253
self.nvfp4_backend
254
)