From c691bd381a9d389bd1d080bdb0d1facc2cf577c7 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 7 May 2026 21:54:38 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp14-=CE=B2):=20wire=20collector-native=20?= =?UTF-8?q?SP14=20producer=20chain=20after=20aux=20forward?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- .../cuda_pipeline/gpu_experience_collector.rs | 233 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 2 + 2 files changed, 235 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index 269dabd1f..274c9964f 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -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::() 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::() 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::() 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. diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 987f46371..60418f7a3 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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).