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 15, 02:15 UTC
[Perf][fp8] Use CustomOp abstraction for fp8 quant for better perf (#19830)
Signed-off-by: Luka Govedic <lgovedic@redhat.com> Co-authored-by: mgoin <mgoin64@gmail.com>
Luka Govedič committed
Jul 11, 2025 at 00:56 UTC
31d5c1797f320b2f407c893673330b3a8766ae47
18 files changed
+368
-104
benchmarks/kernels/bench_per_token_quant_fp8.py
new
+98
@@ -0,0 +1,98 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import itertools
4
+from typing import Callable
5
+
6
+import torch
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config
10
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
11
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
12
+from vllm.triton_utils import triton
13
+
14
+
15
+# TODO(luka): use standalone_compile utility
16
+def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int):
17
+ def inner(*args):
18
+ torch._dynamo.mark_dynamic(args[arg_index], dim_index)
19
+ return fn(*args)
20
+
21
+ return inner
22
+
23
+
24
+torch._dynamo.config.recompile_limit = 8888
25
+compilation_config = CompilationConfig(custom_ops=["none"])
26
+with set_current_vllm_config(VllmConfig(compilation_config=compilation_config)):
27
+ torch_per_token_quant_fp8 = torch.compile(
28
+ QuantFP8(False, GroupShape.PER_TOKEN),
29
+ fullgraph=True,
30
+ dynamic=False, # recompile for different shapes
31
+ )
32
+
33
+ # First dim is explicitly dynamic to simulate vLLM usage
34
+ torch_per_token_quant_fp8 = with_dyn_arg(torch_per_token_quant_fp8, 0, 0)
35
+
36
+
37
+def cuda_per_token_quant_fp8(
38
+ input: torch.Tensor,
39
+) -> tuple[torch.Tensor, torch.Tensor]:
40
+ return ops.scaled_fp8_quant(input)
41
+
42
+
43
+def calculate_diff(batch_size: int, seq_len: int):
44
+ """Calculate difference between Triton and CUDA implementations."""
45
+ device = torch.device("cuda")
46
+ x = torch.rand((batch_size * seq_len, 4096), dtype=torch.float16, device=device)
47
+
48
+ torch_out, torch_scale = torch_per_token_quant_fp8(x)
49
+ cuda_out, cuda_scale = cuda_per_token_quant_fp8(x)
50
+
51
+ if torch.allclose(
52
+ cuda_out.to(torch.float32), torch_out.to(torch.float32), rtol=1e-3, atol=1e-5
53
+ ) and torch.allclose(cuda_scale, torch_scale, rtol=1e-3, atol=1e-5):
54
+ print("✅ All implementations match")
55
+ else:
56
+ print("❌ Implementations differ")
57
+
58
+
59
+batch_size_range = [1, 16, 32, 64, 128]
60
+seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
61
+
62
+configs = list(itertools.product(batch_size_range, seq_len_range))
63
+
64
+
65
+@triton.testing.perf_report(
66
+ triton.testing.Benchmark(
67
+ x_names=["batch_size", "seq_len"],
68
+ x_vals=configs,
69
+ line_arg="provider",
70
+ line_vals=["torch", "cuda"],
71
+ line_names=["Torch", "CUDA"],
72
+ styles=[("blue", "-"), ("green", "-")],
73
+ ylabel="us",
74
+ plot_name="per-token-dynamic-quant-fp8-performance",
75
+ args={},
76
+ )
77
+)
78
+def benchmark_quantization(batch_size, seq_len, provider):
79
+ dtype = torch.float16
80
+ device = torch.device("cuda")
81
+
82
+ x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
83
+
84
+ quantiles = [0.5, 0.2, 0.8]
85
+
86
+ if provider == "torch":
87
+ fn = lambda: torch_per_token_quant_fp8(x.clone())
88
+ elif provider == "cuda":
89
+ fn = lambda: cuda_per_token_quant_fp8(x.clone())
90
+
91
+ ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
92
+
93
+ return 1000 * ms, 1000 * max_ms, 1000 * min_ms
94
+
95
+
96
+if __name__ == "__main__":
97
+ calculate_diff(batch_size=4, seq_len=4096)
98
+ benchmark_quantization.run(print_data=True)
tests/compile/test_fusion.py
+7
-4
@@ -44,7 +44,9 @@ class TestModel(torch.nn.Module):
44
]
45
self.fp8_linear = Fp8LinearOp(
46
cutlass_fp8_supported=cutlass_fp8_enabled,
47
- use_per_token_if_dynamic=True)
47
+ act_quant_static=static,
48
+ act_quant_group_shape=group_shape,
49
+ )
50
51
def forward(self, x):
52
resid = torch.sqrt(x)
@@ -91,9 +93,10 @@ def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
93
maybe_create_device_identity() # needed for certain non-cutlass fp8 paths
94
95
vllm_config = VllmConfig(compilation_config=CompilationConfig(
94
- level=CompilationLevel.PIECEWISE, custom_ops=["+rms_norm"]))
95
- vllm_config.compilation_config.pass_config = \
96
- PassConfig(enable_fusion=True, enable_noop=True)
96
+ level=CompilationLevel.PIECEWISE,
97
+ custom_ops=["+rms_norm", "+quant_fp8"],
98
+ pass_config=PassConfig(enable_fusion=True, enable_noop=True),
99
+ ))
100
with vllm.config.set_current_vllm_config(vllm_config):
101
# Reshape pass is needed for the fusion pass to work
102
noop_pass = NoOpEliminationPass(vllm_config)
tests/compile/test_fusion_attn.py
+2
@@ -50,6 +50,7 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
50
# DYNAMO_ONCE does not properly propagate shapes.
51
level=CompilationLevel.DYNAMO_AS_IS,
52
backend="tests.compile.test_fusion_attn.backend_unfused",
53
+ custom_ops=["+quant_fp8"],
54
)
55
vllm_config = VllmConfig(compilation_config=compile_config)
56
backend_unfused = TestBackend(NoOpEliminationPass(vllm_config))
@@ -73,6 +74,7 @@ def test_attention_fusion(example_prompts, monkeypatch, model: str,
74
# DYNAMO_ONCE does not properly propagate shapes.
75
level=CompilationLevel.DYNAMO_AS_IS,
76
backend="tests.compile.test_fusion_attn.backend",
77
+ custom_ops=["+quant_fp8"],
78
)
79
vllm_config = VllmConfig(compilation_config=compile_config)
80
tests/compile/test_silu_mul_quant_fusion.py
+30
-7
@@ -4,33 +4,56 @@ import pytest
4
import torch
5
6
import vllm.envs as envs
7
-from vllm._custom_ops import scaled_fp8_quant
7
from vllm.compilation.activation_quant_fusion import ActivationQuantFusionPass
8
from vllm.compilation.fx_utils import find_auto_fn, find_auto_fn_maybe
9
+from vllm.compilation.noop_elimination import NoOpEliminationPass
10
from vllm.config import CompilationConfig, PassConfig, VllmConfig
11
from vllm.model_executor.layers.activation import SiluAndMul
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ GroupShape)
14
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15
+ CUTLASS_FP8_SUPPORTED, Fp8LinearOp)
16
+from vllm.platforms import current_platform
17
18
from .backend import TestBackend
19
20
21
class TestModel(torch.nn.Module):
22
18
- def __init__(self, *args, **kwargs):
23
+ def __init__(self, hidden_size: int, cutlass_fp8_enabled: bool, *args,
24
+ **kwargs):
25
super().__init__(*args, **kwargs)
26
self.silu_and_mul = SiluAndMul()
27
+ self.wscale = torch.rand(1, dtype=torch.float32)
28
self.scale = torch.rand(1, dtype=torch.float32)
29
30
+ self.w = (torch.rand(
31
+ hidden_size,
32
+ hidden_size).to(dtype=current_platform.fp8_dtype()).t())
33
+
34
+ self.fp8_linear = Fp8LinearOp(
35
+ cutlass_fp8_supported=cutlass_fp8_enabled,
36
+ act_quant_static=True,
37
+ act_quant_group_shape=GroupShape.PER_TENSOR,
38
+ )
39
+
40
def forward(self, x):
41
y = self.silu_and_mul(x)
25
- x2 = scaled_fp8_quant(y, self.scale)
42
+ x2 = self.fp8_linear.apply(y,
43
+ self.w,
44
+ self.wscale,
45
+ input_scale=self.wscale)
46
return x2
47
48
49
@pytest.mark.parametrize("num_tokens", [256])
50
@pytest.mark.parametrize("hidden_size", [64])
51
+@pytest.mark.parametrize("cutlass_fp8_enabled",
52
+ [True, False] if CUTLASS_FP8_SUPPORTED else [False])
53
@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
54
reason="Only test on CUDA and ROCm")
33
-def test_fusion_silu_and_mul_quant(num_tokens, hidden_size):
55
+def test_fusion_silu_and_mul_quant(num_tokens, hidden_size,
56
+ cutlass_fp8_enabled):
57
torch.set_default_device("cuda")
58
torch.set_default_dtype(torch.float16)
59
@@ -40,11 +63,11 @@ def test_fusion_silu_and_mul_quant(num_tokens, hidden_size):
63
pass_config=PassConfig(enable_fusion=True, enable_noop=True))
64
fusion_pass = ActivationQuantFusionPass(config)
65
43
- backend = TestBackend(fusion_pass)
44
- model = TestModel()
66
+ backend = TestBackend(NoOpEliminationPass(config), fusion_pass)
67
+ model = TestModel(hidden_size, cutlass_fp8_enabled)
68
69
# First dimension dynamic
47
- x = torch.rand(num_tokens, hidden_size)
70
+ x = torch.rand(num_tokens, hidden_size * 2)
71
torch._dynamo.mark_dynamic(x, 0)
72
73
result = model(x)
vllm/attention/backends/abstract.py
+4
-2
@@ -9,6 +9,8 @@ from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional,
9
10
import torch
11
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ GroupShape)
14
from vllm.multimodal import MultiModalPlaceholderMap
15
16
if TYPE_CHECKING:
@@ -289,7 +291,7 @@ class AttentionImpl(ABC, Generic[T]):
291
raise NotImplementedError
292
293
def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
292
- group_shape: tuple[int, int]):
294
+ group_shape: GroupShape):
295
"""
296
Does this attention implementation support fused output quantization.
297
This is used by the AttnFusionPass to only fuse output quantization
@@ -298,7 +300,7 @@ class AttentionImpl(ABC, Generic[T]):
300
TODO(luka) merge parameters into QuantDescriptor
301
:param dtype: quantized dtype
302
:param static: static or dynamic quantization
301
- :param group_shape: quant group shape. (-1, -1) for per-tensor.
303
+ :param group_shape: quant group shape.
304
:return: is fusion supported for this type of quantization
305
"""
306
return False
vllm/attention/backends/rocm_flash_attn.py
+4
-2
@@ -19,6 +19,8 @@ from vllm.attention.ops.paged_attn import (PagedAttention,
19
PagedAttentionMetadata)
20
from vllm.config import get_current_vllm_config
21
from vllm.logger import init_logger
22
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
+ GroupShape)
24
from vllm.platforms import current_platform
25
from vllm.platforms.rocm import use_rocm_custom_paged_attention
26
@@ -598,10 +600,10 @@ class ROCmFlashAttentionImpl(AttentionImpl):
600
head_dim))
601
602
def fused_output_quant_supported(self, dtype: torch.dtype, static: bool,
601
- group_shape: tuple[int, int]):
603
+ group_shape: GroupShape):
604
if self.use_triton_flash_attn:
605
return dtype == current_platform.fp8_dtype(
604
- ) and static and group_shape == (-1, -1) # per-tensor
606
+ ) and static and group_shape == GroupShape.PER_TENSOR
607
608
# Only supported in the Triton backend
609
return False
vllm/compilation/fusion.py
+3
-22
@@ -1,6 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-from typing import Callable, ClassVar, NamedTuple, Optional
3
+from typing import Callable, NamedTuple, Optional
4
5
import torch
6
import torch._inductor.pattern_matcher as pm
@@ -11,6 +11,8 @@ from torch._ops import OpOverload
11
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
15
+ GroupShape)
16
from vllm.platforms import current_platform
17
18
from .fx_utils import find_getitem_maybe
@@ -33,27 +35,6 @@ RMS_OP = torch.ops._C.rms_norm.default
35
RMS_ADD_OP = torch.ops._C.fused_add_rms_norm.default
36
37
36
-# Use proxy as NamedTuple direct subclasses cannot have static members
37
-class _GroupShape(NamedTuple):
38
- row: int
39
- col: int
40
-
41
-
42
-class GroupShape(_GroupShape):
43
- """
44
- This class describes the quantization group shape.
45
- It includes static members for common shapes (per-tensor, per-token).
46
- """
47
-
48
- # Aliases for common quantization group shapes
49
- PER_TENSOR: ClassVar['GroupShape']
50
- PER_TOKEN: ClassVar['GroupShape']
51
-
52
-
53
-GroupShape.PER_TENSOR = GroupShape(-1, -1)
54
-GroupShape.PER_TOKEN = GroupShape(1, -1)
55
-
56
-
38
class QuantKey(NamedTuple):
39
"""
40
Named tuple for identifying the type of quantization.
vllm/model_executor/layers/fused_moe/utils.py
+2
@@ -111,6 +111,8 @@ def _fp8_quantize(
111
is provided, the output will be blocked.
112
"""
113
if block_shape is None:
114
+ # TODO(luka): use QuantFP8 custom op
115
+ # https://github.com/vllm-project/vllm/issues/20711
116
A, A_scale = ops.scaled_fp8_quant(
117
A, A_scale, use_per_token_if_dynamic=per_act_token)
118
else:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_24.py
+16
-8
@@ -15,6 +15,9 @@ from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
15
QKVParallelLinear)
16
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
17
CompressedTensorsScheme)
18
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
19
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
+ GroupShape)
21
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
22
convert_to_channelwise, sparse_cutlass_supported)
23
from vllm.model_executor.parameter import (BasevLLMParameter,
@@ -24,6 +27,8 @@ from vllm.model_executor.parameter import (BasevLLMParameter,
27
28
__all__ = ["CompressedTensors24"]
29
30
+from vllm.platforms import current_platform
31
+
32
33
class CompressedTensors24(CompressedTensorsScheme):
34
@@ -45,6 +50,12 @@ class CompressedTensors24(CompressedTensorsScheme):
50
and self.model_compressor.sparsity_config.format
51
== CompressionFormat.sparse_24_bitmask.value)
52
53
+ if quantized and input_quant is not None and \
54
+ self._get_quant_dtype() == current_platform.fp8_dtype():
55
+ static = not input_quant.dynamic
56
+ g_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN
57
+ self.quant_fp8 = QuantFP8(static, g_shape)
58
+
59
@classmethod
60
def get_min_capability(cls) -> int:
61
# Only cutlass 3.x kernels are implemented so far
@@ -232,9 +243,7 @@ class CompressedTensors24(CompressedTensorsScheme):
243
:return: The output tensor of the layer
244
"""
245
if self.quantized:
235
- scale = None
236
- if hasattr(layer, "input_scale"):
237
- scale = layer.input_scale
246
+ scale = getattr(layer, 'input_scale', None)
247
248
if self.weights_dtype == torch.int8:
249
ops_output = ops.scaled_int8_quant(x, scale=scale)
@@ -242,11 +251,7 @@ class CompressedTensors24(CompressedTensorsScheme):
251
input_scale = ops_output[1]
252
else:
253
assert self.weights_dtype == torch.float8_e4m3fn
245
- if scale is not None:
246
- q_input, input_scale = ops.scaled_fp8_quant(x, scale=scale)
247
- else:
248
- q_input, input_scale = ops.scaled_fp8_quant(
249
- x, use_per_token_if_dynamic=True)
254
+ q_input, input_scale = self.quant_fp8(x, scale=scale)
255
256
else:
257
# Not quantized, nothing to do with the input_scales, use as is
@@ -269,7 +274,10 @@ class CompressedTensors24(CompressedTensorsScheme):
274
def _get_params_dtype(self, params_dtype: torch.dtype) -> torch.dtype:
275
if not self.quantized:
276
return params_dtype
277
+ return self._get_quant_dtype()
278
279
+ def _get_quant_dtype(self) -> torch.dtype:
280
+ assert self.quantized
281
assert self.weight_quant is not None
282
assert self.input_quant is not None
283
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+7
-1
@@ -9,6 +9,8 @@ from torch.nn import Parameter
9
10
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11
CompressedTensorsScheme)
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ GroupShape)
14
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)
@@ -26,7 +28,11 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
28
self.strategy = strategy
29
self.out_dtype = torch.get_default_dtype()
30
self.is_static_input_scheme = is_static_input_scheme
29
- self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
31
+ self.act_q_group_shape = GroupShape.PER_TENSOR \
32
+ if is_static_input_scheme else GroupShape.PER_TOKEN
33
+ self.fp8_linear = Fp8LinearOp(
34
+ act_quant_static=self.is_static_input_scheme,
35
+ act_quant_group_shape=self.act_q_group_shape)
36
37
@classmethod
38
def get_min_capability(cls) -> int:
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+3
-3
@@ -16,7 +16,7 @@ from vllm.model_executor.layers.quantization.base_config import (
16
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
17
apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin)
18
from vllm.model_executor.layers.quantization.utils.quant_utils import (
19
- is_layer_skipped)
19
+ GroupShape, is_layer_skipped)
20
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
21
Fp8LinearOp, maybe_create_device_identity, normalize_e4m3fn_to_e4m3fnuz)
22
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
@@ -37,7 +37,6 @@ class FBGEMMFp8Config(QuantizationConfig):
37
# For GPUs that lack FP8 hardware support, we can leverage the Marlin
38
# kernel for fast weight-only FP8 quantization
39
self.use_marlin = not current_platform.has_device_capability(89)
40
- self.fp8_linear = Fp8LinearOp()
40
41
@classmethod
42
def get_name(cls) -> QuantizationMethods:
@@ -76,7 +75,8 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
75
76
def __init__(self, quant_config: FBGEMMFp8Config):
77
self.quant_config = quant_config
79
- self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
78
+ self.fp8_linear = Fp8LinearOp(
79
+ act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN)
80
self.out_dtype = torch.get_default_dtype()
81
82
def create_weights(
vllm/model_executor/layers/quantization/fp8.py
+11
-3
@@ -29,7 +29,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
29
apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin,
30
prepare_moe_fp8_layer_for_marlin)
31
from vllm.model_executor.layers.quantization.utils.quant_utils import (
32
- is_layer_skipped)
32
+ GroupShape, is_layer_skipped)
33
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
34
Fp8LinearOp, all_close_1d, cutlass_block_fp8_supported,
35
cutlass_fp8_supported, maybe_create_device_identity,
@@ -202,9 +202,17 @@ class Fp8LinearMethod(LinearMethodBase):
202
and current_platform.is_fp8_fnuz())
203
204
self.block_quant = self.quant_config.weight_block_size is not None
205
+ self.act_q_static = self.quant_config.activation_scheme == "static"
206
+ # Use per-token quantization for better perf if dynamic and cutlass
207
+ if not self.act_q_static and cutlass_fp8_supported():
208
+ self.act_q_group_shape = GroupShape.PER_TOKEN
209
+ else:
210
+ self.act_q_group_shape = GroupShape.PER_TENSOR
211
+
212
self.fp8_linear = Fp8LinearOp(
206
- # Default to using per_token quantization if cutlass is supported
207
- use_per_token_if_dynamic=cutlass_fp8_supported())
213
+ act_quant_static=self.act_q_static,
214
+ act_quant_group_shape=self.act_q_group_shape,
215
+ cutlass_fp8_supported=cutlass_fp8_supported())
216
217
def create_weights(
218
self,
vllm/model_executor/layers/quantization/input_quant_fp8.py
new
+103
@@ -0,0 +1,103 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
+import torch
6
+import torch.nn.functional as F
7
+
8
+from vllm import _custom_ops as ops
9
+from vllm.model_executor.custom_op import CustomOp
10
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
11
+ GroupShape)
12
+from vllm.platforms import current_platform
13
+
14
+# Using the default value (240.0) from pytorch will cause accuracy
15
+# issue on dynamic quantization models. Here use 224.0 for fnuz on ROCm.
16
+_FP8_DTYPE = current_platform.fp8_dtype()
17
+_FP8_FINFO = torch.finfo(_FP8_DTYPE)
18
+_FP8_MAX = 224.0 if current_platform.is_fp8_fnuz() else _FP8_FINFO.max
19
+_FP8_MIN = -224.0 if current_platform.is_fp8_fnuz() else _FP8_FINFO.min
20
+_FP8_MIN_SCALING_FACTOR = 1.0 / (_FP8_MAX * 512.0)
21
+
22
+
23
+@CustomOp.register("quant_fp8")
24
+class QuantFP8(CustomOp):
25
+ """
26
+ Quantize input tensor to per-tensor or per-token FP8.
27
+ This CustomOp supports both static and dynamic quantization.
28
+ """
29
+
30
+ def __init__(self,
31
+ static: bool,
32
+ group_shape: GroupShape,
33
+ num_token_padding: Optional[int] = None):
34
+ """
35
+
36
+ :param static: static or dynamic quantization
37
+ :param group_shape: quantization group shape (PER_TOKEN or PER_TENSOR)
38
+ :param num_token_padding: Pad the token dimension of output to this size
39
+ """
40
+ super().__init__()
41
+ self.num_token_padding = num_token_padding
42
+ assert group_shape in {GroupShape.PER_TOKEN, GroupShape.PER_TENSOR}
43
+ assert not static or group_shape == GroupShape.PER_TENSOR, \
44
+ "Only per-tensor scales supported for static quantization."
45
+ self.static = static
46
+ self.group_shape = group_shape
47
+ self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
48
+
49
+ def forward_cuda(
50
+ self,
51
+ x: torch.Tensor,
52
+ scale: Optional[torch.Tensor] = None,
53
+ scale_ub: Optional[torch.Tensor] = None,
54
+ ) -> tuple[torch.Tensor, torch.Tensor]:
55
+ assert (scale is not None) == self.static
56
+ assert scale_ub is None or (not self.static and self.group_shape
57
+ == GroupShape.PER_TOKEN
58
+ and scale_ub.numel() == 1)
59
+
60
+ return ops.scaled_fp8_quant(
61
+ x,
62
+ scale,
63
+ num_token_padding=self.num_token_padding,
64
+ scale_ub=scale_ub,
65
+ use_per_token_if_dynamic=self.use_per_token_if_dynamic)
66
+
67
+ def forward_native(
68
+ self,
69
+ x: torch.Tensor,
70
+ scale: Optional[torch.Tensor] = None,
71
+ scale_ub: Optional[torch.Tensor] = None,
72
+ ):
73
+ assert (scale is not None) == self.static
74
+ assert scale_ub is None or (not self.static and self.group_shape
75
+ == GroupShape.PER_TOKEN
76
+ and scale_ub.numel() == 1)
77
+
78
+ if scale is None:
79
+ if self.group_shape == GroupShape.PER_TOKEN:
80
+ x_max, _ = x.abs().max(dim=-1)
81
+ x_max = x_max.unsqueeze(-1).to(torch.float32)
82
+ if scale_ub is not None:
83
+ x_max = x_max.clamp(max=scale_ub)
84
+ else:
85
+ x_max = x.abs().max().unsqueeze(-1).to(torch.float32)
86
+
87
+ scale = x_max / _FP8_MAX
88
+ scale = scale.clamp(min=_FP8_MIN_SCALING_FACTOR)
89
+
90
+ # Even for dynamic per-token scales,
91
+ # reciprocal performs slightly better than division
92
+ out = x.to(torch.float32) * scale.reciprocal()
93
+ out = out.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
94
+
95
+ # This currently generates an extra Triton kernel in compilation.
96
+ # Fortunately, we don't use padding if compiling.
97
+ # TODO(luka): benchmark torch._scaled_mm to hopefully remove padding
98
+ # in general.
99
+ if self.num_token_padding is not None:
100
+ padding = max(self.num_token_padding - out.size(0), 0)
101
+ out = F.pad(out, (0, 0, 0, padding), "constant", 0.0)
102
+
103
+ return out, scale
vllm/model_executor/layers/quantization/modelopt.py
+3
-2
@@ -22,7 +22,7 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
22
apply_fp4_marlin_linear, is_fp4_marlin_supported,
23
prepare_fp4_layer_for_marlin, prepare_moe_fp4_layer_for_marlin)
24
from vllm.model_executor.layers.quantization.utils.quant_utils import (
25
- is_layer_skipped)
25
+ GroupShape, is_layer_skipped)
26
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27
Fp8LinearOp, requantize_with_max_scale)
28
from vllm.model_executor.parameter import (ModelWeightParameter,
@@ -102,7 +102,8 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
102
103
def __init__(self, quant_config: ModelOptFp8Config):
104
self.quant_config = quant_config
105
- self.fp8_linear = Fp8LinearOp()
105
+ self.fp8_linear = Fp8LinearOp(
106
+ act_quant_static=True, act_quant_group_shape=GroupShape.PER_TENSOR)
107
108
def create_weights(
109
self,
vllm/model_executor/layers/quantization/ptpc_fp8.py
+5
-3
@@ -17,7 +17,7 @@ from vllm.model_executor.layers.quantization.fp8 import (Fp8Config,
17
Fp8KVCacheMethod,
18
Fp8LinearMethod)
19
from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
- is_layer_skipped)
20
+ GroupShape, is_layer_skipped)
21
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
22
Fp8LinearOp)
23
from vllm.platforms import current_platform
@@ -95,8 +95,10 @@ class PTPCFp8LinearMethod(Fp8LinearMethod):
95
super().__init__(quant_config=quant_config)
96
# Force weight quantization
97
self.quant_config.is_checkpoint_fp8_serialized = False
98
- self.fp8_linear = Fp8LinearOp(cutlass_fp8_supported=False,
99
- use_per_token_if_dynamic=True)
98
+ self.fp8_linear = Fp8LinearOp(
99
+ act_quant_static=False,
100
+ cutlass_fp8_supported=False,
101
+ act_quant_group_shape=GroupShape.PER_TOKEN)
102
103
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
104
layer.weight = torch.nn.Parameter(layer.weight.data,
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+11
-5
@@ -7,6 +7,8 @@ import torch
7
from torch.nn import Parameter
8
9
from vllm.model_executor.layers.quantization.quark.schemes import QuarkScheme
10
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
11
+ GroupShape)
12
from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
13
Fp8LinearOp, normalize_e4m3fn_to_e4m3fnuz, requantize_with_max_scale)
14
from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
@@ -28,10 +30,14 @@ class QuarkW8A8Fp8(QuarkScheme):
30
self.is_static_input_scheme = not cast(
31
bool, input_config.get("is_dynamic"))
32
self.input_qscheme = cast(str, input_config.get("qscheme"))
31
- self.use_per_token_if_dynamic = (not self.is_static_input_scheme \
32
- and self.input_qscheme == "per_channel")
33
+
34
+ per_token = (not self.is_static_input_scheme
35
+ and self.input_qscheme == "per_channel")
36
+ self.act_quant_group_shape = GroupShape.PER_TOKEN \
37
+ if per_token else GroupShape.PER_TENSOR
38
self.fp8_linear = Fp8LinearOp(
34
- use_per_token_if_dynamic=self.use_per_token_if_dynamic)
39
+ act_quant_static=self.is_static_input_scheme,
40
+ act_quant_group_shape=self.act_quant_group_shape)
41
self.out_dtype = torch.get_default_dtype()
42
43
@classmethod
@@ -44,7 +50,7 @@ class QuarkW8A8Fp8(QuarkScheme):
50
# tensor scales (thus N scales being passed to the kernel),
51
# requantize so we can always run per tensor
52
if self.weight_qscheme == "per_tensor":
47
- if current_platform.is_rocm():
53
+ if current_platform.is_fp8_fnuz():
54
input_scale = getattr(layer, 'input_scale', None)
55
weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
56
weight=layer.weight,
@@ -82,7 +88,7 @@ class QuarkW8A8Fp8(QuarkScheme):
88
requires_grad=False)
89
else:
90
weight_scale = layer.weight_scale.data
85
- if self.use_per_token_if_dynamic:
91
+ if self.act_quant_group_shape == GroupShape.PER_TOKEN:
92
weight_scale = weight_scale.view(-1, 1)
93
layer.weight = Parameter(weight.t(), requires_grad=False)
94
# required by torch.compile to be torch.nn.Parameter
vllm/model_executor/layers/quantization/utils/quant_utils.py
+28
-7
@@ -3,7 +3,7 @@
3
"""This file is used for /tests and /benchmarks"""
4
from collections.abc import Mapping
5
from types import MappingProxyType
6
-from typing import Optional
6
+from typing import ClassVar, NamedTuple, Optional
7
8
import numpy
9
import torch
@@ -12,13 +12,30 @@ from vllm.model_executor.layers.quantization.qqq import (
12
MARLIN_QQQ_SUPPORTED_NUM_BITS)
13
from vllm.scalar_type import ScalarType, scalar_types
14
15
-SUPPORTED_GPTQ_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
16
-SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
15
+
16
+# Use proxy as NamedTuple direct subclasses cannot have static members
17
+class _GroupShape(NamedTuple):
18
+ row: int
19
+ col: int
20
+
21
+
22
+class GroupShape(_GroupShape):
23
+ """
24
+ This class describes the quantization group shape.
25
+ It includes static members for common shapes (per-tensor, per-token).
26
+ """
27
+
28
+ # Aliases for common quantization group shapes
29
+ PER_TENSOR: ClassVar['GroupShape']
30
+ PER_TOKEN: ClassVar['GroupShape']
31
+
32
+
33
+GroupShape.PER_TENSOR = GroupShape(-1, -1)
34
+GroupShape.PER_TOKEN = GroupShape(1, -1)
35
36
37
# Normalize the group_shape to the full extent for any dims that are -1
20
-def _normalize_quant_group_shape(x: torch.Tensor, group_shape: tuple[int,
21
- int]):
38
+def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape):
39
# -1 means full extent
40
return (group_shape[0] if group_shape[0] > 0 else x.shape[-2],
41
group_shape[1] if group_shape[1] > 0 else x.shape[-1])
@@ -58,7 +75,7 @@ def group_broadcast(t, shape):
75
# (i.e. per-token-per-group)
76
def scaled_quantize(
77
x: torch.Tensor,
61
- group_shape: tuple[int, int],
78
+ group_shape: GroupShape,
79
quant_dtype: torch.dtype,
80
) -> tuple[torch.Tensor, torch.Tensor]:
81
group_shape = _normalize_quant_group_shape(x, group_shape)
@@ -99,7 +116,7 @@ def scaled_quantize(
116
def scaled_dequantize(
117
x_q: torch.Tensor,
118
x_s: torch.Tensor,
102
- group_shape: Optional[tuple[int, int]] = None,
119
+ group_shape: Optional[GroupShape] = None,
120
out_dtype: torch.dtype = torch.float32,
121
) -> tuple[torch.Tensor, torch.Tensor]:
122
if group_shape is not None:
@@ -332,6 +349,10 @@ def quantize_weights(w: torch.Tensor,
349
)
350
351
352
+SUPPORTED_GPTQ_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
353
+SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]
354
+
355
+
356
def gptq_quantize_weights(w: torch.Tensor,
357
quant_type: ScalarType,
358
group_size: int,
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+31
-35
@@ -8,6 +8,9 @@ import torch
8
from vllm import _custom_ops as ops
9
from vllm import envs
10
from vllm.config import CompilationLevel, get_current_vllm_config
11
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ GroupShape)
14
from vllm.platforms import current_platform
15
16
# Input scaling factors are no longer optional in _scaled_mm starting
@@ -271,20 +274,21 @@ def torch_channelwise_w8a8_scaled_mm(*, qinput: torch.Tensor,
274
275
def dispatch_w8a8_scaled_mm(
276
cutlass_fp8_supported: bool, per_tensor_weights: bool,
274
- per_tensor_activations: bool, use_per_token_if_dynamic: Optional[bool]
275
-) -> Callable[..., torch.Tensor]:
277
+ per_tensor_activations: bool) -> Callable[..., torch.Tensor]:
278
279
+ # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
280
if cutlass_fp8_supported:
281
return cutlass_w8a8_scaled_mm
282
if per_tensor_weights and per_tensor_activations:
283
if current_platform.is_rocm():
284
return rocm_per_tensor_w8a8_scaled_mm
285
return torch_per_tensor_w8a8_scaled_mm
283
- # torch.scaled_mm supports per tensor weights + activations only
284
- # so fallback to naive if per channel or per token
285
- if (use_per_token_if_dynamic and not per_tensor_weights
286
- and not per_tensor_activations and USE_ROWWISE_TORCH_SCALED_MM):
286
+ # If torch.scaled_mm supports per-channel (weights) per-token (inputs)
287
+ if not per_tensor_weights and not per_tensor_activations \
288
+ and USE_ROWWISE_TORCH_SCALED_MM:
289
return torch_per_token_w8a8_scaled_mm
290
+ # Normally, torch.scaled_mm supports per tensor weights + activations only
291
+ # so fallback to naive if per channel or per token
292
return torch_channelwise_w8a8_scaled_mm
293
294
@@ -299,11 +303,11 @@ class Fp8LinearOp:
303
"""
304
305
def __init__(self,
306
+ act_quant_static: bool,
307
cutlass_fp8_supported: bool = cutlass_fp8_supported(),
303
- use_per_token_if_dynamic: bool = False,
308
+ act_quant_group_shape: GroupShape = GroupShape.PER_TENSOR,
309
pad_output: Optional[bool] = None):
310
self.cutlass_fp8_supported = cutlass_fp8_supported
306
- self.use_per_token_if_dynamic = use_per_token_if_dynamic
311
312
# Note: we pad the input because torch._scaled_mm is more performant
313
# for matrices with batch dimension > 16.
@@ -312,9 +316,16 @@ class Fp8LinearOp:
316
# as it breaks with dynamic shapes.
317
if pad_output is None:
318
config = get_current_vllm_config().compilation_config
315
- pad_output = config.level < CompilationLevel.PIECEWISE
316
- self.output_padding = 17 if (
317
- pad_output and not current_platform.is_rocm()) else None
319
+ pad_output = config.level < CompilationLevel.PIECEWISE and \
320
+ not cutlass_fp8_supported and \
321
+ not current_platform.is_rocm()
322
+
323
+ self.output_padding = 17 if pad_output else None
324
+ self.act_quant_static = act_quant_static
325
+ self.act_quant_group_shape = act_quant_group_shape
326
+ self.quant_fp8 = QuantFP8(static=act_quant_static,
327
+ group_shape=act_quant_group_shape,
328
+ num_token_padding=self.output_padding)
329
330
def apply(
331
self,
@@ -325,8 +336,6 @@ class Fp8LinearOp:
336
input_scale: Optional[torch.Tensor] = None,
337
input_scale_ub: Optional[torch.Tensor] = None,
338
bias: Optional[torch.Tensor] = None,
328
- # TODO(luka) remove this parameter in favor of __init__
329
- use_per_token_if_dynamic: Optional[bool] = None
339
) -> torch.Tensor:
340
# ops.scaled_fp8_quant supports both dynamic and static quant.
341
# If dynamic, layer.input_scale is None and x_scale computed from x.
@@ -336,40 +345,27 @@ class Fp8LinearOp:
345
input_2d = input.view(-1, input.shape[-1])
346
output_shape = [*input.shape[:-1], weight.shape[1]]
347
339
- # TODO(luka) this is here because currently MLA only decides this
340
- # during the forward method instead of in __init__.
341
- if use_per_token_if_dynamic is None:
342
- use_per_token_if_dynamic = self.use_per_token_if_dynamic
343
-
348
if out_dtype is None:
349
out_dtype = input.dtype
350
347
- # cutlass_scaled_mm supports per tensor/channel W and per tensor/token A
348
- if self.cutlass_fp8_supported:
349
- assert input.dtype != current_platform.fp8_dtype(
350
- ), "FP8 input to cutlass is not currently implemented"
351
- qinput, x_scale = ops.scaled_fp8_quant(
351
+ # If input not quantized
352
+ # TODO(luka) remove this path if not used anymore
353
+ if input.dtype != current_platform.fp8_dtype():
354
+ qinput, x_scale = self.quant_fp8(
355
input_2d,
356
input_scale,
354
- scale_ub=input_scale_ub,
355
- use_per_token_if_dynamic=use_per_token_if_dynamic)
357
+ input_scale_ub,
358
+ )
359
else:
357
- if input.dtype != current_platform.fp8_dtype():
358
- # Maybe apply padding to output, see comment in __init__
359
- qinput, x_scale = ops.scaled_fp8_quant(
360
- input_2d,
361
- input_scale,
362
- num_token_padding=self.output_padding,
363
- use_per_token_if_dynamic=use_per_token_if_dynamic)
364
- else:
365
- qinput, x_scale = input_2d, input_scale
360
+ qinput, x_scale = input_2d, input_scale
361
362
per_tensor_weights = (weight_scale.numel() == 1)
363
per_tensor_activations = (x_scale.numel() == 1)
364
365
+ # TODO(luka) do this dispatch during init (after ScaledMM refactor)
366
w8a8_scaled_mm_func = dispatch_w8a8_scaled_mm(
367
self.cutlass_fp8_supported, per_tensor_weights,
372
- per_tensor_activations, use_per_token_if_dynamic)
368
+ per_tensor_activations)
369
370
return w8a8_scaled_mm_func(qinput=qinput,
371
weight=weight,