gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

diag: pinned memory readback of batch_states after state_gather

Browse Source 
Uses PinnedHostBuf + cuMemcpyDtoHAsync for state_gather diagnostic.
Reads batch_states[0..state_dim] to verify OFI at positions [66..84).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 13:31:16 +02:00
parent 4ddd6e1ffa
commit bd0d90482d
1 changed files with 18 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -2051,20 +2051,28 @@ impl GpuExperienceCollector {
)))?;
}
// One-shot diagnostic: after first state_gather, readback batch_states[0] to verify OFI
// 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 mut state0 = vec![0.0_f32; sd];
if let Ok(()) = self.stream.memcpy_dtoh(&self.batch_states, &mut state0) {
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_ptr=0x{:x}, ofi_dim={}, total_bars={}",
ofi_slice, ofi_nonzero, ofi_slice.len(),
self.ofi_gpu.device_ptr(&self.stream).0,
self.ofi_dim, config.total_bars,
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 ──────────────