From da5e564ccfad43071f579a824ef122dc7017d837 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 14 May 2026 02:13:05 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp22-vnext):=20Phase=20B4b-2=20=E2=80=94?= =?UTF-8?q?=20replay-buffer=20scatter=20+=20trainer=20setter?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second half of Phase B4b. Wires the trade-outcome label column through the replay buffer (struct field, allocation, scatter on insert, gather on sample, direct-to-trainer pointer + setter) and connects the trainer's aux_to_label_buf to receive per-batch sampled labels. Completes the end-to-end producer → ring → trainer i32 path the K=3 sparse CE consumer reads. Replay buffer changes (crates/ml-dqn/src/gpu_replay_buffer.rs): - New struct fields: aux_outcome_labels (capacity-sized ring), sample_aux_outcome_labels (mbs-sized fallback gather), trainer_aux_ outcome_labels_ptr (direct-path destination) - New aux_outcome_labels_ptr field on GpuBatchPtrs - insert_batch signature: new aux_outcome arg between aux_conf_in and bs. Scatters via existing K-generic scatter_insert_i32 (same kernel the K=2 aux_sign_labels uses). - sample_proportional direct gather when trainer_aux_outcome_labels_ptr != 0; fallback gather otherwise. Mirrors aux_sign_labels direct/ fallback semantic exactly. - New setter set_trainer_aux_outcome_labels_ptr mirrors set_trainer_aux_conf_ptr. Collector emission: - New aux_outcome_labels field on GpuExperienceBatch - Populated at end of collect_experiences_gpu via dtod_clone_i32 from Phase B4b-1's per-(env, t) producer scratch. Trainer wireup: - aux_to_label_buf_ptr() accessor on GpuDqnTrainer (mirrors aux_nb_label_buf_ptr) - trainer_aux_to_label_buf_ptr() delegating accessor on FusedTrainingCtx - New set_trainer_aux_outcome_labels_ptr call in training_loop at the same site where set_trainer_buffers + set_trainer_aux_conf_ptr fire. Two call sites updated (lines ~835 + ~2857). - insert_batch call in training_loop passes &gpu_batch.aux_outcome _labels as new arg. Test fixtures updated: 4 smoke test files + 3 unit-test fixtures in gpu_replay_buffer.rs alloc zero-init aux_outcome i32 arg. End-to-end chain complete: trade_outcome_label_kernel (A2) → collector per-step launch (B3) → collector emission (B4b-1) → replay-buffer insert + scatter (B4b-2) → PER sample + direct gather (B4b-2) → trainer aux_heads_forward.loss_reduce (B4) reads sparse {-1,0,1,2} → trainer aux_heads_backward (B4) computes per-sample partials → Adam SAXPY (B1+B4) updates W1, b1, W2, b2 at [163..167) The "degraded predict-Profit-everywhere" cold-start from Phase B4 is resolved. K=3 head trains on real sparse trade-outcome labels. Verification: - cargo check -p ml clean (21 warnings, none new). - cargo test -p ml --lib → 1016/0 on clean runs; pre-existing NoisyLinear flake still surfaces ~30-50% of runs (unrelated to vNext work — see ebc1b1502 / 20e1aea27 commit notes). Remaining vNext work: - B5: input concat 256 → 262 with plan_params - Phase C: 3-slot state assembly (state[121..124]) - Phase D: 12-weight W atom-shift (4 actions × 3 outcomes) - Phase E: dW backward + Adam for W[4, 3] - Phase F: validation smoke at β=0.5 structural prior Audit: docs/dqn-wire-up-audit.md Phase B4b-2 section. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-dqn/src/gpu_replay_buffer.rs | 138 +++++++++++++++++- .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 13 ++ .../cuda_pipeline/gpu_experience_collector.rs | 24 +++ crates/ml/src/trainers/dqn/fused_training.rs | 13 ++ .../trainers/dqn/smoke_tests/gpu_residency.rs | 4 +- .../trainers/dqn/smoke_tests/performance.rs | 3 +- .../dqn/smoke_tests/training_stability.rs | 6 +- .../src/trainers/dqn/trainer/training_loop.rs | 31 ++++ docs/dqn-wire-up-audit.md | 48 ++++++ 9 files changed, 271 insertions(+), 9 deletions(-) diff --git a/crates/ml-dqn/src/gpu_replay_buffer.rs b/crates/ml-dqn/src/gpu_replay_buffer.rs index 0a4b2d6c8..66eecdaf1 100644 --- a/crates/ml-dqn/src/gpu_replay_buffer.rs +++ b/crates/ml-dqn/src/gpu_replay_buffer.rs @@ -40,6 +40,17 @@ pub struct GpuBatchPtrs { /// `sample_aux_sign_labels` (not direct-to-trainer) until B1 lands the /// trainer destination buffer. pub aux_sign_labels_ptr: u64, + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome label per + /// transition (i32). `-1` = mask (no trade close at this bar), `0` + /// = Profit, `1` = Stop, `2` = Timeout. Producer: + /// `trade_outcome_label_kernel`'s per-(env, t) output (Phase B4b-1 + /// addition); rings through replay buffer's `aux_outcome_labels` + /// CudaSlice. Direct-path destination: trainer's `aux_to_label_buf` + /// (set via `set_trainer_aux_outcome_labels_ptr`). When the trainer + /// pointer is 0 (default at construction), gather falls back to + /// `sample_aux_outcome_labels` and this pointer points there. + /// Same direct-path pattern as `aux_sign_labels_ptr`. + pub aux_outcome_labels_ptr: u64, /// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf at the /// SAMPLED bar (f32). Range [0, 0.5] — K=2 peak-softmax-above- /// uniform confidence. 0.0 = no information (sentinel / aux head @@ -178,6 +189,14 @@ pub struct GpuReplayBuffer { /// Written by `scatter_insert_i32` in `insert_batch`, read by /// `gather_i32_scalar` in `sample_proportional`. aux_sign_labels: CudaSlice, + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): per-transition trade- + /// outcome label ring `[capacity]` i32. Stored as i32 because `-1` + /// (mask = no trade close at this bar) must survive the round-trip + /// — u32 would alias it to 4294967295. Mirrors `aux_sign_labels` + /// exactly (same kernel `scatter_insert_i32` for insert, + /// `gather_i32_scalar` for sample). Sparse: most slots ~95-99% are + /// -1 (mask). + aux_outcome_labels: CudaSlice, /// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf ring /// buffer `[capacity]` f32. Range [0, 0.5] — K=2 peak-softmax- /// above-uniform confidence. 0.0 = no information (sentinel). @@ -215,6 +234,12 @@ pub struct GpuReplayBuffer { /// `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, + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): per-batch trade-outcome + /// label gather destination. `gather_i32_scalar` writes here from + /// the ring when the trainer-direct pointer isn't wired; when wired + /// the gather goes directly to the trainer's `aux_to_label_buf` and + /// this buffer goes unused. + sample_aux_outcome_labels: CudaSlice, /// SP20 Phase 3 Task 3.4 (2026-05-10): per-batch aux_conf gather /// destination. `gather_f32_scalar` writes here from the ring; /// `GpuBatchPtrs.aux_conf_ptr` points at this buffer until Phase 5 @@ -268,6 +293,15 @@ pub struct GpuReplayBuffer { /// actions / rewards. Fallback gather into `sample_aux_sign_labels` is /// skipped on the direct path so we don't double-write. trainer_aux_sign_labels_ptr: u64, + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): trainer destination for + /// the trade-outcome label i32 column. When non-zero, + /// `sample_proportional` gathers labels directly into the trainer's + /// `aux_to_label_buf`. When 0 (default at construction), falls back + /// to gathering into `sample_aux_outcome_labels` and + /// `GpuBatchPtrs.aux_outcome_labels_ptr` points there. Same direct- + /// path pattern as `trainer_aux_sign_labels_ptr` — wired via + /// `set_trainer_aux_outcome_labels_ptr` mirror. + trainer_aux_outcome_labels_ptr: u64, /// SP20 Phase 3 Task 3.4 (2026-05-10): trainer destination for the /// aux_conf f32 column. When non-zero, `sample_proportional` /// gathers `aux_conf_at_state` directly into the trainer's @@ -400,6 +434,12 @@ impl GpuReplayBuffer { // 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")?; + // SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome label + // ring + per-batch gather destination. Mirrors aux_sign_labels. + // Zero-init = label 0 (Profit) sentinel; populated by the + // Phase B4b-1 producer kernel's per-(env, t) writes. + let aux_outcome_lbl = a32i(stream, cap, "aux_outcome_labels")?; + let s_aux_outcome = a32i(stream, mbs, "s_aux_outcome_labels")?; // SP20 Phase 3 Task 3.4 (2026-05-10): aux_conf ring + per-batch // gather destination. Zero-init = K=2 uniform-prior sentinel // (no information). Mirrors the aux_sign_labels f32 path. @@ -456,6 +496,7 @@ impl GpuReplayBuffer { states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p, episode_ids: ep_ids, aux_sign_labels: aux_lbl, + aux_outcome_labels: aux_outcome_lbl, aux_conf: aux_conf_ring, write_cursor: 0, size: 0, max_priority: mp, pending_max_priority: None, current_step: 0, @@ -468,6 +509,7 @@ impl GpuReplayBuffer { sample_priorities: sp, sample_weights: sw, sample_episode_ids: s_ep, sample_aux_sign_labels: s_aux, + sample_aux_outcome_labels: s_aux_outcome, sample_aux_conf: s_aux_conf, sample_max_weight: smw, total_sum_buf: tsb, update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare), @@ -487,6 +529,7 @@ impl GpuReplayBuffer { trainer_dones_ptr: 0, trainer_is_weights_ptr: 0, trainer_aux_sign_labels_ptr: 0, + trainer_aux_outcome_labels_ptr: 0, trainer_aux_conf_ptr: 0, trainer_state_dim_padded: 0, learning_health_cache: 1.0, @@ -668,6 +711,11 @@ impl GpuReplayBuffer { dn: &CudaSlice, aux_sign: &CudaSlice, // SP13 Layer B: 30-bar aux sign label, -1/0/1 aux_conf_in: &CudaSlice, // SP20 Phase 3 Task 3.4: per-bar aux_conf, [0, 0.5] + // SP22 H6 vNext Phase B4b-2 (2026-05-14): per-(env, t) trade- + // outcome label column from the collector's + // `exp_aux_to_label_per_sample`. Sparse: ~95-99% bars are -1 + // (mask). Same scatter pattern as `aux_sign`. + aux_outcome: &CudaSlice, bs: usize, ) -> Result<(), MLError> { if bs == 0 { return Ok(()); } @@ -716,6 +764,14 @@ impl GpuReplayBuffer { .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}")))?; } + // SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome label + // scatter (i32, {-1, 0, 1, 2}). Same pattern as aux_sign. + let s_aux_outcome_in = if off > 0 { aux_outcome.slice(off..) } else { aux_outcome.slice(0..) }; + unsafe { + self.stream.launch_builder(&self.kernels.scatter_insert_i32) + .arg(&self.aux_outcome_labels).arg(&s_aux_outcome_in).arg(&ci).arg(&cpi).arg(&bsi) + .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc aux_outcome i32: {e}")))?; + } // SP20 Phase 3 Task 3.4 (2026-05-10): aux_conf scatter // (f32, [0, 0.5]). Mirrors the aux_sign_labels i32 path // immediately above; reuses `scatter_insert_f32` (the same @@ -1118,6 +1174,24 @@ impl GpuReplayBuffer { .launch(lcfg(batch_size)) .map_err(|e| MLError::ModelError(format!("gather aux_sign_labels direct: {e}")))?; } + // SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome labels + // gather i32 scalar → trainer.aux_to_label_buf. Mirror of + // aux_sign_labels direct path. When trainer ptr is 0 (unwired), + // the fallback gather below populates sample_aux_outcome_labels + // and the trainer's CE loss reads stale alloc_zeros labels (all + // Profit, degraded cold-start matching K=2 head's B1.1a state). + let t_aux_outcome = self.trainer_aux_outcome_labels_ptr; + if t_aux_outcome != 0 { + unsafe { + self.stream.launch_builder(&self.kernels.gather_i32_scalar) + .arg(&t_aux_outcome) + .arg(&self.aux_outcome_labels) + .arg(&self.sample_indices_i64) + .arg(&bsi) + .launch(lcfg(batch_size)) + .map_err(|e| MLError::ModelError(format!("gather aux_outcome_labels direct: {e}")))?; + } + } } else { // ── Fallback: gather into intermediate sample_* buffers ── // #30: gather f32 rows (full-precision state storage) @@ -1180,6 +1254,24 @@ impl GpuReplayBuffer { } } + // SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome labels + // fallback gather. Mirrors aux_sign_labels fallback. Only fires + // when direct-to-trainer path didn't gather (trainer ptr 0) OR + // when direct path is active but trainer_aux_outcome_labels_ptr + // is independently 0 (e.g., aux_to_label_buf not yet wired — + // the direct path's i32 gather above is conditional on + // t_aux_outcome != 0). The fallback writes into + // sample_aux_outcome_labels which GpuBatchPtrs.aux_outcome + // _labels_ptr points at when trainer ptr is 0. + if !direct_to_trainer || self.trainer_aux_outcome_labels_ptr == 0 { + unsafe { + self.stream.launch_builder(&self.kernels.gather_i32_scalar) + .arg(&mut self.sample_aux_outcome_labels).arg(&self.aux_outcome_labels) + .arg(&self.sample_indices_i64).arg(&bsi) + .launch(lcfg(batch_size)).map_err(|e| MLError::ModelError(format!("g aux_outcome i32: {e}")))?; + } + } + // Step 3d (SP20 Phase 3 Task 3.4): gather per-bar aux_conf (f32). // Always gathers into `sample_aux_conf` (the GpuBatchPtrs // fallback path) UNLESS the trainer wired @@ -1286,6 +1378,17 @@ impl GpuReplayBuffer { // (already populated by the direct-path gather above); the // fallback `sample_aux_sign_labels` is unused on this path. aux_sign_labels_ptr: self.trainer_aux_sign_labels_ptr, + // SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome + // labels — same direct/fallback semantic as aux_sign. + // When trainer ptr is wired the direct gather populated + // it; when 0 the fallback gather populated + // sample_aux_outcome_labels and that's where the pointer + // points. + aux_outcome_labels_ptr: if self.trainer_aux_outcome_labels_ptr != 0 { + self.trainer_aux_outcome_labels_ptr + } else { + self.sample_aux_outcome_labels.raw_ptr() + }, // SP20 Phase 3 Task 3.4: aux_conf points at the trainer // destination if wired (Phase 5), otherwise at the // fallback `sample_aux_conf` (already populated by the @@ -1306,6 +1409,7 @@ impl GpuReplayBuffer { 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(), + aux_outcome_labels_ptr: self.sample_aux_outcome_labels.raw_ptr(), aux_conf_ptr: aux_conf_dst_ptr, batch_size, state_dim: sd, @@ -1325,6 +1429,25 @@ impl GpuReplayBuffer { self.trainer_aux_conf_ptr = ptr; } + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): wire the trainer's + /// destination pointer for trade-outcome label gathering. When + /// non-zero, `sample_proportional` gathers labels directly into + /// the trainer's `aux_to_label_buf` (mirrors + /// `set_trainer_buffers`'s `aux_sign_labels_ptr` direct-path + /// semantic). When 0 (default at construction), falls back to the + /// intermediate `sample_aux_outcome_labels` buffer. The trainer + /// calls this once at init after its `aux_to_label_buf` is + /// allocated (Phase B2). + pub fn set_trainer_aux_outcome_labels_ptr(&mut self, ptr: u64) { + self.trainer_aux_outcome_labels_ptr = ptr; + } + + /// SP22 H6 vNext Phase B4b-2 accessor: returns the wired trainer + /// trade-outcome labels destination pointer (0 if unwired). + pub const fn trainer_aux_outcome_labels_ptr(&self) -> u64 { + self.trainer_aux_outcome_labels_ptr + } + /// SP20 Phase 3 Task 3.4 accessor: returns the wired trainer /// aux_conf destination pointer (0 if unwired). Used by oracle /// tests to verify the wiring contract. @@ -1692,10 +1815,11 @@ mod tests { let d = stream.alloc_zeros::(n).unwrap(); let aux = stream.alloc_zeros::(n).unwrap(); let aux_conf = stream.alloc_zeros::(n).unwrap(); + let aux_outcome = stream.alloc_zeros::(n).unwrap(); // Non-trivial rewards drive non-zero priorities through `per_insert_pa`. let rewards_host: Vec = (0..n).map(|i| (i as f32 + 1.0) * 0.05).collect(); stream.memcpy_htod(&rewards_host, &mut r).unwrap(); - b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, n).unwrap(); + b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, &aux_outcome, n).unwrap(); assert_eq!(b.len(), n, "post-insert size should be n"); assert!(!b.is_empty(), "post-insert buffer should not be empty"); assert!(b.can_sample(bs), "post-insert sampler should be ready"); @@ -1772,10 +1896,11 @@ mod tests { let d = stream.alloc_zeros::(n).unwrap(); let aux = stream.alloc_zeros::(n).unwrap(); let aux_conf = stream.alloc_zeros::(n).unwrap(); + let aux_outcome = 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, &aux, &aux_conf, n).unwrap(); + b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, &aux_outcome, n).unwrap(); assert_eq!(b.len(), n); // sample_proportional triggers per_prefix_scan — must complete, not hang @@ -1839,7 +1964,8 @@ mod tests { let d = stream.alloc_zeros::(bs).unwrap(); let aux = stream.alloc_zeros::(bs).unwrap(); let aux_conf = stream.alloc_zeros::(bs).unwrap(); - b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, bs).unwrap(); + let aux_outcome = stream.alloc_zeros::(bs).unwrap(); + b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, &aux_outcome, bs).unwrap(); // 2. Fire `update_priorities_gpu` with TD errors spanning [td_min, td_max]. // Health stays at default 1.0, so the standard `per_update_pa` path @@ -1970,8 +2096,9 @@ mod tests { .collect(); let mut aux_conf = stream.alloc_zeros::(n).unwrap(); stream.memcpy_htod(&aux_conf_host, &mut aux_conf).unwrap(); + let aux_outcome = stream.alloc_zeros::(n).unwrap(); - buf.insert_batch(&s, &ns, &a, &r, &d, &aux_sign, &aux_conf, n) + buf.insert_batch(&s, &ns, &a, &r, &d, &aux_sign, &aux_conf, &aux_outcome, n) .expect("insert_batch with aux_conf"); assert_eq!(buf.len(), n, "post-insert size = n"); @@ -2191,8 +2318,9 @@ mod tests { .collect(); let mut aux_conf = stream.alloc_zeros::(n).unwrap(); stream.memcpy_htod(&aux_conf_host, &mut aux_conf).unwrap(); + let aux_outcome = stream.alloc_zeros::(n).unwrap(); - buf.insert_batch(&s, &ns, &a, &r, &d, &aux_sign, &aux_conf, n) + buf.insert_batch(&s, &ns, &a, &r, &d, &aux_sign, &aux_conf, &aux_outcome, n) .expect("insert_batch with aux_conf"); assert_eq!(buf.len(), n, "post-insert size = n"); diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 8ccd7b6c6..9591b9a07 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -34744,6 +34744,19 @@ impl GpuDqnTrainer { /// writes the SAMPLED bar's aux_conf into here on every step. pub(crate) fn aux_conf_at_state_buf_ptr(&self) -> u64 { self.aux_conf_at_state_buf.raw_ptr() } + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): raw pointer to + /// `aux_to_label_buf` for direct-to-trainer PER gather. Mirrors the + /// SP13 B1.1b `aux_nb_label_buf_ptr` direct-to-trainer pattern. The + /// buffer is `CudaSlice` `[batch_size]`; the K=3 trade-outcome + /// sparse CE loss reduce consumer (`aux_trade_outcome_loss_reduce`) + /// reads `{-1, 0, 1, 2}` labels — `-1` is the mask sentinel that + /// must survive signed (i32, not u32). + /// `GpuReplayBuffer::set_trainer_aux_outcome_labels_ptr` is invoked + /// once at trainer init in `training_loop.rs` so PER's + /// `gather_i32_scalar` writes the SAMPLED bar's label into here on + /// every step. + pub(crate) fn aux_to_label_buf_ptr(&self) -> u64 { self.aux_to_label_buf.raw_ptr() } + /// Padded state dimension (128-byte aligned) for direct-to-trainer gather. pub(crate) fn state_dim_padded(&self) -> usize { ml_core::state_layout::STATE_DIM_PADDED } diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 946e1b5c2..2bd7f2b40 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -477,6 +477,18 @@ pub struct GpuExperienceBatch { /// buffer threads this column through `insert_batch` → `gather_i32_scalar` /// → `GpuBatchPtrs.aux_sign_labels_ptr`. pub aux_sign_labels: CudaSlice, + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): per-(env, t) trade-outcome + /// label column. Sized `[n_episodes × timesteps]` i32. Values in + /// `{-1, 0, 1, 2}` where -1 = mask (no trade close at this bar), + /// 0 = Profit, 1 = Stop, 2 = Timeout. Sparse — ~95-99% slots are -1. + /// + /// Producer: `trade_outcome_label_kernel`'s `out_labels_per_sample` + /// output (Phase B4b-1 kernel amendment), written per-rollout-step. + /// Consumer: `GpuReplayBuffer::insert_batch` scatters this column + /// into the replay ring; `sample_proportional` gathers per-batch + /// labels into the trainer's `aux_to_label_buf` via the + /// `trainer_aux_outcome_labels_ptr` direct path. + pub aux_outcome_labels: CudaSlice, /// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf per /// transition `[total]` on GPU (f32) — the K=2 peak-softmax-above- /// uniform `aux_conf ∈ [0, 0.5]` value at the time the bar was @@ -5309,6 +5321,17 @@ impl GpuExperienceCollector { &self.stream, &self.aux_conf_per_sample, total, "aux_conf" )?; + // SP22 H6 vNext Phase B4b-2 (2026-05-14): copy the per-(env, t) + // trade-outcome label tile out of the collector's scratch into + // the emitted batch. Same dtod-clone pattern as `aux_sign_labels` + // (the K=2 sibling label column above). The kernel's per-step + // writes (Phase B4b-1 amendment) accumulate into the scratch + // across the entire rollout — every (env, t) slot has its label + // by the time the rollout completes. + let aux_outcome_labels = dtod_clone_i32( + &self.stream, &self.exp_aux_to_label_per_sample, total, "aux_outcome_labels" + )?; + Ok(GpuExperienceBatch { states, next_states, @@ -5317,6 +5340,7 @@ impl GpuExperienceCollector { dones, episode_ids, aux_sign_labels, + aux_outcome_labels, aux_conf, state_dim: sd, n_episodes, diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index cc2605f34..6d6c2421a 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -4173,6 +4173,19 @@ impl FusedTrainingCtx { self.trainer.aux_conf_at_state_buf_ptr() } + /// SP22 H6 vNext Phase B4b-2 (2026-05-14): raw pointer to trainer's + /// `aux_to_label_buf` for direct-to-trainer PER gather. Wraps the + /// inner `GpuDqnTrainer::aux_to_label_buf_ptr()` accessor — same + /// delegation pattern as `trainer_aux_sign_labels_buf_ptr` / + /// `trainer_aux_conf_at_state_buf_ptr`. Caller (`training_loop.rs`) + /// pipes this into `GpuReplayBuffer::set_trainer_aux_outcome_labels + /// _ptr` so PER's `gather_i32_scalar` writes the per-batch sampled + /// label straight into the trainer's i32 buffer the K=3 trade- + /// outcome CE loss reduce consumer reads. + pub(crate) fn trainer_aux_to_label_buf_ptr(&self) -> u64 { + self.trainer.aux_to_label_buf_ptr() + } + /// Padded state dimension for direct-to-trainer gather. pub(crate) fn trainer_state_dim_padded(&self) -> usize { self.trainer.state_dim_padded() diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index 993ceddf5..bd9563c85 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -74,7 +74,9 @@ fn insert_random_batch( let dones = zeros_cuda_f32(n, stream)?; let aux_sign = stream.alloc_zeros::(n).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?; let aux_conf = stream.alloc_zeros::(n).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?; - buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, n)?; + // SP22 H6 vNext Phase B4b-2 (2026-05-14): trade-outcome labels arg. + let aux_outcome = stream.alloc_zeros::(n).map_err(|e| anyhow::anyhow!("aux_outcome alloc: {e}"))?; + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, &aux_outcome, n)?; Ok(()) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs index f8e13132d..862c218d5 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs @@ -141,7 +141,8 @@ fn test_per_sample_latency() -> anyhow::Result<()> { let dones = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?; let aux_sign = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?; let aux_conf = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?; - buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, batch_insert)?; + let aux_outcome = stream.alloc_zeros::(batch_insert).map_err(|e| anyhow::anyhow!("aux_outcome alloc: {e}"))?; + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, &aux_outcome, batch_insert)?; } assert_eq!(buf.len(), fill_count); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index ccf0825b7..6558cc453 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -151,7 +151,8 @@ fn test_per_weights_valid() -> anyhow::Result<()> { let dones = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; let aux_sign = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?; let aux_conf = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?; - buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, 1)?; + let aux_outcome = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("aux_outcome alloc: {e}"))?; + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, &aux_outcome, 1)?; } buf.sample_proportional(16)?; @@ -197,7 +198,8 @@ fn test_per_indices_valid() -> anyhow::Result<()> { let dones = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("{e}"))?; let aux_sign = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?; let aux_conf = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("aux_conf alloc: {e}"))?; - buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, 1)?; + let aux_outcome = stream.alloc_zeros::(1).map_err(|e| anyhow::anyhow!("aux_outcome alloc: {e}"))?; + buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, &aux_conf, &aux_outcome, 1)?; } buf.sample_proportional(16)?; diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 44a9212ee..9ed4e51ca 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -835,6 +835,20 @@ impl DQNTrainer { agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr( fused.trainer_aux_conf_at_state_buf_ptr(), ); + // SP22 H6 vNext Phase B4b-2 (2026-05-14): wire the + // trainer's `aux_to_label_buf` raw_ptr so PER's + // `gather_i32_scalar` writes per-batch trade-outcome + // labels directly into the K=3 sparse CE consumer's + // i32 buffer. Same direct-to-trainer pattern as + // aux_sign_labels (i32) and aux_conf (f32) above. + // Without this wire-up aux_to_label_buf stays at + // alloc_zeros (all Profit) — the degraded "predict + // Profit everywhere" state matching K=2's B1.1a→ + // B1.1b transition. With it the head trains on the + // real sparse {-1, 0, 1, 2} label distribution. + agent_w.memory_mut().gpu.set_trainer_aux_outcome_labels_ptr( + fused.trainer_aux_to_label_buf_ptr(), + ); fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr()); } } @@ -2762,6 +2776,9 @@ impl DQNTrainer { &gpu_batch.dones, &gpu_batch.aux_sign_labels, &gpu_batch.aux_conf, + // SP22 H6 vNext Phase B4b-2 (2026-05-14): per-(env, t) + // trade-outcome labels for the replay-buffer scatter chain. + &gpu_batch.aux_outcome_labels, total, ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; } @@ -2872,6 +2889,20 @@ impl DQNTrainer { agent_w.memory_mut().gpu.set_trainer_aux_conf_ptr( fused.trainer_aux_conf_at_state_buf_ptr(), ); + // SP22 H6 vNext Phase B4b-2 (2026-05-14): wire the + // trainer's `aux_to_label_buf` raw_ptr so PER's + // `gather_i32_scalar` writes per-batch trade-outcome + // labels directly into the K=3 sparse CE consumer's + // i32 buffer. Same direct-to-trainer pattern as + // aux_sign_labels (i32) and aux_conf (f32) above. + // Without this wire-up aux_to_label_buf stays at + // alloc_zeros (all Profit) — the degraded "predict + // Profit everywhere" state matching K=2's B1.1a→ + // B1.1b transition. With it the head trains on the + // real sparse {-1, 0, 1, 2} label distribution. + agent_w.memory_mut().gpu.set_trainer_aux_outcome_labels_ptr( + fused.trainer_aux_to_label_buf_ptr(), + ); fused.set_per_rng_step_dev_ptr(agent_w.memory().gpu.rng_step_dev_ptr()); } } diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index ac4f1369d..86e608e1a 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -17932,3 +17932,51 @@ First half of Phase B4b (the replay-buffer label scatter chain). Amends the Phas **B4b-1 (this commit)** = producer chain complete. The per-(env, t) column is populated correctly every rollout step. Consumer wiring (replay-buffer scatter + trainer setter) is B4b-2's scope. **Verification**: `cargo check -p ml` clean. `cargo test -p ml --lib` → 1016/0 on clean runs; the pre-existing `test_dqn_checkpoint_round_trip` NoisyLinear flake still surfaces ~50-70% of full-suite runs (varies `pred1=-1 pred2=1` ↔ `pred1=1 pred2=-1` non-deterministically, even with the `bind_to_thread` fix from B4). Flake predates Phase B4b and is unrelated to the trade-outcome head — `disable_noise()` zeros NoisyLinear ε but apparently leaves some other randomness source intact. Worth dedicated investigation separate from vNext work. + +#### Phase B4b-2 — Replay-buffer scatter + trainer setter (2026-05-14) + +Second half of Phase B4b. Wires the trade-outcome label column through the replay buffer (struct field, allocation, scatter on insert, gather on sample, direct-to-trainer pointer + setter) and connects the trainer's `aux_to_label_buf` to receive per-batch sampled labels. Completes the end-to-end producer → ring → trainer i32 path the K=3 sparse CE consumer reads. + +**Replay buffer changes** (`crates/ml-dqn/src/gpu_replay_buffer.rs`): +- New struct fields: `aux_outcome_labels: CudaSlice` (capacity-sized ring), `sample_aux_outcome_labels: CudaSlice` (mbs-sized fallback gather destination), `trainer_aux_outcome_labels_ptr: u64` (direct-path destination, set via setter). +- New `aux_outcome_labels_ptr: u64` field on `GpuBatchPtrs` exposing the per-batch destination to consumers. +- `insert_batch` signature extended: new `aux_outcome: &CudaSlice` arg between `aux_conf_in` and `bs`. Scatters via existing K-generic `scatter_insert_i32` kernel (same kernel the `aux_sign_labels` column uses). +- `sample_proportional` direct gather: when `trainer_aux_outcome_labels_ptr != 0`, `gather_i32_scalar` writes per-batch labels straight into the trainer's `aux_to_label_buf`. Fallback gather populates `sample_aux_outcome_labels` when direct path unwired. +- New setter `set_trainer_aux_outcome_labels_ptr(ptr: u64)` mirrors `set_trainer_aux_conf_ptr` exactly. + +**Collector emission** (`crates/ml/src/cuda_pipeline/gpu_experience_collector.rs`): +- New `aux_outcome_labels: CudaSlice` field on `GpuExperienceBatch`. +- Populated at end of `collect_experiences_gpu` via `dtod_clone_i32(&self.exp_aux_to_label_per_sample, total)` — the per-(env, t) tile populated by the Phase B4b-1 producer kernel writes. + +**Trainer wireup** (`crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + `crates/ml/src/trainers/dqn/fused_training.rs` + `crates/ml/src/trainers/dqn/trainer/training_loop.rs`): +- Added `aux_to_label_buf_ptr()` accessor on `GpuDqnTrainer` (mirrors `aux_nb_label_buf_ptr`). +- Added `trainer_aux_to_label_buf_ptr()` delegating accessor on `FusedTrainingCtx` (mirrors `trainer_aux_sign_labels_buf_ptr`). +- At trainer init in `training_loop.rs` (the same site where `set_trainer_buffers` + `set_trainer_aux_conf_ptr` are called): new `agent_w.memory_mut().gpu.set_trainer_aux_outcome_labels_ptr(fused.trainer_aux_to_label_buf_ptr())` call. The two call sites (lines ~835 and ~2857) are both updated (`replace_all` matched two patterns). +- Updated `insert_batch` call in `training_loop.rs` to pass `&gpu_batch.aux_outcome_labels` as the new arg. + +**Test fixture updates**: 4 smoke test files (`gpu_residency.rs`, `training_stability.rs` ×2, `performance.rs`) and 3 unit-test fixtures inside `gpu_replay_buffer.rs` updated to allocate `aux_outcome` arg (zero-init i32 of the relevant size). + +**End-to-end chain now complete**: +``` +Phase A2 trade_outcome_label_kernel (per-env labels + per-(env, t) labels) + → Phase B3 collector per-step launch (writes exp_aux_to_label_per_sample [N*L] each step) + → Phase B4b-1 collector emission (dtod_clone into GpuExperienceBatch.aux_outcome_labels) + → Phase B4b-2 replay-buffer insert (scatter_insert_i32 into ring) + → Phase B4b-2 PER sample (gather_i32_scalar into trainer aux_to_label_buf — direct path) + → Phase B4 trainer aux_heads_forward (aux_to_fwd.loss_reduce reads sparse {-1,0,1,2} labels) + → Phase B4 trainer aux_heads_backward (aux_to_bwd.backward computes per-sample partials from real labels) + → Phase B1+B4 Adam SAXPY (uniform iteration over [163..167) updates W1, b1, W2, b2) +``` + +The "degraded predict-Profit-everywhere" cold-start state from Phase B4 is now resolved. The K=3 trade-outcome head trains on the real sparse trade-close-outcome label distribution. + +**Verification**: +- `cargo check -p ml` clean (21 warnings, none new). +- `cargo test -p ml --lib` → 1016/0 on clean runs; pre-existing NoisyLinear flake still surfaces ~30-50% of runs (unrelated to vNext work). + +**Remaining vNext work**: +- **B5** (the spec's "Phase B"): Input concat 256 → 262 with `plan_params`. Kernel + W1 shape change. Smaller commit than B4b's full mirror. +- **Phase C**: 3-slot state assembly (state[121..124] = p_Profit, p_Stop, p_Timeout). Replaces Phase 2's single-slot K=2 `prev_aux_dir_prob` with 3 slots from `exp_aux_to_softmax_buf` (currently produced but unused — Phase B3 wired the producer). +- **Phase D**: 12-weight W atom-shift (4 actions × 3 outcomes). Largest remaining commit — extends the existing 4-weight Phase 3 atom-shift mechanism. +- **Phase E**: dW backward + Adam for W[4, 3]. +- **Phase F**: validation smoke at structural prior (β=0.5).