Skip to content

Commit c9e95ba

Browse files
authored
[MRV2][Feature] Support sleep-mode (vllm-project#13005)
### What this PR does / why we need it? This PR adapts 2 modifications into sleep mode in MRV2. #### Problem 1: The missing hook The upstream `gpu_worker.py` (shared by V1 and V2) always refreshes model-runner state after a KV-cache wake-up: ```python # vllm/v1/worker/gpu_worker.py L222-242 def wake_up(self, tags=None): self._get_sleep_mode_backend().resume(tags) # ... buffer restore ... if tags is None or "kv_cache" in tags: self.model_runner.post_kv_cache_wake_up() ``` The Ascend worker (`vllm_ascend/worker/worker.py`) does **not** issue this call. Its `wake_up` performs: 1. `CaMemAllocator.wake_up(tags=tags)` — remaps memory; KV-cache chunks get fresh `data_ptr()` values. 2. MoE expert-weight layout restoration (`w13_weight` / `w2_weight` transpose). 3. Level-2 buffer restore. 4. `SleepWakeupManager.wakeup()` — recaptures ACL graphs. Between steps 3 and 4, `post_kv_cache_wake_up()` is missing. #### Problem 2: AttributeError on extra-cleanup path `SleepWakeupManager.sleep()` accesses `model_runner.use_aclgraph` (`sleep_mem_optimized.py` L50). V1 `NPUModelRunner.__init__` sets this attribute explicitly; V2 `NPUModelRunner.__init__` does **not**. When `enable_sleep_mode_extra_cleanup=True` (not used in the one-card test, but a supported config), V2 workers would crash with `AttributeError`. ### Does this PR introduce _any_ user-facing change? None. ### How was this patch tested? It is tested by tests - vllm\vllm-ascend\tests\e2e\pull_request\one_card\test_camem.py - vllm\vllm-ascend\tests\ut\device_allocator\test_camem.py - vLLM version: v0.25.1 - vLLM main: vllm-project/vllm@fe784ff Signed-off-by: Raining__ <wangruining5@huawei.com>
1 parent c8710f2 commit c9e95ba

2 files changed

Lines changed: 11 additions & 1 deletion

File tree

vllm_ascend/worker/v2/model_runner.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
import numpy as np
2323
import torch
2424
from vllm.config import VllmConfig
25-
from vllm.config.compilation import CUDAGraphMode
25+
from vllm.config.compilation import CompilationMode, CUDAGraphMode
2626
from vllm.v1.core.sched.output import SchedulerOutput
2727
from vllm.v1.kv_cache_interface import KVCacheConfig
2828
from vllm.v1.worker.gpu import model_runner as vllm_model_runner
@@ -70,6 +70,12 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
7070
with torch_cuda_wrapper():
7171
super().__init__(vllm_config, device)
7272

73+
self.use_aclgraph = (
74+
self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE
75+
and self.compilation_config.mode == CompilationMode.VLLM_COMPILE
76+
and not self.model_config.enforce_eager
77+
)
78+
7379
# because we will override these attribute, delete these attribute to
7480
# make sure it's collected by python gc immediately.
7581
del self.req_states

vllm_ascend/worker/worker.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,10 @@ def wake_up(self, tags: list[str] | None = None) -> None:
272272
if name in self._sleep_saved_buffers:
273273
buffer.data.copy_(self._sleep_saved_buffers[name].data)
274274
self._sleep_saved_buffers = {}
275+
276+
if tags is None or "kv_cache" in tags:
277+
self.model_runner.post_kv_cache_wake_up()
278+
275279
cleanup_enabled = getattr(get_ascend_config(), "enable_sleep_mode_extra_cleanup", False)
276280
if cleanup_enabled:
277281
self.sleep_wakeup_manager.wakeup(tags)

0 commit comments

Comments
 (0)