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 9, 07:30 UTC
Support FP8 block quant for CompressedTensorsW8A16Fp8 (#33280)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jan 30, 2026 at 11:15 UTC
fd0e3772441a32e1c5948d62a8702417fac59979
4 files changed
+74
-64
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+2
-2
@@ -651,7 +651,7 @@ class CompressedTensorsConfig(QuantizationConfig):
651
# note: input_quant will be present for converted models;
652
# will be ignored during inference post loading
653
return CompressedTensorsW8A16Fp8(
654
- strategy=weight_quant.strategy,
654
+ weight_quant=weight_quant,
655
is_static_input_scheme=not input_quant.dynamic,
656
)
657
@@ -659,7 +659,7 @@ class CompressedTensorsConfig(QuantizationConfig):
659
if self._is_fp8_w8a16(weight_quant, input_quant):
660
is_static_input_scheme = input_quant and not input_quant.dynamic
661
return CompressedTensorsW8A16Fp8(
662
- strategy=weight_quant.strategy,
662
+ weight_quant=weight_quant,
663
is_static_input_scheme=is_static_input_scheme,
664
)
665
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
+70
-59
@@ -4,11 +4,17 @@
4
from collections.abc import Callable
5
6
import torch
7
-from compressed_tensors.quantization import QuantizationStrategy
7
+from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
8
9
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
10
CompressedTensorsScheme,
11
)
12
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
13
+ create_fp8_scale_parameter,
14
+ create_fp8_weight_parameter,
15
+ process_fp8_weight_block_strategy,
16
+ validate_fp8_block_shape,
17
+)
18
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
19
apply_fp8_marlin_linear,
20
prepare_fp8_layer_for_marlin,
@@ -17,57 +23,40 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
23
convert_to_channelwise,
24
)
25
from vllm.model_executor.parameter import (
26
+ BlockQuantScaleParameter,
27
ChannelQuantScaleParameter,
21
- ModelWeightParameter,
28
PerTensorScaleParameter,
29
)
30
+from vllm.model_executor.utils import replace_parameter
31
32
__all__ = ["CompressedTensorsW8A16Fp8"]
33
27
-SUPPORTED_STRATEGIES = [QuantizationStrategy.CHANNEL, QuantizationStrategy.TENSOR]
34
+strategy_to_parameter_type = {
35
+ QuantizationStrategy.BLOCK: BlockQuantScaleParameter,
36
+ QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter,
37
+ QuantizationStrategy.TENSOR: PerTensorScaleParameter,
38
+}
39
40
41
class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
31
- def __init__(self, strategy: str, is_static_input_scheme: bool):
32
- self.strategy = strategy
42
+ def __init__(self, weight_quant: QuantizationArgs, is_static_input_scheme: bool):
43
+ self.weight_quant = weight_quant
44
+ self.strategy = weight_quant.strategy
45
self.is_static_input_scheme = is_static_input_scheme
46
+ self.weight_block_size = self.weight_quant.block_structure
47
48
@classmethod
49
def get_min_capability(cls) -> int:
50
# turing and up
51
return 75
52
40
- # W8A8-Fp8 kernels support only per-tensor and per-channel cases.
41
- # So if we have a fused module (QKV, MLP) with per tensor scales,
42
- # we expand each scale to its shard's channels.
43
- def process_weights_after_loading(self, layer) -> None:
44
- if self.strategy == QuantizationStrategy.TENSOR:
45
- ws_channelwise = convert_to_channelwise(
46
- layer.weight_scale, layer.logical_widths
47
- )
48
- layer.weight_scale = torch.nn.Parameter(ws_channelwise, requires_grad=False)
49
- else:
50
- # required by torch.compile to be torch.nn.Parameter
51
- layer.weight_scale = torch.nn.Parameter(
52
- layer.weight_scale.data, requires_grad=False
53
- )
54
-
55
- # Weights must be transposed for marlin
56
- layer.weight = torch.nn.Parameter(layer.weight.t(), requires_grad=False)
57
-
58
- if self.is_static_input_scheme:
59
- # required by torch.compile to be torch.nn.Parameter
60
- layer.input_scale = torch.nn.Parameter(
61
- layer.input_scale.data, requires_grad=False
62
- )
63
- prepare_fp8_layer_for_marlin(layer)
64
-
53
def create_weights(
54
self,
55
layer: torch.nn.Module,
68
- input_size: int,
69
- output_partition_sizes: list[int],
56
input_size_per_partition: int,
57
+ output_partition_sizes: list[int],
58
+ input_size: int,
59
+ output_size: int,
60
params_dtype: torch.dtype,
61
weight_loader: Callable,
62
**kwargs,
@@ -79,38 +68,33 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
68
layer.orig_dtype = params_dtype
69
layer.weight_block_size = None
70
82
- # WEIGHT
83
- weight = ModelWeightParameter(
84
- data=torch.empty(
85
- output_size_per_partition,
71
+ if self.strategy == QuantizationStrategy.BLOCK:
72
+ assert self.weight_block_size is not None
73
+ layer.weight_block_size = self.weight_block_size
74
+ # Validate block quantization shapes
75
+ validate_fp8_block_shape(
76
+ layer,
77
+ input_size,
78
+ output_size,
79
input_size_per_partition,
87
- dtype=torch.float8_e4m3fn,
88
- ),
89
- input_dim=1,
90
- output_dim=0,
91
- weight_loader=weight_loader,
80
+ output_partition_sizes,
81
+ self.weight_block_size,
82
+ )
83
+
84
+ # WEIGHT
85
+ weight = create_fp8_weight_parameter(
86
+ output_size_per_partition, input_size_per_partition, weight_loader
87
)
88
layer.register_parameter("weight", weight)
89
90
# WEIGHT SCALE
96
- if self.strategy == QuantizationStrategy.CHANNEL:
97
- weight_scale = ChannelQuantScaleParameter(
98
- data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
99
- output_dim=0,
100
- weight_loader=weight_loader,
101
- )
102
- elif self.strategy == QuantizationStrategy.TENSOR:
103
- weight_scale = PerTensorScaleParameter(
104
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
105
- weight_loader=weight_loader,
106
- )
107
- else:
108
- raise ValueError(
109
- f"Unsupported weight strategy={self.strategy}, "
110
- f"supported strategies are {SUPPORTED_STRATEGIES}"
111
- )
112
-
113
- weight_scale[:] = torch.finfo(torch.float32).min
91
+ weight_scale = create_fp8_scale_parameter(
92
+ strategy_to_parameter_type[self.strategy],
93
+ output_partition_sizes,
94
+ input_size_per_partition,
95
+ layer.weight_block_size,
96
+ weight_loader,
97
+ )
98
layer.register_parameter("weight_scale", weight_scale)
99
100
# INPUT SCALE (to deal with converted checkpoints)
@@ -121,6 +105,33 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
105
)
106
layer.register_parameter("input_scale", input_scale)
107
108
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
109
+ weight = layer.weight
110
+ weight_scale = layer.weight_scale
111
+ size_k_first = True
112
+ # TODO(rob): refactor block quant into separate class.
113
+ if self.strategy == QuantizationStrategy.BLOCK:
114
+ assert self.is_static_input_scheme is False
115
+ size_k_first = False
116
+ weight, weight_scale = process_fp8_weight_block_strategy(
117
+ weight, weight_scale
118
+ )
119
+ else:
120
+ # Weights must be transposed for marlin
121
+ weight = weight.t()
122
+ if self.strategy == QuantizationStrategy.TENSOR:
123
+ # If we have a fused module (QKV, MLP) with per tensor scales,
124
+ # we expand each scale to its shard's channels.
125
+ weight_scale = convert_to_channelwise(
126
+ weight_scale, layer.logical_widths
127
+ )
128
+
129
+ # Update layer with new values
130
+ replace_parameter(layer, "weight", weight.data)
131
+ replace_parameter(layer, "weight_scale", weight_scale.data)
132
+
133
+ prepare_fp8_layer_for_marlin(layer, size_k_first=size_k_first)
134
+
135
def apply_weights(
136
self,
137
layer: torch.nn.Module,
vllm/model_executor/layers/quantization/fp8.py
-2
@@ -400,7 +400,6 @@ class Fp8LinearMethod(LinearMethodBase):
400
None,
401
weight_loader,
402
)
403
- set_weight_attrs(scale, {"scale_type": "weight_scale"})
403
layer.register_parameter("weight_scale", scale)
404
else:
405
assert not self.act_q_static
@@ -412,7 +411,6 @@ class Fp8LinearMethod(LinearMethodBase):
411
self.weight_block_size,
412
weight_loader,
413
)
415
- set_weight_attrs(scale, {"scale_type": "weight_scale"})
414
# The weight_scale_inv name is intentional for deepseekv3
415
layer.register_parameter("weight_scale_inv", scale)
416
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+2
-1
@@ -29,7 +29,7 @@ from vllm.model_executor.parameter import (
29
ChannelQuantScaleParameter,
30
PerTensorScaleParameter,
31
)
32
-from vllm.model_executor.utils import replace_parameter
32
+from vllm.model_executor.utils import replace_parameter, set_weight_attrs
33
from vllm.platforms import current_platform
34
from vllm.triton_utils import tl, triton
35
from vllm.utils.deep_gemm import (
@@ -1520,6 +1520,7 @@ def create_fp8_scale_parameter(
1520
raise ValueError(f"Unknown parameter type: {parameter_type}")
1521
1522
scale[:] = torch.finfo(torch.float32).min
1523
+ set_weight_attrs(scale, {"scale_type": "weight_scale"})
1524
return scale
1525
1526