@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / cef32104b4
Fork 0

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

[FP8] Extend per-token-group quantization support to QuantFP8 (#24342)

Signed-off-by: Tahsin Tunan <tahsintunan@gmail.com> Signed-off-by: Luka Govedič <lgovedic@redhat.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Luka Govedič <lgovedic@redhat.com>

Tahsin Tunan committed Sep 17, 2025 at 07:31 UTC cef32104b4b411e5093581ad77d2e09a50c2837c
5 files changed +444 -61
benchmarks/kernels/bench_per_token_quant_fp8.py
+216 -47
@@ -2,14 +2,25 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import itertools
4 from typing import Callable
5 +from unittest.mock import patch
6
7 +import pandas as pd
8 import torch
9
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 +from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, FlexibleArgumentParser
14 +
15 +
16 +def with_triton_mode(fn):
17 + """Temporarily force the Triton fallback path"""
18 +
19 + def wrapped(*args, **kwargs):
20 + with patch("vllm.platforms.current_platform.is_cuda", return_value=False):
21 + return fn(*args, **kwargs)
22 +
23 + return wrapped
24
25
26 # TODO(luka): use standalone_compile utility
@@ -21,78 +32,236 @@ def with_dyn_arg(fn: Callable, arg_index: int, dim_index: int):
32 return inner
33
34
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 - )
35 +def bench_compile(fn: Callable):
36 + # recompile for different shapes
37 + fwd = torch.compile(fn, fullgraph=True, dynamic=False)
38
39 # 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)
40 + return with_dyn_arg(fwd, 0, 0)
41
42
37 -def cuda_per_token_quant_fp8(
38 - input: torch.Tensor,
39 -) -> tuple[torch.Tensor, torch.Tensor]:
40 - return ops.scaled_fp8_quant(input)
43 +torch._dynamo.config.recompile_limit = 8888
44
45
43 -def calculate_diff(batch_size: int, seq_len: int):
44 - """Calculate difference between Triton and CUDA implementations."""
46 +def calculate_diff(
47 + batch_size: int,
48 + hidden_size: int,
49 + group_shape: GroupShape,
50 + dtype: torch.dtype,
51 +):
52 + """Calculate the difference between Inductor and CUDA implementations."""
53 device = torch.device("cuda")
46 - x = torch.rand((batch_size * seq_len, 4096), dtype=torch.float16, device=device)
54 + x = torch.rand((batch_size * hidden_size, 4096), dtype=dtype, device=device)
55 +
56 + quant_fp8 = QuantFP8(False, group_shape, column_major_scales=False)
57
48 - torch_out, torch_scale = torch_per_token_quant_fp8(x)
49 - cuda_out, cuda_scale = cuda_per_token_quant_fp8(x)
58 + torch_out, torch_scale = bench_compile(quant_fp8.forward_native)(x)
59 + torch_eager_out, torch_eager_scale = quant_fp8.forward_native(x)
60 + cuda_out, cuda_scale = quant_fp8.forward_cuda(x)
61
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):
62 + out_allclose = lambda o1, o2: torch.allclose(
63 + o1.to(torch.float32),
64 + o2.to(torch.float32),
65 + rtol=1e-3,
66 + atol=1e-5,
67 + )
68 + scale_allclose = lambda s1, s2: torch.allclose(s1, s2, rtol=1e-3, atol=1e-5)
69 +
70 + if (
71 + out_allclose(cuda_out, torch_out)
72 + and scale_allclose(cuda_scale, torch_scale)
73 + and out_allclose(cuda_out, torch_eager_out)
74 + and scale_allclose(cuda_scale, torch_eager_scale)
75 + ):
76 print("✅ All implementations match")
77 else:
78 print("❌ Implementations differ")
79
80
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))
81 +configs = []
82
83
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
84 +def benchmark_quantization(
85 + batch_size,
86 + hidden_size,
87 + provider,
88 + group_shape: GroupShape,
89 + col_major: bool,
90 + dtype: torch.dtype,
91 +):
92 device = torch.device("cuda")
93
82 - x = torch.randn(batch_size * seq_len, 4096, device=device, dtype=dtype)
94 + x = torch.randn(batch_size * hidden_size, 4096, device=device, dtype=dtype)
95
96 quantiles = [0.5, 0.2, 0.8]
97 + quant_fp8 = QuantFP8(False, group_shape, column_major_scales=col_major)
98
99 if provider == "torch":
87 - fn = lambda: torch_per_token_quant_fp8(x.clone())
100 + fn = lambda: bench_compile(quant_fp8.forward_native)(x.clone())
101 elif provider == "cuda":
89 - fn = lambda: cuda_per_token_quant_fp8(x.clone())
102 + fn = lambda: quant_fp8.forward_cuda(x.clone())
103 + elif provider == "triton":
104 + if not group_shape.is_per_group():
105 + # Triton only supported for per-group
106 + return 0, 0, 0
107 +
108 + fn = lambda: with_triton_mode(quant_fp8.forward_cuda)(x.clone())
109
110 ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles)
111
112 return 1000 * ms, 1000 * max_ms, 1000 * min_ms
113
114
115 +# TODO(luka) extract to utils
116 +def compute_geomean_speedups(
117 + df: pd.DataFrame,
118 + baseline_col: str,
119 + speedup_cols: list[str],
120 + groupby_cols: list[str] | None = None,
121 +) -> pd.DataFrame:
122 + """
123 + Compute geometric mean speedups over a baseline column.
124 +
125 + Args:
126 + df: Input dataframe
127 + baseline_col: Column to use as baseline
128 + speedup_cols: Columns to compute speedups for
129 + groupby_cols: Columns to group by. If None, compute over entire df.
130 +
131 + Returns:
132 + pd.DataFrame with geometric mean speedups
133 + """
134 + from scipy.stats import gmean
135 +
136 + def geo_speedup(group: pd.DataFrame) -> pd.Series:
137 + ratios = {
138 + col: (group[baseline_col] / group[col]).values for col in speedup_cols
139 + }
140 + return pd.Series({col: gmean(vals) for col, vals in ratios.items()})
141 +
142 + if groupby_cols is None:
143 + result = geo_speedup(df).to_frame().T
144 + else:
145 + result = (
146 + df.groupby(groupby_cols)
147 + .apply(geo_speedup, include_groups=False)
148 + .reset_index()
149 + )
150 +
151 + return result
152 +
153 +
154 if __name__ == "__main__":
97 - calculate_diff(batch_size=4, seq_len=4096)
98 - benchmark_quantization.run(print_data=True)
155 + parser = FlexibleArgumentParser(
156 + description="Benchmark the various implementations of QuantFP8 (dynamic-only)"
157 + )
158 + parser.add_argument("-c", "--check", action="store_true")
159 + parser.add_argument(
160 + "--dtype", type=str, choices=["half", "bfloat16", "float"], default="half"
161 + )
162 + parser.add_argument(
163 + "--hidden-sizes",
164 + type=int,
165 + nargs="+",
166 + default=None,
167 + help="Hidden sizes to benchmark (default: 1,16,64,128,256,512,1024,2048,4096)",
168 + )
169 + parser.add_argument(
170 + "--batch-sizes",
171 + type=int,
172 + nargs="+",
173 + default=None,
174 + help="Batch sizes to benchmark (default: 1,16,32,64,128)",
175 + )
176 + parser.add_argument(
177 + "--group-sizes",
178 + type=int,
179 + nargs="+",
180 + default=None,
181 + help="Group sizes for GroupShape(1,N) to benchmark. "
182 + "Use 0 for PER_TENSOR, -1 for PER_TOKEN (default: 0,-1,64,128)",
183 + )
184 + parser.add_argument(
185 + "--no-column-major",
186 + action="store_true",
187 + help="Disable column-major scales testing",
188 + )
189 +
190 + args = parser.parse_args()
191 + assert args
192 +
193 + dtype = STR_DTYPE_TO_TORCH_DTYPE[args.dtype]
194 +
195 + hidden_sizes = args.hidden_sizes or [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
196 + batch_sizes = args.batch_sizes or [1, 16, 32, 64, 128]
197 +
198 + if args.group_sizes is not None:
199 + group_shapes = []
200 + for size in args.group_sizes:
201 + if size == 0:
202 + group_shapes.append(GroupShape.PER_TENSOR)
203 + elif size == -1:
204 + group_shapes.append(GroupShape.PER_TOKEN)
205 + else:
206 + group_shapes.append(GroupShape(1, size))
207 + else:
208 + group_shapes = [
209 + GroupShape.PER_TENSOR,
210 + GroupShape.PER_TOKEN,
211 + GroupShape(1, 64),
212 + GroupShape(1, 128),
213 + ]
214 +
215 + column_major_scales = [False] if args.no_column_major else [True, False]
216 +
217 + config_gen = itertools.product(
218 + group_shapes,
219 + column_major_scales,
220 + batch_sizes,
221 + hidden_sizes,
222 + )
223 +
224 + # filter out column-major scales for non-group, reverse order
225 + configs.extend(c[::-1] for c in config_gen if (c[0].is_per_group() or not c[1]))
226 +
227 + print(f"Running {len(configs)} configurations:")
228 + print(f" Hidden sizes: {hidden_sizes}")
229 + print(f" Batch sizes: {batch_sizes}")
230 + print(f" Group shapes: {[str(g) for g in group_shapes]}")
231 + print(f" Column major scales: {column_major_scales}")
232 + print()
233 +
234 + if args.check:
235 + for group_shape in group_shapes:
236 + group_size = group_shape[1]
237 + print(f"{group_size=}")
238 + calculate_diff(
239 + batch_size=4, hidden_size=4096, group_shape=group_shape, dtype=dtype
240 + )
241 +
242 + benchmark = triton.testing.perf_report(
243 + triton.testing.Benchmark(
244 + x_names=["hidden_size", "batch_size", "col_major", "group_shape"],
245 + x_vals=configs,
246 + line_arg="provider",
247 + line_vals=["torch", "cuda", "triton"],
248 + line_names=["Torch (Compiled)", "CUDA", "Triton"],
249 + styles=[("blue", "-"), ("green", "-"), ("black", "-")],
250 + ylabel="us",
251 + plot_name="QuantFP8 performance",
252 + args={},
253 + )
254 + )(benchmark_quantization)
255 +
256 + df = benchmark.run(print_data=True, dtype=dtype, return_df=True)
257 +
258 + # Print geomean speedups
259 + geo_table_grouped = compute_geomean_speedups(
260 + df,
261 + baseline_col="Torch (Compiled)",
262 + speedup_cols=["CUDA", "Triton"],
263 + groupby_cols=["col_major", "group_shape"],
264 + )
265 +
266 + print("Speedup over Torch (Compiled)")
267 + print(geo_table_grouped.to_string(index=False))
tests/kernels/quantization/test_fp8_quant_group.py new
+150
@@ -0,0 +1,150 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for QuantFP8 Group Quantization implementation."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
9 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
10 + GroupShape)
11 +from vllm.platforms import current_platform
12 +
13 +
14 +@pytest.mark.parametrize(
15 + "batch_size,hidden_dim,group_size",
16 + [
17 + (16, 256, 32), # Small
18 + (64, 1024, 64), # Medium
19 + (128, 2048, 128), # Large
20 + (8, 513, 64), # Non-divisible (native only)
21 + ])
22 +@pytest.mark.parametrize("seed", [42])
23 +@torch.inference_mode()
24 +def test_quantfp8_group_functionality(batch_size: int, hidden_dim: int,
25 + group_size: int, seed: int) -> None:
26 + """Test QuantFP8 group quantization with various configurations.
27 +
28 + Tests both CUDA and native implementations, column-major scales,
29 + and verifies consistency between implementations.
30 + """
31 + current_platform.seed_everything(seed)
32 +
33 + x = torch.randn(
34 + (batch_size, hidden_dim), dtype=torch.bfloat16, device="cuda") * 8
35 + expected_num_groups = (hidden_dim + group_size - 1) // group_size
36 + is_divisible = hidden_dim % group_size == 0
37 +
38 + group_shape = GroupShape(1, group_size)
39 + quant_op = QuantFP8(static=False,
40 + group_shape=group_shape,
41 + column_major_scales=False)
42 +
43 + # 1. Test native implementation (always available)
44 + x_quant_native, scales_native = quant_op.forward_native(x.clone())
45 + assert x_quant_native.shape == x.shape
46 + assert scales_native.shape == (batch_size, expected_num_groups)
47 +
48 + # 2. Test column-major scales configuration
49 + quant_op_col = QuantFP8(static=False,
50 + group_shape=group_shape,
51 + column_major_scales=True)
52 + _, scales_col = quant_op_col.forward_native(x.clone())
53 + assert scales_col.shape == (expected_num_groups, batch_size)
54 +
55 + # 3. Test CUDA implementation (only for divisible dimensions)
56 + if is_divisible:
57 + x_quant_cuda, scales_cuda = quant_op.forward_cuda(x.clone())
58 + assert x_quant_cuda.shape == x.shape
59 + assert scales_cuda.shape == (batch_size, expected_num_groups)
60 +
61 + # Verify CUDA/native consistency
62 + assert torch.allclose(scales_cuda, scales_native, rtol=1e-9, atol=1e-8)
63 +
64 + # Quantized values should mostly match
65 + diff_count = (x_quant_cuda != x_quant_native).sum().item()
66 + diff_ratio = diff_count / x_quant_cuda.numel()
67 + assert diff_ratio < 0.002, f"Too many differences: {diff_ratio:.4%}"
68 +
69 +
70 +@pytest.mark.parametrize("seed", [42])
71 +@torch.inference_mode()
72 +def test_quantfp8_group_multidimensional(seed: int) -> None:
73 + current_platform.seed_everything(seed)
74 +
75 + group_size = 64
76 +
77 + # Test with 3D input
78 + batch1, batch2, hidden_dim = 4, 8, 512
79 + x_3d = torch.randn(
80 + (batch1, batch2, hidden_dim), dtype=torch.bfloat16, device="cuda") * 8
81 +
82 + group_shape = GroupShape(1, group_size)
83 + quant_op = QuantFP8(static=False,
84 + group_shape=group_shape,
85 + column_major_scales=False)
86 +
87 + x_quant, scales = quant_op.forward_native(x_3d.clone())
88 + assert x_quant.shape == x_3d.shape
89 + assert scales.shape == (batch1, batch2, hidden_dim // group_size)
90 +
91 + # Test column_major_scales with multi-dim
92 + quant_op_col = QuantFP8(static=False,
93 + group_shape=group_shape,
94 + column_major_scales=True)
95 + _, scales_col = quant_op_col.forward_native(x_3d.clone())
96 + assert scales_col.shape == (batch1, hidden_dim // group_size, batch2)
97 +
98 + # Test with 4D input
99 + batch1, batch2, batch3, hidden_dim = 2, 3, 4, 256
100 + x_4d = torch.randn((batch1, batch2, batch3, hidden_dim),
101 + dtype=torch.bfloat16,
102 + device="cuda") * 8
103 +
104 + x_quant_4d, scales_4d = quant_op.forward_native(x_4d.clone())
105 + assert x_quant_4d.shape == x_4d.shape
106 + assert scales_4d.shape == (batch1, batch2, batch3,
107 + hidden_dim // group_size)
108 +
109 + _, scales_4d_col = quant_op_col.forward_native(x_4d.clone())
110 + assert scales_4d_col.shape == (batch1, batch2, hidden_dim // group_size,
111 + batch3)
112 +
113 +
114 +@pytest.mark.parametrize("seed", [42])
115 +@torch.inference_mode()
116 +def test_quantfp8_group_edge_cases(seed: int) -> None:
117 + current_platform.seed_everything(seed)
118 +
119 + batch_size = 16
120 + group_size = 64
121 +
122 + # Test with single group (group_size >= hidden_dim)
123 + x_small = torch.randn(
124 + (batch_size, 32), dtype=torch.bfloat16, device="cuda") * 8
125 + group_shape = GroupShape(1, group_size)
126 + quant_op = QuantFP8(static=False,
127 + group_shape=group_shape,
128 + column_major_scales=False)
129 +
130 + x_quant_small, scales_small = quant_op.forward_native(x_small.clone())
131 + assert x_quant_small.shape == x_small.shape
132 + assert scales_small.shape == (batch_size, 1)
133 +
134 + # Test with zero inputs
135 + x_zero = torch.zeros((batch_size, 256),
136 + dtype=torch.bfloat16,
137 + device="cuda")
138 + x_quant_zero, scales_zero = quant_op.forward_native(x_zero.clone())
139 + assert x_quant_zero.shape == x_zero.shape
140 + assert (scales_zero > 0).all(), "Scales should be clamped to minimum"
141 +
142 + # Test very large values
143 + x_large = torch.full((batch_size, 256),
144 + 1000.0,
145 + dtype=torch.bfloat16,
146 + device="cuda")
147 + x_quant_large, scales_large = quant_op.forward_native(x_large.clone())
148 + assert x_quant_large.shape == x_large.shape
149 + # FP8 max is typically 448 or 224, so scales should be > 1
150 + assert (scales_large > 1.0).all(), "Large values should have scales > 1"
vllm/model_executor/layers/fused_moe/fused_moe.py
+3 -1
@@ -32,9 +32,11 @@ from vllm.model_executor.layers.fused_moe.prepare_finalize import (
32 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
33 TopKWeightAndReduceNoOP)
34 from vllm.model_executor.layers.fused_moe.utils import (
35 - _resize_cache, moe_kernel_quantize_input, per_token_group_quant_fp8)
35 + _resize_cache, moe_kernel_quantize_input)
36 from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (
37 calculate_tile_tokens_dim)
38 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
39 + per_token_group_quant_fp8)
40 from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
41 dequant_mxfp4)
42 from vllm.platforms import current_platform
vllm/model_executor/layers/quantization/input_quant_fp8.py
+66 -13
@@ -23,28 +23,39 @@ _FP8_MIN_SCALING_FACTOR = 1.0 / (_FP8_MAX * 512.0)
23 @CustomOp.register("quant_fp8")
24 class QuantFP8(CustomOp):
25 """
26 - Quantize input tensor to per-tensor or per-token FP8.
26 + Quantize input tensor to FP8 (per-tensor, per-token, or per-group).
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):
33 + num_token_padding: Optional[int] = None,
34 + column_major_scales: bool = False):
35 """
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
37 + :param group_shape: quantization group shape (PER_TOKEN, PER_TENSOR,
38 + or arbitrary block size)
39 + :param num_token_padding: Pad the token dimension of output to this
40 + size
41 + :param column_major_scales: For group quantization, output scales in
42 + column major format
43 """
44 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
47 + self.num_token_padding = num_token_padding
48 + self.column_major_scales = column_major_scales
49 +
50 + self.is_group_quant = group_shape.is_per_group()
51 + if self.is_group_quant:
52 + assert not static, "Group quantization only supports dynamic mode"
53 + self.group_size = group_shape.col
54 + else:
55 + assert group_shape in {GroupShape.PER_TOKEN, GroupShape.PER_TENSOR}
56 + assert not static or group_shape == GroupShape.PER_TENSOR, \
57 + "Only per-tensor scales supported for static quantization."
58 + self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
59
60 def forward_cuda(
61 self,
@@ -52,11 +63,19 @@ class QuantFP8(CustomOp):
63 scale: Optional[torch.Tensor] = None,
64 scale_ub: Optional[torch.Tensor] = None,
65 ) -> tuple[torch.Tensor, torch.Tensor]:
66 + if self.is_group_quant:
67 + assert scale is None, "Group quantization is always dynamic"
68 + from vllm.model_executor.layers.quantization.utils import fp8_utils
69 + return fp8_utils.per_token_group_quant_fp8(
70 + x,
71 + group_size=self.group_size,
72 + column_major_scales=self.column_major_scales,
73 + dtype=_FP8_DTYPE)
74 +
75 assert (scale is not None) == self.static
76 assert scale_ub is None or (not self.static and self.group_shape
77 == GroupShape.PER_TOKEN
78 and scale_ub.numel() == 1)
59 -
79 return ops.scaled_fp8_quant(
80 x,
81 scale,
@@ -70,6 +89,10 @@ class QuantFP8(CustomOp):
89 scale: Optional[torch.Tensor] = None,
90 scale_ub: Optional[torch.Tensor] = None,
91 ):
92 + if self.is_group_quant:
93 + assert scale is None, "Group quantization is always dynamic"
94 + return self._quantize_group_native(x)
95 +
96 assert (scale is not None) == self.static
97 assert scale_ub is None or (not self.static and self.group_shape
98 == GroupShape.PER_TOKEN
@@ -84,8 +107,7 @@ class QuantFP8(CustomOp):
107 else:
108 x_max = x.abs().max().unsqueeze(-1).to(torch.float32)
109
87 - scale = x_max / _FP8_MAX
88 - scale = scale.clamp(min=_FP8_MIN_SCALING_FACTOR)
110 + scale = (x_max / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR)
111
112 # Even for dynamic per-token scales,
113 # reciprocal performs slightly better than division
@@ -101,3 +123,34 @@ class QuantFP8(CustomOp):
123 out = F.pad(out, (0, 0, 0, padding), "constant", 0.0)
124
125 return out, scale
126 +
127 + def _quantize_group_native(
128 + self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
129 + orig_shape = x.shape
130 + hidden_dim = x.shape[-1]
131 + num_groups = (hidden_dim + self.group_size - 1) // self.group_size
132 + padded_dim = num_groups * self.group_size
133 +
134 + if padded_dim != hidden_dim:
135 + padding = padded_dim - hidden_dim
136 + x = F.pad(x, (0, padding), mode='constant', value=0.0)
137 +
138 + x_grouped = x.view(-1, num_groups, self.group_size)
139 + absmax = x_grouped.abs().max(dim=-1, keepdim=True)[0].float()
140 + scales = (absmax / _FP8_MAX).clamp(min=_FP8_MIN_SCALING_FACTOR)
141 +
142 + x_scaled = x_grouped / scales
143 + x_quant = x_scaled.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
144 +
145 + x_quant = x_quant.view(-1, padded_dim)
146 + if padded_dim != hidden_dim:
147 + x_quant = x_quant[..., :hidden_dim]
148 + x_quant = x_quant.view(orig_shape)
149 +
150 + scales = scales.squeeze(-1)
151 + scales = scales.reshape(orig_shape[:-1] + (num_groups, ))
152 +
153 + if self.column_major_scales:
154 + scales = scales.transpose(-2, -1).contiguous()
155 +
156 + return x_quant, scales
vllm/model_executor/layers/quantization/utils/quant_utils.py
+9
@@ -34,6 +34,15 @@ class GroupShape(_GroupShape):
34 PER_TENSOR: ClassVar['GroupShape']
35 PER_TOKEN: ClassVar['GroupShape']
36
37 + def is_per_tensor(self) -> bool:
38 + return self.row == -1 and self.col == -1
39 +
40 + def is_per_token(self) -> bool:
41 + return self.row == 1 and self.col == -1
42 +
43 + def is_per_group(self) -> bool:
44 + return self.row == 1 and self.col >= 1
45 +
46
47 GroupShape.PER_TENSOR = GroupShape(-1, -1)
48 GroupShape.PER_TOKEN = GroupShape(1, -1)