diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 99cf28c90..914038133 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -581,6 +581,7 @@ pub struct GpuDqnTrainer { total_params: usize, params_initialized: bool, target_params_initialized: bool, + attention_initialized: bool, // ── CUDA Graphs (split: forward+backward, then optimizer) ────── // Graph A: zero → cuBLAS forward → C51 loss → C51 grad → cuBLAS backward @@ -1990,11 +1991,15 @@ impl GpuDqnTrainer { let sh2 = self.config.shared_h2; let n_bytes = batch_size * sh2 * f32_size; - // Stream ordering ensures graph replay completed before touching save_h_s2. - // Disable event tracking for all buffer access in this method. // After CUDA Graph capture, cudarc's device_ptr() records write events // that fail with CUDA_ERROR_INVALID_VALUE on stale captured graph events. + // First call syncs to clear stale events; subsequent calls rely on stream ordering. + if !self.attention_initialized { + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let _ = self.stream.context().check_err(); + self.attention_initialized = true; + } let _evt_guard = EventTrackingGuard::new(self.stream.context()); // Run attention forward: save_h_s2 → attention output buffer @@ -2647,6 +2652,7 @@ impl GpuDqnTrainer { total_params, params_initialized: false, target_params_initialized: false, + attention_initialized: false, graph_forward: None, graph_adam: None, upload_staging_buf, @@ -5485,14 +5491,12 @@ impl GpuDqnTrainer { tau: f32, ) -> Result<(), MLError> { // Sync stream and clear any stale errors from graph capture phase. - let _eg = EventTrackingGuard::new(self.stream.context()); - // cudarc stores errors from cuStreamWaitEvent on disabled events during - // graph capture. check_err() consumes them so bind_to_thread() succeeds. - let _ = self.stream.context().check_err(); - - // First call: flatten target weights into flat buffer (20 DtoD copies, once only). - // Subsequent calls reuse the flat buffer which is kept in sync by unflatten at end. + // First call only: clear stale cudarc events from graph capture phase. + // Subsequent calls skip the sync — stream ordering is sufficient. if !self.target_params_initialized { + let _eg = EventTrackingGuard::new(self.stream.context()); + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); } + let _ = self.stream.context().check_err(); self.flatten_target_weights(target_d, target_b)?; self.target_params_initialized = true; }