From 523139aac489aeaad5fbdf4c2d51841edd991620 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 20 Mar 2026 21:33:04 +0100 Subject: [PATCH] fix: synchronize stream before DtoH readback after CUDA Graph launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CUDA Graph capture disables cudarc event tracking, so device_ptr() cannot rely on write events for ordering. Without explicit sync, the DtoH memcpy can get CUDA_ERROR_INVALID_VALUE because the kernel writes haven't completed yet. Add stream.synchronize() between graph launch and scalar readback. This ensures kernel output buffers (total_loss, grad_norm) are valid before the host reads them. The sync cost is negligible — it's 1 sync per training step (already bottlenecked by kernel execution). Fixes H100 CI failure: all DQN pipeline tests returned "DtoH grad_norm: CUDA_ERROR_INVALID_VALUE" consistently. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index fb9eae8cf..60fcde24a 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -864,8 +864,15 @@ impl GpuDqnTrainer { self.submit_training_ops(online_dueling, online_branching)?; } + // ── Synchronize stream before readback ────────────────────── + // CUDA Graph replay and direct kernel launches are async. The stream + // must be synchronized before DtoH to ensure kernel writes are visible. + // Also required because graph capture disables event tracking, so + // cudarc's device_ptr() can't rely on write events for ordering. + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("stream sync before readback: {e}")))?; + // ── Gather 2 scalars to host: total_loss + grad_norm ──────── - // Read each 1-element GPU buffer directly to host (2 × 4B DtoH). let mut loss_host = [0.0_f32; 1]; self.stream.memcpy_dtoh(&self.total_loss_buf, &mut loss_host) .map_err(|e| MLError::ModelError(format!("DtoH total_loss: {e}")))?;