From 021de692353a670b41eb4ab61b8f439d0f7cfe42 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 20 Mar 2026 22:42:50 +0100 Subject: [PATCH] fix: restore scoped event tracking disable during CUDA Graph capture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit removed disable_event_tracking() from graph capture, causing CUDA_ERROR_STREAM_CAPTURE_INVALIDATED on H100 (events recorded during capture invalidate the capture). Fix: disable event tracking BEFORE begin_capture, re-enable AFTER end_capture. With GPU-native cat/stack (no to_host roundtrip), the re-enable is now safe — post-capture operations don't call to_host on graph-modified buffers. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index a6f6d3726..066fa7d82 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -1183,8 +1183,9 @@ impl GpuDqnTrainer { self.stream.synchronize() .map_err(|e| MLError::ModelError(format!("stream sync before capture: {e}")))?; - // Event tracking is already disabled at DQNTrainer construction (single-stream - // context — events are unnecessary with explicit sync barriers). + // Disable event tracking during capture — cudarc's device_ptr/device_ptr_mut + // record CudaEvents which are DISALLOWED inside CUDA Graph capture. + unsafe { self.stream.context().disable_event_tracking(); } // Begin stream capture — only work submitted from this thread on this // stream is captured (THREAD_LOCAL mode, safe for single-stream use). @@ -1205,10 +1206,10 @@ impl GpuDqnTrainer { cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_AUTO_FREE_ON_LAUNCH, ); - // Keep event tracking disabled — graph capture wrote stale events to - // CudaSlices. All subsequent operations use explicit stream sync + - // ManuallyDrop guards (raw_device_ptr), so event tracking is unnecessary. - // Re-enabling would leave stale events that cause CUDA_ERROR_INVALID_VALUE. + // Re-enable event tracking after capture. With GPU-native cat/stack + // (no to_host roundtrip), stale events from capture don't cause issues + // because all post-capture GPU ops use DtoD copies, not DtoH. + unsafe { self.stream.context().enable_event_tracking(); } // Propagate submission error first submit_result?;