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 6, 12:00 UTC
[bugfix][async scheduling] fix extra cuda context in device 0 with EP/DP (#37449)
Signed-off-by: youkaichao <youkaichao@gmail.com>
youkaichao committed
Mar 19, 2026 at 02:32 UTC
70b81c4f3d1a1699303b4b6d82bf4d7373ef0a01
1 file changed
+23
-11
vllm/v1/executor/multiproc_executor.py
+23
-11
@@ -597,17 +597,6 @@ class WorkerProc:
597
wrapper.init_worker(all_kwargs)
598
self.worker = wrapper
599
600
- scheduler_config = vllm_config.scheduler_config
601
- self.use_async_scheduling = scheduler_config.async_scheduling
602
- if self.use_async_scheduling:
603
- self.async_output_queue: queue.Queue = queue.Queue()
604
- self.async_output_copy_thread = Thread(
605
- target=self.async_output_busy_loop,
606
- daemon=True,
607
- name="WorkerAsyncOutputCopy",
608
- )
609
- self.async_output_copy_thread.start()
610
-
600
self.setup_proc_title_and_log_prefix(
601
enable_ep=vllm_config.parallel_config.enable_expert_parallel
602
)
@@ -622,6 +611,17 @@ class WorkerProc:
611
)
612
self.worker.load_model()
613
614
+ scheduler_config = vllm_config.scheduler_config
615
+ self.use_async_scheduling = scheduler_config.async_scheduling
616
+ if self.use_async_scheduling:
617
+ self.async_output_queue: queue.Queue = queue.Queue()
618
+ self.async_output_copy_thread = Thread(
619
+ target=self.async_output_busy_loop,
620
+ daemon=True,
621
+ name="WorkerAsyncOutputCopy",
622
+ )
623
+ self.async_output_copy_thread.start()
624
+
625
# Set block size based on the attention backends
626
current_platform.update_block_size_for_backend(vllm_config)
627
@@ -911,6 +911,18 @@ class WorkerProc:
911
912
def async_output_busy_loop(self):
913
"""Entrypoint for the thread which handles outputs asynchronously."""
914
+
915
+ # set device to the worker device for the thread.
916
+ # a thread will not inherit the context of the main thread.
917
+ # when calling any cuda runtime functions, it will implicitly
918
+ # create a new cuda context on device 0, consuming extra memory.
919
+ # here we set the device to the worker device for the thread,
920
+ # enforcing the context to be the same as the main thread.
921
+ from vllm.platforms import current_platform
922
+
923
+ if hasattr(self.worker, "device"):
924
+ current_platform.set_device(self.worker.device)
925
+
926
while True:
927
output = self.async_output_queue.get()
928
self.enqueue_output(output)