Runtime Draft Weight Update for Speculative Decoding (#46725)
Signed-off-by: vx120 <893600387@qq.com> Signed-off-by: vx120 <57470515+vx120@users.noreply.github.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: crp0128 <191679376@qq.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
vx120 committed
Jul 12, 2026 at 13:51 UTC
fc1c548093029f6487bbdc9c612995dfe7621a75
12 files changed
+129
-1
docs/training/weight_transfer/base.md
+4
@@ -49,6 +49,10 @@ update_request = WeightTransferUpdateRequest(
49
)
50
```
51
52
+At the LLM/API layer, call `start_draft_weight_update()` instead of
53
+`start_weight_update()` to target the speculative draft model;
54
+`update_weights` / `finish_weight_update` are unchanged.
55
+
56
### WeightTransferUpdateInfo
57
58
The base `WeightTransferUpdateInfo` is a marker class for backend-specific update info:
tests/entrypoints/openai/test_openai_schema.py
+1
@@ -145,6 +145,7 @@ def test_openapi_stateless(case: schemathesis.Case):
145
if case.operation.path in (
146
"/init_weight_transfer_engine",
147
"/start_weight_update",
148
+ "/start_draft_weight_update",
149
"/update_weights",
150
"/finish_weight_update",
151
):
tests/v1/worker/test_gpu_worker_weight_transfer.py
+6
@@ -19,6 +19,7 @@ class _RecordingEngine:
19
self.raise_on_update = raise_on_update
20
self.started = False
21
self.finished = False
22
+ self.reset_count = 0
23
self.update_calls: list[dict] = []
24
25
def start_weight_update(self) -> None:
@@ -32,6 +33,9 @@ class _RecordingEngine:
33
def finish_weight_update(self) -> None:
34
self.finished = True
35
36
+ def reset_weight_update_target(self) -> None:
37
+ self.reset_count += 1
38
+
39
40
def _make_worker(engine: _RecordingEngine | None) -> Worker:
41
worker = object.__new__(Worker)
@@ -54,6 +58,7 @@ def test_start_update_finish_delegates_to_engine():
58
59
Worker.finish_weight_update(worker)
60
assert engine.finished is True
61
+ assert engine.reset_count == 1
62
assert worker._weight_update_active is False
63
64
@@ -85,6 +90,7 @@ def test_update_resets_active_on_error():
90
Worker.update_weights(worker, {"names": ["w"]})
91
92
# A failed update ends the session so the next start is clean.
93
+ assert engine.reset_count == 1
94
assert worker._weight_update_active is False
95
96
vllm/distributed/weight_transfer/base.py
+18
@@ -73,6 +73,8 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
73
init_info_cls: type[TInitInfo]
74
update_info_cls: type[TUpdateInfo]
75
76
+ supports_draft_weight_update: bool = True
77
+
78
def __init__(
79
self,
80
config: WeightTransferConfig,
@@ -95,6 +97,22 @@ class WeightTransferEngine(ABC, Generic[TInitInfo, TUpdateInfo]):
97
self.model_config = vllm_config.model_config
98
self.device = device
99
self.model = model
100
+ self._default_model_config = self.model_config
101
+ self._default_model = model
102
+
103
+ def set_weight_update_target(
104
+ self,
105
+ model: torch.nn.Module,
106
+ model_config: Any,
107
+ ) -> None:
108
+ """Set the model that will receive the active weight update."""
109
+ self.model = model
110
+ self.model_config = model_config
111
+
112
+ def reset_weight_update_target(self) -> None:
113
+ """Restore weight updates to the engine's default target model."""
114
+ self.model = self._default_model
115
+ self.model_config = self._default_model_config
116
117
def parse_init_info(self, init_dict: dict[str, Any]) -> TInitInfo:
118
"""
vllm/distributed/weight_transfer/sparse_nccl_engine.py
+1
@@ -99,6 +99,7 @@ class SparseNCCLWeightTransferEngine(
99
100
init_info_cls = NCCLWeightTransferInitInfo
101
update_info_cls = SparseNCCLWeightTransferUpdateInfo
102
+ supports_draft_weight_update = False
103
104
def __init__(
105
self,
vllm/engine/protocol.py
+4
@@ -248,6 +248,10 @@ class EngineClient(ABC):
248
"""Start a new weight update."""
249
raise NotImplementedError
250
251
+ async def start_draft_weight_update(self) -> None:
252
+ """Start a new weight update targeting the speculative draft model."""
253
+ raise NotImplementedError
254
+
255
async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
256
"""Batched weight update for RL training."""
257
raise NotImplementedError
vllm/entrypoints/llm.py
+4
@@ -877,6 +877,10 @@ class LLM(BeamSearchOfflineMixin, PoolingOfflineMixin, OfflineInferenceMixin):
877
"""Start a new weight update."""
878
self.llm_engine.collective_rpc("start_weight_update")
879
880
+ def start_draft_weight_update(self) -> None:
881
+ """Start a new weight update targeting the speculative draft model."""
882
+ self.llm_engine.collective_rpc("start_draft_weight_update")
883
+
884
def update_weights(self, request: WeightTransferUpdateRequest | dict) -> None:
885
"""
886
Update the weights of the model.
vllm/entrypoints/serve/dev/rlhf/api_router.py
+6
@@ -133,6 +133,12 @@ async def start_weight_update(raw_request: Request):
133
return JSONResponse(content={"message": "Weight update started"})
134
135
136
+@router.post("/start_draft_weight_update")
137
+async def start_draft_weight_update(raw_request: Request):
138
+ await engine_client(raw_request).start_draft_weight_update()
139
+ return JSONResponse(content={"message": "Draft weight update started"})
140
+
141
+
142
@router.post("/update_weights")
143
async def update_weights(raw_request: Request):
144
try:
vllm/v1/engine/async_llm.py
+4
@@ -1084,6 +1084,10 @@ class AsyncLLM(EngineClient):
1084
"""Start a new weight update."""
1085
await self.collective_rpc("start_weight_update")
1086
1087
+ async def start_draft_weight_update(self) -> None:
1088
+ """Start a new weight update targeting the speculative draft model."""
1089
+ await self.collective_rpc("start_draft_weight_update")
1090
+
1091
async def update_weights(self, request: WeightTransferUpdateRequest) -> None:
1092
"""
1093
Batched weight update for RL training.
vllm/v1/worker/gpu/model_runner.py
+6
@@ -371,6 +371,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
371
def get_model(self) -> nn.Module:
372
return self.model
373
374
+ def get_draft_model(self) -> nn.Module | None:
375
+ speculator = self.speculator
376
+ if not isinstance(speculator, DraftModelSpeculator):
377
+ return None
378
+ return speculator.model
379
+
380
def reload_weights(self, *args, **kwargs) -> None:
381
# TODO(Wentao): Use full version instead of import when fully migrated to v2
382
from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
vllm/v1/worker/gpu_model_runner.py
+11
@@ -3282,6 +3282,17 @@ class GPUModelRunner(
3282
return self.model.unwrap()
3283
return self.model
3284
3285
+ def get_draft_model(self) -> nn.Module | None:
3286
+ drafter = getattr(self, "drafter", None)
3287
+ if drafter is None:
3288
+ return None
3289
+ model = getattr(drafter, "model", None)
3290
+ if isinstance(
3291
+ model, (CUDAGraphWrapper, UBatchWrapper, BreakableCUDAGraphWrapper)
3292
+ ):
3293
+ return cast(nn.Module, model.unwrap())
3294
+ return cast(nn.Module | None, model)
3295
+
3296
def get_supported_generation_tasks(self) -> list[GenerationTask]:
3297
model = self.get_model()
3298
supported_tasks = list[GenerationTask]()
vllm/v1/worker/gpu_worker.py
+64
-1
@@ -154,6 +154,7 @@ class Worker(WorkerBase):
154
155
# Buffers saved before sleep
156
self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
157
+ self._sleep_rebuild_draft_metadata_buffers = False
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.
@@ -198,6 +199,11 @@ class Worker(WorkerBase):
199
self._sleep_saved_buffers = {
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
+ )
207
208
self._get_sleep_mode_backend().suspend(level)
209
@@ -229,6 +235,14 @@ class Worker(WorkerBase):
235
buffer.data.copy_(self._sleep_saved_buffers[name].data)
236
self._sleep_saved_buffers = {}
237
238
+ if self._sleep_rebuild_draft_metadata_buffers:
239
+ draft = self.get_draft_model()
240
+ 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
245
+
246
if tags is None or "kv_cache" in tags:
247
self.model_runner.post_kv_cache_wake_up()
248
@@ -915,6 +929,29 @@ class Worker(WorkerBase):
929
def get_model(self) -> nn.Module:
930
return self.model_runner.get_model()
931
932
+ def get_draft_model(self) -> nn.Module | None:
933
+ return self.model_runner.get_draft_model()
934
+
935
+ def _set_draft_weight_update_target(self) -> None:
936
+ assert self.weight_transfer_engine is not None
937
+
938
+ draft_model = self.get_draft_model()
939
+ if draft_model is None:
940
+ raise RuntimeError(
941
+ "Draft model weight update requested, but no draft model is configured."
942
+ )
943
+
944
+ speculative_config = self.speculative_config
945
+ if speculative_config is None or speculative_config.draft_model_config is None:
946
+ raise RuntimeError(
947
+ "Draft model weight update requested, but no draft model "
948
+ "config is configured."
949
+ )
950
+
951
+ self.weight_transfer_engine.set_weight_update_target(
952
+ draft_model, speculative_config.draft_model_config
953
+ )
954
+
955
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
956
return self.model_runner.get_supported_tasks()
957
@@ -1179,16 +1216,38 @@ class Worker(WorkerBase):
1216
the configured weight transfer engine. The worker only tracks that a
1217
session is active.
1218
"""
1219
+ self._start_weight_update()
1220
+
1221
+ def start_draft_weight_update(self) -> None:
1222
+ """
1223
+ Like start_weight_update, but retargets the engine at the speculative
1224
+ draft model for this session.
1225
+ """
1226
+ self._start_weight_update(is_draft=True)
1227
+
1228
+ def _start_weight_update(self, is_draft: bool = False) -> None:
1229
self._check_weight_transfer_engine()
1230
assert self.weight_transfer_engine is not None
1231
1232
+ if is_draft and not self.weight_transfer_engine.supports_draft_weight_update:
1233
+ raise RuntimeError(
1234
+ f"{type(self.weight_transfer_engine).__name__} does not support "
1235
+ "draft model weight updates."
1236
+ )
1237
+
1238
if self._weight_update_active:
1239
raise RuntimeError(
1240
"start_weight_update called while a weight update is already "
1241
"active. Call finish_weight_update first."
1242
)
1243
1191
- self.weight_transfer_engine.start_weight_update()
1244
+ try:
1245
+ if is_draft:
1246
+ self._set_draft_weight_update_target()
1247
+ self.weight_transfer_engine.start_weight_update()
1248
+ except BaseException:
1249
+ self.weight_transfer_engine.reset_weight_update_target()
1250
+ raise
1251
self._weight_update_active = True
1252
1253
def update_weights(self, update_info: dict) -> None:
@@ -1197,6 +1256,8 @@ class Worker(WorkerBase):
1256
1257
start_weight_update must be called before update_weights and
1258
finish_weight_update must be called after all chunks have been sent.
1259
+ Every chunk loads into whichever model the session's start_weight_update
1260
+ / start_draft_weight_update call selected.
1261
1262
Args:
1263
update_info: Dictionary containing backend-specific update info
@@ -1213,6 +1274,7 @@ class Worker(WorkerBase):
1274
self.weight_transfer_engine.update_weights(update_info)
1275
except BaseException:
1276
self._weight_update_active = False
1277
+ self.weight_transfer_engine.reset_weight_update_target()
1278
raise
1279
1280
def finish_weight_update(self) -> None:
@@ -1226,6 +1288,7 @@ class Worker(WorkerBase):
1288
)
1289
1290
self.weight_transfer_engine.finish_weight_update()
1291
+ self.weight_transfer_engine.reset_weight_update_target()
1292
self._weight_update_active = False
1293
1294
def shutdown(self) -> None: