use thread local cpature mode (#2850)

awni · web-flow · commit 6c5785bc2fed · 2025-12-01T19:02:47.000-08:00
diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
@@ -87,7 +87,7 @@ CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {
     return;
   }
   CHECK_CUDA_ERROR(
-      cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeGlobal));
+      cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeThreadLocal));
 }
 
 CommandEncoder::CaptureContext::~CaptureContext() {

Original file line number	Diff line number	Diff line change
`@@ -87,7 +87,7 @@ CommandEncoder::CaptureContext::CaptureContext(CommandEncoder& enc) : enc(enc) {`
`87`	`87`	`return;`
`88`	`88`	`}`
`89`	`89`	`CHECK_CUDA_ERROR(`
`90`		`- cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeGlobal));`
	`90`	`+ cudaStreamBeginCapture(enc.stream(), cudaStreamCaptureModeThreadLocal));`
`91`	`91`	`}`
`92`	`92`
`93`	`93`	`CommandEncoder::CaptureContext::~CaptureContext() {`