[Bugfix][Spec Decode] Preserve draft buffers across level-2 sleep (#49774)
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>
aoshen02 committed
Jul 28, 2026 at 15:16 UTC
b09688a6e77f0c061c08c67aeed56ad669c617dc
1 file changed
+12
-11
vllm/v1/worker/gpu_worker.py
+12
-11
@@ -154,7 +154,7 @@ class Worker(WorkerBase):
154
self.worker_sentinel = WorkerSentinel(worker=self)
155
# Buffers saved before sleep
156
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
157
- self._sleep_rebuild_draft_metadata_buffers = False
157
+ self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {}
158
159
# Weight transfer engine is created in `load_model` once the model
160
# is available, since the engine needs a reference to the model.
@@ -200,10 +200,10 @@ class Worker(WorkerBase):
200
name: buffer.cpu().clone() for name, buffer in model.named_buffers()
201
}
202
draft = self.get_draft_model()
203
- inner = getattr(draft, "model", None) if draft is not None else None
204
- self._sleep_rebuild_draft_metadata_buffers = inner is not None and hasattr(
205
- inner, "_build_fused_kv_buffers"
206
- )
203
+ if draft is not None:
204
+ self._sleep_saved_draft_buffers = {
205
+ name: buffer.cpu().clone() for name, buffer in draft.named_buffers()
206
+ }
207
208
self._get_sleep_mode_backend().suspend(level)
209
@@ -228,20 +228,21 @@ class Worker(WorkerBase):
228
self._get_sleep_mode_backend().resume(tags)
229
230
# Restore the buffers after level 2 sleep
231
- if len(self._sleep_saved_buffers):
231
+ wake_weights = tags is None or "weights" in tags
232
+ if wake_weights and len(self._sleep_saved_buffers):
233
model = self.model_runner.model
234
for name, buffer in model.named_buffers():
235
if name in self._sleep_saved_buffers:
236
buffer.data.copy_(self._sleep_saved_buffers[name].data)
237
self._sleep_saved_buffers = {}
238
238
- if self._sleep_rebuild_draft_metadata_buffers:
239
+ if wake_weights and len(self._sleep_saved_draft_buffers):
240
draft = self.get_draft_model()
241
if draft is not None:
241
- inner = getattr(draft, "model", None)
242
- if inner is not None and hasattr(inner, "_build_fused_kv_buffers"):
243
- inner._build_fused_kv_buffers()
244
- self._sleep_rebuild_draft_metadata_buffers = False
242
+ for name, buffer in draft.named_buffers():
243
+ if name in self._sleep_saved_draft_buffers:
244
+ buffer.data.copy_(self._sleep_saved_draft_buffers[name].data)
245
+ self._sleep_saved_draft_buffers = {}
246
247
if tags is None or "kv_cache" in tags:
248
self.model_runner.post_kv_cache_wake_up()