Suppress benign cuBLAS warning when capturing cudagraphs with DBO (#25596)

Signed-off-by: Sage Moore <[email protected]>
Signed-off-by: yewentao256 <[email protected]>

Author: SageMoore

vllm/v1/worker/gpu_ubatch_wrapper.py

@@ -104,6 +104,7 @@ def __init__(self, runnable: Callable, vllm_config: VllmConfig,
             self.graph_pool = current_platform.get_global_graph_pool()
 
         self.sm_control = self._create_sm_control_context(vllm_config)
+        self.device = device
 
     @staticmethod
     def _create_sm_control_context(vllm_config: VllmConfig):
@@ -168,6 +169,7 @@ def _capture_ubatches(self, ubatch_metadata, model) -> torch.Tensor:
 
         @torch.inference_mode()
         def _capture_ubatch_thread(results, ubatch_metadata):
+            torch.cuda.set_device(self.device)
             ubatch_context = ubatch_metadata.context
             with torch.cuda.stream(ubatch_context.compute_stream):
                 _ = torch.cuda.current_blas_handle()