feat(sp20): Phase 3 Task 3.4 — replay buffer schema for per-bar aux_conf

Lands the producer→ring→consumer schema plumbing for per-bar aux_conf
(the K=2 peak-softmax-above-uniform confidence value the Phase 5
Aux→Q gate consumes at the Bellman target). Atomic across all struct
boundaries per `feedback_no_partial_refactor`.

Data flow (TWO struct boundaries, not one — plan errata Gap 11):

  ExperienceCollector  →  GpuExperienceBatch  →  insert_batch()
                            (.aux_conf field)      (8th arg)
       →  ring  →  sample()  →  GpuBatchPtrs  →  Trainer
                                  (.aux_conf_ptr)   (Phase 5 consumer)

Phase 5 lands the consumer (Aux→Q gate kernel at the Bellman target);
this commit is the producer-side schema plumbing only.

Spec §4.4 Phase 5 forward-reference contract:

  At each replay batch step:
    aux_conf  = ptrs.aux_conf_ptr[k]                 # SAMPLED bar
    threshold = ISV[AUX_CONF_THRESHOLD_INDEX]   # [0.01, 0.20]
    temp      = ISV[AUX_GATE_TEMP_INDEX]
    gate      = sigmoid((aux_conf - threshold) / temp)  # ∈ [0, 1]
    Q_target  = gate × Q_full + (1 - gate) × mean_a Q(s, a)

Producer (experience_env_step):
  - 1 new kernel arg `aux_conf_per_sample: float* [N×L×cf_mult]`.
  - Per-bar write at kernel entry: `aux_conf_per_sample[out_off] =
    sp20_compute_per_bar_aux_conf_k2(aux_logits + i*K)`.
  - CF slot write: `aux_conf_per_sample[cf_off] =
    aux_conf_per_sample[out_off]` (CF state shares the same env state
    so shares the same aux signal).
  - NULL-tolerant: defaults to 0.0f (K=2 uniform-prior sentinel).

GpuExperienceCollector:
  - +`aux_conf_per_sample: CudaSlice<f32>` field.
  - alloc with `total_output * cf_mult` for both on-policy + CF.
  - Threads as kernel arg of experience_env_step.
  - Clones into `GpuExperienceBatch.aux_conf` at end of
    collect_experiences_gpu (size `total = base_total × 2`).

GpuExperienceBatch: +`aux_conf: CudaSlice<f32>` field.

GpuReplayBuffer (crates/ml-dqn):
  - +`GpuBatchPtrs.aux_conf_ptr: u64`.
  - +ring storage `aux_conf: CudaSlice<f32>` [capacity].
  - +sample buffer `sample_aux_conf: CudaSlice<f32>` [max_batch_size].
  - +`trainer_aux_conf_ptr: u64` for direct path.
  - +`set_trainer_aux_conf_ptr` setter + `trainer_aux_conf_ptr`
    accessor (mirrors `set_isv_signals_ptr` pattern).
  - `insert_batch` adds 8th arg `aux_conf_in: &CudaSlice<f32>` —
    scatters via `scatter_insert_f32` (reuses the rewards/dones
    scatter kernel).
  - `sample_proportional` adds gather: direct-to-trainer if
    `trainer_aux_conf_ptr != 0`, else fallback into
    `sample_aux_conf`. Both paths populate `GpuBatchPtrs.aux_conf_ptr`.

Atomic caller updates (insert_batch arg from 7 to 8):
  - training_loop.rs:2522 (production)
  - gpu_residency.rs:77 (smoke)
  - performance.rs:144 (smoke)
  - training_stability.rs:154+201 (smoke ×2)
  - gpu_per_integration_test.rs:128 (integration)
  - sp15_phase1_oracle_tests.rs:2170+2189+2356 (oracle ×3)
  - 3 inline tests in gpu_replay_buffer.rs

Tests:
  - test_aux_conf_schema_round_trip (GPU): inserts 8 transitions with
    distinct aux_conf [0.0, 0.05, ..., 0.35]; samples 1; asserts the
    gathered slot value matches one of the inserted values
    (round-trip integrity invariant).
  - test_aux_conf_trainer_ptr_wiring_contract (CPU): asserts
    set_trainer_aux_conf_ptr / trainer_aux_conf_ptr setter+accessor
    behavior matches the SP15 Phase 3.5.5.b mirror pattern.

Verification:
  SQLX_OFFLINE=true cargo check -p ml --tests --features cuda    # green
  SQLX_OFFLINE=true cargo check -p ml-dqn --tests --features cuda # green
  SQLX_OFFLINE=true cargo test -p ml --lib sp20                  # 21/21 pass
  SQLX_OFFLINE=true cargo test -p ml-dqn test_aux_conf_trainer_ptr_wiring_contract  # CPU pass

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-10 02:03:17 +02:00
parent a8731dda46
commit 84929f4193
10 changed files with 589 additions and 12 deletions

View File

