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 5, 02:15 UTC
[Quantization] - Layerwise reloading of Attention/KV quantized models (#38995)
Signed-off-by: Josephasafg <ajgard7@gmail.com>
Asaf Gardin committed
Apr 16, 2026 at 04:03 UTC
19fa90ed0d616f3c0bc0dbb3bf6bfb7772df7f32
4 files changed
+98
-27
tests/model_executor/model_loader/test_reload.py
+28
@@ -164,6 +164,34 @@ def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
164
assert add_perp < mul_perp
165
166
167
+def test_kv_scale_reload(vllm_runner):
168
+ """Test reloading a checkpoint that contains k_scale/v_scale weights."""
169
+ if not current_platform.supports_fp8():
170
+ pytest.skip(reason="Requires FP8 support")
171
+
172
+ model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV"
173
+
174
+ # Load dummy weights, then reload real checkpoint
175
+ with vllm_runner(
176
+ model_name=model,
177
+ load_format="dummy",
178
+ enable_prefix_caching=False,
179
+ max_model_len=16,
180
+ max_num_seqs=1,
181
+ ) as llm:
182
+ llm.collective_rpc(
183
+ "update_config",
184
+ kwargs={"overrides": {"load_config": {"load_format": "auto"}}},
185
+ )
186
+ llm.collective_rpc("reload_weights", kwargs={"weights_path": model})
187
+ reloaded_perp = llm.generate_prompt_perplexity(
188
+ ["The capital of France is the city of Paris"],
189
+ mask=["The capital of France is"],
190
+ )[0]
191
+
192
+ assert reloaded_perp < 10
193
+
194
+
195
@pytest.mark.parametrize(
196
"tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
197
)
vllm/model_executor/model_loader/reload/__init__.py
+2
-3
@@ -8,10 +8,9 @@ which is useful for weight updates without full model reconstruction.
8
9
Limitations:
10
1. Composition with CPU offloading has not been implemented
11
-2. Reloading Attention/MLA weights (q_scale, k_scale, v_scale) has not been implemented
12
-3. Tied parameters will only reflect processing from one of the parent layers (for
11
+2. Tied parameters will only reflect processing from one of the parent layers (for
12
example, only processing from embed_tokens will have an effect)
14
-4. This design assumes that the number of weights loaded from disk is the same as the
13
+3. This design assumes that the number of weights loaded from disk is the same as the
14
number of weights created at model init time. This is not true for quant methods
15
which (1) pad weights or (2) load qkv weights into the same parameter. Both of these
16
cases are non-issues for today's quant methods, but future quantizations may cause
vllm/model_executor/model_loader/reload/layerwise.py
+66
-22
@@ -200,6 +200,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
200
if hasattr(model, "_original_do_torchao_reload"):
201
model._do_torchao_reload = model._original_do_torchao_reload
202
203
+ deferred_attn: list[tuple[torch.nn.Module, LayerReloadingInfo]] = []
204
+
205
for layer in model.modules():
206
info = get_layerwise_info(layer)
207
if not info.can_load():
@@ -208,22 +210,11 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
210
211
# Attention/MLA layers are processed after all other layers
212
if isinstance(layer, (Attention, MLAAttention)):
211
- if info.load_numel > 0:
212
- raise NotImplementedError(
213
- "Layerwise reloading of Q/K/V scale weights is not implemented yet"
214
- )
215
-
216
- elif info.kernel_tensors is None:
217
- raise NotImplementedError(
218
- "Layerwise loading of Q/K/V scale weights is not implemented yet"
219
- )
220
-
221
- else:
222
- _place_kernel_tensors(layer, info)
223
- layer.process_weights_after_loading(model_config.dtype)
213
+ deferred_attn.append((layer, info))
214
+ continue
215
216
# No weights were loaded
226
- elif info.load_numel <= 0:
217
+ if info.load_numel <= 0:
218
# first load: checkpoint did not contain weights for this layer
219
if info.kernel_tensors is None:
220
_layerwise_process(layer, info)
@@ -244,11 +235,58 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon
235
236
info.reset()
237
238
+ # Process attention layers after all other layers are done
239
+ for layer, info in deferred_attn:
240
+ _finalize_attention_layer(layer, info, model_config)
241
+ info.reset()
242
+
243
244
def finalize_layerwise_reload(*args, **kwargs):
245
finalize_layerwise_processing(*args, **kwargs)
246
247
248
+def _finalize_attention_layer(
249
+ layer: torch.nn.Module, info: LayerReloadingInfo, model_config: ModelConfig
250
+) -> None:
251
+ if info.load_numel > 0 and info.kernel_tensors is not None:
252
+ # Reload with new scale weights from checkpoint
253
+ _place_kernel_tensors(layer, info)
254
+ _reload_attention_scales(layer, info)
255
+ elif info.load_numel > 0 or info.kernel_tensors is None:
256
+ raise ValueError(
257
+ "Layerwise loading of attention layers is not supported. "
258
+ "Attention must always process after linears."
259
+ )
260
+ else:
261
+ _place_kernel_tensors(layer, info)
262
+ layer.process_weights_after_loading(model_config.dtype)
263
+
264
+
265
+def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) -> None:
266
+ """Load and process attention scale weights (k_scale, v_scale, etc.)
267
+ during reload.
268
+
269
+ Assumes dtype/shapes of attention tensors do not change during
270
+ processing, since we use .data.copy_() to preserve kernel tensor
271
+ references."""
272
+ quant_method = getattr(layer, "quant_method", None)
273
+ if quant_method is None:
274
+ return
275
+
276
+ # Re-create scale Parameters with sentinel values so unloaded scales
277
+ # are correctly detected by process_weights_after_loading
278
+ quant_method.create_weights(layer)
279
+
280
+ for name, args in info.loaded_weights:
281
+ param = getattr(layer, name)
282
+ args.arguments["param"] = param
283
+ _get_weight_loader(param)(*args.args, **args.kwargs)
284
+
285
+ quant_method.process_weights_after_loading(layer)
286
+
287
+ _copy_and_restore_kernel_tensors(layer, info)
288
+
289
+
290
def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
291
"""
292
Finalize layer loading after all weights have been buffered.
@@ -278,7 +316,6 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
316
param.weight_loader(*args.args, **args.kwargs)
317
318
# Process weights (quantization, repacking, etc.)
281
- # Attention/MLA are processed in `finalize_layerwise_reload`
319
quant_method = getattr(layer, "quant_method", None)
320
if isinstance(quant_method, QuantizeMethodBase):
321
quant_method.process_weights_after_loading(layer)
@@ -286,13 +323,7 @@ def _layerwise_process(layer: torch.nn.Module, info: LayerReloadingInfo):
323
# Copy processed values into original tensor storage (preserves cudagraph refs)
324
# this code is a no-op if not reloading (because kernel tensors is empty)
325
if info.kernel_tensors is not None:
289
- parameters, buffers = info.kernel_tensors
290
- for name, param in parameters.items():
291
- param.data.copy_(getattr(layer, name))
292
- for name, buffer in buffers.items():
293
- buffer.data.copy_(getattr(layer, name))
294
-
295
- _place_kernel_tensors(layer, info)
326
+ _copy_and_restore_kernel_tensors(layer, info)
327
328
info.reset()
329
logger.debug("%s: Processed", layer.__class__.__name__)
@@ -311,6 +342,19 @@ def _get_weight_loader(tensor: torch.Tensor):
342
return getattr(tensor, "weight_loader", default_weight_loader)
343
344
345
+def _copy_and_restore_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
346
+ """Copy processed values into original kernel tensor storage and restore
347
+ kernel tensor references on the layer. Preserves cudagraph references."""
348
+ assert info.kernel_tensors is not None
349
+ parameters, buffers = info.kernel_tensors
350
+ for name, param in parameters.items():
351
+ param.data.copy_(getattr(layer, name))
352
+ for name, buffer in buffers.items():
353
+ buffer.data.copy_(getattr(layer, name))
354
+
355
+ _place_kernel_tensors(layer, info)
356
+
357
+
358
def _place_kernel_tensors(layer: torch.nn.Module, info: LayerReloadingInfo):
359
for name in get_layer_tensors(layer):
360
delattr(layer, name)
vllm/model_executor/model_loader/weight_utils.py
+2
-2
@@ -1364,8 +1364,8 @@ def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> N
1364
if param.numel() == 1 and loaded_weight.numel() == 1:
1365
# Sometimes scalar values aren't considered tensors with shapes
1366
# so if both param and loaded_weight are a scalar,
1367
- # "broadcast" instead of copy
1368
- param.data.fill_(loaded_weight.item())
1367
+ # reshape to match before copying
1368
+ param.data.copy_(loaded_weight.view(param.shape))
1369
else:
1370
assert param.size() == loaded_weight.size(), (
1371
f"Attempted to load weight ({loaded_weight.size()}) "