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

fix: PER insert used 1D kernel for 2D state matrices — model trained on garbage

Browse Source 
scatter_insert_f32 is a 1D scalar kernel (5 args: dst, src, cursor, cap,
batch_size). insert_batch called it with 6 args for state matrices, passing
state_dim as batch_size. CUDA silently dropped the 6th arg (actual batch_size).

Result: only 96 floats (1 state row) inserted per experience batch into PER.
The model was training on ~99.99% uninitialized GPU memory. This bug affected
ALL state features, not just OFI — market features and portfolio were also
garbage in PER-sampled training batches.

Fix: added scatter_insert_f32_rows kernel (2D-aware, 6 args: dst, src,
cursor, cap, state_dim, batch_size) matching the existing scatter_insert
(bf16) pattern. States and next_states now use the row-aware kernel.

Verified locally: OFI_DIAG shows non-zero values through the full chain
(state_gather → env_step → PER insert → PER sample → trainer states_buf).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-20 13:57:23 +02:00
parent bd0d90482d
commit 7af0a228fb
2 changed files with 27 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/ml-dqn/src/gpu_replay_buffer.rs
View File

@@ -44,6 +44,7 @@ pub struct GpuBatchPtrs {
struct ReplayKernels {
scatter_insert_f32: CudaFunction,
scatter_insert_f32_rows: CudaFunction,
scatter_insert_u32: CudaFunction,
gather_f32: CudaFunction,
gather_u32: CudaFunction,
@@ -90,6 +91,7 @@ impl ReplayKernels {
Ok(Self {
scatter_insert_f32: ld("scatter_insert_f32")?,
scatter_insert_f32_rows: ld("scatter_insert_f32_rows")?,
scatter_insert_u32: ld("scatter_insert_u32")?,
gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?,
gather_f32_rows: ld("gather_f32_rows")?,
@@ -406,17 +408,17 @@ impl GpuReplayBuffer {
self.ensure_insert_scratch(eff)?;
let el = eff * sd;
// #30: f32 states — scatter insert directly (full precision storage)
// #30: f32 states — 2D row scatter into ring buffer [cap, state_dim]
let ss = if off > 0 { sf.slice(off * sd..) } else { sf.slice(0..) };
let sn = if off > 0 { nf.slice(off * sd..) } else { nf.slice(0..) };
let (ci, cpi, sdi, bsi) = (self.write_cursor as i32, cap as i32, sd as i32, eff as i32);
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
self.stream.launch_builder(&self.kernels.scatter_insert_f32_rows)
.arg(&self.states).arg(&ss).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc s f32: {e}")))?;
}
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
self.stream.launch_builder(&self.kernels.scatter_insert_f32_rows)
.arg(&self.next_states).arg(&sn).arg(&ci).arg(&cpi).arg(&sdi).arg(&bsi)
.launch(lcfg(el)).map_err(|e| MLError::ModelError(format!("sc n f32: {e}")))?;
}

22
crates/ml-dqn/src/replay_buffer_kernels.cu
View File

@@ -65,6 +65,28 @@ void scatter_insert(
dst[dst_row * state_dim + col] = src[row * state_dim + col];
}
// ── 3b. 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].
extern "C" __global__
void scatter_insert_f32_rows(
float* __restrict__ dst, // [capacity * state_dim]
const float* __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];
}
// ── 4. Gather f32 by indices ────────────────────────────────────────────────
extern "C" __global__