feat(sp14-β): wire collector-native SP14 producer chain after aux forward

Step 4 of β migration: 3 SP14/SP13-EGF producers fire per-rollout-step
in collector using collector-owned kernel handles + collector stream.
Reads rollout-time q_values (post-expected-Q, pre-IQR/ensemble/noise)
+ exp_aux_nb_softmax. Writes to shared ISV.

Producer order preserved (matches trainer submit_aux_ops chain):
  1. SP13 dir-acc reduce → 2 fixed-α EMAs → aux_pred to ISV[375]
  2. SP14 q_disagreement_update (reads aux softmax + q_values)
  3. SP14 alpha_grad_compute (pure ISV state machine)

Same kernel gate (commit 9d0c124ce) preserves EMAs across
no-contribution rollout steps.

q_logits semantic note: collector's q_values buffer is the
post-expected-Q output, BEFORE IQR/ensemble/noise SAXPY bonuses (those
run after this block). The kernel's argmax-over-K=4 finds Q's intended
direction; this matches the trainer's q_out_buf semantic exactly. If a
future audit shows noise-induced argmax flips matter, the launch site
is one indirection from the noise-free expected_q_kernel output.

Gated on isv_signals_dev_ptr != 0 && trainer_params_ptr != 0. No
seed_phase_active_cache gate — EGF Gate 1 needs to observe both
seed-phase scripted-policy and post-seed Q-policy actions across the
curriculum.

Compile clean; sp14_oracle_tests 2/2 non-GPU pass (7 GPU tests
ignored on RTX 3050 Ti host).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 21:54:38 +02:00
parent 296ba92282
commit c691bd381a
2 changed files with 235 additions and 0 deletions

View File

