xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-07 09:49:37 +08:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Sage Moore <sage@neuralmagic.com>
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Sage Moore 2025-09-24 12:02:08 -07:00 committed by yewentao256
parent 72a5101c7a
commit 0e0d51c9c6
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
vllm/v1/worker/gpu_ubatch_wrapper.py
View File

@ -104,6 +104,7 @@ class UBatchWrapper:
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 @@ class UBatchWrapper:
@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()