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] Ensure integrity of reused CPU tensors during async scheduling (#24527)
Signed-off-by: Nick Hill <nhill@redhat.com> Co-authored-by: guoze.lin <guozelin@tencent.com>
Nick Hill committed
Sep 10, 2025 at 06:15 UTC
f4f1a8df22424eac6228629e34b0e8f254175551
1 file changed
+24
-8
vllm/v1/worker/gpu_model_runner.py
+24
-8
@@ -326,6 +326,14 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
326
self.mrope_positions = self._make_buffer(
327
(3, self.max_num_tokens + 1), dtype=torch.int64)
328
329
+ # CUDA event to synchronize use of reused CPU tensors between steps
330
+ # when async scheduling is enabled.
331
+ self.prepare_inputs_event: Optional[torch.cuda.Event] = None
332
+ if self.use_async_scheduling:
333
+ self.prepare_inputs_event = torch.cuda.Event()
334
+ # Start in a completed state.
335
+ self.prepare_inputs_event.record(torch.cuda.default_stream())
336
+
337
# None in the first PP rank. The rest are set after load_model.
338
self.intermediate_tensors: Optional[IntermediateTensors] = None
339
@@ -354,11 +362,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
362
# Cudagraph dispatcher for runtime cudagraph dispatching.
363
self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)
364
357
- self.mm_budget = (MultiModalBudget(
365
+ self.mm_budget = MultiModalBudget(
366
self.model_config,
367
self.scheduler_config,
368
self.mm_registry,
361
- ) if self.supports_mm_inputs else None)
369
+ ) if self.supports_mm_inputs else None
370
371
self.reorder_batch_threshold: Optional[int] = None
372
@@ -991,10 +999,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
999
builder,
1000
)
1001
994
- attn_metadata_i = (builder.build(
1002
+ attn_metadata_i = builder.build(
1003
common_prefix_len=common_prefix_len,
1004
common_attn_metadata=common_attn_metadata,
997
- ))
1005
+ )
1006
1007
for layer_name in attn_group.layer_names:
1008
attn_metadata[layer_name] = attn_metadata_i
@@ -1866,10 +1874,18 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1874
"prompt tokens, tokens, please disable it when the requests"
1875
" need prompt logprobs")
1876
1869
- # Prepare the decoder inputs.
1870
- (attn_metadata, logits_indices, spec_decode_metadata,
1871
- num_scheduled_tokens_np, spec_decode_common_attn_metadata,
1872
- max_query_len) = self._prepare_inputs(scheduler_output)
1877
+ if self.prepare_inputs_event is not None:
1878
+ # Ensure prior step has finished with reused CPU tensors.
1879
+ self.prepare_inputs_event.synchronize()
1880
+ try:
1881
+ # Prepare the decoder inputs.
1882
+ (attn_metadata, logits_indices, spec_decode_metadata,
1883
+ num_scheduled_tokens_np, spec_decode_common_attn_metadata,
1884
+ max_query_len) = self._prepare_inputs(scheduler_output)
1885
+
1886
+ finally:
1887
+ if self.prepare_inputs_event is not None:
1888
+ self.prepare_inputs_event.record()
1889
1890
(
1891
num_scheduled_tokens,