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, 06:30 UTC
[Quantization] fix: overflow with static per-tensor scaling (#29867)
Signed-off-by: Mickael Seznec <mickael@mistral.ai> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Mickaël Seznec committed
Jan 13, 2026 at 13:56 UTC
a5bbbd2f24983f309f915a1de911b9301e1699e2
2 files changed
+71
-56
vllm/model_executor/layers/quantization/utils/quant_utils.py
+61
-2
@@ -5,7 +5,7 @@
5
from collections.abc import Callable, Mapping
6
from dataclasses import dataclass
7
from types import MappingProxyType
8
-from typing import ClassVar, NamedTuple
8
+from typing import TYPE_CHECKING, ClassVar, NamedTuple
9
10
import numpy
11
import torch
@@ -15,6 +15,9 @@ from vllm._custom_ops import cutlass_scaled_mm_supports_fp4
15
from vllm.platforms import current_platform
16
from vllm.scalar_type import ScalarType, scalar_types
17
18
+if TYPE_CHECKING:
19
+ from vllm.model_executor.layers.linear import LinearBase
20
+
21
FP8_DTYPE = current_platform.fp8_dtype()
22
FP4_DTYPE = torch.uint8
23
@@ -239,7 +242,7 @@ def scaled_dequantize(
242
x_s: torch.Tensor,
243
group_shape: GroupShape | None = None,
244
out_dtype: torch.dtype = torch.float32,
242
-) -> tuple[torch.Tensor, torch.Tensor]:
245
+) -> torch.Tensor:
246
if group_shape is not None:
247
group_shape = _normalize_quant_group_shape(x_q, group_shape)
248
@@ -270,6 +273,62 @@ def scaled_dequantize(
273
return (x_q.to(torch.float32) * x_s).to(out_dtype)
274
275
276
+def get_attribute_fallback(obj, attributes: list[str]):
277
+ for attr in attributes:
278
+ if hasattr(obj, attr):
279
+ return getattr(obj, attr)
280
+ raise AttributeError(f"'{obj}' has no recognized attributes: {attributes}.")
281
+
282
+
283
+def get_and_maybe_dequant_weights(
284
+ layer: "LinearBase", out_dtype: torch.dtype = torch.float32
285
+):
286
+ """Return layer's unquantized weights in [out, in] layout"""
287
+ from vllm.model_executor.layers.linear import UnquantizedLinearMethod
288
+ from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
289
+
290
+ weight = get_attribute_fallback(layer, ["weight", "qweight", "weight_packed"])
291
+
292
+ # Unquantized layer: just return base weights
293
+ if layer.quant_method is None or isinstance(
294
+ layer.quant_method, UnquantizedLinearMethod
295
+ ):
296
+ return weight.to(out_dtype)
297
+
298
+ # Simple Fp8 case: rescale with tensor or block weight scales
299
+ if (
300
+ isinstance(layer.quant_method, Fp8LinearMethod)
301
+ and not layer.quant_method.use_marlin
302
+ ):
303
+ weight_scales = get_attribute_fallback(
304
+ layer, ["weight_scale", "weight_scale_inv"]
305
+ )
306
+ dequant_weights = scaled_dequantize(
307
+ weight,
308
+ weight_scales,
309
+ group_shape=layer.weight_block_size,
310
+ out_dtype=out_dtype,
311
+ )
312
+ # per-tensor scaling stores weights in [in, out] layout
313
+ if not layer.quant_method.block_quant:
314
+ dequant_weights = dequant_weights.T
315
+ return dequant_weights
316
+
317
+ # NOTE: Most generic base case
318
+ # - Call the layer with identity matrix which returns unquantized weights.
319
+ # - Must be used with extra care when dealing with static activation quantization:
320
+ # quantizing 1.0 may lead to over/underflows
321
+ # - Should only be used offline, since it's O(N^3)
322
+ assert hasattr(layer, "input_size_per_partition")
323
+ eye = torch.eye(
324
+ layer.input_size_per_partition,
325
+ dtype=out_dtype,
326
+ device=weight.device,
327
+ )
328
+ dequant_weights = layer.quant_method.apply(layer, eye, bias=None).to(out_dtype)
329
+ return dequant_weights.T
330
+
331
+
332
def pack_quantized_values_into_int32(
333
w_q: torch.Tensor, wtype: ScalarType, packed_dim: int = 0
334
):
vllm/v1/attention/backends/mla/common.py
+10
-54
@@ -207,8 +207,9 @@ from vllm.model_executor.layers.batch_invariant import (
207
)
208
from vllm.model_executor.layers.linear import (
209
ColumnParallelLinear,
210
- LinearBase,
211
- UnquantizedLinearMethod,
210
+)
211
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
212
+ get_and_maybe_dequant_weights,
213
)
214
from vllm.platforms import current_platform
215
from vllm.utils.flashinfer import has_nvidia_artifactory
@@ -1184,35 +1185,13 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1185
self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()
1186
1187
def process_weights_after_loading(self, act_dtype: torch.dtype):
1187
- def get_layer_weight(layer):
1188
- WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
1189
- for attr in WEIGHT_NAMES:
1190
- if hasattr(layer, attr):
1191
- return getattr(layer, attr)
1192
- raise AttributeError(
1193
- f"Layer '{layer}' has no recognized weight attribute: {WEIGHT_NAMES}."
1194
- )
1195
-
1196
- def get_and_maybe_dequant_weights(layer: LinearBase):
1197
- if layer.quant_method is not None and not isinstance(
1198
- layer.quant_method, UnquantizedLinearMethod
1199
- ):
1200
- # NOTE: This should only be used offline, since it's O(N^3)
1201
- eye = torch.eye(
1202
- layer.input_size_per_partition,
1203
- dtype=act_dtype,
1204
- device=get_layer_weight(layer).device,
1205
- )
1206
- dequant_weights = layer.quant_method.apply(layer, eye, bias=None)
1207
- del eye
1208
- # standardize to (output, input)
1209
- return dequant_weights.T
1210
- return layer.weight
1211
-
1188
# we currently do not have quantized bmm's which are needed for
1189
# `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
1190
# the bmm's in 16-bit, the extra memory overhead of this is fairly low
1215
- kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
1191
+ kv_b_proj_weight = get_and_maybe_dequant_weights(
1192
+ self.kv_b_proj, out_dtype=act_dtype
1193
+ ).T
1194
+
1195
assert kv_b_proj_weight.shape == (
1196
self.kv_lora_rank,
1197
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
@@ -1600,35 +1579,12 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1579
return attn_out, lse.transpose(0, 1).contiguous()
1580
1581
def process_weights_after_loading(self, act_dtype: torch.dtype):
1603
- def get_layer_weight(layer):
1604
- WEIGHT_NAMES = ("weight", "qweight", "weight_packed")
1605
- for attr in WEIGHT_NAMES:
1606
- if hasattr(layer, attr):
1607
- return getattr(layer, attr)
1608
- raise AttributeError(
1609
- f"Layer '{layer}' has no recognized weight attribute: {WEIGHT_NAMES}."
1610
- )
1611
-
1612
- def get_and_maybe_dequant_weights(layer: LinearBase):
1613
- if layer.quant_method is not None and not isinstance(
1614
- layer.quant_method, UnquantizedLinearMethod
1615
- ):
1616
- # NOTE: This should only be used offline, since it's O(N^3)
1617
- eye = torch.eye(
1618
- layer.input_size_per_partition,
1619
- dtype=act_dtype,
1620
- device=get_layer_weight(layer).device,
1621
- )
1622
- dequant_weights = layer.quant_method.apply(layer, eye, bias=None)
1623
- del eye
1624
- # standardize to (output, input)
1625
- return dequant_weights.T
1626
- return layer.weight
1627
-
1582
# we currently do not have quantized bmm's which are needed for
1583
# `W_UV` and `W_UK_T`, we just store fp16/bf16 copies and perform
1584
# the bmm's in 16-bit, the extra memory overhead of this is fairly low
1631
- kv_b_proj_weight = get_and_maybe_dequant_weights(self.kv_b_proj).T
1585
+ kv_b_proj_weight = get_and_maybe_dequant_weights(
1586
+ self.kv_b_proj, out_dtype=act_dtype
1587
+ ).T
1588
assert kv_b_proj_weight.shape == (
1589
self.kv_lora_rank,
1590
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),