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?;