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 14, 05:30 UTC
[BugFix] Fix KVConnectorOutput TPU breakage (#22598)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Aug 10, 2025 at 19:33 UTC
5898b135abc7b7c0ef7107d21a07d54a84314b7c
3 files changed
+19
-10
tests/v1/kv_connector/unit/utils.py
+8
-4
@@ -179,6 +179,13 @@ def create_model_runner_output(
179
sampled_token = EOS_TOKEN_ID if use_eos else 0
180
sampled_token_ids = [[sampled_token] for _ in req_ids]
181
182
+ kv_connector_output = None if (
183
+ finished_sending is None
184
+ and finished_recving is None) else KVConnectorOutput(
185
+ finished_sending=finished_sending,
186
+ finished_recving=finished_recving,
187
+ )
188
+
189
# Make output data structure.
190
return ModelRunnerOutput(
191
req_ids=req_ids,
@@ -188,10 +195,7 @@ def create_model_runner_output(
195
logprobs=None,
196
prompt_logprobs_dict={},
197
pooler_output=None,
191
- kv_connector_output=KVConnectorOutput(
192
- finished_sending=finished_sending,
193
- finished_recving=finished_recving,
194
- ),
198
+ kv_connector_output=kv_connector_output,
199
)
200
201
vllm/v1/core/sched/scheduler.py
+2
-2
@@ -1151,8 +1151,8 @@ class Scheduler(SchedulerInterface):
1151
scheduler the request during the next step.
1152
"""
1153
1154
- assert self.connector is not None
1155
- self.connector.update_connector_output(kv_connector_output)
1154
+ if self.connector is not None:
1155
+ self.connector.update_connector_output(kv_connector_output)
1156
1157
# KV Connector:: update recv and send status from last step.
1158
for req_id in (kv_connector_output.finished_recving or ()):
vllm/v1/worker/tpu_model_runner.py
+9
-4
@@ -1138,6 +1138,13 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1138
i, target_slice] = valid_sampled_token_ids[i]
1139
req_state.output_token_ids.extend(valid_sampled_token_ids[i])
1140
1141
+ kv_connector_output = None if (
1142
+ finished_sending is None
1143
+ and finished_recving is None) else KVConnectorOutput(
1144
+ finished_sending=finished_sending,
1145
+ finished_recving=finished_recving,
1146
+ )
1147
+
1148
model_runner_output = ModelRunnerOutput(
1149
req_ids=req_ids,
1150
req_id_to_index=self.input_batch.req_id_to_index,
@@ -1146,10 +1153,8 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1153
logprobs=logprobs_lists,
1154
prompt_logprobs_dict=prompt_logprobs_dict,
1155
pooler_output=[],
1149
- kv_connector_output=KVConnectorOutput(
1150
- finished_sending=finished_sending,
1151
- finished_recving=finished_recving,
1152
- ))
1156
+ kv_connector_output=kv_connector_output,
1157
+ )
1158
1159
# Check there are no new graphs compiled - all the graphs should be
1160
# captured and compiled during warm up.