diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index ad5aa8fc3..046e4d400 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -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}")))?; } diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index 593e46ef6..0c3a5f950 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -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__