092bdd6d57 · tduncan/Nvidia-Optimization-Stack-VLLM

docs/features/sleep_mode.md

+21

@@ -74,6 +74,18 @@ llm.wake_up(tags=["weights"])

74 llm.wake_up(tags=["kv_cache"])

75 ```

76

77 +#### Release only KV cache memory

78 +

79 +`LLM.release_kv_cache_memory()` discards KV cache physical memory while keeping model weights resident. It requires a completed pause and all executor memory to be resident: full sleep, partial wake-up, and repeated release without restoring memory are rejected. Requests retained with `mode="keep"` are recomputed after wake-up.

80 +

81 +Use the default `cumem` backend with managed allocations: `enable_cumem_allocator=True` on CUDA/ROCm (also enabled by `enable_sleep_mode=True`), or `enable_sleep_mode=True` on XPU. Other backends must implement selective discard; CPU and unmanaged allocations are unsupported.

82 +

83 +```python

84 +llm.sleep(level=0, mode="keep") # Wait for the pause to complete.

85 +llm.release_kv_cache_memory()

86 +llm.wake_up(tags=["kv_cache"]) # Reallocate KV cache and resume scheduling.

87 +```

88 +

89 ### Online Serving

90

91 To enable sleep mode in a vLLM server you need to initialize it with the flag `VLLM_SERVER_DEV_MODE=1` and pass `--enable-sleep-mode` to the vLLM server.

@@ -107,9 +119,18 @@ curl -X POST 'http://localhost:8000/collective_rpc' -H 'Content-Type: applicatio

119 curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache'

120 ```

121

122 +To release only KV cache memory, wait for the level 0 pause to complete before calling the release endpoint. The same resident-memory and backend requirements as the Python API apply.

123 +

124 +```bash

125 +curl -X POST 'http://localhost:8000/sleep?level=0&mode=keep'

126 +curl -X POST 'http://localhost:8000/release_kv_cache_memory'

127 +curl -X POST 'http://localhost:8000/wake_up?tags=kv_cache'

128 +```

129 +

130 #### HTTP endpoints

131

132 - `POST /sleep?level=1` — Put the model to sleep (`level=1`).

133 +- `POST /release_kv_cache_memory` — Discard KV cache memory after a completed pause, while all executor memory is resident.

134 - `POST /wake_up` — Wake up the model. Supports optional `tags` query parameters for partial wake-up (e.g., `?tags=weights`).

135 - `POST /collective_rpc` — Perform a collective remote procedure call (RPC).

136 - `GET /is_sleeping` — Check if the model is sleeping.

rust/src/engine-core-client/src/client.rs

+13 -4

@@ -911,13 +911,22 @@ impl EngineCoreClient {

911 Ok(())

912 }

913

914 - /// Wake the engine from sleep, optionally limiting the wake-up to specific

915 - /// tags.