@@ -4580,6 +4580,239 @@ impl GpuExperienceCollector {
// q_values now contains expected Q-values [N, q_stride] f32
}
// ── 3a-β. SP14 β-migration step 4 (Option B, 2026-05-07):
// collector-native EGF producer chain.
//
// 3 launches on `self.stream` using the collector-owned kernel
// handles (loaded in step 1 of the β migration). Same kernel
// semantics as the trainer's `submit_aux_ops` chain — the
// ONLY difference is the data the kernels read: this path
// reads `exp_aux_nb_softmax_buf` (rollout-time aux predictions
// from step 3) and the collector's `q_values` (post-expected-Q,
// pre-IQR / pre-ensemble / pre-noise — bit-equivalent to the
// trainer's `q_out_buf` semantic) where Thompson sampling
// forces Short/Long picks every step. The trainer-time path
// sees Q-direction picks dominated by Hold/Flat (the natural
// argmax distribution of the converged 4-way head); rollout-
// time picks via Thompson are the genuine source of aux↔Q
// disagreement signal.
//
// Producer order (must match trainer `submit_aux_ops`):
// 1. SP13 dir-acc reduce + 2 EMAs + aux_pred_to_isv
// 2. SP14 q_disagreement_update (reads aux softmax + q_values)
// 3. SP14 alpha_grad_compute (pure ISV state machine)
//
// q_logits semantic note: the collector's q_values is the
// post-expected-Q output (computed above by `expected_q_kernel`).
// The IQR / ensemble / noise SAXPY bonuses below (steps 3b/3c/3d)
// run AFTER this block so q_values here is noise-free, matching
// the trainer's q_out_buf semantic at the trainer's
// launch_sp14_q_disagreement_update call site (gpu_dqn_trainer.rs:8682).
// q_stride = total_branch_actions = 13 (same as trainer's q_out_buf).
//
// Gated on `isv_signals_dev_ptr != 0` (matches SP13 hold-rate
// pattern at line ~4626 below — without ISV wired the EMAs
// have nowhere to write) AND `trainer_params_ptr != 0` (the
// aux forward in step 3 was skipped if NULL, so the softmax
// tile would be stale `alloc_zeros` data).
//
// No `seed_phase_active_cache` gate here — unlike the SP13
// hold-rate observer (which prices Hold-cost during seed
// phase as well), the SP14 EGF Schmitt-trigger Gate 1 needs
// to observe BOTH seed-phase scripted-policy actions and
// post-seed Q-policy actions to discriminate aux signal
// quality across the curriculum. The kernel's own
// total_cnt > 0 gate (commit 9d0c124ce) handles empty
// contributions correctly.
if self.isv_signals_dev_ptr != 0 && self.trainer_params_ptr != 0 {
use crate::cuda_pipeline::sp13_isv_slots::{
AUX_DIR_ACC_SHORT_EMA_INDEX, AUX_DIR_ACC_LONG_EMA_INDEX,
AUX_DIR_PREDICTION_INDEX, DIR_ACC_EMA_SENTINEL,
};
use super::gpu_aux_heads::AUX_NEXT_BAR_K;
let aux_softmax_dev = self.exp_aux_nb_softmax_buf.raw_ptr();
let aux_label_dev = self.exp_aux_nb_label_buf.raw_ptr();
let dir_acc_out_dev = self.exp_aux_dir_acc_buf.dev_ptr;
let isv_dev = self.isv_signals_dev_ptr;
let n_i32_dir = n as i32;
let k_i32_dir = AUX_NEXT_BAR_K as i32;
// Producer 1a: dir-acc reduce. Writes 6 floats (dir_acc,
// pos_pred_frac, pos_label_frac, n_down, n_up, n_skip)
// to mapped-pinned `exp_aux_dir_acc_buf`. Tree-reduce in
// shmem; smem layout per `aux_dir_acc_reduce_kernel.cu`:
// 6 i32 strips × bdim. Single block.
{
let bdim: u32 = 256;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (bdim, 1, 1),
shared_mem_bytes: 6 * bdim * std::mem::size_of::<i32>() as u32,
};
unsafe {
self.stream
.launch_builder(&self.exp_sp13_aux_dir_acc_reduce_kernel)
.arg(&aux_softmax_dev)
.arg(&aux_label_dev)
.arg(&n_i32_dir)
.arg(&k_i32_dir)
.arg(&dir_acc_out_dev)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!(
"sp14-β: aux_dir_acc_reduce t={t}: {e}"
)))?;
}
}
// Producer 1b: short EMA (α=0.3) into ISV[373]. n=1
// (only out_6[0] = dir_acc feeds the EMA; pos_pred /
// pos_label / n_down / n_up / n_skip are diagnostic-only).
// Sentinel 0.5 = random-baseline per
// `pearl_first_observation_bootstrap`.
{
let n_slots: i32 = 1;
let isv_off: i32 = AUX_DIR_ACC_SHORT_EMA_INDEX as i32;
let alpha: f32 = 0.3;
let sentinel: f32 = DIR_ACC_EMA_SENTINEL;
unsafe {
self.stream
.launch_builder(&self.sp13_apply_fixed_alpha_ema_kernel)
.arg(&dir_acc_out_dev)
.arg(&n_slots)
.arg(&isv_off)
.arg(&alpha)
.arg(&sentinel)
.arg(&isv_dev)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"sp14-β: dir_acc short EMA t={t}: {e}"
)))?;
}
}
// Producer 1c: slow EMA (α=0.05) into ISV[374].
// Stagnation-detector comparator (compares against ISV[373]).
{
let n_slots: i32 = 1;
let isv_off: i32 = AUX_DIR_ACC_LONG_EMA_INDEX as i32;
let alpha: f32 = 0.05;
let sentinel: f32 = DIR_ACC_EMA_SENTINEL;
unsafe {
self.stream
.launch_builder(&self.sp13_apply_fixed_alpha_ema_kernel)
.arg(&dir_acc_out_dev)
.arg(&n_slots)
.arg(&isv_off)
.arg(&alpha)
.arg(&sentinel)
.arg(&isv_dev)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"sp14-β: dir_acc long EMA t={t}: {e}"
)))?;
}
}
// Producer 1d: aux_pred → ISV[375] tanh. Reads softmax
// tile, writes `mean(softmax[1] - softmax[0])` ∈ [-1, +1]
// to the shared scalar slot 375 (per-step overwrite, not
// an EMA). Single block, tree-reduce; smem = bdim × f32.
{
let bdim: u32 = 256;
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (bdim, 1, 1),
shared_mem_bytes: bdim * std::mem::size_of::<f32>() as u32,
};
let isv_slot_off: i32 = AUX_DIR_PREDICTION_INDEX as i32;
unsafe {
self.stream
.launch_builder(&self.exp_sp13_aux_pred_to_isv_tanh_kernel)
.arg(&aux_softmax_dev)
.arg(&n_i32_dir)
.arg(&k_i32_dir)
.arg(&isv_slot_off)
.arg(&isv_dev)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!(
"sp14-β: aux_pred_to_isv_tanh t={t}: {e}"
)))?;
}
}
// Producer 2: SP14 q_disagreement_update. Reads aux
// softmax (K=2) + q_values (K=4 over the first 4 columns
// per row, q_stride=13). Writes ISV[383/384/389]
// (q_dis_short, q_dis_long, var_q EMAs). Single block,
// 256 threads; shmem = 2 × 256 × f32 (numerator + count
// tiles per kernel header).
{
let q_logits_dev = self.q_values.raw_ptr();
let n_i32_q = n as i32;
let q_stride: i32 = (self.branch_sizes[0]
+ self.branch_sizes[1]
+ self.branch_sizes[2]
+ self.branch_sizes[3]) as i32;
let alpha_short: f32 = 0.3;
let alpha_long: f32 = 0.05;
let alpha_var: f32 = 0.05;
let bdim: u32 = 256;
let shmem_bytes: u32 = 2 * bdim * std::mem::size_of::<f32>() as u32;
unsafe {
self.stream
.launch_builder(&self.exp_sp14_q_disagreement_update_kernel)
.arg(&aux_softmax_dev)
.arg(&q_logits_dev)
.arg(&isv_dev)
.arg(&n_i32_q)
.arg(&q_stride)
.arg(&alpha_short)
.arg(&alpha_long)
.arg(&alpha_var)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (bdim, 1, 1),
shared_mem_bytes: shmem_bytes,
})
.map_err(|e| MLError::ModelError(format!(
"sp14-β: q_disagreement_update t={t}: {e}"
)))?;
}
}
// Producer 3: SP14 α_grad compute. Pure ISV-state state
// machine — Schmitt-trigger Gate 1 + adaptive k +
// adaptive β. Reads ISV[372..395) drivers (including
// the q_dis EMAs just written above on the same stream)
// and writes ISV[392/393] (α_grad raw + smoothed) plus
// bookkeeping slots [385..391]. Single thread state
// machine; we launch with 32 threads to mirror the B.4
// oracle test launch shape.
{
let alpha_var: f32 = 0.05;
unsafe {
self.stream
.launch_builder(&self.exp_sp14_alpha_grad_compute_kernel)
.arg(&isv_dev)
.arg(&alpha_var)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (32, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"sp14-β: alpha_grad_compute t={t}: {e}"
)))?;
}
}
}
// ── 3b. Add IQR exploration bonus to Q-values (anti-collapse) ──
// IQR bonus is pre-tiled [N, q_stride] and pre-scaled by alpha.
// SAXPY: q_values[i] += 1.0 * iqr_tiled[i] for all i.