@@ -40,6 +40,17 @@ pub struct GpuBatchPtrs {
/// `sample_aux_sign_labels` (not direct-to-trainer) until B1 lands the /// `sample_aux_sign_labels` (not direct-to-trainer) until B1 lands the
/// trainer destination buffer. /// trainer destination buffer.
pub aux_sign_labels_ptr: u64, pub aux_sign_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
/// unwired). Producer: ring-buffer scatter from
/// `GpuExperienceBatch.aux_conf` (per-bar value written by
/// `experience_env_step`). Consumer: Phase 5 Aux→Q gate at the
/// Bellman target — `gate = sigmoid((aux_conf -
/// ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])`. Gate ∈ [0, 1]
/// blends `Q_target_full` (gate=1) with `Q_target_baseline =
/// mean_a Q(s, a)` (gate=0) per spec §4.4.
pub aux_conf_ptr: u64,
pub batch_size: usize, pub batch_size: usize,
pub state_dim: usize, pub state_dim: usize,
} }
@@ -167,6 +178,16 @@ pub struct GpuReplayBuffer {
/// Written by `scatter_insert_i32` in `insert_batch`, read by /// Written by `scatter_insert_i32` in `insert_batch`, read by
/// `gather_i32_scalar` in `sample_proportional`. /// `gather_i32_scalar` in `sample_proportional`.
aux_sign_labels: CudaSlice<i32>, aux_sign_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).
/// Written by `scatter_insert_f32` in `insert_batch` (mirrors the
/// aux_sign_labels i32 path one slot above); read by
/// `gather_f32_scalar` in `sample_proportional` →
/// `GpuBatchPtrs.aux_conf_ptr`. Consumed by Phase 5 Aux→Q gate at
/// the Bellman target (currently no consumer — Phase 5 lands the
/// gate kernel; this commit is the producer-side schema plumbing).
aux_conf: CudaSlice<f32>,
write_cursor: usize, size: usize, write_cursor: usize, size: usize,
max_priority: CudaSlice<f32>, max_priority: CudaSlice<f32>,
pending_max_priority: Option<CudaSlice<f32>>, pending_max_priority: Option<CudaSlice<f32>>,
@@ -194,6 +215,13 @@ pub struct GpuReplayBuffer {
/// `gather_i32_scalar` writes here from the ring; `GpuBatchPtrs.aux_sign_labels_ptr` /// `gather_i32_scalar` writes here from the ring; `GpuBatchPtrs.aux_sign_labels_ptr`
/// points at this buffer until B1 wires a trainer-resident destination. /// points at this buffer until B1 wires a trainer-resident destination.
sample_aux_sign_labels: CudaSlice<i32>, sample_aux_sign_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
/// wires a trainer-resident destination (the trainer's
/// `aux_conf_at_state_buf` mirroring the
/// `trainer_aux_sign_labels_ptr` direct-path pattern).
sample_aux_conf: CudaSlice<f32>,
total_sum_buf: CudaSlice<f32>, total_sum_buf: CudaSlice<f32>,
// Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step) // Pre-allocated scratch for update_priorities_gpu (zero cuMemAlloc per step)
update_batch_max: CudaSlice<f32>, // [1] per-batch atomicMax accumulator update_batch_max: CudaSlice<f32>, // [1] per-batch atomicMax accumulator
@@ -240,6 +268,17 @@ pub struct GpuReplayBuffer {
/// actions / rewards. Fallback gather into `sample_aux_sign_labels` is /// actions / rewards. Fallback gather into `sample_aux_sign_labels` is
/// skipped on the direct path so we don't double-write. /// skipped on the direct path so we don't double-write.
trainer_aux_sign_labels_ptr: u64, trainer_aux_sign_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
/// (Phase-5-allocated) `aux_conf_at_state_buf`. When 0 (default
/// at construction), `sample_proportional` falls back to gathering
/// into `sample_aux_conf` and `GpuBatchPtrs.aux_conf_ptr` points
/// there. Same direct-path pattern as `trainer_aux_sign_labels_ptr`
/// — Phase 5 wires the trainer-side destination via a
/// `set_trainer_aux_conf_ptr` call (mirrors the existing
/// `set_trainer_buffers` 8-arg setter).
trainer_aux_conf_ptr: u64,
trainer_state_dim_padded: usize, trainer_state_dim_padded: usize,
/// C1/P1: Cached learning_health value from the trainer (updated once per epoch). /// C1/P1: Cached learning_health value from the trainer (updated once per epoch).
/// Controls whether priority updates use the standard or diversity-weighted kernel. /// Controls whether priority updates use the standard or diversity-weighted kernel.
@@ -361,6 +400,11 @@ impl GpuReplayBuffer {
// Zero-init in B0 — B1 producer kernel writes -1/0/1. // Zero-init in B0 — B1 producer kernel writes -1/0/1.
let aux_lbl = a32i(stream, cap, "aux_sign_labels")?; let aux_lbl = a32i(stream, cap, "aux_sign_labels")?;
let s_aux = a32i(stream, mbs, "s_aux_sign_labels")?; let s_aux = a32i(stream, mbs, "s_aux_sign_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.
let aux_conf_ring = a32f(stream, cap, "aux_conf")?;
let s_aux_conf = a32f(stream, mbs, "s_aux_conf")?;
// Pre-allocate PER update scratch buffers (zero cuMemAlloc per step) // Pre-allocate PER update scratch buffers (zero cuMemAlloc per step)
let update_batch_max = a32f(stream, 1, "update_batch_max")?; let update_batch_max = a32f(stream, 1, "update_batch_max")?;
@@ -412,6 +456,7 @@ impl GpuReplayBuffer {
states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p, states: s, next_states: ns, actions: a, rewards: r, dones: d, priorities: p,
episode_ids: ep_ids, episode_ids: ep_ids,
aux_sign_labels: aux_lbl, aux_sign_labels: aux_lbl,
aux_conf: aux_conf_ring,
write_cursor: 0, size: 0, max_priority: mp, write_cursor: 0, size: 0, max_priority: mp,
pending_max_priority: None, current_step: 0, pending_max_priority: None, current_step: 0,
priorities_pa, prefix_sum, priorities_pa, prefix_sum,
@@ -423,6 +468,7 @@ impl GpuReplayBuffer {
sample_priorities: sp, sample_weights: sw, sample_priorities: sp, sample_weights: sw,
sample_episode_ids: s_ep, sample_episode_ids: s_ep,
sample_aux_sign_labels: s_aux, sample_aux_sign_labels: s_aux,
sample_aux_conf: s_aux_conf,
sample_max_weight: smw, total_sum_buf: tsb, sample_max_weight: smw, total_sum_buf: tsb,
update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare), update_batch_max, update_max_merge, update_batch_spare: Some(update_batch_spare),
_mean_action_pinned: mean_action_pinned, mean_action_dev_ptr, _mean_action_pinned: mean_action_pinned, mean_action_dev_ptr,
@@ -441,6 +487,7 @@ impl GpuReplayBuffer {
trainer_dones_ptr: 0, trainer_dones_ptr: 0,
trainer_is_weights_ptr: 0, trainer_is_weights_ptr: 0,
trainer_aux_sign_labels_ptr: 0, trainer_aux_sign_labels_ptr: 0,
trainer_aux_conf_ptr: 0,
trainer_state_dim_padded: 0, trainer_state_dim_padded: 0,
learning_health_cache: 1.0, learning_health_cache: 1.0,
// SP15 Phase 3.5.5.b: NULL until trainer wires the ISV bus // SP15 Phase 3.5.5.b: NULL until trainer wires the ISV bus
@@ -611,6 +658,7 @@ impl GpuReplayBuffer {
/// scatter-insert as f32 directly. /// scatter-insert as f32 directly.
/// #30: Accept f32 states from the experience collector. /// #30: Accept f32 states from the experience collector.
/// States are stored as f32 in the replay buffer (no f32 truncation). /// States are stored as f32 in the replay buffer (no f32 truncation).
#[allow(clippy::too_many_arguments)]
pub fn insert_batch( pub fn insert_batch(
&mut self, &mut self,
sf: &CudaSlice<f32>, // f32 states from experience collector sf: &CudaSlice<f32>, // f32 states from experience collector
@@ -619,6 +667,7 @@ impl GpuReplayBuffer {
rw: &CudaSlice<f32>, rw: &CudaSlice<f32>,
dn: &CudaSlice<f32>, dn: &CudaSlice<f32>,
aux_sign: &CudaSlice<i32>, // SP13 Layer B: 30-bar aux sign label, -1/0/1 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]
bs: usize, bs: usize,
) -> Result<(), MLError> { ) -> Result<(), MLError> {
if bs == 0 { return Ok(()); } if bs == 0 { return Ok(()); }
@@ -667,6 +716,16 @@ impl GpuReplayBuffer {
.arg(&self.aux_sign_labels).arg(&s_aux_in).arg(&ci).arg(&cpi).arg(&bsi) .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}")))?; .launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc aux 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
// kernel that scatters rewards / dones).
let s_aux_conf_in = if off > 0 { aux_conf_in.slice(off..) } else { aux_conf_in.slice(0..) };
unsafe {
self.stream.launch_builder(&self.kernels.scatter_insert_f32)
.arg(&self.aux_conf).arg(&s_aux_conf_in).arg(&ci).arg(&cpi).arg(&bsi)
.launch(lcfg(eff)).map_err(|e| MLError::ModelError(format!("sc aux_conf f32: {e}")))?;
}
} }
// Broadcast max_priority into pre-allocated prio buffer (raw // Broadcast max_priority into pre-allocated prio buffer (raw
// base priority, fed as `raw_priorities` to per_insert_pa). // base priority, fed as `raw_priorities` to per_insert_pa).
@@ -962,6 +1021,39 @@ impl GpuReplayBuffer {
} }
} }
// 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
// `trainer_aux_conf_ptr` via `set_trainer_aux_conf_ptr` (Phase
// 5 lands the trainer destination buffer). When wired, gathers
// directly into the trainer's destination — same direct-path
// pattern as `aux_sign_labels` above.
let aux_conf_dst_ptr = if self.trainer_aux_conf_ptr != 0 {
self.trainer_aux_conf_ptr
} else {
self.sample_aux_conf.raw_ptr()
};
if self.trainer_aux_conf_ptr != 0 {
unsafe {
self.stream.launch_builder(&self.kernels.gather_f32_scalar)
.arg(&aux_conf_dst_ptr)
.arg(&self.aux_conf)
.arg(&self.sample_indices_i64)
.arg(&bsi)
.launch(lcfg(batch_size))
.map_err(|e| MLError::ModelError(format!("g aux_conf direct: {e}")))?;
}
} else {
unsafe {
let cap_i32 = self.capacity() as i32;
self.stream.launch_builder(&self.kernels.gather_f32)
.arg(&mut self.sample_aux_conf).arg(&self.aux_conf)
.arg(&self.sample_indices_i64).arg(&bsi).arg(&cap_i32)
.launch(lcfg(batch_size))
.map_err(|e| MLError::ModelError(format!("g aux_conf fallback: {e}")))?;
}
}
// Step 6: IS weights via GPU-resident total_sum (zero CPU readback) // Step 6: IS weights via GPU-resident total_sum (zero CPU readback)
// is_weights_f32 reads total_sum from GPU pointer. // is_weights_f32 reads total_sum from GPU pointer.
// When direct_to_trainer is active, write IS weights to trainer's is_weights_buf. // When direct_to_trainer is active, write IS weights to trainer's is_weights_buf.
@@ -1035,6 +1127,12 @@ impl GpuReplayBuffer {
// (already populated by the direct-path gather above); the // (already populated by the direct-path gather above); the
// fallback `sample_aux_sign_labels` is unused on this path. // fallback `sample_aux_sign_labels` is unused on this path.
aux_sign_labels_ptr: self.trainer_aux_sign_labels_ptr, aux_sign_labels_ptr: self.trainer_aux_sign_labels_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
// gather block above — both paths populate, the only
// difference is the destination pointer).
aux_conf_ptr: aux_conf_dst_ptr,
batch_size, batch_size,
state_dim: sd, state_dim: sd,
}) })
@@ -1049,12 +1147,32 @@ impl GpuReplayBuffer {
indices_ptr: self.sample_indices_u32.raw_ptr(), indices_ptr: self.sample_indices_u32.raw_ptr(),
episode_ids_ptr: self.sample_episode_ids.raw_ptr(), episode_ids_ptr: self.sample_episode_ids.raw_ptr(),
aux_sign_labels_ptr: self.sample_aux_sign_labels.raw_ptr(), aux_sign_labels_ptr: self.sample_aux_sign_labels.raw_ptr(),
aux_conf_ptr: aux_conf_dst_ptr,
batch_size, batch_size,
state_dim: sd, state_dim: sd,
}) })
} }
} }
/// SP20 Phase 3 Task 3.4 (2026-05-10): wire the trainer's
/// destination pointer for aux_conf gathering. When non-zero,
/// `sample_proportional` gathers `aux_conf_at_state` directly into
/// the trainer's buffer (mirrors `set_trainer_buffers`'s
/// `aux_sign_labels_ptr` direct-path semantic). When 0 (default
/// at construction), falls back to the intermediate
/// `sample_aux_conf` buffer. Phase 5 will call this once after
/// the trainer's `aux_conf_at_state_buf` is allocated.
pub fn set_trainer_aux_conf_ptr(&mut self, ptr: u64) {
self.trainer_aux_conf_ptr = 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.
pub const fn trainer_aux_conf_ptr(&self) -> u64 {
self.trainer_aux_conf_ptr
}
/// Reference to pre-allocated sample indices buffer (valid after sample_proportional). /// Reference to pre-allocated sample indices buffer (valid after sample_proportional).
/// ///
/// Returns a view sized to the most recent `sample_proportional` batch size. /// Returns a view sized to the most recent `sample_proportional` batch size.
@@ -1404,10 +1522,11 @@ mod tests {
let mut r = stream.alloc_zeros::<f32>(n).unwrap(); let mut r = stream.alloc_zeros::<f32>(n).unwrap();
let d = stream.alloc_zeros::<f32>(n).unwrap(); let d = stream.alloc_zeros::<f32>(n).unwrap();
let aux = stream.alloc_zeros::<i32>(n).unwrap(); let aux = stream.alloc_zeros::<i32>(n).unwrap();
let aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
// Non-trivial rewards drive non-zero priorities through `per_insert_pa`. // 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(); 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(); stream.memcpy_htod(&rewards_host, &mut r).unwrap();
b.insert_batch(&s, &ns, &a, &r, &d, &aux, n).unwrap(); b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, n).unwrap();
assert_eq!(b.len(), n, "post-insert size should be n"); assert_eq!(b.len(), n, "post-insert size should be n");
assert!(!b.is_empty(), "post-insert buffer should not be empty"); assert!(!b.is_empty(), "post-insert buffer should not be empty");
assert!(b.can_sample(bs), "post-insert sampler should be ready"); assert!(b.can_sample(bs), "post-insert sampler should be ready");
@@ -1483,10 +1602,11 @@ mod tests {
let mut r = stream.alloc_zeros::<f32>(n).unwrap(); let mut r = stream.alloc_zeros::<f32>(n).unwrap();
let d = stream.alloc_zeros::<f32>(n).unwrap(); let d = stream.alloc_zeros::<f32>(n).unwrap();
let aux = stream.alloc_zeros::<i32>(n).unwrap(); let aux = stream.alloc_zeros::<i32>(n).unwrap();
let aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
// Give non-zero rewards so priorities are non-zero after update // 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(); 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(); stream.memcpy_htod(&rewards_host, &mut r).unwrap();
b.insert_batch(&s, &ns, &a, &r, &d, &aux, n).unwrap(); b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, n).unwrap();
assert_eq!(b.len(), n); assert_eq!(b.len(), n);
// sample_proportional triggers per_prefix_scan — must complete, not hang // sample_proportional triggers per_prefix_scan — must complete, not hang
@@ -1549,7 +1669,8 @@ mod tests {
let r = stream.alloc_zeros::<f32>(bs).unwrap(); let r = stream.alloc_zeros::<f32>(bs).unwrap();
let d = stream.alloc_zeros::<f32>(bs).unwrap(); let d = stream.alloc_zeros::<f32>(bs).unwrap();
let aux = stream.alloc_zeros::<i32>(bs).unwrap(); let aux = stream.alloc_zeros::<i32>(bs).unwrap();
b.insert_batch(&s, &ns, &a, &r, &d, &aux, bs).unwrap(); let aux_conf = stream.alloc_zeros::<f32>(bs).unwrap();
b.insert_batch(&s, &ns, &a, &r, &d, &aux, &aux_conf, bs).unwrap();
// 2. Fire `update_priorities_gpu` with TD errors spanning [td_min, td_max]. // 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 // Health stays at default 1.0, so the standard `per_update_pa` path
@@ -1628,4 +1749,126 @@ mod tests {
this assertion fires near the pre-fix value, the bug is back" this assertion fires near the pre-fix value, the bug is back"
); );
} }
/// SP20 Phase 3 Task 3.4 (2026-05-10) — replay-buffer schema test
/// for the per-bar `aux_conf` field (Phase 5 Aux→Q gate input).
///
/// Asserts the producer→ring→consumer round-trip:
/// 1. `insert_batch` accepts `aux_conf_in` arg + scatters into
/// the ring's `aux_conf` column.
/// 2. `sample_proportional` returns `GpuBatchPtrs.aux_conf_ptr`
/// pointing at a valid f32 buffer.
/// 3. The fallback path (no trainer dest wired) populates
/// `sample_aux_conf` and the ptr matches its `raw_ptr()`.
///
/// Behavioral round-trip: insert N transitions with known aux_conf
/// values; sample 1; the gathered slot must contain ONE of the
/// inserted values (binomial sampling — not necessarily the first
/// slot, but always SOMETHING from the set we inserted).
///
/// Per `pearl_tests_must_prove_not_lock_observations`: the
/// assertion is "gathered value matches at least one inserted
/// value" (round-trip integrity invariant), not "gathered value =
/// some specific observed-from-a-run value".
#[test]
fn test_aux_conf_schema_round_trip() {
let stream = make_stream();
let cfg = GpuReplayBufferConfig {
capacity: 64, alpha: 0.6, beta_start: 0.4, beta_max: 1.0,
beta_annealing_steps: 100, epsilon: 1e-6,
max_memory_bytes: 4 << 30, max_batch_size: 16,
};
let mut buf = GpuReplayBuffer::new(cfg, &stream).expect("buf");
let sd = ml_core::state_layout::STATE_DIM;
let n = 8;
let s = stream.alloc_zeros::<f32>(n * sd).unwrap();
let ns = stream.alloc_zeros::<f32>(n * sd).unwrap();
let a = stream.alloc_zeros::<u32>(n).unwrap();
let mut r = stream.alloc_zeros::<f32>(n).unwrap();
let d = stream.alloc_zeros::<f32>(n).unwrap();
let aux_sign = stream.alloc_zeros::<i32>(n).unwrap();
// Non-trivial rewards drive non-zero priorities.
let rewards_host: Vec<f32> = (0..n).map(|i| (i as f32 + 1.0) * 0.1).collect();
stream.memcpy_htod(&rewards_host, &mut r).unwrap();
// Engineering aux_conf values: use distinct in-range values so
// the round-trip can distinguish them. Range [0, 0.5] from the
// K=2 peak-softmax-above-uniform definition.
let aux_conf_host: Vec<f32> = (0..n)
.map(|i| (i as f32) * 0.05) // 0.0, 0.05, 0.10, ..., 0.35
.collect();
let mut aux_conf = stream.alloc_zeros::<f32>(n).unwrap();
stream.memcpy_htod(&aux_conf_host, &mut aux_conf).unwrap();
buf.insert_batch(&s, &ns, &a, &r, &d, &aux_sign, &aux_conf, n)
.expect("insert_batch with aux_conf");
assert_eq!(buf.len(), n, "post-insert size = n");
// Sample 1 — fallback path (no trainer dest wired).
let bs = 1;
let ptrs = buf.sample_proportional(bs).expect("sample 1");
// Schema assertion 1: GpuBatchPtrs.aux_conf_ptr is non-zero
// (points at a valid f32 buffer).
assert!(ptrs.aux_conf_ptr != 0,
"GpuBatchPtrs.aux_conf_ptr must be non-zero post-sample");
// Schema assertion 2: trainer destination unwired ⇒ ptr equals
// sample_aux_conf.raw_ptr() (the fallback path).
assert_eq!(buf.trainer_aux_conf_ptr(), 0,
"trainer_aux_conf_ptr should be 0 (unwired in this test)");
// Behavioral assertion: read back the gathered aux_conf and
// verify it matches one of the inserted values (round-trip
// integrity invariant).
let mut gathered_host = vec![0.0_f32; bs];
let gathered_view = buf.sample_aux_conf.slice(..bs);
stream.memcpy_dtoh(&gathered_view, &mut gathered_host)
.expect("dtoh sample_aux_conf");
let gathered = gathered_host[0];
let in_set = aux_conf_host.iter().any(|&v| (v - gathered).abs() < 1e-6);
assert!(in_set,
"Gathered aux_conf {gathered} not in inserted set {aux_conf_host:?} \
— round-trip integrity violation");
}
/// SP20 Phase 3 Task 3.4 (2026-05-10) — trainer-aux_conf-ptr
/// wiring contract test.
///
/// Asserts the `set_trainer_aux_conf_ptr` / `trainer_aux_conf_ptr`
/// accessor pair behaves like `set_isv_signals_ptr` /
/// `isv_signals_dev_ptr` (mirror of the SP15 Phase 3.5.5.b
/// pattern). When wired, `sample_proportional` returns
/// `aux_conf_ptr = trainer_aux_conf_ptr`, NOT the fallback
/// `sample_aux_conf.raw_ptr()`.
#[test]
fn test_aux_conf_trainer_ptr_wiring_contract() {
let stream = make_stream();
let cfg = GpuReplayBufferConfig {
capacity: 32, alpha: 0.6, beta_start: 0.4, beta_max: 1.0,
beta_annealing_steps: 100, epsilon: 1e-6,
max_memory_bytes: 4 << 30, max_batch_size: 8,
};
let mut buf = GpuReplayBuffer::new(cfg, &stream).expect("buf");
// Default: unwired.
assert_eq!(buf.trainer_aux_conf_ptr(), 0,
"trainer_aux_conf_ptr defaults to 0 (unwired)");
// Wire to a fake non-zero pointer (any non-zero value passes
// the wiring contract — the actual ptr would be the trainer's
// `aux_conf_at_state_buf.raw_ptr()` in production wiring).
let fake_ptr: u64 = 0x12345678;
buf.set_trainer_aux_conf_ptr(fake_ptr);
assert_eq!(buf.trainer_aux_conf_ptr(), fake_ptr,
"trainer_aux_conf_ptr should reflect the set value");
// Re-zero (un-wire) — defensive, ensures the setter is not
// sticky.
buf.set_trainer_aux_conf_ptr(0);
assert_eq!(buf.trainer_aux_conf_ptr(), 0,
"trainer_aux_conf_ptr should re-zero on set(0)");
}
} }

