@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0a0aa07747
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 10, 10:15 UTC

[Quant] Make static quant support all group shapes (#30833)

Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Lucas Wilkinson committed Jan 9, 2026 at 15:49 UTC 0a0aa07747cf851228598334008366ff2a8fd760
7 files changed +339 -47
csrc/ops.h
+4 -2
@@ -2,6 +2,7 @@
2
3 #include <optional>
4 #include <torch/library.h>
5 +#include <tuple>
6
7 #include "core/scalar_type.hpp"
8
@@ -346,8 +347,9 @@ torch::Tensor gptq_gemm(torch::Tensor a, torch::Tensor b_q_weight,
347
348 void gptq_shuffle(torch::Tensor q_weight, torch::Tensor q_perm, int64_t bit);
349
349 -void static_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
350 - torch::Tensor const& scale);
350 +void static_scaled_fp8_quant(
351 + torch::Tensor& out, torch::Tensor const& input, torch::Tensor const& scale,
352 + std::optional<std::tuple<int64_t, int64_t>> group_shape = std::nullopt);
353
354 void dynamic_scaled_fp8_quant(torch::Tensor& out, torch::Tensor const& input,
355 torch::Tensor& scale);
csrc/quantization/w8a8/fp8/common.cu
+178 -22
@@ -4,28 +4,77 @@
4 #include "quantization/vectorization_utils.cuh"
5 #include <c10/cuda/CUDAGuard.h>
6 #include <ATen/cuda/Exceptions.h>
7 +#include <tuple>
8
9 namespace vllm {
10
10 -template <typename scalar_t, typename fp8_type>
11 -__global__ void scaled_fp8_quant_kernel_strided(
11 +// STRIDE_I_ZERO: true if scale_stride_i == 0 (per-tensor or per-channel)
12 +// STRIDE_J_ZERO: true if scale_stride_j == 0 (per-tensor or per-token)
13 +template <typename scalar_t, typename fp8_type, bool STRIDE_I_ZERO,
14 + bool STRIDE_J_ZERO>
15 +__global__ void scaled_fp8_quant_kernel_strided_group_shape(
16 fp8_type* __restrict__ out, const scalar_t* __restrict__ input,
17 const float* __restrict__ scale, int hidden_size, int64_t in_row_stride,
14 - int64_t out_row_stride) {
15 - const int64_t token_idx = blockIdx.x; // one token per block
18 + int64_t out_row_stride, int group_m, int group_n, int64_t scale_stride_i,
19 + int64_t scale_stride_j) {
20 + const int64_t token_idx = blockIdx.x;
21 const int tid = threadIdx.x;
22
23 const scalar_t* token_in = input + token_idx * in_row_stride;
24 fp8_type* token_out = out + token_idx * out_row_stride;
25
21 - const float inv_scale = 1.0f / (*scale);
22 -
23 - vectorize_with_alignment<16>(
24 - token_in, token_out, hidden_size, tid, blockDim.x,
25 - [=] __device__(fp8_type & dst, const scalar_t& src) {
26 - dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src),
27 - inv_scale);
28 - });
26 + // Precompute row-level base offset for scale access (compile-time eliminated
27 + // when STRIDE_I_ZERO)
28 + const int64_t scale_row_base =
29 + STRIDE_I_ZERO ? 0
30 + : static_cast<int>(token_idx) / group_m * scale_stride_i;
31 +
32 + auto get_inv_scale = [&](int gj) {
33 + return 1.0f / scale[scale_row_base + gj * scale_stride_j];
34 + };
35 +
36 + int cached_gj = -1;
37 + float cached_inv_scale = 0.0f;
38 + auto get_inv_scale_cached = [&](int gj) {
39 + if (gj != cached_gj) {
40 + cached_inv_scale = 1.0f / scale[scale_row_base + gj * scale_stride_j];
41 + cached_gj = gj;
42 + }
43 + return cached_inv_scale;
44 + };
45 +
46 + constexpr int VEC_SIZE = 16; // FP8 so vectorize to 128 bits
47 + auto scaled_fp8_conversion_vectorized = [&](const scalar_t* in, fp8_type* out,
48 + int size, float inv_scale) {
49 + vectorize_with_alignment<VEC_SIZE>(
50 + in, out, size, tid, blockDim.x,
51 + [=] __device__(fp8_type & dst, const scalar_t& src) {
52 + dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src),
53 + inv_scale);
54 + });
55 + };
56 +
57 + if (STRIDE_J_ZERO && hidden_size % VEC_SIZE == 0) {
58 + // Per-tensor or per-token: single scale per row, vectorize full row
59 + scaled_fp8_conversion_vectorized(token_in, token_out, hidden_size,
60 + get_inv_scale(0));
61 + } else if (group_n % VEC_SIZE == 0) {
62 + // Multiple column groups with vectorization
63 + const int num_groups_n = hidden_size / group_n;
64 +
65 + for (int gj = 0; gj < num_groups_n; gj++) {
66 + scaled_fp8_conversion_vectorized(token_in + gj * group_n,
67 + token_out + gj * group_n, group_n,
68 + get_inv_scale(gj));
69 + }
70 + } else {
71 + // Scalar path for small column groups (group_n < VEC_SIZE)
72 + for (int n = tid; n < hidden_size; n += blockDim.x) {
73 + const int gj = n / group_n;
74 + token_out[n] = scaled_fp8_conversion<true, fp8_type>(
75 + static_cast<float>(token_in[n]), get_inv_scale_cached(gj));
76 + }
77 + }
78 }
79
80 template <typename scalar_t, typename fp8_type>
@@ -133,17 +182,116 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided(
182
183 } // namespace vllm
184
136 -void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
137 - torch::Tensor const& input, // [..., d]
138 - torch::Tensor const& scale) // [1]
185 +void static_scaled_fp8_quant(
186 + torch::Tensor& out, // [..., d]
187 + torch::Tensor const& input, // [..., d]
188 + torch::Tensor const& scale, // various shapes
189 + std::optional<std::tuple<int64_t, int64_t>>
190 + opt_group_shape) // optional explicit (group_m, group_n)
191 {
192 TORCH_CHECK(input.stride(-1) == 1,
193 "last dimension of input must be contiguous");
194 TORCH_CHECK(out.stride(-1) == 1,
195 "last dimension of output must be contiguous");
196
145 - const int hidden_size = input.size(-1);
146 - const int num_tokens = input.numel() / hidden_size;
197 + const int hidden_size = input.size(-1); // N (columns)
198 + const int num_tokens = input.numel() / hidden_size; // M (rows)
199 +
200 + // Determine group_m, group_n, and scale strides from scale shape
201 + // Scale indexing: scale[gi * scale_stride_j + gj * scale_stride_i]
202 + // where gi = m / group_m, gj = n / group_n
203 + int group_m, group_n;
204 + int64_t scale_stride_i, scale_stride_j;
205 +
206 + if (scale.dim() == 0 || scale.numel() == 1) {
207 + // Per-tensor: one scale for the entire tensor
208 + group_m = num_tokens;
209 + group_n = hidden_size;
210 + scale_stride_i = 0;
211 + scale_stride_j = 0;
212 + } else if (scale.dim() == 1) {
213 + // 1D scale: require explicit group_shape to disambiguate per-channel vs
214 + // per-token (avoids edge case where num_tokens == hidden_size)
215 + TORCH_CHECK(opt_group_shape.has_value(),
216 + "1D scale requires explicit group_shape to disambiguate "
217 + "per-channel vs per-token quantization. "
218 + "Use group_shape=(-1, 1) for per-channel or group_shape=(1, "
219 + "-1) for per-token.");
220 +
221 + const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
222 + group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
223 + group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
224 +
225 + // Validate the explicit group shape matches the 1D scale
226 + const int64_t scale_len = scale.numel();
227 + const int64_t expected_scale_m = num_tokens / group_m;
228 + const int64_t expected_scale_n = hidden_size / group_n;
229 + const int64_t expected_scale_numel = expected_scale_m * expected_scale_n;
230 +
231 + TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (",
232 + scale_len, ") does not match expected size (",
233 + expected_scale_numel, ") for group_shape (", opt_group_m, ", ",
234 + opt_group_n, ") with input shape (", num_tokens, ", ",
235 + hidden_size, ")");
236 +
237 + // For 1D scale, determine strides based on which dim is trivial
238 + // Scale indexing: scale[gi * scale_stride_i + gj * scale_stride_j]
239 + // where gi = m / group_m (row group), gj = n / group_n (col group)
240 + if (expected_scale_m == 1) {
241 + // Per-channel style: one scale in M dim, scale varies along N
242 + // gi = 0 always, gj varies, so stride_1 traverses the scale
243 + scale_stride_i = 0;
244 + scale_stride_j = scale.stride(0);
245 + } else if (expected_scale_n == 1) {
246 + // Per-token style: one scale in N dim, scale varies along M
247 + // gj = 0 always, gi varies, so stride_0 traverses the scale
248 + scale_stride_i = scale.stride(0);
249 + scale_stride_j = 0;
250 + } else {
251 + TORCH_CHECK(
252 + false,
253 + "1D scale can only be used when one of the scale dimensions is 1. "
254 + "For 2D group scaling, use a 2D scale tensor.");
255 + }
256 + } else if (scale.dim() == 2) {
257 + // 2D scale: infer group sizes from scale dimensions (or use explicit if
258 + // provided)
259 + const int64_t scale_size_0 = scale.size(0);
260 + const int64_t scale_size_1 = scale.size(1);
261 +
262 + TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens,
263 + ") must be divisible by scale.size(0) (", scale_size_0, ")");
264 + TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (", hidden_size,
265 + ") must be divisible by scale.size(1) (", scale_size_1, ")");
266 +
267 + // Infer from 2D scale shape
268 + int inferred_group_m = num_tokens / scale_size_0;
269 + int inferred_group_n = hidden_size / scale_size_1;
270 +
271 + // Use explicit if provided, otherwise use inferred
272 + if (opt_group_shape.has_value()) {
273 + const auto& [opt_group_m, opt_group_n] = opt_group_shape.value();
274 + group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m);
275 + group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n);
276 +
277 + // Validate explicit matches inferred
278 + TORCH_CHECK(group_m == inferred_group_m && group_n == inferred_group_n,
279 + "Explicit group_shape (", opt_group_m, ", ", opt_group_n,
280 + ") does not match inferred group shape (", inferred_group_m,
281 + ", ", inferred_group_n, ") from 2D scale tensor shape (",
282 + scale_size_0, ", ", scale_size_1, ")");
283 + } else {
284 + group_m = inferred_group_m;
285 + group_n = inferred_group_n;
286 + }
287 +
288 + scale_stride_i = scale.stride(0);
289 + scale_stride_j = scale.stride(1);
290 + } else {
291 + TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ",
292 + scale.dim(), "D");
293 + }
294 +
295 const int block_size = 256;
296 dim3 grid(num_tokens);
297 dim3 block(block_size);
@@ -153,15 +301,23 @@ void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
301
302 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
303 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
304 +
305 + // Dispatch to template-specialized kernel based on stride pattern
306 VLLM_DISPATCH_FLOATING_TYPES(
307 input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
308 VLLM_DISPATCH_FP8_TYPES(
309 out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
160 - vllm::scaled_fp8_quant_kernel_strided<scalar_t, fp8_t>
161 - <<<grid, block, 0, stream>>>(
162 - out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
163 - scale.data_ptr<float>(), hidden_size, in_row_stride,
164 - out_row_stride);
310 + VLLM_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] {
311 + VLLM_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] {
312 + vllm::scaled_fp8_quant_kernel_strided_group_shape<
313 + scalar_t, fp8_t, S0_ZERO, S1_ZERO>
314 + <<<grid, block, 0, stream>>>(
315 + out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
316 + scale.data_ptr<float>(), hidden_size, in_row_stride,
317 + out_row_stride, group_m, group_n, scale_stride_i,
318 + scale_stride_j);
319 + });
320 + });
321 });
322 });
323 }
csrc/torch_bindings.cpp
+5 -2
@@ -599,9 +599,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
599 ops.impl("gptq_shuffle", torch::kCUDA, &gptq_shuffle);
600
601 // Compute FP8 quantized tensor for given scaling factor.
602 + // Supports per-tensor, per-channel, per-token, and arbitrary 2D group
603 + // scaling. Optional group_m/group_n specify the group shape explicitly;
604 + // required for 1D scales to disambiguate per-channel vs per-token.
605 ops.def(
603 - "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale) -> "
604 - "()");
606 + "static_scaled_fp8_quant(Tensor! result, Tensor input, Tensor scale, "
607 + "(int, int)? group_shape=None) -> ()");
608 ops.impl("static_scaled_fp8_quant", torch::kCUDA, &static_scaled_fp8_quant);
609
610 // Compute dynamic-per-tensor FP8 quantized tensor and scaling factor.
tests/kernels/quantization/test_fp8_quant.py
+103 -2
@@ -11,6 +11,10 @@ from tests.kernels.quant_utils import (
11 ref_dynamic_per_token_quant,
12 )
13 from tests.kernels.utils import opcheck
14 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 + scaled_quantize,
16 +)
17 +from vllm.platforms import current_platform
18 from vllm.utils.torch_utils import set_random_seed
19
20 DTYPES = [torch.bfloat16, torch.float]
@@ -21,10 +25,18 @@ SEEDS = [0]
25
26
27 def opcheck_fp8_quant(
24 - output, input, scale=None, scale_ub=None, use_per_token_if_dynamic=False
28 + output,
29 + input,
30 + scale=None,
31 + scale_ub=None,
32 + use_per_token_if_dynamic=False,
33 + group_shape=None,
34 ):
35 if scale is not None:
27 - opcheck(torch.ops._C.static_scaled_fp8_quant, (output, input, scale))
36 + opcheck(
37 + torch.ops._C.static_scaled_fp8_quant,
38 + (output, input, scale, group_shape),
39 + )
40 elif use_per_token_if_dynamic:
41 scale = torch.empty(
42 (input.shape[0], 1), device=input.device, dtype=torch.float32
@@ -118,3 +130,92 @@ def test_fp8_quant_large(seed: int) -> None:
130 ops_out = ops_out.to(dtype=dtype)
131
132 torch.testing.assert_close(ref_out, ops_out)
133 +
134 +
135 +# Test static FP8 quantization with 2D group scales
136 +GROUP_SHAPES_2D = [
137 + (-1, -1), # Per-tensor
138 + (-1, 1), # Per-channel
139 + (1, -1), # Per-token
140 + (-1, 128), # Per-head quantization
141 + (1, 128), # DeepSeek-style per-token-per-group (group_m=1, group_n=128)
142 + (128, 128), # DeepSeek-style block quantization
143 + (1, 64), # Smaller group size
144 + (1, 16), # Small group (scalar path in kernel)
145 + (4, 256), # Non-trivial both dimensions
146 +]
147 +# Use sizes divisible by all group shapes
148 +NUM_TOKENS_GROUP = [128, 512]
149 +HIDDEN_SIZES_GROUP = [256, 1024, 2048]
150 +
151 +
152 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS_GROUP)
153 +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES_GROUP)
154 +@pytest.mark.parametrize("group_shape", GROUP_SHAPES_2D)
155 +@pytest.mark.parametrize("dtype", DTYPES)
156 +@pytest.mark.parametrize("seed", SEEDS)
157 +@torch.inference_mode()
158 +def test_static_fp8_quant_group_2d(
159 + num_tokens: int,
160 + hidden_size: int,
161 + group_shape: tuple[int, int],
162 + dtype: torch.dtype,
163 + seed: int,
164 +) -> None:
165 + """Test static FP8 quantization with 2D group scales using scaled_quantize."""
166 + # Normalize group_shape (-1 means full extent)
167 + norm_group_m = num_tokens if group_shape[0] == -1 else group_shape[0]
168 + norm_group_n = hidden_size if group_shape[1] == -1 else group_shape[1]
169 +
170 + # Skip if sizes are not divisible by group shape
171 + if num_tokens % norm_group_m != 0 or hidden_size % norm_group_n != 0:
172 + pytest.skip(
173 + f"Skipping: ({num_tokens}, {hidden_size}) not divisible by "
174 + f"group_shape ({group_shape[0]}, {group_shape[1]})"
175 + )
176 +
177 + current_platform.seed_everything(seed)
178 +
179 + x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
180 + ref_out, scale = scaled_quantize(
181 + x, group_shape, FP8_DTYPE, compute_dtype=torch.float32
182 + )
183 + ops_out, ops_scale = ops.scaled_fp8_quant(x, scale=scale, group_shape=group_shape)
184 +
185 + torch.testing.assert_close(scale, ops_scale)
186 + torch.testing.assert_close(ref_out.float(), ops_out.float(), rtol=0.12, atol=0.0)
187 +
188 + opcheck_fp8_quant(ops_out, x, scale=scale)
189 +
190 +
191 +@pytest.mark.parametrize("num_tokens", NUM_TOKENS_GROUP)
192 +@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES_GROUP)
193 +@pytest.mark.parametrize("dtype", DTYPES)
194 +@pytest.mark.parametrize("seed", SEEDS)
195 +@pytest.mark.parametrize("group_shape", [(1, -1), (-1, 1)]) # per-token, per-channel
196 +@torch.inference_mode()
197 +def test_static_fp8_quant_1d_scale(
198 + num_tokens: int,
199 + hidden_size: int,
200 + dtype: torch.dtype,
201 + seed: int,
202 + group_shape: tuple[int, int],
203 +) -> None:
204 + """Test static FP8 quantization with 1D scale (per-token or per-channel)."""
205 + current_platform.seed_everything(seed)
206 +
207 + x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
208 + ref_out, scale_2d = scaled_quantize(
209 + x, group_shape, FP8_DTYPE, compute_dtype=torch.float32
210 + )
211 +
212 + # Flatten scale to 1D for testing 1D scale path
213 + scale_1d = scale_2d.flatten()
214 + ops_out, ops_scale = ops.scaled_fp8_quant(
215 + x, scale=scale_1d, group_shape=group_shape
216 + )
217 +
218 + torch.testing.assert_close(scale_1d, ops_scale)
219 + torch.testing.assert_close(ref_out.float(), ops_out.float(), rtol=0.12, atol=0.0)
220 +
221 + opcheck_fp8_quant(ops_out, x, scale=scale_1d, group_shape=group_shape)
vllm/_custom_ops.py
+13 -4
@@ -1752,6 +1752,7 @@ def scaled_fp8_quant(
1752 scale_ub: torch.Tensor | None = None,
1753 use_per_token_if_dynamic: bool = False,
1754 output: torch.Tensor | None = None,
1755 + group_shape: tuple[int, int] | None = None,
1756 ) -> tuple[torch.Tensor, torch.Tensor]:
1757 """
1758 Quantize input tensor to FP8 and return quantized tensor and scale.
@@ -1763,14 +1764,23 @@ def scaled_fp8_quant(
1764 will benefit from padding.
1765
1766 Args:
1766 - input: The input tensor to be quantized to FP8
1767 - scale: Optional scaling factor for the FP8 quantization
1767 + input: The input tensor to be quantized to FP8 (must be 2D: [M, N])
1768 + scale: Optional scaling factor for the FP8 quantization. Supports:
1769 + - 0D or [1]: per-tensor scaling
1770 + - 1D: requires explicit group_shape to disambiguate per-channel
1771 + vs per-token (use (-1, 1) for per-channel, (1, -1) for per-token)
1772 + - 2D [M/group_m, N/group_n]: group scaling (e.g. [M, N/128] for
1773 + DeepSeek-style (1,128) groups, or [M/128, N/128] for (128,128))
1774 scale_ub: Optional upper bound for scaling factor in dynamic
1775 per token case
1776 num_token_padding: If specified, pad the first dimension
1777 of the output to at least this value.
1778 use_per_token_if_dynamic: Whether to do per_tensor or per_token
1779 in the dynamic quantization case.
1780 + group_shape: Optional tuple (group_m, group_n) specifying the group
1781 + shape for static quantization. Use -1 for "full extent" (e.g.,
1782 + (-1, -1) for per-tensor, (-1, 1) for per-channel, etc.)
1783 + Required for 1D scales; optional for 2D scales.
1784
1785 Returns:
1786 tuple[torch.Tensor, torch.Tensor]: The output tensor in FP8 and
@@ -1799,8 +1809,7 @@ def scaled_fp8_quant(
1809 scale = torch.empty(1, device=input.device, dtype=torch.float32)
1810 torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)
1811 else:
1802 - assert scale.numel() == 1, f"{scale.shape}"
1803 - torch.ops._C.static_scaled_fp8_quant(output, input, scale)
1812 + torch.ops._C.static_scaled_fp8_quant(output, input, scale, group_shape)
1813
1814 return output, scale
1815
vllm/model_executor/layers/quantization/input_quant_fp8.py
+17 -11
@@ -10,6 +10,7 @@ from vllm.model_executor.custom_op import CustomOp
10 from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 GroupShape,
12 get_fp8_min_max,
13 + group_broadcast,
14 )
15 from vllm.platforms import current_platform
16
@@ -22,7 +23,7 @@ _FP8_MIN_SCALING_FACTOR = 1.0 / (_FP8_MAX * 512.0)
23 @CustomOp.register("quant_fp8")
24 class QuantFP8(CustomOp):
25 """
25 - Quantize input tensor to FP8 (per-tensor, per-token, or per-group).
26 + Quantize input tensor to FP8 (per-tensor, per-token, per-channel, or per-group).
27 This CustomOp supports both static and dynamic quantization.
28 """
29
@@ -57,14 +58,14 @@ class QuantFP8(CustomOp):
58
59 self.is_group_quant = group_shape.is_per_group()
60 if self.is_group_quant:
60 - assert not static, "Group quantization only supports dynamic mode"
61 self.group_size = group_shape.col
62 else:
63 - assert group_shape in {GroupShape.PER_TOKEN, GroupShape.PER_TENSOR}
64 - assert not static or group_shape == GroupShape.PER_TENSOR, (
65 - "Only per-tensor scales supported for static quantization."
66 - )
63 self.use_per_token_if_dynamic = group_shape == GroupShape.PER_TOKEN
64 + if not static:
65 + assert group_shape in (GroupShape.PER_TOKEN, GroupShape.PER_TENSOR), (
66 + "Only per-token or per-tensor scales are supported for dynamic "
67 + "non-group quantization."
68 + )
69
70 def forward_cuda(
71 self,
@@ -72,8 +73,8 @@ class QuantFP8(CustomOp):
73 scale: torch.Tensor | None = None,
74 scale_ub: torch.Tensor | None = None,
75 ) -> tuple[torch.Tensor, torch.Tensor]:
75 - if self.is_group_quant:
76 - assert scale is None, "Group quantization is always dynamic"
76 + if self.is_group_quant and not self.static:
77 + assert scale is None, "Dynamic group quantization does not use scale"
78 from vllm.model_executor.layers.quantization.utils import fp8_utils
79
80 return fp8_utils.per_token_group_quant_fp8(
@@ -90,12 +91,14 @@ class QuantFP8(CustomOp):
91 and self.group_shape == GroupShape.PER_TOKEN
92 and scale_ub.numel() == 1
93 )
94 +
95 return ops.scaled_fp8_quant(
96 x,
97 scale,
98 num_token_padding=self.num_token_padding,
99 scale_ub=scale_ub,
100 use_per_token_if_dynamic=self.use_per_token_if_dynamic,
101 + group_shape=self.group_shape if self.static else None,
102 )
103
104 def forward_hip(
@@ -131,8 +134,8 @@ class QuantFP8(CustomOp):
134 scale: torch.Tensor | None = None,
135 scale_ub: torch.Tensor | None = None,
136 ):
134 - if self.is_group_quant:
135 - assert scale is None, "Group quantization is always dynamic"
137 + if self.is_group_quant and not self.static:
138 + assert scale is None, "Dynamic group quantization does not use scale"
139 return self._quantize_group_native(x)
140
141 assert (scale is not None) == self.static
@@ -155,7 +158,10 @@ class QuantFP8(CustomOp):
158
159 # Even for dynamic per-token scales,
160 # reciprocal performs slightly better than division
158 - out = x.to(torch.float32) * scale.reciprocal()
161 + out = (
162 + x.to(torch.float32)
163 + * group_broadcast(scale.to(torch.float32), x.shape[-2:]).reciprocal()
164 + )
165 out = out.clamp(_FP8_MIN, _FP8_MAX).to(_FP8_DTYPE)
166
167 # This currently generates an extra Triton kernel in compilation.
vllm/model_executor/layers/quantization/utils/quant_utils.py
+19 -4
@@ -158,11 +158,14 @@ def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape):
158 # with an extent of 1, since this can be done implicitly by pytorch
159 def group_broadcast(t, shape):
160 for i, s in enumerate(shape):
161 - if t.shape[i] != s and t.shape[i] != 1:
162 - assert s % t.shape[i] == 0
161 + # If tensor has fewer dimensions than target shape, treat missing
162 + # dimensions as size 1 (standard PyTorch broadcasting behavior)
163 + t_dim_size = t.shape[i] if i < t.ndim else 1
164 + if t_dim_size != s and t_dim_size != 1:
165 + assert s % t_dim_size == 0
166 t = (
167 t.unsqueeze(i + 1)
165 - .expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
168 + .expand(*t.shape[: i + 1], s // t_dim_size, *t.shape[i + 1 :])
169 .flatten(i, i + 1)
170 )
171 return t
@@ -180,7 +183,16 @@ def scaled_quantize(
183 x: torch.Tensor,
184 group_shape: GroupShape,
185 quant_dtype: torch.dtype,
186 + compute_dtype: torch.dtype | None = None,
187 ) -> tuple[torch.Tensor, torch.Tensor]:
188 + """
189 + Args:
190 + x: Input tensor to quantize
191 + group_shape: Shape of quantization groups
192 + quant_dtype: Target quantized dtype (e.g., torch.float8_e4m3fn)
193 + compute_dtype: Optional dtype for intermediate computations.
194 + If None, uses input dtype. Use torch.float32 for higher precision.
195 + """
196 group_shape = _normalize_quant_group_shape(x, group_shape)
197 assert quant_dtype.is_floating_point, (
198 "currently `scaled_quantize` only supports floating point dtypes "
@@ -189,11 +201,14 @@ def scaled_quantize(
201
202 finfo = torch.finfo(quant_dtype)
203
204 + # Convert to compute dtype if specified
205 + x_compute = x if compute_dtype is None else x.to(compute_dtype)
206 +
207 # Reshape (M, N) into (BLK_M, BLOCK_SIZE_M, BLK_N, BLOCK_SIZE_N)
208 assert x.ndim == 2
209 assert x.shape[0] % group_shape[0] == 0 and x.shape[1] % group_shape[1] == 0
210 blk_m, blk_n = x.shape[0] // group_shape[0], x.shape[1] // group_shape[1]
196 - x_blkd = x.reshape(blk_m, group_shape[0], blk_n, group_shape[1])
211 + x_blkd = x_compute.reshape(blk_m, group_shape[0], blk_n, group_shape[1])
212
213 # Permute to (BLK_M, BLK_N, BLOCK_SIZE_M, BLOCK_SIZE_N)
214 x_blkd_permd = x_blkd.permute(0, 2, 1, 3)