From 49444c19d83ef547fa8796e20d7bfe33602b43f1 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sat, 21 Mar 2026 23:58:01 +0100
Subject: [PATCH] fix: memcpy_dtoh buffer size mismatch in Q-value readback

q_out_buf is [config.batch_size, total_actions] but host readback buffers
were sized for [sample_size, total_actions]. When sample_size < batch_size,
cudarc's memcpy_dtoh assertion (dst.len >= src.len) panics with SIGABRT.

Fix: allocate host buffer to match q_out.len(), then truncate to sample_size.
Affects: compute_epoch_q_diagnostics, compute_validation_loss, PER refresh.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml/src/trainers/dqn/trainer/metrics.rs    | 16 ++++++++--------
 .../ml/src/trainers/dqn/trainer/training_loop.rs |  6 +++---
 2 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/crates/ml/src/trainers/dqn/trainer/metrics.rs b/crates/ml/src/trainers/dqn/trainer/metrics.rs
index e46d9c930..79ac17dc4 100644
--- a/crates/ml/src/trainers/dqn/trainer/metrics.rs
+++ b/crates/ml/src/trainers/dqn/trainer/metrics.rs
@@ -224,13 +224,13 @@ impl DQNTrainer {
             Err(_) => return None,
         };
 
-        // Download Q-values to host (cold path, epoch boundary only)
-        // host_q has exactly sample_size * total_actions elements;
-        // cudarc memcpy_dtoh copies dst.len() elements from device buffer.
+        // Download Q-values to host (cold path, epoch boundary only).
+        // q_out is [config.batch_size, total_actions] but we only need [sample_size, total_actions].
+        // Allocate host buffer to match GPU buffer size, then truncate.
         let stream = self.cuda_stream.as_ref()?;
-        let total_elems = sample_size * total_actions;
-        let mut host_q = vec![0.0_f32; total_elems];
+        let mut host_q = vec![0.0_f32; q_out.len()];
         stream.memcpy_dtoh(q_out, &mut host_q).ok()?;
+        host_q.truncate(sample_size * total_actions);
 
         // Compute gap (best - second_best) and per-action averages on host
         let cols = total_actions;
@@ -474,11 +474,11 @@ impl DQNTrainer {
             .map_err(|e| anyhow::anyhow!("cuBLAS validation forward: {e}"))?;
 
         // Download Q-values to CPU for action selection + Sharpe computation.
-        // All remaining computation stays on CPU — epoch boundary cold path.
-        let q_total_elems = sample_size * total_actions;
-        let mut host_q = vec![0.0_f32; q_total_elems];
+        // q_out is [config.batch_size, total_actions]; allocate to match, then truncate.
+        let mut host_q = vec![0.0_f32; q_out.len()];
         stream.memcpy_dtoh(q_out, &mut host_q)
             .map_err(|e| anyhow::anyhow!("Q-val DtoH: {e}"))?;
+        host_q.truncate(sample_size * total_actions);
 
         // ── CPU greedy action selection ──
         // Direction LUT maps exposure head index → position size direction.
diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs
index 5ee2d21d5..7a5cb067a 100644
--- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs
+++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs
@@ -1392,10 +1392,10 @@ impl DQNTrainer {
                 let q_result = fused.compute_q_values(&states_gpu, batch_size_refresh);
 
                 if let Ok(q_out) = q_result {
-                    // Download Q-values for max-per-row computation (small batch, epoch boundary)
-                    let q_total = batch_size_refresh * total_actions;
-                    let mut host_q = vec![0.0_f32; q_total];
+                    // Download Q-values (q_out may be larger than needed — match GPU buffer size)
+                    let mut host_q = vec![0.0_f32; q_out.len()];
                     if refresh_stream.memcpy_dtoh(q_out, &mut host_q).is_ok() {
+                        host_q.truncate(batch_size_refresh * total_actions);
                         let mut max_vals = Vec::with_capacity(batch_size_refresh);
                         for row in 0..batch_size_refresh {
                             let offset = row * total_actions;