set capture mode thread_local (#2560)

InternLM · Oct 21, 2024 · a465e60 · a465e60
1 parent d009335
commit a465e60
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/lmdeploy/pytorch/backends/cuda/graph_runner.py b/lmdeploy/pytorch/backends/cuda/graph_runner.py
@@ -77,9 +77,12 @@ def capture(self, **kwargs):
         self.model(**padded_kwargs)
 
         self._graph = torch.cuda.CUDAGraph()
+        # unsafe kernel call in other thread might invalid the capture
+        # so we set thread_safe capture mode here.
         with torch.cuda.graph(self._graph,
                               pool=self.pool,
-                              stream=current_stream):
+                              stream=current_stream,
+                              capture_error_mode='thread_local'):
             output = self.model(**padded_kwargs)
 
         output_buffers = dict(logits=output)