916 - pub async fn wake_up(&self, tags: Option<Vec<String>>) -> Result<()> {

917 - self.call_utility::<(), _>("wake_up", (tags,)).await?;

914 + /// Release KV cache memory while keeping model weights resident.

915 + pub async fn release_kv_cache_memory(&self) -> Result<()> {

916 + self.call_utility::<(), _>("release_kv_cache_memory", ()).await?;

917 Ok(())

918 }

919

920 + /// Wake the engine from sleep, optionally limiting the wake-up to specific

921 + /// tags, and return whether every engine is fully awake.

922 + pub async fn wake_up(&self, tags: Option<Vec<String>>) -> Result<bool> {

923 + Ok(self

924 + .call_utility::<bool, _>("wake_up", (tags,))

925 + .await?

926 + .into_iter()

927 + .all(|fully_awake| fully_awake))

928 + }

929 +

930 /// Pause the scheduler so generation can be halted

931 pub async fn pause_scheduler(&self, mode: PauseMode, clear_cache: bool) -> Result<()> {

932 self.call_utility::<(), _>("pause_scheduler", (mode, clear_cache)).await?;

rust/src/mock-engine/src/engine.rs

+4 -6

@@ -95,12 +95,10 @@ fn utility_response(

95 "get_supported_tasks" => utility_envelope(vec!["generate"]),

96 "is_sleeping" => utility_envelope(false),

97 "reset_prefix_cache" => utility_envelope(true),

98 - "reset_mm_cache"

99 - | "reset_encoder_cache"

100 - | "profile"

101 - | "sleep"

102 - | "wake_up"

103 - | "execute_dummy_batch" => utility_envelope(()),

98 + "wake_up" => utility_envelope(true),

99 + "reset_mm_cache" | "reset_encoder_cache" | "profile" | "sleep" | "execute_dummy_batch" => {

100 + utility_envelope(())

101 + }

102 _ => utility_envelope(Value::Nil),

103 }?;

104

rust/src/mock-engine/src/tests.rs

+1

@@ -195,6 +195,7 @@ async fn utility_requests_return_minimal_success_responses() {

195 assert!(client.reset_prefix_cache(false, false).await.expect("reset prefix cache"));

196 client.reset_mm_cache().await.expect("reset mm cache");

197 client.reset_encoder_cache().await.expect("reset encoder cache");

198 + assert!(client.wake_up(Some(vec!["kv_cache".to_string()])).await.expect("wake up"));

199

200 shutdown_mock(client, shutdown, task).await;

201 }

rust/src/server/src/middleware/metrics.rs

+1

@@ -26,6 +26,7 @@ const EXCLUDED_HANDLERS: &[&str] = &[

26 "/reset_encoder_cache",

27 "/collective_rpc",

28 "/sleep",

29 + "/release_kv_cache_memory",

30 "/wake_up",

31 "/is_sleeping",

32 ];

rust/src/server/src/routes.rs

+4

@@ -142,6 +142,10 @@ fn build_router_with_options(

142 .route("/weight_info", get(weight_transfer::weight_info))

143 .route("/abort_requests", post(abort_requests::abort_requests))

144 .route("/sleep", post(sleep::sleep))

145 + .route(

146 + "/release_kv_cache_memory",

147 + post(sleep::release_kv_cache_memory),

148 + )

149 .route("/wake_up", post(sleep::wake_up))

150 .route("/is_sleeping", get(sleep::is_sleeping))

151 .route("/pause", post(pause::pause))

rust/src/server/src/routes/sleep.rs

+24 -9

@@ -5,7 +5,7 @@ use std::sync::Arc;

5

6 use axum::Json;

7 use axum::extract::rejection::QueryRejection;

8 -use axum::extract::{Query, State};

8 +use axum::extract::{Query, RawQuery, State};

9 use axum::http::StatusCode;

10 use serde::{Deserialize, Serialize};

11 use vllm_engine_core_client::protocol::utility::PauseMode;

@@ -27,12 +27,6 @@ pub(crate) struct SleepParams {

27 mode: PauseMode,

28 }

29

30 -#[derive(Debug, Default, Deserialize)]

31 -pub(crate) struct WakeUpParams {

32 - #[serde(default)]

33 - tags: Option<Vec<String>>,

34 -}

35 -

30 const fn default_sleep_level() -> u32 {

31 1

32 }

@@ -57,14 +51,35 @@ pub async fn sleep(

51 Ok(StatusCode::OK)

52 }

53

54 +/// Release KV cache memory while keeping model weights resident.

55 +pub async fn release_kv_cache_memory(

56 + State(state): State<Arc<AppState>>,

57 +) -> Result<StatusCode, ApiError> {

58 + state

59 + .engine_core_client()

60 + .release_kv_cache_memory()

61 + .await

62 + .map_err(|error| utility_call_error("release_kv_cache_memory", error))?;

63 +

64 + Ok(StatusCode::OK)

65 +}

66 +

67 /// Wake the engine from sleep mode.

68 pub async fn wake_up(

69 State(state): State<Arc<AppState>>,

63 - Query(params): Query<WakeUpParams>,

70 + RawQuery(query): RawQuery,

71 ) -> Result<StatusCode, ApiError> {

72 + let tags = query.and_then(|query| {

73 + let tags = url::form_urlencoded::parse(query.as_bytes())

74 + .filter(|(key, _)| key == "tags")

75 + .map(|(_, value)| value.into_owned())

76 + .collect::<Vec<_>>();

77 + (!tags.is_empty()).then_some(tags)

78 + });

79 +

80 state

81 .engine_core_client()

67 - .wake_up(params.tags)

82 + .wake_up(tags)

83 .await

84 .map_err(|error| utility_call_error("wake_up", error))?;

85

rust/src/server/src/routes/tests.rs

+87 -2

@@ -5805,7 +5805,46 @@ async fn sleep_route_uses_python_compatible_default_query_values() {

5805

5806 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5807 #[serial]

5808 -async fn wake_up_route_without_tags_sends_none() {

5808 +async fn release_kv_cache_memory_route_sends_expected_utility_call() {

5809 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {

5810 + boxed_test_future(async move {

5811 + let utility = recv_engine_message(dealer).await;

5812 + assert_eq!(utility[0].as_ref(), &[0x03]);

5813 +

5814 + let payload = decode_value(&utility[1]).expect("decode utility payload");

5815 + let array = payload.as_array().expect("utility payload array");

5816 + let call_id = array[1].as_u64().expect("call id");

5817 +

5818 + assert_eq!(array[2], Value::from("release_kv_cache_memory"));

5819 + assert_eq!(array[3], Value::Array(Vec::new()));

5820 +

5821 + send_outputs(push, utility_outputs(call_id, utility_none_result())).await;

5822 + })

5823 + })

5824 + .await;

5825 +

5826 + let response = app

5827 + .clone()

5828 + .call(

5829 + Request::builder()

5830 + .method("POST")

5831 + .uri("/release_kv_cache_memory")

5832 + .body(Body::empty())

5833 + .expect("build request"),

5834 + )

5835 + .await

5836 + .expect("call app");

5837 +

5838 + let status = response.status();

5839 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

5840 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));

5841 + assert!(body.is_empty());

5842 + engine_task.await.expect("mock engine task");

5843 +}

5844 +

5845 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5846 +#[serial]

5847 +async fn wake_up_route_without_tags_accepts_bool_result() {

5848 let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {

5849 boxed_test_future(async move {

5850 let utility = recv_engine_message(dealer).await;

@@ -5818,7 +5857,7 @@ async fn wake_up_route_without_tags_sends_none() {

5857 assert_eq!(array[2], Value::from("wake_up"));

5858 assert_eq!(array[3], Value::Array(vec![Value::Nil]));

5859

5821 - send_outputs(push, utility_outputs(call_id, utility_none_result())).await;

5860 + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;

5861 })

5862 })

5863 .await;

@@ -5842,6 +5881,51 @@ async fn wake_up_route_without_tags_sends_none() {

5881 engine_task.await.expect("mock engine task");

5882 }

5883

5884 +#[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5885 +#[serial]

5886 +async fn wake_up_route_accepts_repeated_tags() {

5887 + let (app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {

5888 + boxed_test_future(async move {

5889 + let utility = recv_engine_message(dealer).await;

5890 + assert_eq!(utility[0].as_ref(), &[0x03]);

5891 +

5892 + let payload = decode_value(&utility[1]).expect("decode utility payload");

5893 + let array = payload.as_array().expect("utility payload array");

5894 + let call_id = array[1].as_u64().expect("call id");

5895 +

5896 + assert_eq!(array[2], Value::from("wake_up"));

5897 + assert_eq!(

5898 + array[3],

5899 + Value::Array(vec![Value::Array(vec![

5900 + Value::from("weights"),

5901 + Value::from("kv_cache"),

5902 + ])])

5903 + );

5904 +

5905 + send_outputs(push, utility_outputs(call_id, utility_result_value(true))).await;

5906 + })

5907 + })

5908 + .await;

5909 +

5910 + let response = app

5911 + .clone()

5912 + .call(

5913 + Request::builder()

5914 + .method("POST")

5915 + .uri("/wake_up?tags=weights&tags%5B%5D=ignored&tags=kv_cache")

5916 + .body(Body::empty())

5917 + .expect("build request"),

5918 + )

5919 + .await

5920 + .expect("call app");

5921 +

5922 + let status = response.status();

5923 + let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");

5924 + assert_eq!(status, StatusCode::OK, "{}", String::from_utf8_lossy(&body));

5925 + assert!(body.is_empty());

5926 + engine_task.await.expect("mock engine task");

5927 +}

5928 +

5929 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

5930 #[serial]

5931 async fn is_sleeping_route_returns_json_payload() {

@@ -6160,6 +6244,7 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {

6244 for (method, uri) in [

6245 ("GET", "/is_sleeping"),

6246 ("POST", "/sleep"),

6247 + ("POST", "/release_kv_cache_memory"),

6248 ("POST", "/wake_up"),

6249 ("GET", "/is_paused"),

6250 ("POST", "/pause"),

tests/basic_correctness/test_mem.py

+85

@@ -109,6 +109,91 @@ def test_basic_cumem():

109 assert torch.allclose(output, torch.ones_like(output) * 3)

110

111

112 +@pytest.mark.parametrize("full_sleep", [0, 1, 2], ids=["kv-only", "sleep-1", "sleep-2"])

113 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

114 +def test_release_kv_cache_memory_preserves_generation(full_sleep, monkeypatch):

115 + """Rejected release and KV-only/full sleep cycles preserve greedy tokens."""

116 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

117 + kv_cache_memory_bytes = 256 * 1024 * 1024

118 + llm = LLM(

119 + "Qwen/Qwen3-0.6B",

120 + enable_sleep_mode=True,

121 + enforce_eager=True,

122 + kv_cache_memory_bytes=kv_cache_memory_bytes,

123 + max_model_len=1024,

124 + max_num_seqs=4,

125 + gpu_memory_utilization=0.05,

126 + )

127 + prompt = "How are you?"

128 + sampling_params = SamplingParams(temperature=0, max_tokens=10)

129 + expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids

130 +

131 + def get_mapped_bytes(worker):

132 + return mapped_usage(get_mem_allocator_instance())

133 +

134 + mapped_before = llm.collective_rpc(get_mapped_bytes)[0]

135 + # Utility RPC transports the engine error as a plain Exception.

136 + with pytest.raises(Exception, match="requires a completed pause"):

137 + llm.release_kv_cache_memory()

138 + assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before

139 + assert not llm.llm_engine.is_sleeping()

140 + assert llm.generate(prompt, sampling_params)[0].outputs[0].token_ids == expected

141 +

142 + llm.sleep(level=0)

143 + llm.release_kv_cache_memory()

144 + mapped_after = llm.collective_rpc(get_mapped_bytes)[0]

145 + assert mapped_before - mapped_after >= kv_cache_memory_bytes * 0.99

146 + assert mapped_after > 0

147 + assert llm.llm_engine.is_sleeping()

148 +

149 + if full_sleep:

150 + llm.sleep(level=full_sleep)

151 + assert llm.collective_rpc(get_mapped_bytes)[0] == 0

152 + llm.wake_up(tags=None if full_sleep else ["kv_cache"])

153 + if full_sleep == 2:

154 + llm.collective_rpc("reload_weights")

155 + assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before

156 + assert not llm.llm_engine.is_sleeping()

157 + actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids

158 + assert actual == expected

159 +

160 +

161 +@pytest.mark.parametrize("first_level", [1, 2])

162 +@pytest.mark.parametrize("second_level", [1, 2])

163 +@create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

164 +def test_sleep_with_only_weights_asleep(first_level, second_level, monkeypatch):

165 + """Repeated sleep after KV-only wake preserves mappings and recoverability."""

166 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")

167 + llm = LLM(

168 + "Qwen/Qwen3-0.6B",

169 + enable_sleep_mode=True,

170 + enforce_eager=True,

171 + kv_cache_memory_bytes=256 * 1024 * 1024,

172 + max_model_len=1024,

173 + max_num_seqs=4,

174 + gpu_memory_utilization=0.05,

175 + )

176 + prompt = "How are you?"

177 + sampling_params = SamplingParams(temperature=0, max_tokens=10)

178 + expected = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids

179 +

180 + def get_mapped_bytes(worker):

181 + return mapped_usage(get_mem_allocator_instance())

182 +

183 + llm.sleep(level=first_level)

184 + llm.wake_up(tags=["kv_cache"])

185 + mapped_before = llm.collective_rpc(get_mapped_bytes)[0]

186 + assert mapped_before > 0

187 + llm.sleep(level=second_level)

188 + assert llm.collective_rpc(get_mapped_bytes)[0] == mapped_before

189 + llm.wake_up(tags=["weights"])

190 + if first_level == 2:

191 + llm.collective_rpc("reload_weights")

192 + assert not llm.llm_engine.is_sleeping()

193 + actual = llm.generate(prompt, sampling_params)[0].outputs[0].token_ids

194 + assert actual == expected

195 +

196 +

197 @create_new_process_for_each_test("fork" if current_platform.is_cuda() else "spawn")

198 def test_discard_tags():

199 """Test that discard(tags) selectively frees GPU memory for specific

tests/entrypoints/serve/dev/test_sleep.py

+24 -1

@@ -1,14 +1,37 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3

4 +from unittest.mock import AsyncMock

5 +

6 +import pytest

7 import requests

8 +from fastapi import FastAPI

9 +from fastapi.testclient import TestClient

10 from prometheus_client.parser import text_string_to_metric_families

11

12 from tests.utils import RemoteOpenAIServer

13 +from vllm.entrypoints.serve.dev.sleep.api_router import attach_router

14

15 MODEL_NAME = "meta-llama/Llama-3.2-1B"

16

17

18 +@pytest.mark.cpu_test

19 +@pytest.mark.parametrize("fails", [False, True])

20 +def test_release_kv_cache_memory_route(fails):

21 + app = FastAPI()

22 + app.state.engine_client = AsyncMock()

23 + release = app.state.engine_client.release_kv_cache_memory

24 + if fails:

25 + release.side_effect = RuntimeError("requires a completed pause first")

26 + attach_router(app)

27 +

28 + with TestClient(app, raise_server_exceptions=False) as client:

29 + response = client.post("/release_kv_cache_memory")

30 +

31 + assert response.status_code == (500 if fails else 200)

32 + release.assert_awaited_once_with()

33 +

34 +

35 def test_sleep_mode():

36 # dtype, max-len etc set so that this can run in CI

37 args = [

@@ -63,7 +86,7 @@ def test_sleep_mode():

86 )

87 assert response.status_code == 200

88

66 - # is sleeping should be false after waking up any part of the engine

89 + # Partial wake keeps the engine sleeping.

90 response = requests.get(remote_server.url_for("is_sleeping"))

91 assert response.status_code == 200

92 assert response.json().get("is_sleeping") is True

tests/multimodal/test_cache.py

+75

@@ -1,7 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import multiprocessing as mp

4 +from concurrent.futures import ThreadPoolExecutor

5 +from functools import partial

6 from types import SimpleNamespace

7 +from unittest.mock import Mock

8

9 import numpy as np

10 import pytest

@@ -34,6 +37,8 @@ from vllm.multimodal.inputs import (

37 PlaceholderRange,

38 )

39 from vllm.multimodal.processing import PromptInsertion

40 +from vllm.renderers.base import BaseRenderer

41 +from vllm.utils.async_utils import make_async

42 from vllm.utils.mem_constants import GiB_bytes, MiB_bytes

43

44 from ..utils import create_new_process_for_each_test

@@ -755,6 +760,76 @@ def test_processor_cache_shared_across_loras():

760 assert feature_lora_b.data == item_data

761

762

763 +@pytest.mark.parametrize("use_async", [False, True])

764 +@pytest.mark.parametrize(

765 + "release_error",

766 + [

767 + None,

768 + "requires a completed pause first",

769 + "requires all executor memory to be resident",

770 + ],

771 + ids=["released", "not-paused", "nonresident-memory"],

772 +)

773 +@pytest.mark.asyncio

774 +async def test_release_kv_cache_resends_mm_payload(use_async, release_error):

775 + """Release must not leave a sender hit pointing at a cleared receiver."""

776 + from vllm.v1.engine.async_llm import AsyncLLM

777 + from vllm.v1.engine.llm_engine import LLMEngine

778 +

779 + model_config = SimpleNamespace(

780 + get_multimodal_config=lambda: MultiModalConfig(mm_processor_cache_gb=1)

781 + )

782 + sender = MultiModalProcessorSenderCache(model_config)

783 + receiver = MultiModalReceiverCache(model_config)

784 + item = _dummy_item({"pixel_values": 16})

785 + mm_hash = "image_A"

786 + payload, _ = sender.get_and_update_item((item, []), mm_hash)

787 + assert receiver.get_and_update_item(payload, mm_hash) is item

788 + assert sender.get_and_update_item(None, mm_hash)[0] is None

789 +

790 + renderer = SimpleNamespace(mm_processor_cache=sender, _mm_cache_stats=None)

791 + renderer.clear_mm_cache = partial(BaseRenderer.clear_mm_cache, renderer)

792 +

793 + def release():

794 + if release_error:

795 + raise RuntimeError(release_error)

796 + receiver.clear_cache()

797 +

798 + with ThreadPoolExecutor(max_workers=1) as executor:

799 + renderer._clear_mm_cache_async = make_async(

800 + renderer.clear_mm_cache, executor=executor

801 + )

802 + renderer.clear_mm_cache_async = partial(

803 + BaseRenderer.clear_mm_cache_async, renderer

804 + )

805 + engine = SimpleNamespace(

806 + renderer=renderer,

807 + engine_core=SimpleNamespace(

808 + release_kv_cache_memory=release,

809 + release_kv_cache_memory_async=make_async(release, executor=executor),

810 + ),

811 + logger_manager=Mock(),

812 + )

813 +

814 + async def call_release():

815 + if use_async:

816 + await AsyncLLM.release_kv_cache_memory(engine)

817 + else:

818 + LLMEngine.release_kv_cache_memory(engine)

819 +

820 + if release_error:

821 + with pytest.raises(RuntimeError, match=release_error):

822 + await call_release()

823 + engine.logger_manager.record_sleep_state.assert_not_called()

824 + else:

825 + await call_release()

826 + engine.logger_manager.record_sleep_state.assert_called_once_with(1, 0)

827 +

828 + payload, _ = sender.get_and_update_item((item, []), mm_hash)

829 + assert payload is item

830 + assert receiver.get_and_update_item(payload, mm_hash) is item

831 +

832 +

833 _SLEEP_VISION_PROMPT = (

834 "<|im_start|>system\nYou are a helpful assistant.<|im_end|>"

835 "\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"

tests/v1/core/test_scheduler.py

+76

@@ -36,10 +36,13 @@ from vllm.utils.hashing import sha256

36 from vllm.v1.core.encoder_cache_manager import EncoderCacheManager

37 from vllm.v1.core.kv_cache_coordinator import HybridKVCacheCoordinator

38 from vllm.v1.core.kv_cache_utils import get_request_block_hasher, init_none_hash

39 +from vllm.v1.core.sched.interface import PauseState

40 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput

41 from vllm.v1.core.sched.scheduler import Scheduler

42 from vllm.v1.core.single_type_kv_cache_manager import register_all_kvcache_specs

43 from vllm.v1.engine import FinishReason

44 +from vllm.v1.engine.core import EngineCore

45 +from vllm.v1.executor.uniproc_executor import UniProcExecutor

46 from vllm.v1.hisparse.coordinator import get_hisparse_coordinator

47 from vllm.v1.kv_cache_interface import (

48 FullAttentionSpec,

@@ -1487,6 +1490,79 @@ def test_scheduler_reset_prefix_cache():

1490 assert scheduler.waiting[i] == request

1491

1492

1493 +def test_kv_cache_release_after_keep_pause_preserves_requests():

1494 + """A completed keep pause permits release and recomputation of live requests."""

1495 + scheduler = create_scheduler(max_num_seqs=1, enable_prefix_caching=True)

1496 + running, waiting = create_requests(num_requests=2)

1497 + for request in (running, waiting):

1498 + scheduler.add_request(request)

1499 + scheduler_output = scheduler.schedule()

1500 + scheduler.update_from_output(

1501 + scheduler_output,

1502 + ModelRunnerOutput(

1503 + req_ids=[running.request_id],

1504 + req_id_to_index={running.request_id: 0},

1505 + sampled_token_ids=[[1]],

1506 + ),

1507 + )

1508 + assert scheduler.kv_cache_manager.get_block_ids(running.request_id)[0]

1509 +

1510 + core = object.__new__(EngineCore)

1511 + core.scheduler = scheduler

1512 + core.model_executor = Mock(is_sleeping=False)

1513 + core.mm_receiver_cache = None

1514 + core.batch_queue = None

1515 +

1516 + assert core.pause_scheduler(mode="keep", clear_cache=False) is None

1517 + assert scheduler.running == [running]

1518 + assert list(scheduler.waiting) == [waiting]

1519 + assert not scheduler.has_requests()

1520 +

1521 + core.release_kv_cache_memory()

1522 +

1523 + core.model_executor.discard.assert_called_once_with(("kv_cache",))

1524 + assert scheduler.requests == {

1525 + request.request_id: request for request in (running, waiting)

1526 + }

1527 + assert not scheduler.running

1528 + assert list(scheduler.waiting) == [running, waiting]

1529 + assert running.status == RequestStatus.PREEMPTED

1530 + assert running.num_computed_tokens == 0

1531 + assert list(running.output_token_ids) == [1]

1532 + assert not scheduler.kv_cache_manager.get_block_ids(running.request_id)[0]

1533 + assert scheduler.schedule().total_num_scheduled_tokens == 0

1534 +

1535 + assert core.wake_up(tags=["kv_cache"])

1536 + resumed = scheduler.schedule()

1537 + assert resumed.num_scheduled_tokens == {running.request_id: running.num_tokens}

1538 +

1539 +

1540 +@pytest.mark.parametrize(

1541 + "sleeping_tags",

1542 + [{"weights", "kv_cache"}, {"weights"}, {"kv_cache"}],

1543 + ids=["full-sleep", "weights-only", "kv-only-or-repeated-release"],

1544 +)

1545 +def test_kv_cache_release_rejects_nonresident_memory(sleeping_tags):

1546 + """Reject release without resetting caches or discarding more memory."""

1547 + core = object.__new__(EngineCore)

1548 + core.scheduler = create_scheduler()

1549 + core.scheduler.set_pause_state(PauseState.PAUSED_ALL)

1550 + core.batch_queue = None

1551 + core._reset_caches = Mock()

1552 + core.model_executor = object.__new__(UniProcExecutor)

1553 + core.model_executor.sleeping_tags = sleeping_tags.copy()

1554 + core.model_executor.collective_rpc = Mock()

1555 +

1556 + with pytest.raises(

1557 + RuntimeError, match="requires all executor memory to be resident"

1558 + ):

1559 + core.release_kv_cache_memory()

1560 +

1561 + core._reset_caches.assert_not_called()

1562 + core.model_executor.collective_rpc.assert_not_called()

1563 + assert core.model_executor.sleeping_tags == sleeping_tags

1564 +

1565 +

1566 def test_reset_connector_cache_no_connector_is_no_op_success():

1567 """``reset_connector_cache`` must return True when no connector is

1568 configured.

tests/v1/engine/test_engine_core.py

+24

@@ -23,6 +23,7 @@ from vllm.config import (

23 from vllm.engine.arg_utils import EngineArgs

24 from vllm.platforms import current_platform

25 from vllm.utils.torch_utils import set_default_torch_num_threads

26 +from vllm.v1.core.sched.interface import PauseState

27 from vllm.v1.engine import EngineCoreRequest

28 from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc

29 from vllm.v1.executor.abstract import Executor

@@ -667,6 +668,29 @@ def _pausable_engine_core_proc() -> EngineCoreProc:

668 return core

669

670

671 +@pytest.mark.parametrize(

672 + "pause_state,has_requests,has_batches",

673 + [

674 + pytest.param(PauseState.UNPAUSED, False, False, id="not-paused"),

675 + pytest.param(PauseState.PAUSED_ALL, True, False, id="pending-requests"),

676 + pytest.param(PauseState.PAUSED_ALL, False, True, id="pending-batches"),

677 + ],

678 +)

679 +def test_kv_cache_release_rejects_unsafe_state(pause_state, has_requests, has_batches):

680 + """Reject release before touching caches or memory if work can still use KV."""

681 + core = _pausable_engine_core_proc()

682 + core.scheduler.pause_state = pause_state

683 + core.scheduler.has_requests.return_value = has_requests

684 + core.batch_queue = [object()] if has_batches else None

685 + core._reset_caches = MagicMock()

686 +

687 + with pytest.raises(RuntimeError, match="requires a completed pause"):

688 + core.release_kv_cache_memory()

689 +

690 + core._reset_caches.assert_not_called()

691 + core.model_executor.discard.assert_not_called()

692 +

693 +

694 @pytest.mark.parametrize("deferred", [False, True])

695 def test_pause_synchronizes_device_before_cache_reset(deferred: bool):

696 """A resolved pause promises an idle device: the barrier must run before

vllm/device_allocator/sleep_mode_backend.py

+14

@@ -73,6 +73,13 @@ class SleepModeBackend(ABC):

73 (suspended) engine from a healthy-serving one (see RFC #34303)."""

74 return self._state

75

76 + def discard(self, tags: tuple[str, ...]) -> None:

77 + """Discard selected tagged allocations without preserving contents."""

78 + raise RuntimeError(

79 + f"Sleep-mode backend {type(self).__name__!r} does not support "

80 + "selective discard."

81 + )

82 +

83 # -- Capability introspection (no instance required) --

84

85 @classmethod

@@ -132,6 +139,13 @@ class CuMemBackend(SleepModeBackend):

139 allocator.wake_up(tags)

140 self._state = "RUNNING"

141

142 + def discard(self, tags: tuple[str, ...]) -> None:

143 + from vllm.device_allocator import get_mem_allocator_instance

144 +

145 + allocator = get_mem_allocator_instance()

146 + allocator.discard(tags)

147 + self._state = "SUSPENDED"

148 +

149 @classmethod

150 def preserves_communicators(cls) -> bool:

151 # Comm identity/topology survive memory suspension; no reinit needed.

vllm/engine/protocol.py

+5

@@ -186,6 +186,11 @@ class EngineClient(ABC):

186 """Sleep the engine."""

187 ...

188

189 + @abstractmethod

190 + async def release_kv_cache_memory(self) -> None:

191 + """Discard KV cache physical GPU memory. Requires a completed pause."""

192 + ...

193 +

194 @abstractmethod

195 async def wake_up(self, tags: list[str] | None = None) -> None:

196 """Wake up the engine."""

vllm/entrypoints/llm.py

+8

@@ -830,6 +830,14 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):

830 """

831 self.llm_engine.sleep(level=level, mode=mode)

832

833 + def release_kv_cache_memory(self) -> None:

834 + """Release the GPU physical memory backing the KV cache.

835 +

836 + Requires a completed ``sleep(level=0)`` and resident executor memory.

837 + Restore with ``wake_up(tags=["kv_cache"])``; kept requests are recomputed.

838 + """

839 + self.llm_engine.release_kv_cache_memory()

840 +

841 def wake_up(self, tags: list[str] | None = None):

842 """Wake up the engine from sleep mode. See the [sleep][vllm.LLM.sleep]

843 method for more details.

vllm/entrypoints/serve/dev/sleep/api_router.py

+6

@@ -27,6 +27,12 @@ async def sleep(raw_request: Request):

27 return Response(status_code=200)

28

29

30 +@router.post("/release_kv_cache_memory")

31 +async def release_kv_cache_memory(raw_request: Request):

32 + await engine_client(raw_request).release_kv_cache_memory()

33 + return Response(status_code=200)

34 +

35 +

36 @router.post("/wake_up")

37 async def wake_up(raw_request: Request):

38 tags = raw_request.query_params.getlist("tags")

vllm/v1/engine/async_llm.py

+9 -2

@@ -1079,10 +1079,17 @@ class AsyncLLM(EngineClient):

1079 if self.logger_manager is not None:

1080 self.logger_manager.record_sleep_state(1, level)

1081

1082 - async def wake_up(self, tags: list[str] | None = None) -> None:

1083 - await self.engine_core.wake_up_async(tags)

1082 + async def release_kv_cache_memory(self) -> None:

1083 + await self.renderer.clear_mm_cache_async()

1084 + await self.engine_core.release_kv_cache_memory_async()

1085

1086 if self.logger_manager is not None:

1087 + self.logger_manager.record_sleep_state(1, 0)

1088 +

1089 + async def wake_up(self, tags: list[str] | None = None) -> None:

1090 + fully_awake = await self.engine_core.wake_up_async(tags)

1091 +

1092 + if self.logger_manager is not None and fully_awake:

1093 self.logger_manager.record_sleep_state(0, 0)

1094

1095 async def checkpoint_prepare(self) -> None:

vllm/v1/engine/core.py

+27 -2

@@ -957,12 +957,15 @@ class EngineCore:

957 pause_future.add_done_callback(pause_complete)

958 return future

959

960 - def wake_up(self, tags: list[str] | None = None):

960 + def wake_up(self, tags: list[str] | None = None) -> bool:

961 """Wake up the engine from sleep.

962

963 Args:

964 tags: Tags to wake up. Use ["scheduling"] for level 0 wake up.

965

966 + Returns:

967 + Whether all executor memory is resident again (fully awake).

968 +

969 """

970 if tags is not None and "scheduling" in tags:

971 # Remove "scheduling" from tags if there are other tags to process.

@@ -973,8 +976,30 @@ class EngineCore:

976

977 # Partial wakes intentionally keep the remaining allocations asleep.

978 # Resume scheduling only once all executor memory is resident again.

976 - if not self.model_executor.is_sleeping:

979 + fully_awake = not self.model_executor.is_sleeping

980 + if fully_awake:

981 self.resume_scheduler()

982 + return fully_awake

983 +

984 + def release_kv_cache_memory(self) -> None:

985 + """Discard KV cache physical memory. Requires a completed pause

986 + and all executor memory to be resident. Kept requests are recomputed

987 + after wake-up.

988 + """

989 + if not (

990 + self.is_scheduler_paused()

991 + and not self.scheduler.has_requests()

992 + and not self.batch_queue

993 + ):

994 + raise RuntimeError(

995 + "release_kv_cache_memory() requires a completed pause first"

996 + )

997 + if self.model_executor.is_sleeping:

998 + raise RuntimeError(

999 + "release_kv_cache_memory() requires all executor memory to be resident"

1000 + )

1001 + self._reset_caches()

1002 + self.model_executor.discard(("kv_cache",))

1003

1004 def is_sleeping(self) -> bool:

1005 """Check if engine is sleeping at any level."""

vllm/v1/engine/core_client.py

+23 -8

@@ -193,7 +193,10 @@ class EngineCoreClient(ABC):

193 def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:

194 raise NotImplementedError

195

196 - def wake_up(self, tags: list[str] | None = None) -> None:

196 + def release_kv_cache_memory(self) -> None:

197 + raise NotImplementedError

198 +

199 + def wake_up(self, tags: list[str] | None = None) -> bool:

200 raise NotImplementedError

201

202 def is_sleeping(self) -> bool:

@@ -285,7 +288,10 @@ class EngineCoreClient(ABC):

288 async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:

289 raise NotImplementedError

290

288 - async def wake_up_async(self, tags: list[str] | None = None) -> None:

291 + async def release_kv_cache_memory_async(self) -> None:

292 + raise NotImplementedError

293 +

294 + async def wake_up_async(self, tags: list[str] | None = None) -> bool:

295 raise NotImplementedError

296

297 async def is_sleeping_async(self) -> bool:

@@ -397,8 +403,11 @@ class InprocClient(EngineCoreClient):

403 result = self.engine_core.sleep(level, mode)

404 assert result is None

405

400 - def wake_up(self, tags: list[str] | None = None) -> None:

401 - self.engine_core.wake_up(tags)

406 + def release_kv_cache_memory(self) -> None:

407 + self.engine_core.release_kv_cache_memory()

408 +

409 + def wake_up(self, tags: list[str] | None = None) -> bool:

410 + return self.engine_core.wake_up(tags)

411

412 def is_sleeping(self) -> bool:

413 return self.engine_core.is_sleeping()

@@ -1021,8 +1030,11 @@ class SyncMPClient(MPClient):

1030 def sleep(self, level: int = 1, mode: PauseMode = "abort") -> None:

1031 self.call_utility("sleep", level, mode)

1032

1024 - def wake_up(self, tags: list[str] | None = None) -> None:

1025 - self.call_utility("wake_up", tags)

1033 + def release_kv_cache_memory(self) -> None:

1034 + self.call_utility("release_kv_cache_memory")

1035 +

1036 + def wake_up(self, tags: list[str] | None = None) -> bool:

1037 + return self.call_utility("wake_up", tags)

1038

1039 def is_sleeping(self) -> bool:

1040 return self.call_utility("is_sleeping")

@@ -1265,8 +1277,11 @@ class AsyncMPClient(MPClient):

1277 async def sleep_async(self, level: int = 1, mode: PauseMode = "abort") -> None:

1278 await self.call_utility_async("sleep", level, mode)

1279

1268 - async def wake_up_async(self, tags: list[str] | None = None) -> None:

1269 - await self.call_utility_async("wake_up", tags)

1280 + async def release_kv_cache_memory_async(self) -> None:

1281 + await self.call_utility_async("release_kv_cache_memory")

1282 +

1283 + async def wake_up_async(self, tags: list[str] | None = None) -> bool:

1284 + return await self.call_utility_async("wake_up", tags)

1285

1286 async def is_sleeping_async(self) -> bool:

1287 return await self.call_utility_async("is_sleeping")

vllm/v1/engine/llm_engine.py

+9 -2

@@ -375,10 +375,17 @@ class LLMEngine:

375 if self.logger_manager is not None:

376 self.logger_manager.record_sleep_state(1, level)

377

378 - def wake_up(self, tags: list[str] | None = None):

379 - self.engine_core.wake_up(tags)

378 + def release_kv_cache_memory(self) -> None:

379 + self.renderer.clear_mm_cache()

380 + self.engine_core.release_kv_cache_memory()

381

382 if self.logger_manager is not None:

383 + self.logger_manager.record_sleep_state(1, 0)

384 +

385 + def wake_up(self, tags: list[str] | None = None):

386 + fully_awake = self.engine_core.wake_up(tags)

387 +

388 + if self.logger_manager is not None and fully_awake:

389 self.logger_manager.record_sleep_state(0, 0)

390

391 def is_sleeping(self) -> bool:

vllm/v1/executor/abstract.py

+26 -6

@@ -32,6 +32,9 @@ logger = init_logger(__name__)

32

33 _R = TypeVar("_R")

34

35 +# Memory-pool tags managed by sleep/wake_up/discard.

36 +SLEEP_TAGS = frozenset(("weights", "kv_cache"))

37 +

38 FailureCallback = Callable[[], None]

39

40

@@ -108,7 +111,6 @@ class Executor(ABC):

111 self.speculative_config = vllm_config.speculative_config

112 self.observability_config = vllm_config.observability_config

113 self._init_executor()

111 - self.is_sleeping = False

114 self.sleeping_tags: set[str] = set()

115 self.kv_output_aggregator: KVOutputAggregator | None = None

116 self.ec_output_aggregator: ECOutputAggregator | None = None

@@ -338,15 +340,18 @@ class Executor(ABC):

340 """Reset the encoder cache in each worker to clear cached encoder outputs."""

341 self.collective_rpc("reset_encoder_cache")

342

343 + @property

344 + def is_sleeping(self) -> bool:

345 + return bool(self.sleeping_tags)

346 +

347 def sleep(self, level: int = 1):

342 - if self.is_sleeping:

348 + if "weights" in self.sleeping_tags:

349 logger.warning("Executor is already sleeping.")

350 return

351 time_before_sleep = time.perf_counter()

352 self.collective_rpc("sleep", kwargs=dict(level=level))

353 time_after_sleep = time.perf_counter()

348 - self.sleeping_tags = {"weights", "kv_cache"}

349 - self.is_sleeping = True

354 + self.sleeping_tags |= SLEEP_TAGS

355 logger.info(

356 "It took %.6f seconds to fall asleep.", time_after_sleep - time_before_sleep

357 )

@@ -375,8 +380,23 @@ class Executor(ABC):

380 self.sleeping_tags.remove(tag)

381 else:

382 self.sleeping_tags.clear()

378 - if not self.sleeping_tags:

379 - self.is_sleeping = False

383 +

384 + def discard(self, tags: tuple[str, ...]) -> None:

385 + tags_to_discard = set(tags) - self.sleeping_tags

386 + if not tags_to_discard:

387 + logger.warning("Tags %s are already sleeping.", tags)

388 + return

389 + time_before_discard = time.perf_counter()

390 + try:

391 + self.collective_rpc("discard", args=(tuple(tags_to_discard),))

392 + finally:

393 + self.sleeping_tags |= tags_to_discard

394 + time_after_discard = time.perf_counter()

395 + logger.info(

396 + "It took %.6f seconds to discard tags %s.",

397 + time_after_discard - time_before_discard,

398 + tags_to_discard,

399 + )

400

401 def reinitialize_distributed(

402 self, reconfig_request: ReconfigureDistributedRequest

vllm/v1/worker/cpu_worker.py

+3

@@ -198,6 +198,9 @@ class CPUWorker(Worker):

198 logger.warning("sleep mode is not supported on CPU, ignore it.")

199 pass

200

201 + def discard(self, tags: tuple[str, ...]) -> None:

202 + raise RuntimeError("Selective discard is not supported on CPU.")

203 +

204 def _should_warm_up_model(self) -> bool:

205 # VLLM_CPU_CI_ENV always skips warmup to save CI time.

206 if envs.VLLM_CPU_CI_ENV:

vllm/v1/worker/gpu_worker.py

+3

@@ -304,6 +304,9 @@ class Worker(WorkerBase):

304

305 self.synchronize_device()

306

307 + def discard(self, tags: tuple[str, ...]) -> None:

308 + self.sleep_mode_backend.discard(tags)

309 +

310 def checkpoint_prepare(self) -> None:

311 checkpoint_prepare_distributed_state()

312

[Frontend][Core] Add release_kv_cache_memory() API (#44890)