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
[Minor] Group async_scheduling related fields in model runner init (#26736)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Oct 14, 2025 at 14:46 UTC
ff4810ba73168237f090b0195fbe0ae4d64e0730
1 file changed
+9
-11
vllm/v1/worker/gpu_model_runner.py
+9
-11
@@ -375,9 +375,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
375
)
376
377
self.use_async_scheduling = self.scheduler_config.async_scheduling
378
- self.async_output_copy_stream = (
379
- torch.cuda.Stream() if self.use_async_scheduling else None
380
- )
378
+ # Separate cuda stream for overlapping transfer of sampled token ids from
379
+ # GPU to CPU when async scheduling is enabled.
380
+ self.async_output_copy_stream: torch.cuda.Stream | None = None
381
+ # cuda event to synchronize use of reused CPU tensors between steps
382
+ # when async scheduling is enabled.
383
+ self.prepare_inputs_event: torch.cuda.Event | None = None
384
+ if self.use_async_scheduling:
385
+ self.async_output_copy_stream = torch.cuda.Stream()
386
+ self.prepare_inputs_event = torch.cuda.Event()
387
388
# TODO(woosuk): Provide an option to tune the max cudagraph batch size.
389
# The convention is different.
@@ -444,14 +450,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
450
(3, self.max_num_tokens + 1), dtype=torch.int64
451
)
452
447
- # CUDA event to synchronize use of reused CPU tensors between steps
448
- # when async scheduling is enabled.
449
- self.prepare_inputs_event: torch.cuda.Event | None = None
450
- if self.use_async_scheduling:
451
- self.prepare_inputs_event = torch.cuda.Event()
452
- # Start in a completed state.
453
- self.prepare_inputs_event.record(torch.cuda.default_stream())
454
-
453
# None in the first PP rank. The rest are set after load_model.
454
self.intermediate_tensors: IntermediateTensors | None = None
455