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 1, 23:45 UTC
Refactor CT NVFP4 linear to use a single class (#42443)
Dipika Sikka committed
Jun 4, 2026 at 08:25 UTC
e68988a24807c9dfb2bf6936eb17425ce7812c5f
6 files changed
+55
-162
tests/quantization/test_compressed_tensors.py
+5
-17
@@ -26,7 +26,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
26
CompressedTensorsW4A4Fp4,
27
CompressedTensorsW4A4Mxfp4,
28
CompressedTensorsW4A8Fp8,
29
- CompressedTensorsW4A16Fp4,
29
CompressedTensorsW8A8Fp8,
30
CompressedTensorsW8A8Int8,
31
CompressedTensorsW8A8Mxfp8,
@@ -37,9 +36,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
36
find_matched_target,
37
)
38
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
40
-from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
41
- cutlass_fp4_supported,
42
-)
39
from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
40
from vllm.platforms import current_platform
41
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
@@ -376,13 +372,12 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
372
@pytest.mark.parametrize(
373
"args",
374
[
379
- # TODO: Enable once model is available again
380
- # ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16", CompressedTensorsW4A16Fp4),
381
- ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", CompressedTensorsW4A4Fp4),
375
+ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16", True),
376
+ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", False),
377
],
378
)
379
def test_compressed_tensors_nvfp4(vllm_runner, args):
385
- model, scheme = args
380
+ model, use_a16 = args
381
with vllm_runner(model, enforce_eager=True) as llm:
382
383
def check_model(model):
@@ -390,15 +385,8 @@ def test_compressed_tensors_nvfp4(vllm_runner, args):
385
386
qkv_proj = layer.self_attn.qkv_proj
387
assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
393
- if (
394
- isinstance(qkv_proj.scheme, scheme)
395
- or isinstance(qkv_proj.scheme, CompressedTensorsW4A16Fp4)
396
- and not cutlass_fp4_supported()
397
- ):
398
- assert True
399
- else:
400
- raise AssertionError("FP4 Scheme Mismatch")
401
-
388
+ assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4)
389
+ assert qkv_proj.scheme.use_a16 == use_a16
390
assert qkv_proj.scheme.group_size == 16
391
392
llm.apply_model(check_model)
vllm/model_executor/kernels/linear/__init__.py
+4
-2
@@ -842,7 +842,7 @@ _NVFP4_BACKEND_TO_KERNEL: dict[str, type[NvFp4LinearKernel]] = {
842
}
843
844
845
-def init_nvfp4_linear_kernel() -> NvFp4LinearKernel:
845
+def init_nvfp4_linear_kernel(use_a16: bool = False) -> NvFp4LinearKernel:
846
"""Select and instantiate the best NVFP4 linear kernel for the
847
current platform."""
848
config = NvFp4LinearLayerConfig()
@@ -885,7 +885,9 @@ def init_nvfp4_linear_kernel() -> NvFp4LinearKernel:
885
elif linear_backend == "auto":
886
# Deprecated env-var overrides — only honoured when --linear-backend
887
# is "auto". Deprecation warnings are emitted from vllm/envs.py.
888
- if envs.VLLM_USE_FBGEMM:
888
+ if use_a16: # force a16 if running weight-only quantization
889
+ force_kernel = MarlinNvFp4LinearKernel
890
+ elif envs.VLLM_USE_FBGEMM:
891
force_kernel = FbgemmNvFp4LinearKernel
892
elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:
893
force_kernel = EmulationNvFp4LinearKernel
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+10
-8
@@ -40,7 +40,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
40
CompressedTensorsW4A4Mxfp4,
41
CompressedTensorsW4A8Fp8,
42
CompressedTensorsW4A8Int,
43
- CompressedTensorsW4A16Fp4,
43
CompressedTensorsW8A8Fp8,
44
CompressedTensorsW8A8Int8,
45
CompressedTensorsW8A8Mxfp8,
@@ -616,8 +615,16 @@ class CompressedTensorsConfig(QuantizationConfig):
615
format = format if format is not None else self.quant_format
616
617
# Detect If Mixed Precision
619
- if self._is_nvfp4_format(weight_quant) and input_quant is None:
620
- return CompressedTensorsW4A16Fp4()
618
+ if self._is_nvfp4_format(weight_quant):
619
+ if input_quant is None:
620
+ return CompressedTensorsW4A4Fp4(use_a16=True)
621
+
622
+ if not self._is_nvfp4_format(input_quant):
623
+ raise ValueError(
624
+ "For NVFP4 weights, input quantization must also be NVFP4 format, ",
625
+ "None for NVFP4A16",
626
+ )
627
+ return CompressedTensorsW4A4Fp4()
628
629
if self._is_mxfp4(weight_quant):
630
return CompressedTensorsW4A4Mxfp4()
@@ -650,11 +657,6 @@ class CompressedTensorsConfig(QuantizationConfig):
657
658
act_quant_format = is_activation_quantization_format(format)
659
if act_quant_format:
653
- if self._is_nvfp4_format(weight_quant) and self._is_nvfp4_format(
654
- input_quant
655
- ):
656
- return CompressedTensorsW4A4Fp4()
657
-
660
if self._is_fp8_w8a8(weight_quant, input_quant):
661
is_fp8_w8a8_supported = self._check_scheme_supported(
662
CompressedTensorsW8A8Fp8.get_min_capability(), error=False
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
-2
@@ -6,7 +6,6 @@ from .compressed_tensors_w4a4_mxfp4 import CompressedTensorsW4A4Mxfp4
6
from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
7
from .compressed_tensors_w4a8_fp8 import CompressedTensorsW4A8Fp8
8
from .compressed_tensors_w4a8_int import CompressedTensorsW4A8Int
9
-from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
9
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
10
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
11
from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8
@@ -20,7 +19,6 @@ __all__ = [
19
"CompressedTensorsW8A8Int8",
20
"CompressedTensorsW8A8Fp8",
21
"WNA16_SUPPORTED_BITS",
23
- "CompressedTensorsW4A16Fp4",
22
"CompressedTensorsW4A4Mxfp4",
23
"CompressedTensorsW4A4Fp4",
24
"CompressedTensorsW4A8Int",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a16_nvfp4.py
deleted
-109
@@ -1,109 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from collections.abc import Callable
4
-
5
-import torch
6
-from torch.nn.parameter import Parameter
7
-
8
-from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
9
- CompressedTensorsScheme,
10
-)
11
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
12
- apply_fp4_marlin_linear,
13
- prepare_fp4_layer_for_marlin,
14
-)
15
-from vllm.model_executor.parameter import (
16
- GroupQuantScaleParameter,
17
- ModelWeightParameter,
18
- PerTensorScaleParameter,
19
-)
20
-
21
-__all__ = ["CompressedTensorsW4A16Fp4"]
22
-
23
-
24
-class CompressedTensorsW4A16Fp4(CompressedTensorsScheme):
25
- def __init__(self):
26
- self.group_size = 16
27
-
28
- @classmethod
29
- def get_min_capability(cls) -> int:
30
- # don't restrict as emulations
31
- return 75
32
-
33
- def create_weights(
34
- self,
35
- layer: torch.nn.Module,
36
- output_partition_sizes: list[int],
37
- input_size_per_partition: int,
38
- params_dtype: torch.dtype,
39
- weight_loader: Callable,
40
- **kwargs,
41
- ):
42
- output_size_per_partition = sum(output_partition_sizes)
43
- layer.logical_widths = output_partition_sizes
44
- layer.input_size_per_partition = input_size_per_partition
45
- layer.output_size_per_partition = output_size_per_partition
46
-
47
- # Weight
48
- weight = ModelWeightParameter(
49
- data=torch.empty(
50
- sum(output_partition_sizes),
51
- input_size_per_partition // 2,
52
- dtype=torch.uint8,
53
- ),
54
- input_dim=1,
55
- output_dim=0,
56
- weight_loader=weight_loader,
57
- )
58
- layer.register_parameter("weight_packed", weight)
59
-
60
- # Global Weight Scale
61
- weight_global_scale = PerTensorScaleParameter(
62
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
63
- weight_loader=weight_loader,
64
- )
65
- layer.register_parameter("weight_global_scale", weight_global_scale)
66
-
67
- # Per Group Weight Scale
68
- weight_scale = GroupQuantScaleParameter(
69
- data=torch.empty(
70
- sum(output_partition_sizes),
71
- input_size_per_partition // self.group_size,
72
- dtype=torch.float8_e4m3fn,
73
- ),
74
- input_dim=1,
75
- output_dim=0,
76
- weight_loader=weight_loader,
77
- )
78
-
79
- layer.register_parameter("weight_scale", weight_scale)
80
-
81
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
82
- # Process parameters for marlin repacking
83
-
84
- # Rename weight_packed to weight that marlin expects
85
- layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
86
- del layer.weight_packed
87
- # ct stores the inverse of what is expected by the marlin kernel
88
- layer.weight_global_scale = Parameter(
89
- 1.0 / layer.weight_global_scale.max().to(torch.float32), requires_grad=False
90
- )
91
-
92
- prepare_fp4_layer_for_marlin(layer)
93
-
94
- def apply_weights(
95
- self,
96
- layer: torch.nn.Module,
97
- x: torch.Tensor,
98
- bias: torch.Tensor | None = None,
99
- ) -> torch.Tensor:
100
- return apply_fp4_marlin_linear(
101
- input=x,
102
- weight=layer.weight,
103
- weight_scale=layer.weight_scale,
104
- weight_global_scale=layer.weight_global_scale,
105
- workspace=layer.workspace,
106
- size_n=layer.output_size_per_partition,
107
- size_k=layer.input_size_per_partition,
108
- bias=bias,
109
- )
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+36
-24
@@ -23,8 +23,9 @@ __all__ = ["CompressedTensorsW4A4Fp4"]
23
24
25
class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
26
- def __init__(self):
27
- self.kernel = init_nvfp4_linear_kernel()
26
+ def __init__(self, use_a16: bool = False):
27
+ self.use_a16 = use_a16
28
+ self.kernel = init_nvfp4_linear_kernel(use_a16=use_a16)
29
self.group_size = 16
30
31
@classmethod
@@ -79,46 +80,57 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
80
81
layer.register_parameter("weight_scale", weight_scale)
82
82
- input_global_scale = PerTensorScaleParameter(
83
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
84
- weight_loader=weight_loader,
85
- )
86
- layer.register_parameter("input_global_scale", input_global_scale)
83
+ if not self.use_a16:
84
+ input_global_scale = PerTensorScaleParameter(
85
+ data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
86
+ weight_loader=weight_loader,
87
+ )
88
+ layer.register_parameter("input_global_scale", input_global_scale)
89
90
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
91
# Rename CT checkpoint names to standardized names
92
layer.weight = layer.weight_packed
93
del layer.weight_packed
94
93
- if (
94
- torch.unique(layer.input_global_scale).numel() != 1
95
- or torch.unique(layer.weight_global_scale).numel() != 1
96
- ):
95
+ # Check for mismatched weight global scales
96
+ if torch.unique(layer.weight_global_scale).numel() != 1:
97
logger.warning_once(
98
- "In NVFP4 linear, the global scale for input or weight are different"
98
+ "In NVFP4 linear, the weight global scale is different"
99
" for parallel layers (e.g. q_proj, k_proj, v_proj). This "
100
" will likely result in reduced accuracy. Please verify the model"
101
" accuracy. Consider using a checkpoint with a shared global NVFP4"
102
" scale for fused layers."
103
)
104
105
- # Process global scales (CT stores as divisors, i.e. 1/scale)
106
- input_global_scale_inv = layer.input_global_scale.max().to(torch.float32)
107
- layer.input_global_scale = Parameter(
108
- (1.0 / input_global_scale_inv).to(torch.float32), requires_grad=False
109
- )
105
+ # Process weight global scale (CT stores as divisors, i.e. 1/scale)
106
weight_global_scale = layer.weight_global_scale.max().to(torch.float32)
107
layer.weight_global_scale = Parameter(
108
1.0 / weight_global_scale, requires_grad=False
109
)
110
115
- # Pre-compute alpha and inverse for runtime quantization
116
- layer.input_global_scale_inv = Parameter(
117
- input_global_scale_inv, requires_grad=False
118
- )
119
- layer.alpha = Parameter(
120
- layer.input_global_scale * layer.weight_global_scale, requires_grad=False
121
- )
111
+ if not self.use_a16:
112
+ if torch.unique(layer.input_global_scale).numel() != 1:
113
+ logger.warning_once(
114
+ "In NVFP4 linear, the input global scale is different"
115
+ " for parallel layers (e.g. q_proj, k_proj, v_proj). This "
116
+ " will likely result in reduced accuracy. Please verify the model"
117
+ " accuracy. Consider using a checkpoint with a shared global NVFP4"
118
+ " scale for fused layers."
119
+ )
120
+ # Process input global scale and pre-compute alpha for W4A4 mode
121
+ input_global_scale_inv = layer.input_global_scale.max().to(torch.float32)
122
+ layer.input_global_scale = Parameter(
123
+ (1.0 / input_global_scale_inv).to(torch.float32), requires_grad=False
124
+ )
125
+
126
+ # Pre-compute alpha and inverse for runtime quantization
127
+ layer.input_global_scale_inv = Parameter(
128
+ input_global_scale_inv, requires_grad=False
129
+ )
130
+ layer.alpha = Parameter(
131
+ layer.input_global_scale * layer.weight_global_scale,
132
+ requires_grad=False,
133
+ )
134
135
# Convert layer to NVFP4 linear kernel format
136
self.kernel.process_weights_after_loading(layer)