[Bugfix][Quantization] Stop sleep(level=2) from zeroing compressed-tensors KV scales

tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py

@@ -2,13 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import asyncio
+import math
 
 import pytest
 import torch
 
 import vllm.device_allocator.cumem as cumem
 import vllm.envs as envs
 from vllm import LLM, AsyncEngineArgs, AsyncLLMEngine, SamplingParams
+from vllm.model_executor.model_loader import get_model_loader
 from vllm.platforms import current_platform
 from vllm.utils.mem_constants import GiB_bytes
 
@@ -232,6 +234,58 @@ async def test():
     asyncio.run(test())
 
 
+def _load_weights(worker) -> None:
+    # Not reload_weights(): its layerwise finalize re-derives the KV scales.
+    model = worker.model_runner.get_model()
+    loader = get_model_loader(worker.load_config)
+    model.load_weights(loader.get_all_weights(worker.model_config, model))
+
+
+def _attn_kv_scales(model) -> dict[str, list[float]]:
+    return {
+        f"{name}.{attr}": getattr(module, attr).flatten().tolist()
+        for name, module in model.named_modules()
+        for attr in ("_q_scale", "_k_scale", "_v_scale")
+        if isinstance(getattr(module, attr, None), torch.Tensor)
+    }
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Segfaults on ROCm.")
+@requires_fp8
+@create_new_process_for_each_test()
+def test_deep_sleep_compressed_tensors_kv_scales(
+    monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared
+):
+    """KV scales must survive level-2 sleep plus a weights reload (else NaN)."""
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    llm = LLM(
+        "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor",
+        enable_sleep_mode=True,
+        kv_cache_dtype="fp8",
+        enforce_eager=True,
+        gpu_memory_utilization=0.6,
+    )
+    sampling_params = SamplingParams(temperature=0, max_tokens=10, logprobs=0)
+    scales = llm.apply_model(_attn_kv_scales)[0]
+    assert any(v != 1.0 for values in scales.values() for v in values)
+    expected = llm.generate("How are you?", sampling_params)[0].outputs[0]
+
+    llm.sleep(level=2)
+    llm.wake_up(tags=["weights"])
+    llm.collective_rpc(_load_weights)
+    llm.wake_up(tags=["kv_cache"])
+
+    assert llm.apply_model(_attn_kv_scales)[0] == scales
+    actual = llm.generate("How are you?", sampling_params)[0].outputs[0]
+    assert actual.token_ids == expected.token_ids
+    logprobs = [
+        [lp[t].logprob for lp, t in zip(o.logprobs, o.token_ids)]
+        for o in (expected, actual)
+    ]
+    assert all(math.isfinite(x) for x in logprobs[1])
+    assert logprobs[0] == logprobs[1]
+
+
 @requires_fp8
 def test_deep_sleep_fp8_kvcache_mrv1(
     monkeypatch: pytest.MonkeyPatch, gpu_memory_cleared

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -1208,9 +1208,10 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         """Override the default vLLM placeholder scales with the llm-compressor loaded
         scales. Zero points are not used as only symmetric quantization is supported.
         """
-        layer._k_scale = layer.k_scale
-        layer._v_scale = layer.v_scale
-        layer._q_scale = layer.q_scale
+        # Buffers (not Parameters) so sleep snapshots them; ATTN_HEAD changes shape.
+        layer.register_buffer("_k_scale", layer.k_scale.detach())
+        layer.register_buffer("_v_scale", layer.v_scale.detach())
+        layer.register_buffer("_q_scale", layer.q_scale.detach())
 
         # Set the _float variants that the attention backend uses.
         def _to_scalar(tensor: torch.Tensor) -> float: