feat(sp22-vnext): Phase B4b-2 — replay-buffer scatter + trainer setter
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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<i32>,
|
||||
/// 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<i32>,
|
||||
/// 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<i32>,
|
||||
/// 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<i32>,
|
||||
/// 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<f32>,
|
||||
aux_sign: &CudaSlice<i32>, // SP13 Layer B: 30-bar aux sign label, -1/0/1
|
||||
aux_conf_in: &CudaSlice<f32>, // 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<i32>,
|
||||
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::<f32>(n).unwrap();
|
||||
let aux = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
let aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let aux_outcome = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
// Non-trivial rewards drive non-zero priorities through `per_insert_pa`.
|
||||
let rewards_host: Vec<f32> = (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::<f32>(n).unwrap();
|
||||
let aux = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
let aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
|
||||
let aux_outcome = stream.alloc_zeros::<i32>(n).unwrap();
|
||||
// Give non-zero rewards so priorities are non-zero after update
|
||||
let rewards_host: Vec<f32> = (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::<f32>(bs).unwrap();
|
||||
let aux = stream.alloc_zeros::<i32>(bs).unwrap();
|
||||
let aux_conf = stream.alloc_zeros::<f32>(bs).unwrap();
|
||||
b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, bs).unwrap();
|
||||
let aux_outcome = stream.alloc_zeros::<i32>(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::<f32>(n).unwrap();
|
||||
stream.memcpy_htod(&aux_conf_host, &mut aux_conf).unwrap();
|
||||
let aux_outcome = stream.alloc_zeros::<i32>(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::<f32>(n).unwrap();
|
||||
stream.memcpy_htod(&aux_conf_host, &mut aux_conf).unwrap();
|
||||
let aux_outcome = stream.alloc_zeros::<i32>(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");
|
||||
|
||||
|
||||
@@ -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<i32>` `[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 }
|
||||
|
||||
|
||||
@@ -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<i32>,
|
||||
/// 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<i32>,
|
||||
/// 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,
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -74,7 +74,9 @@ fn insert_random_batch(
|
||||
let dones = zeros_cuda_f32(n, stream)?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(n).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(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::<i32>(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(())
|
||||
}
|
||||
|
||||
|
||||
@@ -141,7 +141,8 @@ fn test_per_sample_latency() -> anyhow::Result<()> {
|
||||
let dones = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(batch_insert).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(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::<i32>(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);
|
||||
|
||||
|
||||
@@ -151,7 +151,8 @@ fn test_per_weights_valid() -> anyhow::Result<()> {
|
||||
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(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::<i32>(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::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
|
||||
let aux_sign = stream.alloc_zeros::<i32>(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
|
||||
let aux_conf = stream.alloc_zeros::<f32>(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::<i32>(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)?;
|
||||
|
||||
@@ -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());
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<i32>` (capacity-sized ring), `sample_aux_outcome_labels: CudaSlice<i32>` (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<i32>` 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<i32>` 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).
|
||||
|
||||
Reference in New Issue
Block a user