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 6, 04:30 UTC
[FP8]add FP8 WoQ kernel abstraction. (#32929)
Signed-off-by: Zhu, Zufang <zufang.zhu@intel.com>
Kunshang Ji committed
Mar 23, 2026 at 17:47 UTC
27d5ee3e6f392ed8eb4d2ceb956f6129547c0e6b
5 files changed
+177
-83
vllm/model_executor/kernels/linear/__init__.py
+4
@@ -72,6 +72,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
72
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
73
FlashInferFP8ScaledMMLinearKernel,
74
)
75
+from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
76
+ MarlinFP8ScaledMMLinearKernel,
77
+)
78
from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
79
ChannelWiseTorchFP8ScaledMMLinearKernel,
80
PerTensorTorchFP8ScaledMMLinearKernel,
@@ -104,6 +107,7 @@ _POSSIBLE_INT8_KERNELS: dict[PlatformEnum, list[type[Int8ScaledMMLinearKernel]]]
107
# in priority/performance order (when available)
108
_POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] = {
109
PlatformEnum.CUDA: [
110
+ MarlinFP8ScaledMMLinearKernel,
111
FlashInferFP8ScaledMMLinearKernel,
112
CutlassFP8ScaledMMLinearKernel,
113
PerTensorTorchFP8ScaledMMLinearKernel,
vllm/model_executor/kernels/linear/scaled_mm/__init__.py
+4
@@ -14,6 +14,9 @@ from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
14
from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
15
FlashInferFP8ScaledMMLinearKernel,
16
)
17
+from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
18
+ MarlinFP8ScaledMMLinearKernel,
19
+)
20
from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
21
ChannelWiseTorchFP8ScaledMMLinearKernel,
22
PerTensorTorchFP8ScaledMMLinearKernel,
@@ -46,6 +49,7 @@ __all__ = [
49
"CutlassFP8ScaledMMLinearKernel",
50
"CutlassInt8ScaledMMLinearKernel",
51
"FlashInferFP8ScaledMMLinearKernel",
52
+ "MarlinFP8ScaledMMLinearKernel",
53
"ChannelWiseTorchFP8ScaledMMLinearKernel",
54
"PerTensorTorchFP8ScaledMMLinearKernel",
55
"RowWiseTorchFP8ScaledMMLinearKernel",
vllm/model_executor/kernels/linear/scaled_mm/marlin.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Sequence
5
+
6
+import torch
7
+
8
+import vllm.envs as envs
9
+from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
10
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
11
+ process_fp8_weight_block_strategy,
12
+)
13
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
14
+ apply_fp8_marlin_linear,
15
+ is_fp8_marlin_supported,
16
+ prepare_fp8_layer_for_marlin,
17
+)
18
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
+ kFp8Static128BlockSym,
20
+)
21
+from vllm.model_executor.utils import replace_parameter
22
+from vllm.platforms import current_platform
23
+
24
+from .ScaledMMLinearKernel import (
25
+ FP8ScaledMMLinearKernel,
26
+ FP8ScaledMMLinearLayerConfig,
27
+)
28
+
29
+
30
+class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
31
+ """
32
+ FP8 Marlin kernel for GPUs that lack FP8 hardware support.
33
+ Leverages the Marlin kernel for fast weight-only FP8 quantization.
34
+ """
35
+
36
+ @classmethod
37
+ def is_supported(
38
+ cls, compute_capability: int | None = None
39
+ ) -> tuple[bool, str | None]:
40
+ if not current_platform.is_cuda():
41
+ return False, "requires CUDA."
42
+ # Check if platform supports FP8 Marlin
43
+ if not is_fp8_marlin_supported():
44
+ return False, "FP8 Marlin requires compute capability 7.5 or higher"
45
+ if vllm_is_batch_invariant():
46
+ return False, "FP8 Marlin not supported for batch invariant execution."
47
+ if (
48
+ compute_capability is not None
49
+ and compute_capability >= 89
50
+ and not envs.VLLM_TEST_FORCE_FP8_MARLIN
51
+ ):
52
+ return (
53
+ False,
54
+ "To apply FP8 Marlin on high-capability GPUs, please set "
55
+ "VLLM_TEST_FORCE_FP8_MARLIN=1",
56
+ )
57
+ return True, None
58
+
59
+ @classmethod
60
+ def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
61
+ return True, None
62
+
63
+ def __init__(
64
+ self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
65
+ ) -> None:
66
+ super().__init__(c, layer_param_names)
67
+ self.marlin_input_dtype = None
68
+ self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym}
69
+ self.size_k_first = not self.block_quant
70
+
71
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
72
+ if self.block_quant:
73
+ weight, weight_scale_inv = process_fp8_weight_block_strategy(
74
+ layer.weight, layer.weight_scale_inv
75
+ )
76
+ # Update layer with new values
77
+ replace_parameter(layer, "weight", weight.data)
78
+ replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
79
+ else:
80
+ weight = layer.weight.t()
81
+ replace_parameter(layer, "weight", weight.data)
82
+ layer.input_scale = None
83
+ prepare_fp8_layer_for_marlin(
84
+ layer, self.size_k_first, input_dtype=self.marlin_input_dtype
85
+ )
86
+ del layer.input_scale
87
+
88
+ def apply_weights(
89
+ self,
90
+ layer: torch.nn.Module,
91
+ x: torch.Tensor,
92
+ bias: torch.Tensor | None = None,
93
+ ) -> torch.Tensor:
94
+ if self.block_quant:
95
+ weight_scale = layer.weight_scale_inv
96
+ else:
97
+ weight_scale = layer.weight_scale
98
+ return apply_fp8_marlin_linear(
99
+ input=x,
100
+ weight=layer.weight,
101
+ weight_scale=weight_scale,
102
+ workspace=layer.workspace,
103
+ size_n=layer.output_size_per_partition,
104
+ size_k=layer.input_size_per_partition,
105
+ input_dtype=self.marlin_input_dtype,
106
+ bias=bias,
107
+ )
108
+
109
+ def apply_scaled_mm(
110
+ self,
111
+ *,
112
+ A: torch.Tensor,
113
+ B: torch.Tensor,
114
+ out_dtype: torch.dtype,
115
+ As: torch.Tensor,
116
+ Bs: torch.Tensor,
117
+ bias: torch.Tensor | None,
118
+ output_shape: list,
119
+ ) -> torch.Tensor:
120
+ pass
vllm/model_executor/layers/quantization/fbgemm_fp8.py
-12
@@ -22,7 +22,6 @@ from vllm.model_executor.layers.quantization.base_config import (
22
QuantizeMethodBase,
23
)
24
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
25
- apply_fp8_marlin_linear,
25
prepare_fp8_layer_for_marlin,
26
)
27
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -177,15 +176,4 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
176
x: torch.Tensor,
177
bias: torch.Tensor | None = None,
178
) -> torch.Tensor:
180
- if self.quant_config.use_marlin:
181
- return apply_fp8_marlin_linear(
182
- input=x,
183
- weight=layer.weight,
184
- weight_scale=layer.weight_scale,
185
- workspace=layer.workspace,
186
- size_n=layer.output_size_per_partition,
187
- size_k=layer.input_size_per_partition,
188
- bias=bias,
189
- )
190
-
179
return self.fp8_linear.apply_weights(layer, x, bias)
vllm/model_executor/layers/quantization/fp8.py
+49
-71
@@ -7,7 +7,6 @@ import torch
7
from torch.nn import Module
8
from torch.utils._python_dispatch import TorchDispatchMode
9
10
-import vllm.envs as envs
10
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm import _custom_ops as ops
12
from vllm._aiter_ops import rocm_aiter_ops
@@ -16,6 +15,7 @@ from vllm.logger import init_logger
15
from vllm.model_executor.kernels.linear import (
16
init_fp8_linear_kernel,
17
)
18
+from vllm.model_executor.kernels.linear.scaled_mm import MarlinFP8ScaledMMLinearKernel
19
from vllm.model_executor.layers.attention import Attention
20
from vllm.model_executor.layers.batch_invariant import (
21
vllm_is_batch_invariant,
@@ -61,10 +61,6 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
61
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
62
get_marlin_input_dtype,
63
)
64
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
65
- apply_fp8_marlin_linear,
66
- prepare_fp8_layer_for_marlin,
67
-)
64
from vllm.model_executor.layers.quantization.utils.quant_utils import (
65
GroupShape,
66
is_layer_skipped,
@@ -280,15 +276,6 @@ class Fp8LinearMethod(LinearMethodBase):
276
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
277
# kernel for fast weight-only FP8 quantization
278
self.marlin_input_dtype = None
283
- self.use_marlin = (
284
- not current_platform.has_device_capability(89)
285
- or envs.VLLM_TEST_FORCE_FP8_MARLIN
286
- )
287
- # Disable marlin for rocm
288
- if current_platform.is_rocm() or current_platform.is_xpu():
289
- self.use_marlin = False
290
- if vllm_is_batch_invariant():
291
- self.use_marlin = False
279
280
self.use_aiter_and_is_supported = rocm_aiter_ops.is_linear_fp8_enabled()
281
self.use_deep_gemm = is_deep_gemm_supported()
@@ -297,7 +284,28 @@ class Fp8LinearMethod(LinearMethodBase):
284
self.block_quant = self.weight_block_size is not None
285
self.act_q_static = self.quant_config.activation_scheme == "static"
286
287
+ # Use per-token quantization for better perf if dynamic and cutlass
288
+ if self.act_q_static:
289
+ activation_quant_key = kFp8StaticTensorSym
290
+ elif cutlass_fp8_supported():
291
+ activation_quant_key = kFp8DynamicTokenSym
292
+ else:
293
+ activation_quant_key = kFp8DynamicTensorSym
294
+
295
if self.block_quant:
296
+ weight_quant_key = kFp8Static128BlockSym
297
+ else:
298
+ weight_quant_key = kFp8StaticTensorSym
299
+
300
+ self.fp8_linear = init_fp8_linear_kernel(
301
+ activation_quant_key=activation_quant_key,
302
+ weight_quant_key=weight_quant_key,
303
+ out_dtype=torch.get_default_dtype(),
304
+ module_name=self.__class__.__name__,
305
+ )
306
+ self.use_marlin = isinstance(self.fp8_linear, MarlinFP8ScaledMMLinearKernel)
307
+
308
+ if self.block_quant and not self.use_marlin:
309
assert not self.act_q_static
310
assert self.weight_block_size is not None
311
self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
@@ -306,21 +314,6 @@ class Fp8LinearMethod(LinearMethodBase):
314
cutlass_block_fp8_supported=self.cutlass_block_fp8_supported,
315
use_aiter_and_is_supported=self.use_aiter_and_is_supported,
316
)
309
- else:
310
- # Use per-token quantization for better perf if dynamic and cutlass
311
- if self.act_q_static:
312
- activation_quant_key = kFp8StaticTensorSym
313
- elif cutlass_fp8_supported():
314
- activation_quant_key = kFp8DynamicTokenSym
315
- else:
316
- activation_quant_key = kFp8DynamicTensorSym
317
-
318
- self.fp8_linear = init_fp8_linear_kernel(
319
- activation_quant_key=activation_quant_key,
320
- weight_quant_key=kFp8StaticTensorSym,
321
- out_dtype=torch.get_default_dtype(),
322
- module_name=self.__class__.__name__,
323
- )
317
318
def create_weights(
319
self,
@@ -387,12 +380,18 @@ class Fp8LinearMethod(LinearMethodBase):
380
layer.register_parameter("input_scale", scale)
381
382
def process_weights_after_loading(self, layer: Module) -> None:
390
- size_k_first = True
383
+ if self.use_marlin:
384
+ # Only Marlin kernels support `marlin_input_dtype`; guard to avoid
385
+ # AttributeError if backend selection changes.
386
+ if hasattr(self.fp8_linear, "marlin_input_dtype"):
387
+ self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
388
+ self.fp8_linear.process_weights_after_loading(layer)
389
+ return
390
+
391
input_scale = None
392
# TODO(rob): refactor block quant into separate class.
393
if self.block_quant:
394
assert not self.act_q_static
395
- size_k_first = False
395
396
weight, weight_scale_inv = process_fp8_weight_block_strategy(
397
layer.weight, layer.weight_scale_inv
@@ -411,16 +410,15 @@ class Fp8LinearMethod(LinearMethodBase):
410
411
# If using w8a8, torch._scaled_mm needs per tensor, so
412
# requantize the logical shards as a single weight.
414
- if not self.use_marlin:
415
- weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
416
- weight,
417
- weight_scale,
418
- layer.logical_widths,
419
- getattr(layer, "input_scale", None),
420
- )
421
- if self.act_q_static:
422
- assert input_scale is not None
423
- input_scale = input_scale.max()
413
+ weight, weight_scale, input_scale = process_fp8_weight_tensor_strategy(
414
+ weight,
415
+ weight_scale,
416
+ layer.logical_widths,
417
+ getattr(layer, "input_scale", None),
418
+ )
419
+ if self.act_q_static:
420
+ assert input_scale is not None
421
+ input_scale = input_scale.max()
422
weight = weight.t()
423
424
# Update layer with new values.
@@ -432,14 +430,6 @@ class Fp8LinearMethod(LinearMethodBase):
430
else:
431
layer.input_scale = None
432
435
- if self.use_marlin:
436
- prepare_fp8_layer_for_marlin(
437
- layer, size_k_first, input_dtype=self.marlin_input_dtype
438
- )
439
- # Activations not quantized for marlin.
440
- del layer.input_scale
441
- return
442
-
433
if self.block_quant:
434
maybe_post_process_fp8_weight_block(layer)
435
@@ -486,21 +476,7 @@ class Fp8LinearMethod(LinearMethodBase):
476
return torch.nn.functional.linear(x, weight_bf16.t(), bias)
477
478
if self.use_marlin:
489
- if self.block_quant:
490
- weight_scale = layer.weight_scale_inv
491
- else:
492
- weight_scale = layer.weight_scale
493
-
494
- return apply_fp8_marlin_linear(
495
- input=x,
496
- weight=layer.weight,
497
- weight_scale=weight_scale,
498
- workspace=layer.workspace,
499
- size_n=layer.output_size_per_partition,
500
- size_k=layer.input_size_per_partition,
501
- input_dtype=self.marlin_input_dtype,
502
- bias=bias,
503
- )
479
+ return self.fp8_linear.apply_weights(layer, x, bias)
480
481
if self.block_quant:
482
assert self.weight_block_size is not None
@@ -623,18 +599,20 @@ class Fp8OnlineLinearMethod(Fp8LinearMethod):
599
600
layer.input_scale = None
601
qweight, weight_scale = ops.scaled_fp8_quant(layer.weight, scale=None)
626
- weight = qweight.t()
602
603
# Update layer with new values.
629
- replace_parameter(layer, "weight", weight.data)
604
+ replace_parameter(layer, "weight", qweight.data)
605
replace_parameter(layer, "weight_scale", weight_scale.data)
606
607
if self.use_marlin:
633
- size_k_first = True
634
- prepare_fp8_layer_for_marlin(
635
- layer, size_k_first, input_dtype=self.marlin_input_dtype
636
- )
637
- # Activations not quantized for marlin.
608
+ # Only Marlin kernels support `marlin_input_dtype`; guard to avoid
609
+ # AttributeError if backend selection changes.
610
+ if hasattr(self.fp8_linear, "marlin_input_dtype"):
611
+ self.fp8_linear.marlin_input_dtype = self.marlin_input_dtype
612
+ self.fp8_linear.process_weights_after_loading(layer)
613
+ else:
614
+ weight = qweight.t()
615
+ replace_parameter(layer, "weight", weight.data)
616
617
# Prevent duplicate processing (e.g., during weight reload)
618
layer._already_called_process_weights_after_loading = True