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 5, 08:30 UTC
[CT][FP8][Marlin] refactor CompressedTensorsW8A16Fp8 to use kernel abstraction (#38244)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com> Signed-off-by: Kunshang Ji <jikunshang95@gmail.com>
Kunshang Ji committed
Apr 10, 2026 at 09:58 UTC
55d037e2e5cc56c38a1a4a77a15c347fee380c50
4 files changed
+127
-56
vllm/model_executor/kernels/linear/__init__.py
+51
@@ -177,6 +177,21 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
177
],
178
}
179
180
+_POSSIBLE_WFP8A16_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = {
181
+ PlatformEnum.CUDA: [
182
+ MarlinFP8ScaledMMLinearKernel,
183
+ ],
184
+ PlatformEnum.ROCM: [
185
+ # To be added
186
+ ],
187
+ PlatformEnum.CPU: [
188
+ # To be added
189
+ ],
190
+ PlatformEnum.XPU: [
191
+ # To be added
192
+ ],
193
+}
194
+
195
# in priority/performance order (when available)
196
_POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
197
PlatformEnum.CUDA: [
@@ -463,6 +478,41 @@ def choose_mp_linear_kernel(
478
)
479
480
481
+def init_wfp8_a16_linear_kernel(
482
+ weight_quant_key: QuantKey,
483
+ activation_quant_key: QuantKey,
484
+ weight_shape: tuple[int, int],
485
+ input_dtype: torch.dtype,
486
+ out_dtype: torch.dtype,
487
+ force_kernel: type[FP8ScaledMMLinearKernel] | None = None,
488
+ module_name: str | None = None,
489
+) -> FP8ScaledMMLinearKernel:
490
+ config = FP8ScaledMMLinearLayerConfig(
491
+ weight_quant_key=weight_quant_key,
492
+ activation_quant_key=activation_quant_key,
493
+ weight_shape=weight_shape,
494
+ input_dtype=input_dtype,
495
+ out_dtype=out_dtype,
496
+ )
497
+
498
+ kernel_type = choose_scaled_mm_linear_kernel(
499
+ config, _POSSIBLE_WFP8A16_KERNELS, force_kernel=force_kernel
500
+ )
501
+
502
+ if module_name:
503
+ logger.info_once(
504
+ "Selected %s for %s",
505
+ kernel_type.__name__,
506
+ module_name,
507
+ scope="global",
508
+ )
509
+
510
+ return kernel_type(
511
+ config,
512
+ layer_param_names=["weight", "weight_scale", "input_scale", "input_scale_ub"],
513
+ )
514
+
515
+
516
# Maps VLLM_NVFP4_GEMM_BACKEND env var values to kernel classes.
517
_NVFP4_BACKEND_TO_KERNEL: dict[str, type[NvFp4LinearKernel]] = {
518
"flashinfer-cutlass": FlashInferCutlassNvFp4LinearKernel,
@@ -588,6 +638,7 @@ __all__ = [
638
"init_nvfp4_linear_kernel",
639
"choose_mp_linear_kernel",
640
"register_linear_kernel",
641
+ "init_wfp8_a16_linear_kernel",
642
"FP8ScaledMMLinearKernel",
643
"Int8ScaledMMLinearKernel",
644
"ScaledMMLinearKernel",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a16_fp8.py
+44
-44
@@ -6,45 +6,49 @@ from collections.abc import Callable
6
import torch
7
from compressed_tensors.quantization import QuantizationArgs, QuantizationStrategy
8
9
+from vllm.config import get_current_vllm_config
10
+from vllm.model_executor.kernels.linear import (
11
+ init_wfp8_a16_linear_kernel,
12
+)
13
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
14
CompressedTensorsScheme,
15
)
16
+from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
17
+ STRATEGY_TO_PARAMETER_TYPE,
18
+ STRATEGY_TO_WEIGHT_QUANT_KEY,
19
+)
20
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
21
create_fp8_scale_parameter,
22
create_fp8_weight_parameter,
15
- process_fp8_weight_block_strategy,
23
validate_fp8_block_shape,
24
)
18
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
19
- apply_fp8_marlin_linear,
20
- prepare_fp8_layer_for_marlin,
25
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
26
+ kFp8DynamicTensorSym,
27
+ kFp8StaticTensorSym,
28
)
29
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
30
convert_to_channelwise,
31
)
25
-from vllm.model_executor.parameter import (
26
- BlockQuantScaleParameter,
27
- ChannelQuantScaleParameter,
28
- PerTensorScaleParameter,
29
-)
32
+from vllm.model_executor.parameter import PerTensorScaleParameter
33
from vllm.model_executor.utils import replace_parameter
34
35
__all__ = ["CompressedTensorsW8A16Fp8"]
36
34
-strategy_to_parameter_type = {
35
- QuantizationStrategy.BLOCK: BlockQuantScaleParameter,
36
- QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter,
37
- QuantizationStrategy.TENSOR: PerTensorScaleParameter,
38
-}
39
-
37
38
class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
39
def __init__(self, weight_quant: QuantizationArgs, 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.input_dtype = get_current_vllm_config().model_config.dtype
44
self.is_static_input_scheme = is_static_input_scheme
45
self.weight_block_size = self.weight_quant.block_structure
46
47
+ self.weight_quant_key = STRATEGY_TO_WEIGHT_QUANT_KEY[self.strategy]
48
+ self.activation_quant_key = (
49
+ kFp8StaticTensorSym if is_static_input_scheme else kFp8DynamicTensorSym
50
+ )
51
+
52
@classmethod
53
def get_min_capability(cls) -> int:
54
# turing and up
@@ -89,7 +93,7 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
93
94
# WEIGHT SCALE
95
weight_scale = create_fp8_scale_parameter(
92
- strategy_to_parameter_type[self.strategy],
96
+ STRATEGY_TO_PARAMETER_TYPE[self.strategy],
97
output_partition_sizes,
98
input_size_per_partition,
99
layer.weight_block_size,
@@ -105,32 +109,36 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
109
)
110
layer.register_parameter("input_scale", input_scale)
111
112
+ self.linear_kernel = init_wfp8_a16_linear_kernel(
113
+ weight_quant_key=self.weight_quant_key,
114
+ activation_quant_key=self.activation_quant_key,
115
+ weight_shape=layer.weight.shape,
116
+ input_dtype=self.input_dtype,
117
+ out_dtype=self.out_dtype,
118
+ )
119
+
120
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.
121
if self.strategy == QuantizationStrategy.BLOCK:
122
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
- )
123
+ # MarlinFP8ScaledMMLinearKernel uses "weight_scale_inv" for block
124
+ # quant, while CT registers the scale as "weight_scale".
125
+ # Rename by deleting the old parameter and adding the new one so
126
+ # that prepare_fp8_layer_for_marlin (which prefers "weight_scale"
127
+ # over "weight_scale_inv") picks up "weight_scale_inv" correctly.
128
+ weight_scale_data = layer.weight_scale.data
129
+ del layer._parameters["weight_scale"]
130
+ replace_parameter(layer, "weight_scale_inv", weight_scale_data)
131
else:
120
- # Weights must be transposed for marlin
121
- weight = weight.t()
132
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
133
+ # For fused modules with per-tensor scales, expand each scale
134
+ # to its shard's channels.
135
+ replace_parameter(
136
+ layer,
137
+ "weight_scale",
138
+ convert_to_channelwise(layer.weight_scale, layer.logical_widths),
139
)
140
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)
141
+ self.linear_kernel.process_weights_after_loading(layer)
142
143
def apply_weights(
144
self,
@@ -138,12 +146,4 @@ class CompressedTensorsW8A16Fp8(CompressedTensorsScheme):
146
x: torch.Tensor,
147
bias: torch.Tensor | None = None,
148
) -> torch.Tensor:
141
- return apply_fp8_marlin_linear(
142
- input=x,
143
- weight=layer.weight,
144
- weight_scale=layer.weight_scale,
145
- workspace=layer.workspace,
146
- size_n=layer.output_size_per_partition,
147
- size_k=layer.input_size_per_partition,
148
- bias=bias,
149
- )
149
+ return self.linear_kernel.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+4
-12
@@ -16,6 +16,9 @@ from vllm.model_executor.kernels.linear import (
16
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
17
CompressedTensorsScheme,
18
)
19
+from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
20
+ STRATEGY_TO_PARAMETER_TYPE,
21
+)
22
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
23
create_fp8_input_scale,
24
create_fp8_scale_parameter,
@@ -34,20 +37,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
37
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
38
cutlass_block_fp8_supported,
39
)
37
-from vllm.model_executor.parameter import (
38
- BlockQuantScaleParameter,
39
- ChannelQuantScaleParameter,
40
- PerTensorScaleParameter,
41
-)
40
41
__all__ = ["CompressedTensorsW8A8Fp8"]
42
45
-strategy_to_parameter_type = {
46
- QuantizationStrategy.BLOCK: BlockQuantScaleParameter,
47
- QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter,
48
- QuantizationStrategy.TENSOR: PerTensorScaleParameter,
49
-}
50
-
43
STATIC_QUANT = True
44
DYNAMIC_QUANT = False
45
activation_quant_key_mapping = {
@@ -130,7 +122,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
122
123
# WEIGHT SCALE
124
weight_scale = create_fp8_scale_parameter(
133
- strategy_to_parameter_type[self.strategy],
125
+ STRATEGY_TO_PARAMETER_TYPE[self.strategy],
126
output_partition_sizes,
127
input_size_per_partition,
128
layer.weight_block_size,
vllm/model_executor/layers/quantization/compressed_tensors/utils.py
+28
@@ -6,8 +6,36 @@ from types import MappingProxyType
6
7
import regex as re
8
from compressed_tensors import CompressionFormat
9
+from compressed_tensors.quantization import QuantizationStrategy
10
from torch.nn import Module
11
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ kFp8Static128BlockSym,
14
+ kFp8StaticChannelSym,
15
+ kFp8StaticTensorSym,
16
+)
17
+from vllm.model_executor.parameter import (
18
+ BlockQuantScaleParameter,
19
+ ChannelQuantScaleParameter,
20
+ PerTensorScaleParameter,
21
+)
22
+
23
+# Maps quantization strategy to the corresponding scale parameter type.
24
+# Shared across compressed-tensor scheme classes (w8a16_fp8, w8a8_fp8, …).
25
+STRATEGY_TO_PARAMETER_TYPE = {
26
+ QuantizationStrategy.BLOCK: BlockQuantScaleParameter,
27
+ QuantizationStrategy.CHANNEL: ChannelQuantScaleParameter,
28
+ QuantizationStrategy.TENSOR: PerTensorScaleParameter,
29
+}
30
+
31
+# Maps quantization strategy to the vLLM weight-quant key used for
32
+# kernel selection. Shared across compressed-tensor scheme classes.
33
+STRATEGY_TO_WEIGHT_QUANT_KEY = {
34
+ QuantizationStrategy.BLOCK: kFp8Static128BlockSym,
35
+ QuantizationStrategy.CHANNEL: kFp8StaticChannelSym,
36
+ QuantizationStrategy.TENSOR: kFp8StaticTensorSym,
37
+}
38
+
39
40
def is_activation_quantization_format(format: str) -> bool:
41
_ACTIVATION_QUANTIZATION_FORMATS = [