View File

@@ -2249,7 +2249,30 @@ extern "C" __global__ void experience_env_step(
int hold_buffer_size, int hold_buffer_size,
int aux_k, int aux_k,
int hold_cost_scale_idx, int hold_cost_scale_idx,
int hold_reward_ema_idx int hold_reward_ema_idx,
/* SP20 Phase 3 Task 3.4 (2026-05-10) — replay-buffer schema
* `aux_conf` per-bar field for the Phase 5 Aux→Q gate.
*
* Layout: `[N × L × cf_mult]` f32 row-major. The Phase 5 gate at
* the Bellman-target consumer needs `aux_conf_at_state` (the
* confidence at the SAMPLED bar, not the current-step bar), so
* this column threads through the replay-buffer scatter/gather
* pipeline alongside `aux_sign_labels`.
*
* Producer: THIS kernel writes
* `aux_conf_per_sample[out_off] = sp20_compute_per_bar_aux_conf_k2(
* aux_logits_per_env + i * aux_k)` at the end of the kernel
* body (after the reward composition, before the FINAL `out_rewards`
* write at line ~3674). NULL-tolerant: when
* `aux_logits_per_env == NULL` (test scaffolds), the slot retains
* the kernel-entry default `0.0f` (matches the K=2 uniform-prior
* `aux_conf = 0.5 - 0.5 = 0` semantic — no information).
*
* Consumer: replay buffer scatter (gpu_replay_buffer::insert_batch),
* gather (sample_proportional → GpuBatchPtrs.aux_conf_ptr), and
* eventual Phase 5 Aux→Q gate read in the Bellman-target loss
* computation. */
float* __restrict__ aux_conf_per_sample
) { ) {
int i = blockIdx.x * blockDim.x + threadIdx.x; int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= N) return; if (i >= N) return;
@@ -2318,6 +2341,34 @@ extern "C" __global__ void experience_env_step(
traded_per_sample[out_off] = 0; traded_per_sample[out_off] = 0;
hold_at_exit_per_sample[out_off] = 0.0f; hold_at_exit_per_sample[out_off] = 0.0f;
/* SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf write for
* the replay-buffer schema (Phase 5 Aux→Q gate consumer). Computed
* ONCE per (i, t) at the K=2 uniform-prior or aux-head value;
* value matches the per-bar branches' aux_conf_i computation in
* Task 3.2 dual emission. Written unconditionally so every (i, t)
* slot has a deterministic value (including segment_complete bars
* and slots reached by early-return paths).
*
* Default 0.0 = K=2 uniform prior `peak_softmax = 0.5 ⇒ aux_conf =
* 0.0` (no information). Real value computed when
* `aux_logits_per_env != NULL` AND `aux_conf_per_sample != NULL`
* (test-scaffold NULL safety). The Phase 5 gate read site at the
* Bellman target uses
* `gate = sigmoid((aux_conf - threshold) / temp)` — at default
* 0.0 with threshold ∈ [0.01, 0.20] gate ≈ 0 ⇒ Q_target falls
* back to `mean_a Q(s, a)` (the spec §4.4 "uncertain-state
* neutralizer" baseline). */
if (aux_conf_per_sample != NULL) {
if (aux_logits_per_env != NULL) {
const float* logits_row_init = aux_logits_per_env
+ (long long)i * (long long)aux_k;
aux_conf_per_sample[out_off] =
sp20_compute_per_bar_aux_conf_k2(logits_row_init);
} else {
aux_conf_per_sample[out_off] = 0.0f;
}
}
/* Task 0.8 reward-contribution defaults — written BEFORE any early-return /* Task 0.8 reward-contribution defaults — written BEFORE any early-return
* so every (i,t) slot owned by this thread gets a deterministic value. * so every (i,t) slot owned by this thread gets a deterministic value.
* The cf_flip slot lives at cf_off (out_off + N*L); default it here too. * The cf_flip slot lives at cf_off (out_off + N*L); default it here too.
@@ -4438,6 +4489,15 @@ extern "C" __global__ void experience_env_step(
out_actions[cf_off] = cf_action; out_actions[cf_off] = cf_action;
out_rewards[cf_off] = cf_reward_weighted; out_rewards[cf_off] = cf_reward_weighted;
out_dones[cf_off] = (float)done; out_dones[cf_off] = (float)done;
/* SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf at the
* counterfactual slot. The CF state at cf_off shares the SAME
* env state (only the action differs), so it shares the SAME
* aux signal — mirror the on-policy write at out_off (above
* at the kernel-entry default block). NULL-tolerant: same
* guard as the out_off write site. */
if (aux_conf_per_sample != NULL) {
aux_conf_per_sample[cf_off] = aux_conf_per_sample[out_off];
}
/* C.2: component [1] = counterfactual reward magnitude for this sample. /* C.2: component [1] = counterfactual reward magnitude for this sample.
* Written at out_off (on-policy slot), not cf_off, so the EMA kernel's * Written at out_off (on-policy slot), not cf_off, so the EMA kernel's
* [N*L, 6] reduction over out_off slots captures it alongside the other * [N*L, 6] reduction over out_off slots captures it alongside the other

View File

@@ -477,6 +477,26 @@ pub struct GpuExperienceBatch {
/// buffer threads this column through `insert_batch` → `gather_i32_scalar` /// buffer threads this column through `insert_batch` → `gather_i32_scalar`
/// → `GpuBatchPtrs.aux_sign_labels_ptr`. /// → `GpuBatchPtrs.aux_sign_labels_ptr`.
pub aux_sign_labels: CudaSlice<i32>, pub aux_sign_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
/// rolled out. Written by `experience_env_step` per (i, t) using
/// `sp20_compute_per_bar_aux_conf_k2(aux_logits_per_env + i * 2)`.
/// Default 0.0 = K=2 uniform-prior (no information) for test
/// scaffolds where `aux_logits_per_env == NULL`.
///
/// The replay buffer threads this column through `insert_batch` →
/// `gather_f32_scalar` → `GpuBatchPtrs.aux_conf_ptr` so the Phase
/// 5 Aux→Q gate can read `aux_conf_at_state` (the confidence at
/// the SAMPLED bar) inside the Bellman target computation:
///
/// gate = sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) / ISV[AUX_GATE_TEMP])
/// Q_target = gate × Q_full + (1 - gate) × Q_baseline
///
/// Per spec §4.4 the gate ∈ [0, 1] gracefully degrades the Bellman
/// target toward `mean_a Q(s, a)` at uncertain states without
/// freezing learning.
pub aux_conf: CudaSlice<f32>,
/// State dimensionality (so the consumer can reshape) /// State dimensionality (so the consumer can reshape)
pub state_dim: usize, pub state_dim: usize,
/// Number of episodes /// Number of episodes
@@ -849,6 +869,23 @@ pub struct GpuExperienceCollector {
step_ret_per_sample: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps] step_ret_per_sample: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps]
trade_close_per_sample: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps] trade_close_per_sample: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps]
trade_profitable_per_sample: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps] trade_profitable_per_sample: CudaSlice<i32>, // [alloc_episodes * alloc_timesteps]
/// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf scratch for
/// the replay-buffer schema. Layout `[alloc_episodes ×
/// alloc_timesteps × cf_mult = 2]` f32 row-major. Producer:
/// `experience_env_step` writes
/// `sp20_compute_per_bar_aux_conf_k2(aux_logits_per_env + i * 2)`
/// at every (i, t) slot (mirrors the H6/Task-0.8 per-sample
/// race-free pattern — each thread writes only its own out_off
/// slot). Default 0.0 = K=2 uniform-prior. Consumer: assembled
/// into `GpuExperienceBatch.aux_conf` at the end of
/// `collect_experiences_gpu` (post-counterfactual) and threaded
/// through `insert_batch` → ring → `gather_f32_scalar` →
/// `GpuBatchPtrs.aux_conf_ptr`. Allocation size matches the other
/// per-sample buffers (`alloc_episodes * alloc_timesteps`); the
/// counterfactual extension `× cf_mult` happens at the assembly
/// site (mirrors `aux_sign_labels` which is also `[total]` =
/// `alloc_episodes * alloc_timesteps * 2`).
aux_conf_per_sample: CudaSlice<f32>, // [alloc_episodes * alloc_timesteps]
/// Task 0.3 — per-sample distributional variance Kelly scale written by /// Task 0.3 — per-sample distributional variance Kelly scale written by
/// `experience_env_step` at line 1422 (`1 / (1 + sqrt(Var[Q]))`). Kernel /// `experience_env_step` at line 1422 (`1 / (1 + sqrt(Var[Q]))`). Kernel
/// writes only when `q_variance` is non-NULL AND `total_actions_for_var > 0`, /// writes only when `q_variance` is non-NULL AND `total_actions_for_var > 0`,
@@ -1972,6 +2009,23 @@ impl GpuExperienceCollector {
let trade_profitable_per_sample = stream let trade_profitable_per_sample = stream
.alloc_zeros::<i32>(total_output) .alloc_zeros::<i32>(total_output)
.map_err(|e| MLError::ModelError(format!("alloc trade_profitable_per_sample: {e}")))?; .map_err(|e| MLError::ModelError(format!("alloc trade_profitable_per_sample: {e}")))?;
// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf scratch.
// Sized `[total_output × cf_mult]` so the kernel writes both
// on-policy (out_off) AND counterfactual (cf_off) slots — the
// CF state at cf_off shares the SAME aux signal (same env
// state, just different action), so both slots get the same
// value, mirroring `cf_flip_per_sample`'s `total_output ×
// cf_mult` sizing pattern (`gpu_experience_collector.rs` line
// ~1945). alloc_zeros → uninitialised slots read as 0.0f =
// K=2 uniform-prior (no information). Kernel overwrites at
// every reached (i, t) with `sp20_compute_per_bar_aux_conf_k2(
// aux_logits + i*K)`; un-reached slots stay 0.0f. Used
// downstream by the Phase 5 Aux→Q gate at `gate =
// sigmoid((aux_conf - ISV[AUX_CONF_THRESHOLD]) /
// ISV[AUX_GATE_TEMP])`.
let aux_conf_per_sample = stream
.alloc_zeros::<f32>(total_output * cf_mult)
.map_err(|e| MLError::ModelError(format!("alloc aux_conf_per_sample: {e}")))?;
// Task 0.3: per-sample Kelly-from-atoms variance scale, written by // Task 0.3: per-sample Kelly-from-atoms variance scale, written by
// experience_env_step (~line 1422). Alloc_zeros → uninitialised slots // experience_env_step (~line 1422). Alloc_zeros → uninitialised slots
// read as 0.0f, which the host reducer filters out (only positive // read as 0.0f, which the host reducer filters out (only positive
@@ -2900,6 +2954,7 @@ impl GpuExperienceCollector {
action_mag_per_sample, action_mag_per_sample,
step_ret_per_sample, step_ret_per_sample,
trade_close_per_sample, trade_close_per_sample,
aux_conf_per_sample,
trade_profitable_per_sample, trade_profitable_per_sample,
var_scale_per_sample, var_scale_per_sample,
slot_live_per_sample, slot_live_per_sample,
@@ -4899,6 +4954,20 @@ impl GpuExperienceCollector {
self.last_experience_count = total; self.last_experience_count = total;
// SP20 Phase 3 Task 3.4 (2026-05-10): clone the per-bar
// aux_conf scratch into a fresh `[total]` CudaSlice for the
// batch struct. The kernel writes at `out_off` + `cf_off =
// out_off + N*L` so the first `total = N*L*2` slots of the
// scratch contain on-policy + CF; same dtod_clone pattern as
// states/rewards/actions/dones above. NULL-tolerant
// downstream: if aux_logits_per_env was NULL during the
// rollout, the kernel left these slots at 0.0f (K=2 uniform
// prior, no information) — Phase 5 gate will fall back to
// `mean_a Q(s, a)` baseline at sample time.
let aux_conf = dtod_clone_f32_native(
&self.stream, &self.aux_conf_per_sample, total, "aux_conf"
)?;
Ok(GpuExperienceBatch { Ok(GpuExperienceBatch {
states, states,
next_states, next_states,
@@ -4907,6 +4976,7 @@ impl GpuExperienceCollector {
dones, dones,
episode_ids, episode_ids,
aux_sign_labels, aux_sign_labels,
aux_conf,
state_dim: sd, state_dim: sd,
n_episodes, n_episodes,
timesteps, timesteps,
@@ -6331,6 +6401,13 @@ impl GpuExperienceCollector {
use crate::cuda_pipeline::sp14_isv_slots::HOLD_REWARD_EMA_INDEX; use crate::cuda_pipeline::sp14_isv_slots::HOLD_REWARD_EMA_INDEX;
HOLD_REWARD_EMA_INDEX as i32 HOLD_REWARD_EMA_INDEX as i32
}) })
// SP20 Phase 3 Task 3.4 (2026-05-10): per-bar aux_conf
// output for the replay-buffer schema. Threaded as
// a `[N × L]` f32 buffer (counterfactual extension
// happens at the `aux_conf` assembly site at the
// end of `collect_experiences_gpu`, mirroring the
// `aux_sign_labels` `[N × L × 2]` final shape).
.arg(&mut self.aux_conf_per_sample)
.launch(launch_cfg) .launch(launch_cfg)
.map_err(|e| MLError::ModelError(format!( .map_err(|e| MLError::ModelError(format!(
"experience_env_step t={t}: {e}" "experience_env_step t={t}: {e}"

View File

@@ -73,7 +73,8 @@ fn insert_random_batch(
let rewards = random_cuda_f32(n, stream)?; let rewards = random_cuda_f32(n, stream)?;
let dones = zeros_cuda_f32(n, stream)?; 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_sign = stream.alloc_zeros::<i32>(n).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, n)?; 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)?;
Ok(()) Ok(())
} }

View File

@@ -140,7 +140,8 @@ fn test_per_sample_latency() -> anyhow::Result<()> {
let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?; let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?;
let dones = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?; 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_sign = stream.alloc_zeros::<i32>(batch_insert).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, batch_insert)?; 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)?;
} }
assert_eq!(buf.len(), fill_count); assert_eq!(buf.len(), fill_count);

View File

@@ -150,7 +150,8 @@ fn test_per_weights_valid() -> anyhow::Result<()> {
let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?; let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?;
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?; 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_sign = stream.alloc_zeros::<i32>(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, 1)?; 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)?;
} }
buf.sample_proportional(16)?; buf.sample_proportional(16)?;
@@ -195,7 +196,8 @@ fn test_per_indices_valid() -> anyhow::Result<()> {
let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?; let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?;
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?; 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_sign = stream.alloc_zeros::<i32>(1).map_err(|e| anyhow::anyhow!("aux_sign alloc: {e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, &aux_sign, 1)?; 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)?;
} }
buf.sample_proportional(16)?; buf.sample_proportional(16)?;

View File

@@ -2526,6 +2526,7 @@ impl DQNTrainer {
&gpu_batch.rewards, &gpu_batch.rewards,
&gpu_batch.dones, &gpu_batch.dones,
&gpu_batch.aux_sign_labels, &gpu_batch.aux_sign_labels,
&gpu_batch.aux_conf,
total, total,
).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?; ).map_err(|e| anyhow::anyhow!("GPU PER insert_batch: {e}"))?;
} }

View File

@@ -122,8 +122,9 @@ fn fill_buffer(buf: &mut GpuReplayBuffer, n: usize, state_dim: usize) {
let rf = stream.clone_htod(&rewards_host).unwrap(); let rf = stream.clone_htod(&rewards_host).unwrap();
let df = stream.clone_htod(&dones_host).unwrap(); let df = stream.clone_htod(&dones_host).unwrap();
let aux_sign: CudaSlice<i32> = stream.alloc_zeros::<i32>(n).unwrap(); let aux_sign: CudaSlice<i32> = stream.alloc_zeros::<i32>(n).unwrap();
let aux_conf: CudaSlice<f32> = stream.alloc_zeros::<f32>(n).unwrap();
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &aux_sign, n).unwrap(); buf.insert_batch(&sf, &nsf, &af, &rf, &df, &aux_sign, &aux_conf, n).unwrap();
} }
#[test] #[test]

View File

@@ -2165,8 +2165,9 @@ mod gpu {
let rf: CudaSlice<f32> = stream.clone_htod(&rewards_host).expect("htod rewards"); let rf: CudaSlice<f32> = stream.clone_htod(&rewards_host).expect("htod rewards");
let df: CudaSlice<f32> = stream.clone_htod(&dones_host).expect("htod dones"); let df: CudaSlice<f32> = stream.clone_htod(&dones_host).expect("htod dones");
let ax: CudaSlice<i32> = stream.clone_htod(&aux_sign_host).expect("htod aux_sign"); let ax: CudaSlice<i32> = stream.clone_htod(&aux_sign_host).expect("htod aux_sign");
let ac: CudaSlice<f32> = stream.alloc_zeros::<f32>(n1).expect("aux_conf alloc");
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, n1) buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, &ac, n1)
.expect("insert_batch (recovery=1)"); .expect("insert_batch (recovery=1)");
stream.synchronize().expect("sync after insert_batch (recovery=1)"); stream.synchronize().expect("sync after insert_batch (recovery=1)");
@@ -2183,8 +2184,9 @@ mod gpu {
let rf2: CudaSlice<f32> = stream.clone_htod(&vec![0.0_f32; n2]).expect("htod rf2"); let rf2: CudaSlice<f32> = stream.clone_htod(&vec![0.0_f32; n2]).expect("htod rf2");
let df2: CudaSlice<f32> = stream.clone_htod(&vec![0.0_f32; n2]).expect("htod df2"); let df2: CudaSlice<f32> = stream.clone_htod(&vec![0.0_f32; n2]).expect("htod df2");
let ax2: CudaSlice<i32> = stream.clone_htod(&vec![0_i32; n2]).expect("htod ax2"); let ax2: CudaSlice<i32> = stream.clone_htod(&vec![0_i32; n2]).expect("htod ax2");
let ac2: CudaSlice<f32> = stream.alloc_zeros::<f32>(n2).expect("aux_conf alloc 2");
buf.insert_batch(&sf2, &nsf2, &af2, &rf2, &df2, &ax2, n2) buf.insert_batch(&sf2, &nsf2, &af2, &rf2, &df2, &ax2, &ac2, n2)
.expect("insert_batch (recovery=0)"); .expect("insert_batch (recovery=0)");
stream.synchronize().expect("sync after insert_batch (recovery=0)"); stream.synchronize().expect("sync after insert_batch (recovery=0)");
@@ -2349,8 +2351,9 @@ mod gpu {
let rf: CudaSlice<f32> = stream.clone_htod(&rewards_host).expect("htod rewards"); let rf: CudaSlice<f32> = stream.clone_htod(&rewards_host).expect("htod rewards");
let df: CudaSlice<f32> = stream.clone_htod(&dones_host).expect("htod dones"); let df: CudaSlice<f32> = stream.clone_htod(&dones_host).expect("htod dones");
let ax: CudaSlice<i32> = stream.clone_htod(&aux_sign_host).expect("htod aux_sign"); let ax: CudaSlice<i32> = stream.clone_htod(&aux_sign_host).expect("htod aux_sign");
let ac: CudaSlice<f32> = stream.alloc_zeros::<f32>(batch_size).expect("aux_conf alloc");
buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, batch_size) buf.insert_batch(&sf, &nsf, &af, &rf, &df, &ax, &ac, batch_size)
.expect("insert_batch (mixed-env)"); .expect("insert_batch (mixed-env)");
stream.synchronize().expect("sync after insert_batch"); stream.synchronize().expect("sync after insert_batch");

View File

@@ -2,6 +2,194 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
## 2026-05-10 — SP20 Phase 3 Task 3.4: replay-buffer schema for per-bar aux_conf (atomic)
Lands the producer→ring→consumer schema plumbing for per-bar
`aux_conf` (the K=2 peak-softmax-above-uniform confidence value the
Phase 5 Aux→Q gate consumes at the Bellman target). Atomic across all
struct boundaries per `feedback_no_partial_refactor`:
```
ExperienceCollector → GpuExperienceBatch → insert_batch() → ring → sample() → GpuBatchPtrs → Trainer
(writes per (i,t)) (.aux_conf field) (8th arg) (ring (.aux_conf_ptr (Phase 5 gate
column) field) consumer)
```
Phase 5 lands the consumer (Bellman target gate kernel); this commit
is the producer-side schema plumbing only.
### Spec §4.4 consumer contract (Phase 5 forward-reference)
```
At each replay batch step:
aux_conf = ptrs.aux_conf_ptr[k] // SAMPLED bar
threshold = ISV[AUX_CONF_THRESHOLD_INDEX] // [0.01, 0.20]
temp = ISV[AUX_GATE_TEMP_INDEX] // adaptive
gate = sigmoid((aux_conf - threshold) / temp) // ∈ [0, 1]
Q_target = gate × Q_target_full + (1 - gate) × Q_target_baseline
where Q_target_baseline = mean_a Q(s, a)
```
The gate ∈ [0, 1] gracefully degrades the Bellman target toward
`mean_a Q(s, a)` at uncertain states (low aux_conf), preventing the
all-actions punishment that drives Hold-everywhere. Spec §4.4 calls
this the "uncertain-state neutralizer" semantic.
### Producer site (`experience_env_step` per (i, t))
New kernel arg appended at the end of the signature (after the Task
3.2 `aux_logits_per_env` / `hold_baseline_buffer` / `hold_buffer_size`
/ `aux_k` / `hold_cost_scale_idx` / `hold_reward_ema_idx` block):
```cuda
float* __restrict__ aux_conf_per_sample /* [N × L] f32 */
```
Written ONCE per (i, t) at the kernel-entry default block (just after
the H6 diagnostic defaults at line ~2270), using the same
`sp20_compute_per_bar_aux_conf_k2` helper Task 3.2 added in
`sp20_hold_baseline.cuh`:
```cuda
if (aux_conf_per_sample != NULL) {
if (aux_logits_per_env != NULL) {
const float* logits_row_init = aux_logits_per_env + i * aux_k;
aux_conf_per_sample[out_off] = sp20_compute_per_bar_aux_conf_k2(logits_row_init);
} else {
aux_conf_per_sample[out_off] = 0.0f;
}
}
```
NULL-tolerant: when `aux_logits_per_env == NULL` (test scaffolds), the
slot retains the K=2 uniform-prior sentinel `0.0f` — no information.
The Phase 5 gate at default 0.0 with threshold ∈ [0.01, 0.20] gives
`gate ≈ 0` ⇒ Q_target falls back to `mean_a Q(s, a)` (graceful
degradation).
### Counterfactual slot
The kernel ALSO writes `aux_conf_per_sample[cf_off]` inside the
existing CF block (line ~4485-ish, after `out_actions[cf_off]` etc.):
```cuda
if (aux_conf_per_sample != NULL) {
aux_conf_per_sample[cf_off] = aux_conf_per_sample[out_off];
}
```
The CF state at `cf_off` shares the SAME env state (only the action
differs), so it shares the SAME aux signal — mirror the on-policy
write rather than recomputing.
### New collector buffer
`GpuExperienceCollector.aux_conf_per_sample: CudaSlice<f32>`
`[alloc_episodes × alloc_timesteps × cf_mult]` f32 row-major. Sized
to match the kernel's `out_off` + `cf_off` layout (out_off ∈
`[0, total_output)`, cf_off ∈ `[total_output, total_output × cf_mult)`).
alloc_zeros default = K=2 uniform-prior sentinel.
### New batch struct field
`GpuExperienceBatch.aux_conf: CudaSlice<f32>` — `[total = base_total
* cf_mult]` f32. Cloned from `aux_conf_per_sample` via
`dtod_clone_f32_native(buf, total)` at the end of
`collect_experiences_gpu` (mirrors states/rewards/actions/dones
clone pattern).
### New `GpuBatchPtrs` field (`crates/ml-dqn/src/gpu_replay_buffer.rs`)
`GpuBatchPtrs.aux_conf_ptr: u64` — points at either the trainer's
destination buffer (when `set_trainer_aux_conf_ptr` was called by
Phase 5) or the fallback `sample_aux_conf` ring-gather destination
(default).
### New `GpuReplayBuffer` ring + sample fields
- `aux_conf: CudaSlice<f32>``[capacity]` ring storage. Written
by `scatter_insert_f32` in `insert_batch` (mirrors the
`aux_sign_labels` i32 path immediately above; reuses the existing
f32 scatter kernel that scatters rewards / dones).
- `sample_aux_conf: CudaSlice<f32>``[max_batch_size]` per-batch
gather destination. Written by `gather_f32` (fallback) or
`gather_f32_scalar` (direct-to-trainer path).
- `trainer_aux_conf_ptr: u64` — destination for the direct path
(Phase 5 wires the trainer's `aux_conf_at_state_buf` via
`set_trainer_aux_conf_ptr`).
### `insert_batch` signature change (atomic across all callers)
Added 7th arg `aux_conf_in: &CudaSlice<f32>` between `aux_sign` and
`bs`. All 6 callers updated atomically:
- `crates/ml/src/trainers/dqn/trainer/training_loop.rs:2522` (production)
- `crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs:77`
- `crates/ml/src/trainers/dqn/smoke_tests/performance.rs:144`
- `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs:154 + :201`
- `crates/ml/tests/gpu_per_integration_test.rs:128`
- `crates/ml/tests/sp15_phase1_oracle_tests.rs:2170 + :2189 + :2356`
- 3 inline tests in `gpu_replay_buffer.rs::tests` (1528, 1608, 1672)
Each smoke / oracle test scaffold passes a fresh `alloc_zeros::<f32>(n)`
for aux_conf — matches the K=2 uniform-prior semantic when the aux
head isn't exercised by that scaffold.
### Tests
- `gpu_replay_buffer::tests::test_aux_conf_schema_round_trip` (GPU
behavioral test): inserts 8 transitions with distinct in-range
aux_conf values [0.0, 0.05, ..., 0.35]; samples 1; asserts the
gathered slot value matches at least one inserted value
(round-trip integrity invariant per
`pearl_tests_must_prove_not_lock_observations`).
- `gpu_replay_buffer::tests::test_aux_conf_trainer_ptr_wiring_contract`
(CPU-only): asserts `set_trainer_aux_conf_ptr` /
`trainer_aux_conf_ptr` setter/accessor pair behavior matches
`set_isv_signals_ptr` / `isv_signals_dev_ptr` mirror.
### Files modified
| File | Status | Purpose |
|------|--------|---------|
| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | +1 kernel arg, +2 write sites | Per-bar aux_conf production |
| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | +field, +alloc, +arg pass, +clone, +batch field | Producer-side schema |
| `crates/ml-dqn/src/gpu_replay_buffer.rs` | +`GpuBatchPtrs.aux_conf_ptr`, +ring/sample fields, +scatter/gather, +setter/accessor, +insert_batch arg | Replay buffer schema |
| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | +1 arg in insert_batch call | Production caller |
| `crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs` | +1 arg | Smoke test |
| `crates/ml/src/trainers/dqn/smoke_tests/performance.rs` | +1 arg | Smoke test |
| `crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs` | +1 arg ×2 | Smoke tests |
| `crates/ml/tests/gpu_per_integration_test.rs` | +1 arg | Integration test |
| `crates/ml/tests/sp15_phase1_oracle_tests.rs` | +1 arg ×3 | Oracle tests |
| `docs/dqn-wire-up-audit.md` | This entry | Audit log |
### Verification
```
SQLX_OFFLINE=true cargo check -p ml --tests --features cuda # green
SQLX_OFFLINE=true cargo check -p ml-dqn --tests --features cuda # green
SQLX_OFFLINE=true cargo test -p ml --lib sp20 # 21/21 pass
SQLX_OFFLINE=true cargo test -p ml-dqn test_aux_conf_trainer_ptr # CPU pass
SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml-dqn --features cuda \
test_aux_conf_schema_round_trip # GPU host pass
```
### Plan accuracy errata
See `docs/superpowers/plans/2026-05-09-sp19-20-wr-first.md` Phase 3
Errata Gap 11 for the `GpuBatchPtrs` vs `GpuExperienceBatch` naming
clarification (both structs got the new field; the data flow has TWO
struct boundaries, not one).
### Phase 3 → Phase 5 forward references
- Phase 5: Aux→Q gate kernel reads `GpuBatchPtrs.aux_conf_ptr` at
the Bellman target loss site. `set_trainer_aux_conf_ptr` is
called once after the trainer's `aux_conf_at_state_buf` is
allocated, switching the gather to direct-to-trainer mode.
## 2026-05-10 — SP20 Phase 3 Task 3.2: Hold opportunity-cost dual emission (Component 2, atomic) ## 2026-05-10 — SP20 Phase 3 Task 3.2: Hold opportunity-cost dual emission (Component 2, atomic)
Lands the per-bar Hold opportunity-cost dual emission at the Lands the per-bar Hold opportunity-cost dual emission at the