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, 22:45 UTC
[BugFix] register quant scale tensors as buffer (#31395)
Signed-off-by: Boyuan Feng <boyuan@meta.com>
Boyuan Feng committed
Dec 27, 2025 at 19:20 UTC
62def07d6786b39cd51b896b254fb2a40cc0f503
1 file changed
+82
-16
vllm/attention/layer.py
+82
-16
@@ -28,6 +28,7 @@ from vllm.model_executor.layers.linear import (
28
UnquantizedLinearMethod,
29
)
30
from vllm.model_executor.layers.quantization import QuantizationConfig
31
+from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
32
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
33
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
34
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
@@ -46,6 +47,35 @@ from vllm.v1.kv_cache_interface import (
47
logger = init_logger(__name__)
48
49
50
+def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
51
+ """Returns whether the quantization method should load quantized weights."""
52
+ return quant_method is not None and not isinstance(
53
+ quant_method, UnquantizedLinearMethod
54
+ )
55
+
56
+
57
+def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) -> None:
58
+ """Sets default quantization scales for the layer."""
59
+ if register_buffer:
60
+ layer.register_buffer("_k_scale", torch.tensor(1.0, dtype=torch.float32))
61
+ layer.register_buffer("_v_scale", torch.tensor(1.0, dtype=torch.float32))
62
+ layer.register_buffer("_q_scale", torch.tensor(1.0, dtype=torch.float32))
63
+ layer.register_buffer("_prob_scale", torch.tensor(1.0, dtype=torch.float32))
64
+ else:
65
+ layer._k_scale.fill_(1.0)
66
+ layer._v_scale.fill_(1.0)
67
+ layer._q_scale.fill_(1.0)
68
+ layer._prob_scale.fill_(1.0)
69
+
70
+ # We also keep q/k/v_scale on host (cpu) memory for attention
71
+ # backends that require the scales to be on host instead of on device.
72
+ # e.g. Flashinfer
73
+ layer._q_scale_float = 1.0
74
+ layer._k_scale_float = 1.0
75
+ layer._v_scale_float = 1.0
76
+ layer._prob_scale_float = 1.0
77
+
78
+
79
def _init_kv_cache_quant(
80
layer: nn.Module,
81
quant_config: QuantizationConfig | None,
@@ -74,17 +104,21 @@ def _init_kv_cache_quant(
104
# with the model weights.
105
layer.kv_cache_dtype = kv_cache_dtype
106
layer.calculate_kv_scales = calculate_kv_scales
77
- layer._k_scale = torch.tensor(1.0, dtype=torch.float32)
78
- layer._v_scale = torch.tensor(1.0, dtype=torch.float32)
79
- layer._q_scale = torch.tensor(1.0, dtype=torch.float32)
80
- layer._prob_scale = torch.tensor(1.0, dtype=torch.float32)
107
82
- # We also keep q/k/v_scale on host (cpu) memory for attention
83
- # backends that require the scales to be on host instead of on device.
84
- # e.g. Flashinfer
85
- layer._q_scale_float = 1.0
86
- layer._k_scale_float = 1.0
87
- layer._v_scale_float = 1.0
108
+ # Note [Register q/k/v/prob scales in state dict]
109
+ # When calling model.to(device), only parameters/buffers in state dict are
110
+ # moved. If not registering q/k/v/prob scales in state dict, there would
111
+ # be an IMA error when a cuda kernel (e.g., quant_fp8) accesses the tensor
112
+ # on cpu.
113
+ # Registering in state dict means it interacts with weight loading. One edge
114
+ # case is when quant_method is None, or quant_method is UnquantizedLinearMethod
115
+ # (i.e., should_load_quant_weights(quant_method) == False).
116
+ # In this case, the checkpoint does not have the scales. We need to
117
+ # initialize the scales to 1.0 and update the scales after weight loading.
118
+ # This is espectially important when we load dummy weights first (providing
119
+ # wrong scales) and then load real weights (which misses scales and keeps the
120
+ # wrong scales from dummy load).
121
+ set_default_quant_scales(layer, register_buffer=True)
122
123
# The output scale on host memory. This should be the input scale of
124
# the quant op after this attention layer.
@@ -93,9 +127,9 @@ def _init_kv_cache_quant(
127
quant_method = (
128
quant_config.get_quant_method(layer, prefix=prefix) if quant_config else None
129
)
96
- if quant_method is not None and not isinstance(
97
- quant_method, UnquantizedLinearMethod
98
- ):
130
+
131
+ # See [Note: Register q/k/v/prob scales in state dict]
132
+ if should_load_quant_weights(quant_method):
133
assert isinstance(quant_method, BaseKVCacheMethod)
134
# TODO (mgoin): kv cache dtype should be specified in the FP8
135
# checkpoint config and become the "auto" behavior
@@ -169,10 +203,16 @@ class Attention(nn.Module, AttentionLayerBase):
203
assert num_heads % num_kv_heads == 0, (
204
f"num_heads ({num_heads}) is not divisible by num_kv_heads ({num_kv_heads})"
205
)
206
+ self.quant_config = quant_config
207
+ self.layer_name = prefix
208
209
# Initialize KV cache quantization attributes
210
_init_kv_cache_quant(
175
- self, quant_config, prefix, kv_cache_dtype, calculate_kv_scales
211
+ self,
212
+ self.quant_config,
213
+ self.layer_name,
214
+ kv_cache_dtype,
215
+ calculate_kv_scales,
216
)
217
218
self.num_heads = num_heads
@@ -249,7 +289,6 @@ class Attention(nn.Module, AttentionLayerBase):
289
if prefix in compilation_config.static_forward_context:
290
raise ValueError(f"Duplicate layer name: {prefix}")
291
compilation_config.static_forward_context[prefix] = self
252
- self.layer_name = prefix
292
self.attn_type = attn_type
293
294
if kv_sharing_target_layer_name is not None:
@@ -378,6 +417,17 @@ class Attention(nn.Module, AttentionLayerBase):
417
def process_weights_after_loading(self, act_dtype: torch.dtype):
418
self.impl.process_weights_after_loading(act_dtype)
419
420
+ # If we should not load quant weights, we initialize the scales to 1.0
421
+ # as the default value. See [Note: Register q/k/v/prob scales in state dict]
422
+ # for more details.
423
+ quant_method = (
424
+ self.quant_config.get_quant_method(self, prefix=self.layer_name)
425
+ if self.quant_config
426
+ else None
427
+ )
428
+ if not should_load_quant_weights(quant_method):
429
+ set_default_quant_scales(self, register_buffer=False)
430
+
431
def get_attn_backend(self) -> type[AttentionBackend]:
432
return self.attn_backend
433
@@ -453,10 +503,15 @@ class MLAAttention(nn.Module, AttentionLayerBase):
503
kv_cache_dtype = "auto"
504
block_size = 16
505
calculate_kv_scales = False
506
+ self.quant_config = quant_config
507
508
# Initialize KV cache quantization attributes
509
_init_kv_cache_quant(
459
- self, quant_config, prefix, kv_cache_dtype, calculate_kv_scales
510
+ self,
511
+ self.quant_config,
512
+ self.layer_name,
513
+ kv_cache_dtype,
514
+ calculate_kv_scales,
515
)
516
517
dtype = torch.get_default_dtype()
@@ -586,6 +641,17 @@ class MLAAttention(nn.Module, AttentionLayerBase):
641
if hasattr(self.impl, "process_weights_after_loading"):
642
self.impl.process_weights_after_loading(act_dtype)
643
644
+ # If we should not load quant weights, we initialize the scales to 1.0
645
+ # as the default value. See [Note: Register q/k/v/prob scales in state dict]
646
+ # for more details.
647
+ quant_method = (
648
+ self.quant_config.get_quant_method(self, prefix=self.layer_name)
649
+ if self.quant_config
650
+ else None
651
+ )
652
+ if not should_load_quant_weights(quant_method):
653
+ set_default_quant_scales(self, register_buffer=False)
654
+
655
def calc_kv_scales(
656
self, q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor
657
) -> None: