diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index dda408b9a..629bb59eb 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -33,6 +33,13 @@ pub struct GpuBatchPtrs { pub weights_ptr: u64, pub indices_ptr: u64, pub episode_ids_ptr: u64, + /// SP13 Layer B: 30-bar aux sign label per transition (i32). + /// `-1` = skip (no lookahead window available), `0` = down, `1` = up. + /// B0 plumbing: written by ring buffer, no consumer yet — B1 wires the + /// aux head's CrossEntropy loss to read this. Always points at + /// `sample_aux_sign_labels` (not direct-to-trainer) until B1 lands the + /// trainer destination buffer. + pub aux_sign_labels_ptr: u64, pub batch_size: usize, pub state_dim: usize, } @@ -46,6 +53,10 @@ struct ReplayKernels { scatter_insert_f32: CudaFunction, scatter_insert_f32_rows: CudaFunction, scatter_insert_u32: CudaFunction, + /// SP13 Layer B: i32 scatter for the 30-bar aux sign label ring buffer. + /// `-1` is a "skip" sentinel that must survive signed; u32 would alias it + /// to 4294967295. Consumer is `gather_i32_scalar` at #18c. + scatter_insert_i32: CudaFunction, gather_f32: CudaFunction, gather_u32: CudaFunction, /// #30 F32 row gather for f32 state storage. @@ -97,6 +108,7 @@ impl ReplayKernels { scatter_insert_f32: ld("scatter_insert_f32")?, scatter_insert_f32_rows: ld("scatter_insert_f32_rows")?, scatter_insert_u32: ld("scatter_insert_u32")?, + scatter_insert_i32: ld("scatter_insert_i32")?, gather_f32: ld("gather_f32")?, gather_u32: ld("gather_u32")?, gather_f32_rows: ld("gather_f32_rows")?, gather_f32_rows_padded: ld("gather_f32_rows_padded")?, @@ -149,6 +161,12 @@ pub struct GpuReplayBuffer { /// u32 chosen to match the gather_u32 / scatter_insert_u32 kernel signatures — /// episode-id values are always non-negative counters. episode_ids: CudaSlice, + /// SP13 Layer B: 30-bar aux sign label ring buffer `[capacity]` i32. + /// Stored as i32 because `-1` is a valid "skip" sentinel that must + /// survive the round-trip — u32 would alias it to 4294967295. + /// Written by `scatter_insert_i32` in `insert_batch`, read by + /// `gather_i32_scalar` in `sample_proportional`. + aux_sign_labels: CudaSlice, write_cursor: usize, size: usize, max_priority: CudaSlice, pending_max_priority: Option>, @@ -172,6 +190,10 @@ pub struct GpuReplayBuffer { sample_weights: CudaSlice, sample_max_weight: CudaSlice, sample_episode_ids: CudaSlice, + /// SP13 Layer B: per-batch aux sign label gather destination. + /// `gather_i32_scalar` writes here from the ring; `GpuBatchPtrs.aux_sign_labels_ptr` + /// points at this buffer until B1 wires a trainer-resident destination. + sample_aux_sign_labels: CudaSlice, total_sum_buf: CudaSlice, // Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step) update_batch_max: CudaSlice, // [1] per-batch atomicMax accumulator @@ -277,6 +299,10 @@ impl GpuReplayBuffer { let tsb = a32f(stream, 1, "ts_buf")?; let ep_ids = a32u(stream, cap, "episode_ids")?; let s_ep = a32u(stream, mbs, "s_episode_ids")?; + // SP13 Layer B: aux sign label ring + per-batch gather destination. + // Zero-init in B0 — B1 producer kernel writes -1/0/1. + let aux_lbl = a32i(stream, cap, "aux_sign_labels")?; + let s_aux = a32i(stream, mbs, "s_aux_sign_labels")?; // Pre-allocate PER update scratch buffers (zero cuMemAlloc per step) let update_batch_max = a32f(stream, 1, "update_batch_max")?; @@ -327,6 +353,7 @@ impl GpuReplayBuffer { config, stream: Arc::clone(stream), kernels: k, states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p, episode_ids: ep_ids, + aux_sign_labels: aux_lbl, write_cursor: 0, size: 0, max_priority: mp, pending_max_priority: None, current_step: 0, priorities_pa, prefix_sum, @@ -337,6 +364,7 @@ impl GpuReplayBuffer { sample_rewards: sr, sample_dones: sdn, sample_priorities: sp, sample_weights: sw, sample_episode_ids: s_ep, + sample_aux_sign_labels: s_aux, sample_max_weight: smw, total_sum_buf: tsb, update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare), _mean_action_pinned: mean_action_pinned, mean_action_dev_ptr, @@ -438,6 +466,7 @@ impl GpuReplayBuffer { ac: &CudaSlice, rw: &CudaSlice, dn: &CudaSlice, + aux_sign: &CudaSlice, // SP13 Layer B: 30-bar aux sign label, -1/0/1 bs: usize, ) -> Result<(), MLError> { if bs == 0 { return Ok(()); } @@ -479,6 +508,13 @@ impl GpuReplayBuffer { .arg(&self.dones).arg(&sd2).arg(&ci).arg(&cpi).arg(&bsi) .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc d f32: {e}")))?; } + // SP13 Layer B: aux sign label scatter (i32, -1/0/1). + let s_aux_in = if off > 0 { aux_sign.slice(off..) } else { aux_sign.slice(0..) }; + unsafe { + self.stream.launch_builder(&self.kernels.scatter_insert_i32) + .arg(&self.aux_sign_labels).arg(&s_aux_in).arg(&ci).arg(&cpi).arg(&bsi) + .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc aux i32: {e}")))?; + } } // Broadcast max_priority into pre-allocated prio buffer unsafe { @@ -715,6 +751,16 @@ impl GpuReplayBuffer { .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g ep: {e}")))?; } + // Step 3c (SP13 Layer B): gather aux sign labels (i32). gather_i32_scalar + // takes [dst, src, indices, batch_size] — no capacity arg, mirrors the + // direct-to-trainer gather pattern at line 651 above. + unsafe { + self.stream.launch_builder(&self.kernels.gather_i32_scalar) + .arg(&mut self.sample_aux_sign_labels).arg(&self.aux_sign_labels) + .arg(&self.sample_indices_i64).arg(&bsi) + .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g aux i32: {e}")))?; + } + // Step 6: IS weights via GPU-resident total_sum (zero CPU readback) // is_weights_f32 reads total_sum from GPU pointer. // When direct_to_trainer is active, write IS weights to trainer's is_weights_buf. @@ -784,6 +830,7 @@ impl GpuReplayBuffer { weights_ptr: self.trainer_is_weights_ptr, indices_ptr: self.sample_indices_u32.raw_ptr(), episode_ids_ptr: self.sample_episode_ids.raw_ptr(), + aux_sign_labels_ptr: self.sample_aux_sign_labels.raw_ptr(), batch_size, state_dim: sd, }) @@ -797,6 +844,7 @@ impl GpuReplayBuffer { weights_ptr: self.sample_weights.raw_ptr(), indices_ptr: self.sample_indices_u32.raw_ptr(), episode_ids_ptr: self.sample_episode_ids.raw_ptr(), + aux_sign_labels_ptr: self.sample_aux_sign_labels.raw_ptr(), batch_size, state_dim: sd, }) @@ -1048,6 +1096,9 @@ fn a32f(s: &Arc, n: usize, nm: &str) -> Result, MLErr fn a32u(s: &Arc, n: usize, nm: &str) -> Result, MLError> { s.alloc_zeros::(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}"))) } +fn a32i(s: &Arc, n: usize, nm: &str) -> Result, MLError> { + s.alloc_zeros::(n).map_err(|e| MLError::ModelError(format!("alloc {nm}: {e}"))) +} #[cfg(test)] mod tests { @@ -1109,10 +1160,11 @@ mod tests { let a = stream.alloc_zeros::(n).unwrap(); let mut r = stream.alloc_zeros::(n).unwrap(); let d = stream.alloc_zeros::(n).unwrap(); + let aux = stream.alloc_zeros::(n).unwrap(); // Give non-zero rewards so priorities are non-zero after update let rewards_host: Vec = (0..n).map(|i| (i as f32 + 1.0) * 0.01).collect(); stream.memcpy_htod(&rewards_host, &mut r).unwrap(); - b.insert_batch(&s, &ns, &a, &r, &d, n).unwrap(); + b.insert_batch(&s, &ns, &a, &r, &d, &aux, n).unwrap(); assert_eq!(b.len(), n); // sample_proportional triggers per_prefix_scan — must complete, not hang diff --git a/crates/ml-dqn/src/replay_buffer_kernels.cu b/crates/ml-dqn/src/replay_buffer_kernels.cu index 2c3b8ecd3..2f793c0f6 100644 --- a/crates/ml-dqn/src/replay_buffer_kernels.cu +++ b/crates/ml-dqn/src/replay_buffer_kernels.cu @@ -3,6 +3,7 @@ // Kernels: // 1. scatter_insert_f32 — ring buffer insert for f32 arrays (rewards, dones, priorities) // 2. scatter_insert_u32 — ring buffer insert for u32 arrays (actions) +// 2b. scatter_insert_i32 — ring buffer insert for i32 arrays (SP13 aux sign labels with -1 sentinel) // 3. scatter_insert — ring buffer insert for state matrices // 4. gather_f32 — index_select for f32 arrays // 5. gather_u32 — index_select for u32 arrays @@ -45,6 +46,26 @@ void scatter_insert_u32( dst[dst_idx] = src[i]; } +// ── 2b. Scatter insert for i32 (1D, signed) ──────────────────────────────── +// Mirror of scatter_insert_u32 with i32 storage. SP13 Layer B uses this +// for the 30-bar aux sign label where -1 is a "skip" sentinel that must +// survive the ring-buffer round-trip as a signed value (gather_i32_scalar +// at #18c is its consumer). + +extern "C" __global__ +void scatter_insert_i32( + int* __restrict__ dst, + const int* __restrict__ src, + int start_idx, + int capacity, + int batch_size +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= batch_size) return; + int dst_idx = (start_idx + i) % capacity; + dst[dst_idx] = src[i]; +} + // ── 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 d8d511b74..41ed6ef2d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -437,6 +437,14 @@ pub struct GpuExperienceBatch { /// Episode `ep` occupies flat indices `[ep * timesteps, (ep+1) * timesteps)`. /// Required by HER Future and Final strategies for GPU-native donor selection. pub episode_ids: CudaSlice, + /// SP13 Layer B: 30-bar aux sign label per transition `[total]` on GPU (i32). + /// `-1` = skip (no 30-bar lookahead window available, e.g. last L-30 bars of + /// each episode), `0` = down (price_t+30 < price_t), `1` = up otherwise. + /// B0 plumbing: zero-initialized — B1 wires the producer kernel that fills + /// real labels and the consumer (CrossEntropy aux loss head). The replay + /// buffer threads this column through `insert_batch` → `gather_i32_scalar` + /// → `GpuBatchPtrs.aux_sign_labels_ptr`. + pub aux_sign_labels: CudaSlice, /// State dimensionality (so the consumer can reshape) pub state_dim: usize, /// Number of episodes @@ -3426,6 +3434,14 @@ impl GpuExperienceCollector { self.last_experience_count = total; + // SP13 Layer B (B0 plumbing): zero-initialized aux sign label column. + // B1 will replace this `alloc_zeros` with a producer kernel that + // computes -1/0/1 from the price trajectory; for B0, all labels = 0 + // (no consumer reads them yet). + let aux_sign_labels = self.stream + .alloc_zeros::(total) + .map_err(|e| MLError::ModelError(format!("alloc aux_sign_labels: {e}")))?; + Ok(GpuExperienceBatch { states, next_states, @@ -3433,6 +3449,7 @@ impl GpuExperienceCollector { rewards, dones, episode_ids, + aux_sign_labels, state_dim: sd, n_episodes, timesteps, diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 6741d0ebe..9c9579bb2 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -2035,6 +2035,7 @@ impl DQNTrainer { actions_u32, &gpu_batch.rewards, &gpu_batch.dones, + &gpu_batch.aux_sign_labels, total, ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index c1eec8c0f..8f8968d5e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -5861,3 +5861,53 @@ The P0a Hold-pricing controller wired the per-bar Hold cost in `experience_kerne ### Files (atomic P0b commit) 3 files: `sp13_isv_slots.rs` (constants), `training_loop.rs` (controller + helper extraction), `tests.rs` (3 unit tests). No new kernels, no new ISV slots, no fingerprint change. + +## SP13 Layer B — Commit B0: replay buffer i32 ring + GpuBatchPtrs plumbing (2026-05-05) + +**Why split B0 from B1**: the original Layer B Commit 1 bundled replay-buffer plumbing (i32 ring + GpuBatchPtrs field) with the aux head contract change (1→2 dim, MSE→CE, softmax-driven Q-head input wire, slot 117 retire, fingerprint bump). Four implementer subagents in a row hit `NEEDS_CONTEXT` on brief-accuracy issues — the bundle was too large for a single dispatch. Splitting at this seam isolates the mechanical column-by-column migration (B0) from the contract-changing aux-head work (B1) so B1 can be re-dispatched to a fresh implementer with a precise audit-derived brief while B0 lands as pure plumbing with zero behavior change. + +**Behavior contract**: B0 is pure plumbing — no consumer reads the new column, all `aux_sign_labels` slots are zero-initialized in the producer, and the round-trip preserves the column through scatter→ring→gather→`GpuBatchPtrs.aux_sign_labels_ptr`. A smoke between B0 and B1 should be bit-identical to the parent commit `0ad5b6fa4` modulo allocator entropy. B1 will replace the `alloc_zeros` producer with a kernel that computes -1/0/1 from the price trajectory and wire the aux head's CrossEntropy loss to consume the column. + +### Wiring (data flow) + + 1. **`scatter_insert_i32` kernel** (`replay_buffer_kernels.cu:49-67`) — mirrors `scatter_insert_u32` but preserves -1 sentinels (u32 would alias them to 4294967295). Symmetric to existing `gather_i32_scalar` consumer (which actions/direct-to-trainer path already uses). + 2. **`ReplayKernels.scatter_insert_i32: CudaFunction`** (`gpu_replay_buffer.rs`) — field + `ld("scatter_insert_i32")` loader entry. + 3. **`GpuReplayBuffer.aux_sign_labels: CudaSlice`** (ring `[capacity]`) + **`sample_aux_sign_labels: CudaSlice`** (per-batch gather destination). + 4. **`insert_batch` signature** gains `aux_sign: &CudaSlice` parameter; body launches `scatter_insert_i32` alongside actions/rewards/dones scatters using the same `(ci, cpi, bsi)` tuple. + 5. **`sample_proportional` Step 3c** launches `gather_i32_scalar` from ring into `sample_aux_sign_labels`. Both `GpuBatchPtrs` return paths (direct-to-trainer + fallback) point `aux_sign_labels_ptr` at the sample buffer's `raw_ptr()`. B1 will add a trainer-resident destination + direct gather for the direct-to-trainer fast path. + 6. **`GpuBatchPtrs.aux_sign_labels_ptr: u64`** publicly exposed. + 7. **`GpuExperienceBatch.aux_sign_labels: CudaSlice`** (`gpu_experience_collector.rs`) field; `collect_experiences_gpu` allocates `total`-sized zero-init slice (B1 replaces with kernel-driven producer). + 8. **Single production caller** (`training_loop.rs:2032` — `agent.primary_dqn_mut().memory.gpu.insert_batch`) threads `&gpu_batch.aux_sign_labels` through. + +### Audit-verified call site cardinality (per implementer 4 finding) + +| Site | Count | Action | +|---|---|---| +| `insert_batch` production callers | 1 | `training_loop.rs:2032` updated | +| `insert_batch` test callers | 1 | in-file `gpu_replay_buffer::tests` updated | +| `GpuBatchPtrs` construction sites | 2 | both inside `sample_proportional` (direct-to-trainer + fallback paths) | +| `GpuExperienceBatch` construction sites | 1 | `collect_experiences_gpu` end | + +Implementer 4 specifically caught the over-counted "4 sites" claim from the original brief — the 2 extra were `insert_batch_with_episode_ids` references in doc comments only (no separate function exists; episode_ids are passed via the same `insert_batch` signature). + +### Hard rules upheld + +- `feedback_no_partial_refactor`: every consumer of the changed `insert_batch` signature migrates atomically (1 production + 1 test call site, both updated) +- `feedback_no_stubs`: B0 is not a return-zero stub — the column carries real data through a real kernel; the field is allocated, written, scatter-launched, gathered, and exposed via `GpuBatchPtrs`. Zero values in B0 are valid I32 values that B1 will overwrite with -1/0/1 in the producer kernel +- `feedback_isv_for_adaptive_bounds`: no new ISV slots introduced (ISV[117]=AUX_LABEL_SCALE_EMA_INDEX retirement deferred to B1 atomic when MSE→CE flips and the EMA loses its consumer) +- `feedback_no_atomicadd`: no new producers/reductions added +- `feedback_cpu_is_read_only`: no host computation introduced; alloc_zeros + scatter_insert_i32 + gather_i32_scalar are all GPU-resident + +Build: `cargo check --workspace` clean in 21s. Only pre-existing warnings (DevicePtrMut unused imports — orthogonal). + +### Files (atomic B0 commit) + +4 files, 92 LOC net additions: +- `crates/ml-dqn/src/replay_buffer_kernels.cu` (+21) — scatter_insert_i32 kernel +- `crates/ml-dqn/src/gpu_replay_buffer.rs` (+54) — kernel field/loader, ring + sample buffers, insert_batch param/launch, sample_proportional gather, both GpuBatchPtrs returns, in-file test caller, a32i helper +- `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+17) — GpuExperienceBatch field + alloc_zeros producer +- `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (+1) — caller threads through + +### Next: B1 (separate atomic commit, fresh-implementer dispatch) + +B1 covers: aux head 1→2 dim, MSE→CE loss, `aux_dir_acc` reads softmax, `aux_pred_to_isv_tanh` rewrite as logit-diff, ISV[117]=AUX_LABEL_SCALE_EMA_INDEX retirement, `dqn_param_layout` fingerprint bump, kernel that populates `aux_sign_labels` with real -1/0/1 labels from the 30-bar price trajectory, `aux_b1_diag` HEALTH_DIAG metric, 17+ GPU oracle unit tests. Brief written from B0's audit findings (this section) — no implementer should re-derive call site cardinality.