View File

@@ -2,6 +2,8 @@
**Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7.
SP14 β-migration step 4 — wire collector-native SP14 EGF producer chain after aux forward (2026-05-07): Step 4 of 5. Inserts the 3 SP14/SP13-EGF producer launches into the collector's per-rollout-step body, immediately after the `expected_q_kernel` (which produces `q_values`) and BEFORE the IQR/ensemble/noise SAXPY blocks (so the SP14 q_disagreement consumer sees noise-free `q_values` matching the trainer's `q_out_buf` semantic). Producer order matches the trainer's `submit_aux_ops` chain (preserved per `feedback_no_partial_refactor.md`): (1) SP13 dir-acc reduce + 2 fixed-α EMAs (α=0.3 short → ISV[373], α=0.05 long → ISV[374]) + aux_pred → ISV[375] tanh — same 4 launches the trainer's `launch_sp13_aux_dir_metrics` orchestrates inline (gpu_dqn_trainer.rs:15270), reading the collector's `exp_aux_nb_softmax_buf [n_episodes, 2]` (from step 3) + `exp_aux_nb_label_buf [n_episodes]` (from step 3's per-step label producer) + writing to the shared ISV via `isv_signals_dev_ptr`. Sentinel 0.5 = random-guessing baseline per `pearl_first_observation_bootstrap`. (2) SP14 `q_disagreement_update_kernel` — reads `exp_aux_nb_softmax_buf [n_episodes, K=2]` + collector's `q_values [n_episodes, q_stride=13]` (post-`expected_q_kernel` output); writes ISV[383/384/389] (q_dis_short, q_dis_long, var_q EMAs). q_stride=13 matches the trainer's q_out_buf stride; the kernel's argmax-over-K_DIR=4 finds the direction-head pick from the first 4 columns of each row. Single block, 256 threads, smem = 2 × 256 × f32. (3) SP14 `alpha_grad_compute_kernel` — pure ISV-state state machine reading the q_dis EMAs just written above (same-stream serial ordering enforces the dependency) plus ISV[372..395) drivers; writes ISV[392/393] (α_grad raw + smoothed) plus the Schmitt-trigger / k_q / k_aux state machine bookkeeping at ISV[385..391]. Single thread (we launch with 32 to mirror the B.4 oracle test launch shape). The kernel's empty-batch decay gate (commit 9d0c124ce) preserves prior EMA values when `total_cnt == 0` for a given step — same protection extends to the rollout path. **Gating**: `if self.isv_signals_dev_ptr != 0 && self.trainer_params_ptr != 0` — without ISV the EMAs have nowhere to write; without trainer params the aux forward in step 3 would have been skipped and the softmax tile would still be `alloc_zeros`. **No `seed_phase_active_cache` gate** here (unlike SP13 hold-rate which prices Hold-cost only post-seed): the EGF Schmitt-trigger Gate 1 needs to observe BOTH seed-phase scripted-policy actions and post-seed Q-policy actions to discriminate aux signal quality across the curriculum. **q_logits placement note**: the SP14 chain runs AFTER `expected_q_kernel` (line 4581) and BEFORE the IQR (3b) / ensemble (3c) / noise (3d) SAXPY bonuses — so `q_values` here is bit-equivalent to the trainer's `q_out_buf` semantic at the trainer's `launch_sp14_q_disagreement_update` call site. The IQR/ensemble bonuses below are exploration noise added for action selection, not gradient flow. Per `pearl_no_host_branches_in_captured_graph` the launches are OUTSIDE any captured graph (the captured `forward` graph has already ended). The trainer's existing private `launch_sp14_q_disagreement_update` / `launch_sp14_alpha_grad_compute` / `launch_sp13_aux_dir_metrics` methods at gpu_dqn_trainer.rs:8671 / 8733 / 15270 are still defined and still pass kernel correctness via the 7 sp14_oracle GPU tests — production has migrated to the collector path; step 5 will delete the now-unused `submit_aux_ops` invocations (commit `200f05fce` lines 2031-2040). Files modified: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+200 lines: 3 producer launch blocks at line ~4583 with comprehensive doc-comments). Verification: `SQLX_OFFLINE=true cargo check -p ml --lib` clean (dev profile); `cargo test -p ml --test sp14_oracle_tests` 2/2 non-GPU pass (7 GPU tests ignored on RTX 3050 Ti host — they verify kernel correctness, which is unchanged by this commit).
SP14 β-migration step 3 — wire collector-side aux-head forward + per-step label producer (2026-05-07): Step 3 of 5. Inserts the aux next-bar forward + per-step sign-label producer into the collector's per-rollout-step body, immediately after the captured forward graph completes (`cuStreamEndCapture` / `cuGraphLaunch` block at line 4444) and before the expected_q kernel + the SP14 EGF chain (step 4). Two launches per rollout step, both on `self.stream`: (1) `aux_sign_label_per_step_kernel` — new thin variant of `aux_sign_label_kernel`. The trajectory kernel writes the entire `[total]` ring at the end of `collect_experiences_gpu`; the per-step kernel writes only the current step's `[n_episodes]` slice using `bar = episode_starts[ep] + t` (host-scalar `t` passed as kernel arg). Same per-thread O(1) map: read `targets[bar*6+2]` (raw_close column) at `bar` and `bar+lookahead`, classify by strict greater-than. Skip sentinel -1 fires when the lookahead window runs past `total_bars`; the CE / dir-acc / q_disagreement consumers all mask -1 rows out of denominators. Lookahead matches the trajectory kernel's `config.hindsight_lookahead.max(1)` so both encode the same K=2 "up vs not-up" boundary. Pure GPU map; no atomicAdd; per `feedback_no_atomicadd.md`. (2) `AuxHeadsForwardOps::forward_next_bar` — same orchestrator method the trainer's `aux_heads_forward` (gpu_dqn_trainer.rs:17208) uses, but bound to `self.stream` and operating on `n_episodes` rows of `exp_h_s2_f32` (the rollout-time h_s2). Param tensor pointers resolved via the existing `f32_weight_ptrs_from_base(self.trainer_params_ptr, &self.param_sizes)` path used elsewhere in the collector — same indices [119..123) the trainer uses (`nb_w1`, `nb_b1`, `nb_w2`, `nb_b2`). Writes hidden + logits + softmax tiles; the `exp_aux_nb_softmax_buf [n_episodes, 2]` is the production output for step 4's SP14 EGF chain. **Cold-start gate**: the entire block is gated on `self.trainer_params_ptr != 0`. The collector's constructor sets `trainer_params_ptr: 0` initially, then the trainer's `set_trainer_params_ptr()` (called after `FusedTrainingCtx::new`) wires the real pointer. Without that wire `f32_weight_ptrs_from_base` would return offsets from a NULL base and the aux forward would dereference garbage. Same defensive-NULL pattern the existing `forward_online_f32` call at line 4160 uses (the captured graph itself fails to capture if `trainer_params_ptr == 0`, so the test scaffold path that doesn't allocate params doesn't reach this block in production code paths either). **Placement**: AFTER captured graph end (`cuStreamEndCapture` for capture path, `cuGraphLaunch` for replay path) and BEFORE expected_q. The aux forward reads `exp_h_s2_f32` populated by `forward_online_f32` inside the captured graph; same-stream serial ordering within the per-step body suffices (no event sync needed). Per `pearl_no_host_branches_in_captured_graph` the launches are OUTSIDE the captured graph — the `f32_weight_ptrs_from_base` call returns a fixed `[u64; NUM_WEIGHT_TENSORS]` array post-capture, the host-scalar `t` is passed as kernel arg outside capture, and the cudarc `launch_builder` calls don't allocate inside any active capture (the captured `forward` graph already ended). Step 4 will wire the SP14 EGF producers (3 launches) AFTER this block but in the same per-step body. Files added: `crates/ml/src/cuda_pipeline/aux_sign_label_per_step_kernel.cu` (per-thread O(1) map kernel, +66 lines). Files modified: `crates/ml/build.rs` (+8 lines: register new cubin); `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+1 struct field for the new kernel handle, +1 cubin static include, +14 lines for the load in `new()`, +1 wire in struct literal, +90 lines for the per-step launch block at line ~4445). Verification: `SQLX_OFFLINE=true cargo check -p ml --lib` clean (dev profile); `cargo test -p ml --test sp14_oracle_tests` 2/2 non-GPU pass.
SP14 β-migration step 2 — allocate rollout-sized aux buffers + AuxHeadsForwardOps in collector (2026-05-07): Step 2 of 5. Adds 5 collector-owned buffers and 1 collector-owned `AuxHeadsForwardOps` instance, all sized to `alloc_episodes` (rollout batch dim) instead of the trainer's `batch_size` (replay batch dim). Buffers: `exp_aux_nb_hidden_buf [alloc_episodes × AUX_HIDDEN_DIM=32] f32` (post-ELU hidden, written by `aux_next_bar_forward`), `exp_aux_nb_logits_buf [alloc_episodes × 2] f32` (saved logits, parity with trainer), `exp_aux_nb_softmax_buf [alloc_episodes × 2] f32` (the production consumer surface — read by all 3 SP14 EGF producers in step 4), `exp_aux_nb_label_buf [alloc_episodes] i32` (per-step `{-1, 0, 1}` filled in step 3), `exp_aux_dir_acc_buf` mapped-pinned 6 floats (matches trainer's post-B1.1a layout `[dir_acc, pos_pred_frac, pos_label_frac, n_down, n_up, n_skip]`). Orchestrator: `exp_aux_heads_fwd: AuxHeadsForwardOps` constructed via `AuxHeadsForwardOps::new(&stream)` — same loader path the trainer uses (gpu_aux_heads.rs:95), but bound to `self.stream` (collector). The collector uses ONLY the `forward_next_bar` method from this orchestrator; the regime forward + CE reduce kernels are loaded inside the AuxHeadsForwardOps constructor as a unit (cudarc Module returns the cubin's full symbol set) but never launched on this stream — the rollout path doesn't compute a regime classification loss. Per `feedback_no_stubs.md` this is acceptable because the unused handles cost nothing at runtime and keeping `AuxHeadsForwardOps` as the unit-of-construction preserves trainer-collector symmetry; deconstructing the orchestrator into next-bar-only would either require a new `AuxNextBarForwardOps` API (orphan-by-design, conflicts with `feedback_no_legacy_aliases.md`) or duplicate the loader logic (conflicts with `feedback_no_partial_refactor.md`'s spirit). Param tensor pointers (the trainer's `params_buf[119..123)` for nb_w1/b1/w2/b2) flow through the existing `f32_weight_ptrs_from_base(self.trainer_params_ptr, &self.param_sizes)` path used by the existing `forward_online_f32` call at line 4160 — no new param-pointer plumbing. No forward calls yet — additive only; step 3 wires the per-step launches. Verification: `SQLX_OFFLINE=true cargo check -p ml --lib` clean (dev profile); `cargo test -p ml --test sp14_oracle_tests` 2/2 non-GPU pass. Files modified: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+6 struct fields, +6 allocations in `new()`, +6 field placements in struct literal).