fix: pinned host memory for async DtoH readback

cuMemcpyDtoHAsync with non-pinned host pointers blocks the CPU on
CUDA 13 / H100 driver 580+ until ALL pending GPU work completes.
With graph_forward + aux_ops + adam queued, this caused multi-minute
hangs between training steps.

Fix: allocate 12 bytes of pinned (page-locked) host memory via
cuMemHostAlloc(DEVICEMAP) for the 3 scalar readbacks (loss, mse_loss,
grad_norm). Pinned memory enables true async DtoH — CPU returns
immediately, GPU copies when it reaches that point in the stream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs

@@ -636,12 +636,10 @@ pub struct GpuDqnTrainer {
 
     /// CUDA event for async scalar readback (replaces per-step stream.synchronize)
     readback_event: Option<CudaEvent>,
-    /// Host-side buffer for deferred loss readback from previous step (blended MSE+C51)
-    pending_loss: f32,
-    /// Host-side buffer for deferred MSE loss readback from previous step
-    pending_mse_loss: f32,
-    /// Host-side buffer for deferred grad_norm readback from previous step (sum-of-squares)
-    pending_grad_norm: f32,
+    /// Pinned host buffer [3 × f32] for async DtoH scalar readback.
+    /// Layout: [0]=loss, [1]=mse_loss, [2]=grad_norm_sq.
+    /// Pinned memory enables true async cuMemcpyDtoHAsync without CPU blocking.
+    readback_pinned: *mut f32,
     /// Whether there's an in-flight readback to collect
     readback_pending: bool,
 
@@ -918,6 +916,10 @@ impl Drop for GpuDqnTrainer {
         self.graph_forward_mse = None;
         self.graph_forward_ddqn = None;
         self.graph_adam = None;
+        // Free pinned host memory
+        if !self.readback_pinned.is_null() {
+            let _ = unsafe { cudarc::driver::result::free_host(self.readback_pinned.cast()) };
+        }
     }
 }
 
@@ -2804,9 +2806,14 @@ impl GpuDqnTrainer {
             scalars_readback_buf,
             scalars_readback_host,
             readback_event: None,
-            pending_loss: 0.0,
-            pending_mse_loss: 0.0,
-            pending_grad_norm: 0.0,
+            readback_pinned: {
+                let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP;
+                unsafe {
+                    cudarc::driver::result::malloc_host(3 * std::mem::size_of::<f32>(), flags)
+                        .map_err(|e| MLError::ModelError(format!("pinned readback alloc: {e}")))?
+                        as *mut f32
+                }
+            },
             readback_pending: false,
             bf16_states_buf,
             bf16_next_states_buf,
@@ -3004,11 +3011,14 @@ impl GpuDqnTrainer {
                 }
             }
             self.readback_pending = false;
-            // Blend MSE + C51 loss: guard against 0.0 * NaN = NaN (IEEE 754)
-            let blended = Self::blend_loss(self.c51_alpha, self.pending_mse_loss, self.pending_loss);
+            // Read from pinned buffer: [0]=loss, [1]=mse_loss, [2]=grad_norm_sq
+            let (loss, mse_loss, grad_norm_sq) = unsafe {
+                (*self.readback_pinned, *self.readback_pinned.add(1), *self.readback_pinned.add(2))
+            };
+            let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss);
             FusedTrainScalars {
                 total_loss: blended,
-                grad_norm: self.pending_grad_norm.sqrt(),
+                grad_norm: grad_norm_sq.sqrt(),
             }
         } else {
             // First step — no previous data
@@ -3022,26 +3032,28 @@ impl GpuDqnTrainer {
         self.replay_adam()?;
         if adam_step < 3 { eprintln!("adam_readback: step {adam_step} — async readback start"); }
 
-        // 3. Sync stream then DtoH readback (synchronous — non-pinned host memory)
-        // cuMemcpyDtoHAsync with non-pinned host pointers degrades to sync on
-        // CUDA 13+ / H100 driver 580+, blocking for ALL pending GPU work.
-        // Explicit sync + sync copy is faster and predictable.
-        unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream()); }
+        // 3. Async DtoH into pinned host buffer (truly async — no CPU blocking)
         unsafe {
-            cudarc::driver::sys::cuMemcpyDtoH_v2(
-                (&mut self.pending_loss as *mut f32).cast(),
+            // [0] = loss
+            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
+                self.readback_pinned.cast(),
                 self.total_loss_buf.raw_ptr(),
                 std::mem::size_of::<f32>(),
+                self.stream.cu_stream(),
             );
-            cudarc::driver::sys::cuMemcpyDtoH_v2(
-                (&mut self.pending_mse_loss as *mut f32).cast(),
+            // [1] = mse_loss
+            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
+                self.readback_pinned.add(1).cast(),
                 self.mse_loss_buf.raw_ptr(),
                 std::mem::size_of::<f32>(),
+                self.stream.cu_stream(),
             );
-            cudarc::driver::sys::cuMemcpyDtoH_v2(
-                (&mut self.pending_grad_norm as *mut f32).cast(),
+            // [2] = grad_norm_sq
+            cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
+                self.readback_pinned.add(2).cast(),
                 self.grad_norm_f32_buf.raw_ptr(),
                 std::mem::size_of::<f32>(),
+                self.stream.cu_stream(),
             );
         }
 
@@ -3075,10 +3087,13 @@ impl GpuDqnTrainer {
             }
         }
         self.readback_pending = false;
-        let blended = Self::blend_loss(self.c51_alpha, self.pending_mse_loss, self.pending_loss);
+        let (loss, mse_loss, grad_norm_sq) = unsafe {
+            (*self.readback_pinned, *self.readback_pinned.add(1), *self.readback_pinned.add(2))
+        };
+        let blended = Self::blend_loss(self.c51_alpha, mse_loss, loss);
         Ok(FusedTrainScalars {
             total_loss: blended,
-            grad_norm: self.pending_grad_norm.sqrt(),
+            grad_norm: grad_norm_sq.sqrt(),
         })
     }