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
Suppress benign cuBLAS warning when capturing cudagraphs with DBO (#25596)
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Sage Moore committed
Sep 24, 2025 at 12:02 UTC
f84a472a03cf86cc7362eedaf6c2f09cfa90982c
1 file changed
+2
vllm/v1/worker/gpu_ubatch_wrapper.py
+2
@@ -104,6 +104,7 @@ class UBatchWrapper:
104
self.graph_pool = current_platform.get_global_graph_pool()
105
106
self.sm_control = self._create_sm_control_context(vllm_config)
107
+ self.device = device
108
109
@staticmethod
110
def _create_sm_control_context(vllm_config: VllmConfig):
@@ -168,6 +169,7 @@ class UBatchWrapper:
169
170
@torch.inference_mode()
171
def _capture_ubatch_thread(results, ubatch_metadata):
172
+ torch.cuda.set_device(self.device)
173
ubatch_context = ubatch_metadata.context
174
with torch.cuda.stream(ubatch_context.compute_stream):
175
_ = torch.cuda.current_blas_handle()