iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Bugfix][Spec Decode] Preserve draft buffers across level-2 sleep (#49774)

Browse Source 
Signed-off-by: aoshen02 <aoshen02@users.noreply.github.com>
Signed-off-by: vx120 <893600387@qq.com>
Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com>
Co-authored-by: vx120 <893600387@qq.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
This commit is contained in:
aoshen02
2026-07-28 00:16:21 -07:00
committed by GitHub
co-authored by aoshen02 vx120 Claude Opus 5 Roger Wang
parent 03a2d03367
commit b09688a6e7
1 changed files with 12 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/v1/worker/gpu_worker.py
+12 -11
View File
@@ -154,7 +154,7 @@ class Worker(WorkerBase):
self.worker_sentinel = WorkerSentinel(worker=self)
# Buffers saved before sleep
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
self._sleep_rebuild_draft_metadata_buffers = False
self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {}
# Weight transfer engine is created in `load_model` once the model
# is available, since the engine needs a reference to the model.
@@ -200,10 +200,10 @@ class Worker(WorkerBase):
name: buffer.cpu().clone() for name, buffer in model.named_buffers()
}
draft = self.get_draft_model()
inner = getattr(draft, "model", None) if draft is not None else None
self._sleep_rebuild_draft_metadata_buffers = inner is not None and hasattr(
inner, "_build_fused_kv_buffers"
)
if draft is not None:
self._sleep_saved_draft_buffers = {
name: buffer.cpu().clone() for name, buffer in draft.named_buffers()
}
self._get_sleep_mode_backend().suspend(level)
@@ -228,20 +228,21 @@ class Worker(WorkerBase):
self._get_sleep_mode_backend().resume(tags)
# Restore the buffers after level 2 sleep
if len(self._sleep_saved_buffers):
wake_weights = tags is None or "weights" in tags
if wake_weights and len(self._sleep_saved_buffers):
model = self.model_runner.model
for name, buffer in model.named_buffers():
if name in self._sleep_saved_buffers:
buffer.data.copy_(self._sleep_saved_buffers[name].data)
self._sleep_saved_buffers = {}
if self._sleep_rebuild_draft_metadata_buffers:
if wake_weights and len(self._sleep_saved_draft_buffers):
draft = self.get_draft_model()
if draft is not None:
inner = getattr(draft, "model", None)
if inner is not None and hasattr(inner, "_build_fused_kv_buffers"):
inner._build_fused_kv_buffers()
self._sleep_rebuild_draft_metadata_buffers = False
for name, buffer in draft.named_buffers():
if name in self._sleep_saved_draft_buffers:
buffer.data.copy_(self._sleep_saved_draft_buffers[name].data)
self._sleep_saved_draft_buffers = {}
if tags is None or "kv_cache" in tags:
self.model_runner.post_kv_cache_wake_up()