A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 18:15 UTC
[MRV2][BugFix] Fix KV connector handling in spec decode case (#43719)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Nick Hill committed
May 26, 2026 at 23:37 UTC
8c94938cfb92cc00b244ae4a933c5f60dbc1139f
3 files changed
+22
-18
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+1
-1
@@ -168,7 +168,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
168
slot_mappings_by_layer=None,
169
hidden_states=None,
170
aux_hidden_states=None,
171
- kv_connector_output=None,
171
+ finished_req_ids=set(),
172
num_tokens_across_dp=None,
173
)
174
runner.postprocess = lambda *args, **kwargs: events.append("postprocess")
vllm/v1/worker/gpu/kv_connector.py
+6
-10
@@ -34,7 +34,7 @@ class KVConnector:
34
pass
35
36
def post_forward(
37
- self, scheduler_output: "SchedulerOutput", wait_for_save: bool = True
37
+ self, finished_req_ids: set[str], wait_for_save: bool = True
38
) -> KVConnectorOutput | None:
39
return None
40
@@ -76,10 +76,7 @@ class ActiveKVConnector(KVConnector):
76
self.kv_connector.start_load_kv(get_forward_context())
77
78
def post_forward(
79
- self,
80
- scheduler_output: "SchedulerOutput",
81
- wait_for_save: bool = True,
82
- clear_metadata: bool = True,
79
+ self, finished_req_ids: set[str], wait_for_save: bool = True
80
) -> KVConnectorOutput | None:
81
if self._disabled:
82
return None
@@ -88,7 +85,7 @@ class ActiveKVConnector(KVConnector):
85
if wait_for_save:
86
self.kv_connector.wait_for_save()
87
output.finished_sending, output.finished_recving = (
91
- self.kv_connector.get_finished(scheduler_output.finished_req_ids)
88
+ self.kv_connector.get_finished(finished_req_ids)
89
)
90
output.invalid_block_ids = self.kv_connector.get_block_ids_with_load_errors()
91
output.kv_connector_stats = self.kv_connector.get_kv_connector_stats()
@@ -96,9 +93,7 @@ class ActiveKVConnector(KVConnector):
93
output.kv_connector_worker_meta = (
94
self.kv_connector.build_connector_worker_meta()
95
)
99
-
100
- if clear_metadata:
101
- self.kv_connector.clear_connector_metadata()
96
+ self.kv_connector.clear_connector_metadata()
97
return output
98
99
def no_forward(self, scheduler_output: "SchedulerOutput") -> ModelRunnerOutput:
@@ -106,7 +101,8 @@ class ActiveKVConnector(KVConnector):
101
return EMPTY_MODEL_RUNNER_OUTPUT
102
103
self.pre_forward(scheduler_output)
109
- kv_connector_output = self.post_forward(scheduler_output, wait_for_save=False)
104
+ finished_req_ids = scheduler_output.finished_req_ids
105
+ kv_connector_output = self.post_forward(finished_req_ids, wait_for_save=False)
106
if kv_connector_output is None or kv_connector_output.is_empty():
107
return EMPTY_MODEL_RUNNER_OUTPUT
108
output = copy.copy(EMPTY_MODEL_RUNNER_OUTPUT)
vllm/v1/worker/gpu/model_runner.py
+15
-7
@@ -50,7 +50,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
50
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
51
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
52
from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
53
-from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
53
+from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
54
from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
55
from vllm.v1.worker.gpu.async_utils import AsyncOutput, AsyncPoolingOutput
56
from vllm.v1.worker.gpu.attn_utils import (
@@ -1207,19 +1207,20 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1207
aux_hidden_states = None
1208
output_intermediate_tensors = model_output
1209
1210
- kv_connector_output = self.kv_connector.post_forward(scheduler_output)
1210
+ finished_req_ids = scheduler_output.finished_req_ids
1211
self.execute_model_state = ExecuteModelState(
1212
input_batch=input_batch,
1213
attn_metadata=attn_metadata,
1214
slot_mappings_by_layer=slot_mappings_by_layer,
1215
hidden_states=hidden_states,
1216
aux_hidden_states=aux_hidden_states,
1217
- kv_connector_output=kv_connector_output,
1217
+ finished_req_ids=finished_req_ids,
1218
)
1219
1220
if not self.is_last_pp_rank:
1221
# Non-last PP rank: return IntermediateTensors for sending.
1222
assert output_intermediate_tensors is not None
1223
+ kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
1224
output_intermediate_tensors.kv_connector_output = kv_connector_output
1225
return output_intermediate_tensors
1226
return None
@@ -1238,7 +1239,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1239
slot_mappings_by_layer = self.execute_model_state.slot_mappings_by_layer
1240
hidden_states = self.execute_model_state.hidden_states
1241
aux_hidden_states = self.execute_model_state.aux_hidden_states
1241
- kv_connector_output = self.execute_model_state.kv_connector_output
1242
+ finished_req_ids = self.execute_model_state.finished_req_ids
1243
self.execute_model_state = None
1244
1245
if not self.is_last_pp_rank:
@@ -1280,8 +1281,8 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1281
req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
1282
sampled_token_ids=None, # type: ignore
1283
prompt_logprobs_dict=prompt_logprobs_dict, # type: ignore[arg-type]
1283
- kv_connector_output=kv_connector_output,
1284
)
1285
+ # Start async output copy here so that it can overlap with speculator proposal.
1286
async_output = AsyncOutput(
1287
model_runner_output=model_runner_output,
1288
sampler_output=sampler_output,
@@ -1344,6 +1345,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1345
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
1346
self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
1347
1348
+ # Post-step KV connector related operations.
1349
+ kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
1350
+ model_runner_output.kv_connector_output = kv_connector_output
1351
+
1352
if self.use_async_scheduling:
1353
return async_output
1354
return async_output.get_output()
@@ -1360,7 +1365,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1365
1366
input_batch = self.execute_model_state.input_batch
1367
hidden_states = self.execute_model_state.hidden_states
1363
- kv_connector_output = self.execute_model_state.kv_connector_output
1368
+ finished_req_ids = self.execute_model_state.finished_req_ids
1369
self.execute_model_state = None
1370
1371
if not self.is_last_pp_rank:
@@ -1372,6 +1377,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1377
hidden_states, input_batch, self.req_states
1378
)
1379
1380
+ # Post-step KV connector related operations.
1381
+ kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
1382
+
1383
# Build the model runner output.
1384
model_runner_output = ModelRunnerOutput(
1385
req_ids=input_batch.req_ids,
@@ -1455,4 +1463,4 @@ class ExecuteModelState(NamedTuple):
1463
slot_mappings_by_layer: dict[str, torch.Tensor] | None
1464
hidden_states: torch.Tensor | None
1465
aux_hidden_states: list[torch.Tensor] | None
1458
- kv_connector_output: KVConnectorOutput | None
1466
+ finished_req_ids: set[str]