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 14, 05:30 UTC
Refactor dense FP8 tensor/channel/block utils and add CT FP8 block (#21404)
Michael Goin committed
Sep 18, 2025 at 08:53 UTC
fbd6523ac00082c398dc8126434cede595169609
5 files changed
+442
-318
vllm/model_executor/layers/linear.py
+7
-7
@@ -805,12 +805,10 @@ class MergedColumnParallelLinear(ColumnParallelLinear):
805
assert loaded_shard_id < len(self.output_sizes)
806
807
if isinstance(param, BlockQuantScaleParameter):
808
- from vllm.model_executor.layers.quantization.fp8 import (
809
- Fp8LinearMethod, Fp8MoEMethod)
808
assert self.quant_method is not None
811
- assert isinstance(self.quant_method,
812
- (Fp8LinearMethod, Fp8MoEMethod))
813
- weight_block_size = self.quant_method.quant_config.weight_block_size
809
+ # Assume the weight block size has been set by quant method
810
+ assert hasattr(self, "weight_block_size")
811
+ weight_block_size = self.weight_block_size
812
assert weight_block_size is not None
813
block_n, _ = weight_block_size[0], weight_block_size[1]
814
shard_offset = (
@@ -989,8 +987,10 @@ class QKVParallelLinear(ColumnParallelLinear):
987
# Note(simon): This is needed for Qwen3's fp8 quantization.
988
if isinstance(param, BlockQuantScaleParameter):
989
assert self.quant_method is not None
992
- assert hasattr(self.quant_method, "quant_config")
993
- weight_block_size = self.quant_method.quant_config.weight_block_size
990
+ # Assume the weight block size has been set by quant method
991
+ assert hasattr(self, "weight_block_size")
992
+ weight_block_size = self.weight_block_size
993
+ assert weight_block_size is not None
994
block_n, _ = weight_block_size[0], weight_block_size[1]
995
shard_offset = (shard_offset + block_n - 1) // block_n
996
shard_size = (shard_size + block_n - 1) // block_n
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+35
-33
@@ -12,7 +12,6 @@ from compressed_tensors.quantization import (QuantizationArgs,
12
QuantizationStrategy,
13
QuantizationType)
14
from compressed_tensors.transform import TransformConfig
15
-from pydantic import BaseModel
15
16
import vllm.envs as envs
17
from vllm.logger import init_logger
@@ -268,7 +267,8 @@ class CompressedTensorsConfig(QuantizationConfig):
267
else:
268
return False
269
271
- def _is_fp4a4_nvfp4(self, weight_quant: BaseModel, input_quant: BaseModel):
270
+ def _is_fp4a4_nvfp4(self, weight_quant: QuantizationArgs,
271
+ input_quant: QuantizationArgs):
272
273
if weight_quant is None or input_quant is None:
274
return False
@@ -288,8 +288,8 @@ class CompressedTensorsConfig(QuantizationConfig):
288
return (is_tensor_group_quant and is_float_type and is_4_bits
289
and is_group_size_16 and is_symmetric)
290
291
- def _is_fp4a16_nvfp4(self, weight_quant: BaseModel,
292
- input_quant: BaseModel):
291
+ def _is_fp4a16_nvfp4(self, weight_quant: QuantizationArgs,
292
+ input_quant: QuantizationArgs):
293
294
is_weight_only = weight_quant is not None and input_quant is None
295
is_tensor_group_quant = (
@@ -303,8 +303,8 @@ class CompressedTensorsConfig(QuantizationConfig):
303
return (is_weight_only and is_tensor_group_quant and is_float_type
304
and is_4_bits and is_group_size_16 and is_symmetric)
305
306
- def _is_static_tensor_w8a8(self, weight_quant: BaseModel,
307
- input_quant: BaseModel) -> bool:
306
+ def _is_static_tensor_w8a8(self, weight_quant: QuantizationArgs,
307
+ input_quant: QuantizationArgs) -> bool:
308
is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8
309
weight_strategy = (
310
weight_quant.strategy == QuantizationStrategy.TENSOR.value
@@ -317,8 +317,8 @@ class CompressedTensorsConfig(QuantizationConfig):
317
# Only symmetric weight quantization supported.
318
return is_8_bits and is_tensor and weight_quant.symmetric and is_static
319
320
- def _is_dynamic_token_w8a8(self, weight_quant: BaseModel,
321
- input_quant: BaseModel) -> bool:
320
+ def _is_dynamic_token_w8a8(self, weight_quant: QuantizationArgs,
321
+ input_quant: QuantizationArgs) -> bool:
322
is_8_bits = weight_quant.num_bits == input_quant.num_bits == 8
323
weight_strategy = (
324
weight_quant.strategy == QuantizationStrategy.TENSOR.value
@@ -331,8 +331,8 @@ class CompressedTensorsConfig(QuantizationConfig):
331
# Only symmetric weight quantization supported.
332
return is_8_bits and is_token and weight_quant.symmetric and is_dynamic
333
334
- def _is_dynamic_token_w4a8_int(self, weight_quant: BaseModel,
335
- input_quant: BaseModel) -> bool:
334
+ def _is_dynamic_token_w4a8_int(self, weight_quant: QuantizationArgs,
335
+ input_quant: QuantizationArgs) -> bool:
336
is_weight_4_bits = weight_quant.num_bits == 4
337
is_activation_8_bits = input_quant.num_bits == 8
338
weight_strategy = (
@@ -347,8 +347,8 @@ class CompressedTensorsConfig(QuantizationConfig):
347
return (is_weight_4_bits and is_activation_8_bits and is_token
348
and weight_quant.symmetric and is_dynamic)
349
350
- def _is_fp8_w8a8(self, weight_quant: BaseModel,
351
- input_quant: BaseModel) -> bool:
350
+ def _is_fp8_w8a8(self, weight_quant: QuantizationArgs,
351
+ input_quant: QuantizationArgs) -> bool:
352
# Confirm weights and activations quantized.
353
if weight_quant is None or input_quant is None:
354
return False
@@ -358,11 +358,12 @@ class CompressedTensorsConfig(QuantizationConfig):
358
and input_quant.type == QuantizationType.FLOAT)
359
is_symmetric_weight = weight_quant.symmetric
360
is_static_weight = not weight_quant.dynamic
361
- is_per_tensor_or_channel_weight = (weight_quant.strategy in [
362
- QuantizationStrategy.TENSOR, QuantizationStrategy.CHANNEL
361
+ is_tensor_or_channel_or_block_weight = (weight_quant.strategy in [
362
+ QuantizationStrategy.TENSOR, QuantizationStrategy.CHANNEL,
363
+ QuantizationStrategy.BLOCK
364
])
365
if not (is_floating_point and is_symmetric_weight and is_static_weight
365
- and is_per_tensor_or_channel_weight):
366
+ and is_tensor_or_channel_or_block_weight):
367
return False
368
369
# Dynamic quantization is always supported if weights supported.
@@ -375,8 +376,8 @@ class CompressedTensorsConfig(QuantizationConfig):
376
input_quant.strategy == QuantizationStrategy.TENSOR)
377
return is_symmetric_activation and is_per_tensor_activation
378
378
- def _is_fp8_w4a8(self, weight_quant: BaseModel,
379
- input_quant: BaseModel) -> bool:
379
+ def _is_fp8_w4a8(self, weight_quant: QuantizationArgs,
380
+ input_quant: QuantizationArgs) -> bool:
381
if not weight_quant or not input_quant:
382
return False
383
is_weight_4_bits = weight_quant.num_bits == 4
@@ -392,24 +393,24 @@ class CompressedTensorsConfig(QuantizationConfig):
393
return (is_weight_4_bits and is_activation_8_bits and is_token
394
and is_symmetric and is_dynamic)
395
395
- def _is_fp8_w4a8_sm90(self, weight_quant: BaseModel,
396
- input_quant: BaseModel) -> bool:
396
+ def _is_fp8_w4a8_sm90(self, weight_quant: QuantizationArgs,
397
+ input_quant: QuantizationArgs) -> bool:
398
return (self._check_scheme_supported(90, error=False, match_exact=True)
399
and self._is_fp8_w4a8(weight_quant, input_quant))
400
400
- def _is_fp8_w8a8_sm90(self, weight_quant: BaseModel,
401
- input_quant: BaseModel) -> bool:
401
+ def _is_fp8_w8a8_sm90(self, weight_quant: QuantizationArgs,
402
+ input_quant: QuantizationArgs) -> bool:
403
return (self._check_scheme_supported(90, error=False, match_exact=True)
404
and self._is_fp8_w8a8(weight_quant, input_quant))
405
405
- def _is_fp8_w8a8_sm100(self, weight_quant: BaseModel,
406
- input_quant: BaseModel) -> bool:
406
+ def _is_fp8_w8a8_sm100(self, weight_quant: QuantizationArgs,
407
+ input_quant: QuantizationArgs) -> bool:
408
return (self._check_scheme_supported(
409
100, error=False, match_exact=True)
410
and self._is_fp8_w8a8(weight_quant, input_quant))
411
411
- def _is_fp8_w8a16(self, weight_quant: BaseModel,
412
- input_quant: BaseModel) -> bool:
412
+ def _is_fp8_w8a16(self, weight_quant: QuantizationArgs,
413
+ input_quant: QuantizationArgs) -> bool:
414
# Confirm weights quantized.
415
if weight_quant is None:
416
return False
@@ -421,18 +422,19 @@ class CompressedTensorsConfig(QuantizationConfig):
422
# Confirm weight scheme is supported.
423
is_symmetric_weight = weight_quant.symmetric
424
is_static_weight = not weight_quant.dynamic
424
- is_per_tensor_or_channel_weight = (weight_quant.strategy in [
425
- QuantizationStrategy.TENSOR, QuantizationStrategy.CHANNEL
425
+ is_tensor_or_channel_or_block_weight = (weight_quant.strategy in [
426
+ QuantizationStrategy.TENSOR, QuantizationStrategy.CHANNEL,
427
+ QuantizationStrategy.BLOCK
428
])
429
if not (is_symmetric_weight and is_static_weight # noqa: SIM103
428
- and is_per_tensor_or_channel_weight):
430
+ and is_tensor_or_channel_or_block_weight):
431
return False
432
433
# All conditions satisfied.
434
return True
435
434
- def _is_wNa16_group_channel(self, weight_quant: BaseModel,
435
- input_quant: BaseModel) -> bool:
436
+ def _is_wNa16_group_channel(self, weight_quant: QuantizationArgs,
437
+ input_quant: QuantizationArgs) -> bool:
438
input_quant_none = input_quant is None
439
is_channel_group = (
440
weight_quant.strategy == QuantizationStrategy.CHANNEL.value
@@ -443,8 +445,8 @@ class CompressedTensorsConfig(QuantizationConfig):
445
446
def _get_scheme_from_parts(
447
self,
446
- weight_quant: BaseModel,
447
- input_quant: BaseModel,
448
+ weight_quant: QuantizationArgs,
449
+ input_quant: QuantizationArgs,
450
format: Optional[str] = None) -> "CompressedTensorsScheme":
451
452
# use the per-layer format if defined, otherwise, use global format
@@ -496,7 +498,7 @@ class CompressedTensorsConfig(QuantizationConfig):
498
CompressedTensorsW8A8Fp8.get_min_capability(), error=False)
499
if is_fp8_w8a8_supported:
500
return CompressedTensorsW8A8Fp8(
499
- strategy=weight_quant.strategy,
501
+ weight_quant=weight_quant,
502
is_static_input_scheme=(input_quant
503
and not input_quant.dynamic))
504
else:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+96
-95
@@ -4,28 +4,41 @@
4
from typing import Callable, Optional
5
6
import torch
7
-from compressed_tensors.quantization import QuantizationStrategy
7
+from compressed_tensors.quantization import (QuantizationArgs,
8
+ QuantizationStrategy)
9
from torch.nn import Parameter
10
11
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
12
CompressedTensorsScheme)
13
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14
+ apply_fp8_block_linear, check_aiter_fp8_linear_support,
15
+ create_fp8_input_scale, create_fp8_scale_parameter,
16
+ create_fp8_weight_parameter, maybe_post_process_fp8_weight_block,
17
+ process_fp8_weight_block_strategy, process_fp8_weight_channel_strategy,
18
+ process_fp8_weight_tensor_strategy, validate_fp8_block_shape)
19
from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
GroupShape)
21
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15
- Fp8LinearOp, maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz,
16
- requantize_with_max_scale)
17
-from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
18
- ModelWeightParameter,
22
+ Fp8LinearOp, cutlass_block_fp8_supported, maybe_create_device_identity)
23
+from vllm.model_executor.parameter import (BlockQuantScaleParameter,
24
+ ChannelQuantScaleParameter,
25
PerTensorScaleParameter)
20
-from vllm.platforms import current_platform
26
27
__all__ = ["CompressedTensorsW8A8Fp8"]
28
29
+strategy_to_parameter_type = {
30
+ QuantizationStrategy.BLOCK: BlockQuantScaleParameter,
31
+ QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter,
32
+ QuantizationStrategy.TENSOR: PerTensorScaleParameter,
33
+}
34
+
35
36
class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
37
27
- def __init__(self, strategy: str, is_static_input_scheme: bool):
28
- self.strategy = strategy
38
+ def __init__(self, weight_quant: QuantizationArgs,
39
+ is_static_input_scheme: bool):
40
+ self.weight_quant = weight_quant
41
+ self.strategy = weight_quant.strategy
42
self.out_dtype = torch.get_default_dtype()
43
self.is_static_input_scheme = is_static_input_scheme
44
self.act_q_group_shape = GroupShape.PER_TENSOR \
@@ -34,120 +47,108 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
47
act_quant_static=self.is_static_input_scheme,
48
act_quant_group_shape=self.act_q_group_shape)
49
50
+ self.weight_block_size = self.weight_quant.block_structure
51
+ self.cutlass_block_fp8_supported = cutlass_block_fp8_supported()
52
+ self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
53
+
54
@classmethod
55
def get_min_capability(cls) -> int:
56
# lovelace and up
57
return 89
58
42
- def process_weights_after_loading(self, layer) -> None:
43
- # If per tensor, when we have a fused module (e.g. QKV) with per
44
- # tensor scales (thus N scales being passed to the kernel),
45
- # requantize so we can always run per tensor
46
- if self.strategy == QuantizationStrategy.TENSOR:
47
- max_w_scale, weight = requantize_with_max_scale(
48
- weight=layer.weight,
49
- weight_scale=layer.weight_scale,
50
- logical_widths=layer.logical_widths,
51
- )
52
-
53
- if current_platform.is_fp8_fnuz():
54
- input_scale = getattr(layer, 'input_scale', None)
55
-
56
- weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
57
- weight=weight,
58
- weight_scale=max_w_scale,
59
- input_scale=input_scale)
60
- if input_scale is not None:
61
- layer.input_scale = Parameter(input_scale,
62
- requires_grad=False)
63
-
64
- layer.weight = Parameter(weight.t(), requires_grad=False)
65
- layer.weight_scale = Parameter(max_w_scale, requires_grad=False)
66
-
67
- # If channelwise, scales are already lined up, so just transpose.
68
- elif self.strategy == QuantizationStrategy.CHANNEL:
69
- weight = layer.weight
70
-
71
- if current_platform.is_fp8_fnuz():
72
- input_scale = getattr(layer, 'input_scale', None)
73
-
74
- weight, weight_scale, input_scale = \
75
- normalize_e4m3fn_to_e4m3fnuz(
76
- weight=weight,
77
- weight_scale=layer.weight_scale,
78
- input_scale=input_scale)
79
- if input_scale is not None:
80
- layer.input_scale = Parameter(input_scale,
81
- requires_grad=False)
82
- else:
83
- weight_scale = layer.weight_scale.data
84
-
85
- layer.weight = Parameter(weight.t(), requires_grad=False)
86
- # required by torch.compile to be torch.nn.Parameter
87
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
88
-
89
- else:
90
- raise ValueError(f"Unknown quantization strategy {self.strategy}")
91
-
92
- # INPUT SCALE
93
- if self.is_static_input_scheme and hasattr(layer, 'input_scale'):
94
- layer.input_scale = Parameter(layer.input_scale.max(),
95
- requires_grad=False)
96
- else:
97
- layer.input_scale = None
98
-
59
def create_weights(self, layer: torch.nn.Module,
100
- output_partition_sizes: list[int],
60
input_size_per_partition: int,
102
- params_dtype: torch.dtype, weight_loader: Callable,
103
- **kwargs):
61
+ output_partition_sizes: list[int], input_size: int,
62
+ output_size: int, params_dtype: torch.dtype,
63
+ weight_loader: Callable, **kwargs):
64
maybe_create_device_identity()
65
66
output_size_per_partition = sum(output_partition_sizes)
67
layer.logical_widths = output_partition_sizes
68
+ layer.weight_block_size = None
69
+
70
+ if self.strategy == QuantizationStrategy.BLOCK:
71
+ assert self.weight_block_size is not None
72
+ layer.weight_block_size = self.weight_block_size
73
+ # Validate block quantization shapes
74
+ validate_fp8_block_shape(layer, input_size, output_size,
75
+ input_size_per_partition,
76
+ output_partition_sizes,
77
+ self.weight_block_size)
78
79
# WEIGHT
110
- weight = ModelWeightParameter(data=torch.empty(
111
- output_size_per_partition,
112
- input_size_per_partition,
113
- dtype=torch.float8_e4m3fn),
114
- input_dim=1,
115
- output_dim=0,
116
- weight_loader=weight_loader)
80
+ weight = create_fp8_weight_parameter(output_size_per_partition,
81
+ input_size_per_partition,
82
+ weight_loader)
83
layer.register_parameter("weight", weight)
84
85
# WEIGHT SCALE
120
- # TODO: update create_xxx_parameter functions to return
121
- # the newly added parameters
122
- if self.strategy == QuantizationStrategy.CHANNEL:
123
- weight_scale = ChannelQuantScaleParameter(
124
- data=torch.empty((sum(output_partition_sizes), 1),
125
- dtype=torch.float32),
126
- output_dim=0,
127
- weight_loader=weight_loader)
128
- else:
129
- assert self.strategy == QuantizationStrategy.TENSOR
130
- weight_scale = PerTensorScaleParameter(data=torch.empty(
131
- len(output_partition_sizes), dtype=torch.float32),
132
- weight_loader=weight_loader)
133
-
134
- # min requirement for fp8 kernels
135
- weight_scale[:] = torch.finfo(torch.float32).min
86
+ weight_scale = create_fp8_scale_parameter(
87
+ strategy_to_parameter_type[self.strategy], output_partition_sizes,
88
+ input_size_per_partition, layer.weight_block_size, weight_loader)
89
layer.register_parameter("weight_scale", weight_scale)
90
91
# INPUT SCALE
92
if self.is_static_input_scheme:
140
- input_scale = PerTensorScaleParameter(data=torch.empty(
141
- len(output_partition_sizes), dtype=torch.float32),
142
- weight_loader=weight_loader)
143
- input_scale[:] = torch.finfo(torch.float32).min
93
+ input_scale = create_fp8_input_scale(output_partition_sizes,
94
+ weight_loader)
95
layer.register_parameter("input_scale", input_scale)
96
97
+ def process_weights_after_loading(self, layer) -> None:
98
+ if self.strategy == QuantizationStrategy.TENSOR:
99
+ weight, weight_scale, input_scale = (
100
+ process_fp8_weight_tensor_strategy(
101
+ layer.weight, layer.weight_scale, layer.logical_widths,
102
+ getattr(layer, 'input_scale', None)))
103
+ weight = weight.t()
104
+
105
+ elif self.strategy == QuantizationStrategy.CHANNEL:
106
+ weight, weight_scale, input_scale = (
107
+ process_fp8_weight_channel_strategy(
108
+ layer.weight, layer.weight_scale,
109
+ getattr(layer, 'input_scale', None)))
110
+ weight = weight.t()
111
+
112
+ elif self.strategy == QuantizationStrategy.BLOCK:
113
+ assert self.is_static_input_scheme is False
114
+ weight, weight_scale = process_fp8_weight_block_strategy(
115
+ layer.weight, layer.weight_scale)
116
+ input_scale = None
117
+
118
+ else:
119
+ raise ValueError(f"Unknown quantization strategy {self.strategy}")
120
+
121
+ # required by torch.compile to be torch.nn.Parameter
122
+ layer.weight = Parameter(weight.data, requires_grad=False)
123
+ layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)
124
+ if input_scale is not None:
125
+ layer.input_scale = Parameter(input_scale.data,
126
+ requires_grad=False)
127
+
128
+ # INPUT SCALE
129
+ if self.is_static_input_scheme and hasattr(layer, 'input_scale'):
130
+ layer.input_scale = Parameter(layer.input_scale.max(),
131
+ requires_grad=False)
132
+ else:
133
+ layer.input_scale = None
134
+
135
+ if self.strategy == QuantizationStrategy.BLOCK:
136
+ maybe_post_process_fp8_weight_block(
137
+ layer, self.cutlass_block_fp8_supported)
138
+
139
def apply_weights(self,
140
layer: torch.nn.Module,
141
x: torch.Tensor,
142
bias: Optional[torch.Tensor] = None) -> torch.Tensor:
143
144
+ if layer.weight_block_size is not None:
145
+ return apply_fp8_block_linear(
146
+ layer,
147
+ input=x,
148
+ bias=bias,
149
+ cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
150
+ use_aiter_and_is_supported=self.use_aiter_and_is_supported)
151
+
152
return self.fp8_linear.apply(input=x,
153
weight=layer.weight,
154
weight_scale=layer.weight_scale,
vllm/model_executor/layers/quantization/fp8.py
+84
-183
@@ -4,7 +4,6 @@
4
from typing import TYPE_CHECKING, Any, Callable, Optional, Union
5
6
import torch
7
-import torch.nn.functional as F
7
from torch.nn import Module
8
from torch.nn.parameter import Parameter
9
@@ -32,8 +31,12 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
31
register_moe_scaling_factors, rotate_flashinfer_fp8_moe_weights,
32
select_cutlass_fp8_gemm_impl, swap_w13_to_w31)
33
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
35
- get_col_major_tma_aligned_tensor, requant_weight_ue8m0_inplace,
36
- should_use_deepgemm_for_fp8_linear)
34
+ apply_fp8_block_linear, check_aiter_fp8_linear_support,
35
+ create_fp8_input_scale, create_fp8_scale_parameter,
36
+ create_fp8_weight_parameter, get_col_major_tma_aligned_tensor,
37
+ maybe_post_process_fp8_weight_block, process_fp8_weight_block_strategy,
38
+ process_fp8_weight_tensor_strategy, requant_weight_ue8m0_inplace,
39
+ validate_fp8_block_shape)
40
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
41
apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin,
42
prepare_moe_fp8_layer_for_marlin)
@@ -42,8 +45,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
46
Fp8LinearOp, all_close_1d, cutlass_block_fp8_supported,
47
cutlass_fp8_supported, maybe_create_device_identity,
45
- normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize,
46
- requantize_with_max_scale)
48
+ normalize_e4m3fn_to_e4m3fnuz, per_tensor_dequantize)
49
from vllm.model_executor.parameter import (BlockQuantScaleParameter,
50
ModelWeightParameter,
51
PerTensorScaleParameter)
@@ -233,14 +235,10 @@ class Fp8LinearMethod(LinearMethodBase):
235
if current_platform.is_rocm():
236
self.use_marlin = False
237
236
- # AITER is only supported on ROCm and only for FP8_FNUZ
237
- # and at the moment are MI300 series
238
- self.use_aiter_and_is_supported = (current_platform.is_rocm()
239
- and envs.VLLM_ROCM_USE_AITER
240
- and envs.VLLM_ROCM_USE_AITER_LINEAR
241
- and current_platform.is_fp8_fnuz())
238
+ self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
239
243
- self.block_quant = self.quant_config.weight_block_size is not None
240
+ self.weight_block_size = self.quant_config.weight_block_size
241
+ self.block_quant = self.weight_block_size is not None
242
self.act_q_static = self.quant_config.activation_scheme == "static"
243
# Use per-token quantization for better perf if dynamic and cutlass
244
if not self.act_q_static and cutlass_fp8_supported():
@@ -273,51 +271,27 @@ class Fp8LinearMethod(LinearMethodBase):
271
layer.weight_block_size = None
272
273
if self.block_quant:
276
- tp_size = getattr(layer, "tp_size",
277
- get_tensor_model_parallel_world_size())
278
- assert self.quant_config.weight_block_size is not None
279
- layer.weight_block_size = self.quant_config.weight_block_size
280
- block_n, block_k = (
281
- self.quant_config.weight_block_size[0],
282
- self.quant_config.weight_block_size[1],
283
- )
284
- # Required by row parallel
285
- if (tp_size > 1
286
- and input_size // input_size_per_partition == tp_size
287
- and input_size_per_partition % block_k != 0):
288
- raise ValueError(
289
- f"Weight input_size_per_partition = "
290
- f"{input_size_per_partition} is not divisible by "
291
- f"weight quantization block_k = {block_k}.")
292
- # Required by column parallel or enabling merged weights
293
- is_tp_split = (tp_size > 1 and
294
- output_size // output_size_per_partition == tp_size)
295
- is_merged_gemm = len(output_partition_sizes) > 1
296
- if is_tp_split or is_merged_gemm:
297
- sizes_to_check = output_partition_sizes
298
- if not is_tp_split and is_merged_gemm:
299
- # In case of merged matrices, we allow the last
300
- # matrix to not be a multiple of block size
301
- sizes_to_check = output_partition_sizes[:-1]
302
- for output_partition_size in sizes_to_check:
303
- if output_partition_size % block_n != 0:
304
- raise ValueError(
305
- f"Weight output_partition_size = "
306
- f"{output_partition_size} is not divisible by "
307
- f"weight quantization block_n = {block_n}.")
274
+ assert self.weight_block_size is not None
275
+ layer.weight_block_size = self.weight_block_size
276
+ validate_fp8_block_shape(layer, input_size, output_size,
277
+ input_size_per_partition,
278
+ output_partition_sizes,
279
+ self.weight_block_size)
280
281
# WEIGHT
310
- weight_dtype = (torch.float8_e4m3fn
311
- if self.quant_config.is_checkpoint_fp8_serialized else
312
- params_dtype)
313
-
314
- weight = ModelWeightParameter(data=torch.empty(
315
- output_size_per_partition,
316
- input_size_per_partition,
317
- dtype=weight_dtype),
318
- input_dim=1,
319
- output_dim=0,
320
- weight_loader=weight_loader)
282
+ if self.quant_config.is_checkpoint_fp8_serialized:
283
+ weight = create_fp8_weight_parameter(output_size_per_partition,
284
+ input_size_per_partition,
285
+ weight_loader)
286
+ else:
287
+ # For non-serialized checkpoints, use original dtype
288
+ weight = ModelWeightParameter(data=torch.empty(
289
+ output_size_per_partition,
290
+ input_size_per_partition,
291
+ dtype=params_dtype),
292
+ input_dim=1,
293
+ output_dim=0,
294
+ weight_loader=weight_loader)
295
layer.register_parameter("weight", weight)
296
297
# If checkpoint is serialized fp8, load them.
@@ -325,154 +299,87 @@ class Fp8LinearMethod(LinearMethodBase):
299
if self.quant_config.is_checkpoint_fp8_serialized:
300
# WEIGHT SCALE
301
if not self.block_quant:
328
- scale = PerTensorScaleParameter(
329
- data=torch.empty(len(output_partition_sizes),
330
- dtype=torch.float32),
331
- weight_loader=weight_loader,
332
- )
333
- scale[:] = torch.finfo(torch.float32).min
302
+ scale = create_fp8_scale_parameter(PerTensorScaleParameter,
303
+ output_partition_sizes,
304
+ input_size_per_partition,
305
+ None, weight_loader)
306
set_weight_attrs(scale, {"scale_type": "weight_scale"})
307
layer.register_parameter("weight_scale", scale)
308
else:
337
- assert self.quant_config.activation_scheme == "dynamic"
338
- scale = BlockQuantScaleParameter(
339
- data=torch.empty(
340
- (output_size_per_partition + block_n - 1) // block_n,
341
- (input_size_per_partition + block_k - 1) // block_k,
342
- dtype=torch.float32,
343
- ),
344
- input_dim=1,
345
- output_dim=0,
346
- weight_loader=weight_loader,
347
- )
348
- scale[:] = torch.finfo(torch.float32).min
309
+ assert not self.act_q_static
310
+ assert self.weight_block_size is not None
311
+ scale = create_fp8_scale_parameter(BlockQuantScaleParameter,
312
+ output_partition_sizes,
313
+ input_size_per_partition,
314
+ self.weight_block_size,
315
+ weight_loader)
316
set_weight_attrs(scale, {"scale_type": "weight_scale"})
317
# The weight_scale_inv name is intentional for deepseekv3
318
layer.register_parameter("weight_scale_inv", scale)
319
320
# INPUT ACTIVATION SCALE
354
- if self.quant_config.activation_scheme == "static":
355
- scale = PerTensorScaleParameter(data=torch.empty(
356
- len(output_partition_sizes), dtype=torch.float32),
357
- weight_loader=weight_loader)
358
-
359
- scale[:] = torch.finfo(torch.float32).min
321
+ if self.act_q_static:
322
+ scale = create_fp8_input_scale(output_partition_sizes,
323
+ weight_loader)
324
set_weight_attrs(scale, {"scale_type": "input_scale"})
325
layer.register_parameter("input_scale", scale)
326
else:
327
layer.register_parameter("input_scale", None)
328
365
- def _maybe_pad_weight(self, weight: torch.Tensor) -> torch.Tensor:
366
- # Pad the weight tensor. This is an optimization on ROCm platform, which
367
- # can benefit from tensors located far enough from one another in memory
368
- if (envs.VLLM_ROCM_FP8_PADDING and current_platform.is_rocm()
369
- and weight.stride(-1) == 1
370
- and (weight.stride(-2) * weight.element_size()) % 512 == 0):
371
- num_pad = 256 // weight.element_size()
372
- weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
373
- torch.cuda.empty_cache()
374
- return weight
375
-
329
def process_weights_after_loading(self, layer: Module) -> None:
330
size_k_first = True
331
+ input_scale = None
332
# TODO(rob): refactor block quant into separate class.
333
if self.block_quant:
380
- assert self.quant_config.activation_scheme == "dynamic"
334
+ assert not self.act_q_static
335
size_k_first = False
382
- if current_platform.is_fp8_fnuz():
383
- weight, weight_scale_inv, _ = \
384
- normalize_e4m3fn_to_e4m3fnuz(
385
- weight=layer.weight,
386
- weight_scale=layer.weight_scale_inv)
387
- else:
388
- weight = layer.weight.data
389
- weight_scale_inv = layer.weight_scale_inv.data
336
391
- weight = self._maybe_pad_weight(weight)
392
-
393
- # Torch.compile cannot use Parameter subclasses.
394
- layer.weight = Parameter(weight, requires_grad=False)
395
- layer.weight_scale_inv = Parameter(weight_scale_inv,
396
- requires_grad=False)
337
+ weight, weight_scale = process_fp8_weight_block_strategy(
338
+ layer.weight, layer.weight_scale_inv)
339
+ # Delete the weight_scale_inv parameter to avoid confusion
340
+ # with the weight_scale parameter
341
+ del layer.weight_scale_inv
342
343
# If checkpoint not serialized fp8, quantize the weights.
344
elif not self.quant_config.is_checkpoint_fp8_serialized:
345
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight,
346
scale=None)
347
+ weight = qweight.t()
348
403
- # Update the layer with the new values.
404
- layer.weight = Parameter(qweight.t(), requires_grad=False)
405
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
406
- # layer.input_scale is None indicates dynamic quant and scale is
407
- # computed from input.
408
- layer.input_scale = None
409
-
410
- # If checkpoint is fp8, handle that there are N scales for N
349
+ # If checkpoint is fp8 per-tensor, handle that there are N scales for N
350
# shards in a fused module
351
else:
413
- layer.weight_scale = torch.nn.Parameter(layer.weight_scale.data,
414
- requires_grad=False)
415
- if self.quant_config.activation_scheme == "static":
416
- layer.input_scale = torch.nn.Parameter(layer.input_scale.data,
417
- requires_grad=False)
418
-
352
weight = layer.weight
353
weight_scale = layer.weight_scale
354
355
# If using w8a8, torch._scaled_mm needs per tensor, so
356
# requantize the logical shards as a single weight.
357
if not self.use_marlin:
425
- # Dequant -> Quant with max scale so we can run per tensor.
426
- if current_platform.is_fp8_fnuz():
427
- weight, weight_scale, input_scale = \
428
- normalize_e4m3fn_to_e4m3fnuz(
429
- weight=weight,
430
- weight_scale=weight_scale,
431
- input_scale=layer.input_scale)
432
- if input_scale is not None:
433
- layer.input_scale = Parameter(input_scale,
434
- requires_grad=False)
435
-
436
- weight_scale, weight = requantize_with_max_scale(
437
- weight=weight,
438
- weight_scale=weight_scale,
439
- logical_widths=layer.logical_widths,
440
- )
441
-
442
- weight = self._maybe_pad_weight(weight)
443
- # Update layer with new values.
444
- layer.weight = Parameter(weight.t(), requires_grad=False)
445
- layer.weight_scale = Parameter(weight_scale, requires_grad=False)
446
- if self.quant_config.activation_scheme == "static":
447
- layer.input_scale = Parameter(layer.input_scale.max(),
448
- requires_grad=False)
358
+ weight, weight_scale, input_scale = (
359
+ process_fp8_weight_tensor_strategy(
360
+ weight, weight_scale, layer.logical_widths,
361
+ getattr(layer, 'input_scale', None)))
362
+ if self.act_q_static:
363
+ assert input_scale is not None
364
+ input_scale = input_scale.max()
365
+ weight = weight.t()
366
+
367
+ # Update layer with new values.
368
+ layer.weight = Parameter(weight.data, requires_grad=False)
369
+ layer.weight_scale = Parameter(weight_scale.data, requires_grad=False)
370
+ layer.input_scale = Parameter(
371
+ input_scale,
372
+ requires_grad=False) if input_scale is not None else None
373
374
if self.use_marlin:
375
prepare_fp8_layer_for_marlin(layer, size_k_first)
376
# Activations not quantized for marlin.
377
del layer.input_scale
378
+ return
379
455
- # On Blackwell or Hopper, if E8M0 for DeepGemm is used, we need to
456
- # requantize the weight and input to the specific scale
457
- # at the same time.
458
- if is_deep_gemm_e8m0_used() and self.block_quant:
459
- assert layer.weight_block_size is not None
460
- block_sz = tuple(layer.weight_block_size)
461
- requant_weight_ue8m0_inplace(
462
- layer.weight.data,
463
- layer.weight_scale_inv.data if hasattr(
464
- layer, "weight_scale_inv") else layer.weight_scale.data,
465
- block_sz,
466
- )
467
-
468
- # SM90 Block FP8 CUTLASS requires row-major weight scales
469
- if (self.block_quant and current_platform.is_device_capability(90)
470
- and self.cutlass_block_fp8_supported
471
- and not should_use_deepgemm_for_fp8_linear(
472
- torch.bfloat16, layer.weight)):
473
- layer.weight_scale_inv = Parameter(
474
- layer.weight_scale_inv.data.T.contiguous(),
475
- requires_grad=False)
380
+ if self.block_quant:
381
+ maybe_post_process_fp8_weight_block(
382
+ layer, self.cutlass_block_fp8_supported)
383
384
def apply(self,
385
layer: torch.nn.Module,
@@ -490,18 +397,12 @@ class Fp8LinearMethod(LinearMethodBase):
397
bias=bias)
398
399
if self.block_quant:
493
- assert self.quant_config.weight_block_size is not None
494
-
495
- return torch.ops.vllm.apply_w8a8_block_fp8_linear(
400
+ return apply_fp8_block_linear(
401
+ layer,
402
input=x,
497
- weight=layer.weight,
498
- block_size=self.quant_config.weight_block_size,
499
- weight_scale=layer.weight_scale_inv,
500
- input_scale=layer.input_scale,
403
bias=bias,
404
cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
503
- use_aiter_and_is_supported=self.use_aiter_and_is_supported,
504
- )
405
+ use_aiter_and_is_supported=self.use_aiter_and_is_supported)
406
407
return self.fp8_linear.apply(input=x,
408
weight=layer.weight,
@@ -528,7 +429,8 @@ class Fp8MoEMethod(FusedMoEMethodBase):
429
super().__init__(layer.moe_config)
430
self.layer = layer
431
self.quant_config = quant_config
531
- self.block_quant = self.quant_config.weight_block_size is not None
432
+ self.weight_block_size = self.quant_config.weight_block_size
433
+ self.block_quant = self.weight_block_size is not None
434
435
self.flashinfer_moe_backend: Optional[FlashinferMoeBackend] = None
436
self.fused_experts: Optional[
@@ -590,12 +492,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
492
if self.quant_config.is_checkpoint_fp8_serialized:
493
params_dtype = torch.float8_e4m3fn
494
if self.block_quant:
593
- assert self.quant_config.weight_block_size is not None
594
- layer.weight_block_size = self.quant_config.weight_block_size
495
+ assert self.weight_block_size is not None
496
+ layer.weight_block_size = self.weight_block_size
497
tp_size = get_tensor_model_parallel_world_size()
498
block_n, block_k = (
597
- self.quant_config.weight_block_size[0],
598
- self.quant_config.weight_block_size[1],
499
+ self.weight_block_size[0],
500
+ self.weight_block_size[1],
501
)
502
# NOTE: To ensure proper alignment of the block-wise quantization
503
# scales, the output_size of the weights for both the gate and up
@@ -952,7 +854,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
854
"BatchedTritonOrDeepGemmExperts(%s): "
855
"max_tokens_per_rank=%s, block_size=%s, per_act_token=%s",
856
self.__class__.__name__, max_num_tokens_per_rank,
955
- self.quant_config.weight_block_size, False)
857
+ self.weight_block_size, False)
858
return BatchedTritonOrDeepGemmExperts(
859
max_num_tokens=max_num_tokens_per_rank,
860
num_dispatchers=prepare_finalize.num_dispatchers(),
@@ -969,8 +871,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
871
else:
872
logger.debug(
873
"TritonOrDeepGemmExperts(%s): block_size=%s, per_act_token=%s",
972
- self.__class__.__name__, self.quant_config.weight_block_size,
973
- False)
874
+ self.__class__.__name__, self.weight_block_size, False)
875
return TritonOrDeepGemmExperts(
876
quant_config=self.moe_quant_config,
877
allow_deep_gemm=self.allow_deep_gemm,
@@ -988,7 +889,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
889
if self.block_quant else layer.w2_weight_scale),
890
a1_scale=layer.w13_input_scale,
891
a2_scale=layer.w2_input_scale,
991
- block_shape=self.quant_config.weight_block_size,
892
+ block_shape=self.weight_block_size,
893
)
894
895
def apply(
@@ -1046,7 +947,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
947
intermediate_size=layer.intermediate_size_per_partition,
948
expert_offset=layer.ep_rank * layer.local_num_experts,
949
local_num_experts=layer.local_num_experts,
1049
- block_shape=self.quant_config.weight_block_size,
950
+ block_shape=self.weight_block_size,
951
routed_scaling=routed_scaling_factor,
952
)
953
else:
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+220
@@ -17,6 +17,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
group_broadcast)
18
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
19
CUTLASS_BLOCK_FP8_SUPPORTED)
20
+from vllm.model_executor.parameter import (BlockQuantScaleParameter,
21
+ ChannelQuantScaleParameter,
22
+ PerTensorScaleParameter)
23
from vllm.platforms import current_platform
24
from vllm.triton_utils import tl, triton
25
from vllm.utils import cdiv, direct_register_custom_op
@@ -794,3 +797,220 @@ def requant_weight_ue8m0_inplace(
797
# Write back the results in-place.
798
w_q.copy_(w_requant)
799
s_old.copy_(s_requant)
800
+
801
+
802
+def check_aiter_fp8_linear_support() -> bool:
803
+ """AITER is only supported on ROCm and only for FP8_FNUZ
804
+ and at the moment are MI300 series"""
805
+ return (current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER
806
+ and envs.VLLM_ROCM_USE_AITER_LINEAR
807
+ and current_platform.is_fp8_fnuz())
808
+
809
+
810
+def _maybe_pad_fp8_weight(weight: torch.Tensor) -> torch.Tensor:
811
+ """Pad the weight tensor. This is an optimization on ROCm platform, which
812
+ can benefit from tensors located far enough from one another in memory"""
813
+ if (envs.VLLM_ROCM_FP8_PADDING and current_platform.is_rocm()
814
+ and weight.stride(-1) == 1
815
+ and (weight.stride(-2) * weight.element_size()) % 512 == 0):
816
+ num_pad = 256 // weight.element_size()
817
+ import torch.nn.functional as F
818
+ weight = F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
819
+ torch.cuda.empty_cache()
820
+ return weight
821
+
822
+
823
+def validate_fp8_block_shape(layer: torch.nn.Module, input_size: int,
824
+ output_size: int, input_size_per_partition: int,
825
+ output_partition_sizes: list[int],
826
+ block_size: list[int]) -> None:
827
+ """Validate block quantization shapes for tensor parallelism."""
828
+ from vllm.distributed import get_tensor_model_parallel_world_size
829
+
830
+ tp_size = getattr(layer, "tp_size", get_tensor_model_parallel_world_size())
831
+ block_n, block_k = block_size[0], block_size[1]
832
+
833
+ # Required by row parallel
834
+ if (tp_size > 1 and input_size // input_size_per_partition == tp_size
835
+ and input_size_per_partition % block_k != 0):
836
+ raise ValueError(
837
+ f"Weight input_size_per_partition = {input_size_per_partition} "
838
+ f"is not divisible by weight quantization block_k = {block_k}.")
839
+
840
+ # Required by column parallel or enabling merged weights
841
+ is_tp_split = (tp_size > 1
842
+ and output_size // sum(output_partition_sizes) == tp_size)
843
+ is_merged_gemm = len(output_partition_sizes) > 1
844
+ if is_tp_split or is_merged_gemm:
845
+ sizes_to_check = output_partition_sizes
846
+ if not is_tp_split and is_merged_gemm:
847
+ # In case of merged matrices, we allow the last
848
+ # matrix to not be a multiple of block size
849
+ sizes_to_check = output_partition_sizes[:-1]
850
+ for output_partition_size in sizes_to_check:
851
+ if output_partition_size % block_n != 0:
852
+ raise ValueError(
853
+ f"Weight output_partition_size = "
854
+ f"{output_partition_size} is not divisible by "
855
+ f"weight quantization block_n = {block_n}.")
856
+
857
+
858
+def create_fp8_weight_parameter(
859
+ output_size_per_partition: int, input_size_per_partition: int,
860
+ weight_loader: Optional[Callable]) -> torch.nn.Parameter:
861
+ """Create FP8 weight parameter."""
862
+ from vllm.model_executor.parameter import ModelWeightParameter
863
+
864
+ return ModelWeightParameter(data=torch.empty(output_size_per_partition,
865
+ input_size_per_partition,
866
+ dtype=torch.float8_e4m3fn),
867
+ input_dim=1,
868
+ output_dim=0,
869
+ weight_loader=weight_loader)
870
+
871
+
872
+def create_fp8_scale_parameter(
873
+ parameter_type: torch.nn.Parameter, output_partition_sizes: list[int],
874
+ input_size_per_partition: int, block_size: Optional[list[int]],
875
+ weight_loader: Optional[Callable]) -> torch.nn.Parameter:
876
+ """Create scale parameter based on quantization strategy."""
877
+ if parameter_type == ChannelQuantScaleParameter:
878
+ scale = parameter_type(data=torch.empty(
879
+ (sum(output_partition_sizes), 1), dtype=torch.float32),
880
+ output_dim=0,
881
+ weight_loader=weight_loader)
882
+ elif parameter_type == BlockQuantScaleParameter:
883
+ assert block_size is not None
884
+ block_n, block_k = block_size[0], block_size[1]
885
+ output_size_per_partition = sum(output_partition_sizes)
886
+ scale = parameter_type(
887
+ data=torch.empty(
888
+ (output_size_per_partition + block_n - 1) // block_n,
889
+ (input_size_per_partition + block_k - 1) // block_k,
890
+ dtype=torch.float32,
891
+ ),
892
+ input_dim=1,
893
+ output_dim=0,
894
+ weight_loader=weight_loader,
895
+ )
896
+ elif parameter_type == PerTensorScaleParameter:
897
+ scale = parameter_type(data=torch.empty(len(output_partition_sizes),
898
+ dtype=torch.float32),
899
+ weight_loader=weight_loader)
900
+ else:
901
+ raise ValueError(f"Unknown parameter type: {parameter_type}")
902
+
903
+ scale[:] = torch.finfo(torch.float32).min
904
+ return scale
905
+
906
+
907
+def create_fp8_input_scale(
908
+ output_partition_sizes: list[int],
909
+ weight_loader: Optional[Callable]) -> torch.nn.Parameter:
910
+ """Create input scale parameter for static activation quantization."""
911
+ from vllm.model_executor.parameter import PerTensorScaleParameter
912
+
913
+ scale = PerTensorScaleParameter(data=torch.empty(
914
+ len(output_partition_sizes), dtype=torch.float32),
915
+ weight_loader=weight_loader)
916
+ scale[:] = torch.finfo(torch.float32).min
917
+ return scale
918
+
919
+
920
+def process_fp8_weight_tensor_strategy(
921
+ weight: torch.Tensor,
922
+ weight_scale: torch.Tensor,
923
+ logical_widths: list[int],
924
+ input_scale: Optional[torch.Tensor] = None
925
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
926
+ """Process weights for tensor-wise quantization strategy."""
927
+ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
928
+ normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale)
929
+
930
+ if current_platform.is_fp8_fnuz():
931
+ weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
932
+ weight=weight, weight_scale=weight_scale, input_scale=input_scale)
933
+
934
+ # Requantize with max scale
935
+ weight_scale, weight = requantize_with_max_scale(
936
+ weight=weight,
937
+ weight_scale=weight_scale,
938
+ logical_widths=logical_widths,
939
+ )
940
+
941
+ weight = _maybe_pad_fp8_weight(weight)
942
+ return weight, weight_scale, input_scale
943
+
944
+
945
+def process_fp8_weight_channel_strategy(
946
+ weight: torch.Tensor,
947
+ weight_scale: torch.Tensor,
948
+ input_scale: Optional[torch.Tensor] = None
949
+) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
950
+ """Process weights for channel-wise quantization strategy."""
951
+ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
952
+ normalize_e4m3fn_to_e4m3fnuz)
953
+
954
+ if current_platform.is_fp8_fnuz():
955
+ weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
956
+ weight=weight, weight_scale=weight_scale, input_scale=input_scale)
957
+
958
+ return weight, weight_scale, input_scale
959
+
960
+
961
+def process_fp8_weight_block_strategy(
962
+ weight: torch.Tensor,
963
+ weight_scale: torch.Tensor,
964
+) -> tuple[torch.Tensor, torch.Tensor]:
965
+ """Process weights for block-wise quantization strategy."""
966
+ from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
967
+ normalize_e4m3fn_to_e4m3fnuz)
968
+
969
+ if current_platform.is_fp8_fnuz():
970
+ weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz(
971
+ weight=weight, weight_scale=weight_scale)
972
+
973
+ weight = _maybe_pad_fp8_weight(weight)
974
+ return weight, weight_scale
975
+
976
+
977
+def maybe_post_process_fp8_weight_block(layer: torch.nn.Module,
978
+ cutlass_block_fp8_supported: bool):
979
+ assert layer.weight_block_size is not None
980
+
981
+ from vllm.utils.deep_gemm import (is_deep_gemm_e8m0_used,
982
+ should_use_deepgemm_for_fp8_linear)
983
+
984
+ # On Blackwell or Hopper, if E8M0 for DeepGemm is used, we need to
985
+ # requantize the weight and input to the specific scale
986
+ # at the same time.
987
+ if is_deep_gemm_e8m0_used():
988
+ block_sz = tuple(layer.weight_block_size)
989
+ requant_weight_ue8m0_inplace(layer.weight.data,
990
+ layer.weight_scale.data, block_sz)
991
+ # SM90 Block FP8 CUTLASS requires row-major weight scales
992
+ elif (current_platform.is_device_capability(90)
993
+ and cutlass_block_fp8_supported
994
+ and not should_use_deepgemm_for_fp8_linear(torch.bfloat16,
995
+ layer.weight)):
996
+ layer.weight_scale = torch.nn.Parameter(
997
+ layer.weight_scale.data.T.contiguous(), requires_grad=False)
998
+
999
+
1000
+def apply_fp8_block_linear(layer: torch.nn.Module, input: torch.Tensor,
1001
+ bias: Optional[torch.Tensor],
1002
+ cutlass_block_fp8_supported: bool,
1003
+ use_aiter_and_is_supported: bool) -> torch.Tensor:
1004
+ """Apply block-wise FP8 linear operation."""
1005
+ assert layer.weight_block_size is not None
1006
+
1007
+ return torch.ops.vllm.apply_w8a8_block_fp8_linear(
1008
+ input=input,
1009
+ weight=layer.weight,
1010
+ block_size=layer.weight_block_size,
1011
+ weight_scale=layer.weight_scale,
1012
+ input_scale=layer.input_scale,
1013
+ bias=bias,
1014
+ cutlass_block_fp8_supported=cutlass_block_fp8_supported,
1015
+ use_aiter_and_is_supported=use_aiter_and_is_supported,
1016
+ )