diff --git a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py index 903fc7d3cbe1..3ed4703c8b64 100644 --- a/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py +++ b/tests/basic_correctness/memory/sleep_mode/test_sleep_mode.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import asyncio +import math import pytest import torch @@ -9,6 +10,7 @@ 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 diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index cb69a1d66525..a3fc11be3746 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/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: