[Bugfix][Quantization] Stop sleep(level=2) from zeroing compressed-tensors KV scales (#57163)
Signed-off-by: AlanFokCo <alanfok2868@gmail.com> Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai>
AlanFokCo committed
Sep 29, 2026 at 23:34 UTC
f4917dadc8b44b75e880d870588f9ddbd874f412
2 files changed
+58
-3
tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py
+54
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import asyncio
5
+import math
6
7
import pytest
8
import torch
@@ -9,6 +10,7 @@ import torch
10
import vllm.device_allocator.cumem as cumem
11
import vllm.envs as envs
12
from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
13
+from vllm.model_executor.model_loader import get_model_loader
14
from vllm.platforms import current_platform
15
from vllm.utils.mem_constants import GiB_bytes
16
@@ -232,6 +234,58 @@ def test_deep_sleep_async(gpu_memory_cleared):
234
asyncio.run(test())
235
236
237
+def _load_weights(worker) -> None:
238
+ # Not reload_weights(): its layerwise finalize re-derives the KV scales.
239
+ model = worker.model_runner.get_model()
240
+ loader = get_model_loader(worker.load_config)
241
+ model.load_weights(loader.get_all_weights(worker.model_config, model))
242
+
243
+
244
+def _attn_kv_scales(model) -> dict[str, list[float]]:
245
+ return {
246
+ f"{name}.{attr}": getattr(module, attr).flatten().tolist()
247
+ for name, module in model.named_modules()
248
+ for attr in ("_q_scale", "_k_scale", "_v_scale")
249
+ if isinstance(getattr(module, attr, None), torch.Tensor)
250
+ }
251
+
252
+
253
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Segfaults on ROCm.")
254
+@requires_fp8
255
+@create_new_process_for_each_test()
256
+def test_deep_sleep_compressed_tensors_kv_scales(
257
+ monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared
258
+):
259
+ """KV scales must survive level-2 sleep plus a weights reload (else NaN)."""
260
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
261
+ llm = LLM(
262
+ "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor",
263
+ enable_sleep_mode=True,
264
+ kv_cache_dtype="fp8",
265
+ enforce_eager=True,
266
+ gpu_memory_utilization=0.6,
267
+ )
268
+ sampling_params = SamplingParams(temperature=0, max_tokens=10, logprobs=0)
269
+ scales = llm.apply_model(_attn_kv_scales)[0]
270
+ assert any(v != 1.0 for values in scales.values() for v in values)
271
+ expected = llm.generate("How are you?", sampling_params)[0].outputs[0]
272
+
273
+ llm.sleep(level=2)
274
+ llm.wake_up(tags=["weights"])
275
+ llm.collective_rpc(_load_weights)
276
+ llm.wake_up(tags=["kv_cache"])
277
+
278
+ assert llm.apply_model(_attn_kv_scales)[0] == scales
279
+ actual = llm.generate("How are you?", sampling_params)[0].outputs[0]
280
+ assert actual.token_ids == expected.token_ids
281
+ logprobs = [
282
+ [lp[t].logprob for lp, t in zip(o.logprobs, o.token_ids)]
283
+ for o in (expected, actual)
284
+ ]
285
+ assert all(math.isfinite(x) for x in logprobs[1])
286
+ assert logprobs[0] == logprobs[1]
287
+
288
+
289
@requires_fp8
290
def test_deep_sleep_fp8_kvcache_mrv1(
291
monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+4
-3
@@ -1208,9 +1208,10 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
1208
"""Override the default vLLM placeholder scales with the llm-compressor loaded
1209
scales. Zero points are not used as only symmetric quantization is supported.
1210
"""
1211
- layer._k_scale = layer.k_scale
1212
- layer._v_scale = layer.v_scale
1213
- layer._q_scale = layer.q_scale
1211
+ # Buffers (not Parameters) so sleep snapshots them; ATTN_HEAD changes shape.
1212
+ layer.register_buffer("_k_scale", layer.k_scale.detach())
1213
+ layer.register_buffer("_v_scale", layer.v_scale.detach())
1214
+ layer.register_buffer("_q_scale", layer.q_scale.detach())
1215
1216
# Set the _float variants that the attention backend uses.
1217
def _to_scalar(tensor: torch.Tensor) -> float: