From a7ef84792273197351dc331cb74fac7c3b3d0a8e Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Mon, 20 Apr 2026 14:18:04 +0200
Subject: [PATCH] cleanup: remove OFI diagnostics + dead bf16 scatter_insert
 kernel
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Diagnostics served their purpose — confirmed OFI flows through full
pipeline on H100 (state_gather → PER → trainer). Remove:
- OFI host verify readback (upload_ofi_features)
- STATE_GATHER_DIAG pinned memory readback (timestep loop)
- Dead bf16 scatter_insert kernel (states are f32, was never called)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 crates/ml-dqn/src/replay_buffer_kernels.cu    | 22 +-----------
 .../cuda_pipeline/gpu_experience_collector.rs | 34 -------------------
 2 files changed, 1 insertion(+), 55 deletions(-)

diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu
index 0c3a5f950..04800c927 100644
--- a/crates/ml-dqn/src/replay_buffer_kernels.cu
+++ b/crates/ml-dqn/src/replay_buffer_kernels.cu
@@ -45,27 +45,7 @@ void scatter_insert_u32(
     dst[dst_idx] = src[i];
 }
 
-// ── 3. Scatter insert for rows [batch, dim] → ring buffer [cap, dim] ──
-
-extern "C" __global__
-void scatter_insert(
-    unsigned short* __restrict__ dst,     // [capacity * state_dim]
-    const unsigned short* __restrict__ src, // [batch_size * state_dim]
-    int start_idx,
-    int capacity,
-    int state_dim,
-    int batch_size
-) {
-    int tid = blockIdx.x * blockDim.x + threadIdx.x;
-    int total = batch_size * state_dim;
-    if (tid >= total) return;
-    int row = tid / state_dim;
-    int col = tid % state_dim;
-    int dst_row = (start_idx + row) % capacity;
-    dst[dst_row * state_dim + col] = src[row * state_dim + col];
-}
-
-// ── 3b. Scatter insert for f32 rows [batch, dim] → ring buffer [cap, dim] ──
+// ── 3. Scatter insert for f32 rows [batch, dim] → ring buffer [cap, dim] ──
 //    2D-aware version of scatter_insert_f32 for state matrices.
 //    Each thread writes one element: dst[(start + row) % cap, col] = src[row, col].
 
diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
index 89af4227b..cb5f76901 100644
--- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
+++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
@@ -2051,30 +2051,6 @@ impl GpuExperienceCollector {
                     )))?;
             }
 
-            // One-shot diagnostic: readback batch_states[0..state_dim] via pinned memory
-            if t == 0 && self.last_experience_count == 0 {
-                let sd = self.state_dim;
-                let pinned = unsafe { PinnedHostBuf::<f32>::new(sd)
-                    .map_err(|e| MLError::ModelError(format!("diag pinned alloc: {e}")))? };
-                unsafe {
-                    cudarc::driver::sys::cuMemcpyDtoHAsync_v2(
-                        pinned.ptr as *mut std::ffi::c_void,
-                        self.batch_states.device_ptr(&self.stream).0,
-                        (sd * 4) as usize,
-                        self.stream.cu_stream(),
-                    );
-                    cudarc::driver::sys::cuStreamSynchronize(self.stream.cu_stream());
-                }
-                let state0 = unsafe { std::slice::from_raw_parts(pinned.ptr, sd) };
-                let ofi_slice = &state0[66..84.min(sd)];
-                let ofi_nonzero = ofi_slice.iter().filter(|&&v| v != 0.0).count();
-                info!(
-                    "STATE_GATHER_DIAG: state[66..84]={:?}, nonzero={}/{}, ofi_dim={}, total_bars={}",
-                    ofi_slice, ofi_nonzero, ofi_slice.len(),
-                    self.ofi_dim, config.total_bars,
-                );
-            }
-
             // ── 2. cuBLAS Q-forward: batch_states → logits ──────────────
             // Graph the ENTIRE forward (bottleneck + Q-network) for deterministic cuBLAS.
             // Ungraphed cublasLtMatmul contaminates cuBLAS internal state, breaking
@@ -2946,16 +2922,6 @@ impl GpuExperienceCollector {
         if ofi_flat.is_empty() {
             return Err(MLError::ModelError("OFI data empty — cannot upload empty features to GPU".into()));
         }
-        // Verify host data before upload — zero-cost check
-        {
-            let first20 = &ofi_flat[..20.min(ofi_flat.len())];
-            let nonzero = first20.iter().filter(|&&v| v != 0.0).count();
-            let total_nonzero = ofi_flat.iter().filter(|&&v| v != 0.0).count();
-            info!(
-                "OFI host verify: bar0={:?}, bar0_nonzero={}/20, total_nonzero={}/{}",
-                &first20[..8.min(first20.len())], nonzero, total_nonzero, ofi_flat.len(),
-            );
-        }
         let gpu_buf = super::clone_htod_f32(&self.stream, ofi_flat)?;
         info!(
             "OFI features uploaded to GPU: {} bars x {} dims ({:.1} KB)",