gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp14-β): wire collector-side aux-head forward + label producer per-rollout-step

Browse Source 
Step 3 of β migration: collector now runs aux_next_bar_forward on
rollout state every step. Label producer (new thin variant
aux_sign_label_per_step_kernel) derives sign(price[t+1] - price[t])
per env using bar = episode_starts[ep] + t. Aux predictions feed
the EGF kernel chain (step 4), NOT the Q-head's input (rollout
Q-head still sees raw h_s2, dir_qaux_concat_ptr remains 0u64).

Placement: AFTER captured forward graph, BEFORE expected_q kernel.
Same-stream serial ordering reads exp_h_s2_f32 populated by
forward_online_f32 inside the captured graph. Cold-start gated on
trainer_params_ptr != 0 to skip the test-scaffold path where the
trainer hasn't wired its params yet.

Files added: aux_sign_label_per_step_kernel.cu (66 lines).
Files modified: build.rs (+8 lines, register cubin),
gpu_experience_collector.rs (+106 lines: struct field, cubin static,
load in new(), per-step launch block).

Compile clean; sp14_oracle_tests 2/2 non-GPU pass.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-07 21:51:51 +02:00
parent 88eb7aa241
commit 296ba92282
4 changed files with 219 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
crates/ml/build.rs
View File

@@ -737,6 +737,14 @@ fn main() {
// known degraded behavior between B1.1a and B1.1b that this
// kernel resolves.
"aux_sign_label_kernel.cu",
// SP14 β-migration step 3 (Option B, 2026-05-07): per-rollout-step
// variant of `aux_sign_label_kernel`. The trajectory kernel writes
// the entire `[total]` ring at the end of `collect_experiences_gpu`;
// this kernel writes the current step's `[n_episodes]` slice using
// `bar = episode_starts[ep] + t` so the collector-native EGF
// producer chain has fresh labels every rollout step. Same per-thread
// O(1) map; same lookahead window semantics; same skip sentinel.
"aux_sign_label_per_step_kernel.cu",
// SP14 Layer B Task B.3 (2026-05-05): Earned Gradient Flow
// producer kernel. Per-step computes the K=4↔K=2 mapped argmax
// mismatch between the Q-head's 4-way direction action

65
crates/ml/src/cuda_pipeline/aux_sign_label_per_step_kernel.cu Normal file
View File

@@ -0,0 +1,65 @@
/* ══════════════════════════════════════════════════════════════════════════
* SP14 β-migration step 3 — per-rollout-step variant of aux_sign_label_kernel
* (2026-05-07).
*
* The full collector-end `aux_sign_label_kernel` (this directory) writes
* the entire trajectory's labels at the end of `collect_experiences_gpu`
* (after the timestep loop). The β migration's collector-native EGF
* producer chain needs the labels per rollout step — only the current
* step's slice of `n_episodes` envs, derived from
* `bar = episode_starts[ep] + t` per env.
*
* Same per-thread O(1) map as the trajectory kernel: read `targets[bar*6+2]`
* (raw_close column — see `dt_kernels.cu:803` for the canonical column
* layout: 0=preproc_close, 1=preproc_next, 2=raw_close, 3=raw_next,
* 4=raw_open, 5=mid_open) at `bar` and `bar+lookahead`, classify by
* strict greater-than. Skip sentinel -1 fires when the lookahead window
* runs past `total_bars`.
*
* Pure map; no atomicAdd, no reduction, no shared memory — block
* tree-reduce N/A here per `feedback_no_atomicadd.md`. GPU-only data
* read/write per `feedback_cpu_is_read_only.md`. Inputs are mapped-pinned
* (targets) or device-allocated (episode_starts, out_labels) per
* `feedback_no_htod_htoh_only_mapped_pinned.md`.
*
* Launch config: grid=(ceil(n_episodes / 256)), block=(256). Mirrors the
* trajectory variant.
*
* Args:
* targets — `[total_bars, 6]` mapped-pinned f32.
* episode_starts — `[n_episodes]` device i32, base bar index per env.
* t — current rollout step (0..timesteps).
* out_labels — `[n_episodes]` device i32; receives -1/0/1 per env.
* n_episodes — number of envs in this rollout slice.
* total_bars — total bars in the data series.
* lookahead — bars to look ahead (must match the trajectory
* kernel's lookahead — both encode the same K=2
* "up" boundary).
* ══════════════════════════════════════════════════════════════════════════ */
extern "C" __global__
void aux_sign_label_per_step_kernel(
const float* __restrict__ targets,
const int* __restrict__ episode_starts,
int t,
int* __restrict__ out_labels,
int n_episodes,
int total_bars,
int lookahead)
{
int ep = blockIdx.x * blockDim.x + threadIdx.x;
if (ep >= n_episodes) return;
int bar = episode_starts[ep] + t;
if (bar < 0 || bar + lookahead >= total_bars) {
out_labels[ep] = -1; // skip sentinel — no lookahead window available
return;
}
float p_now = targets[bar * 6 + 2];
float p_fut = targets[(bar + lookahead) * 6 + 2];
// Strict greater-than tie-break: matches trajectory kernel — flat/down
// both map to "down"=0 because the K=2 head can only distinguish
// "up" from "not-up".
out_labels[ep] = (p_fut > p_now) ? 1 : 0;
}

144
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -995,6 +995,16 @@ pub struct GpuExperienceCollector {
/// observability of the rollout-time label distribution).
exp_aux_dir_acc_buf: super::mapped_pinned::MappedF32Buffer,
/// SP14 β-migration step 3 (2026-05-07): per-rollout-step label
/// producer kernel handle. Sibling of the trajectory-wide
/// `aux_sign_label_kernel` (line 3776 above) — same per-thread O(1)
/// map but indexed by `episode_starts[ep] + t` for the current
/// rollout step instead of the full `bar_indices` ring. Loaded
/// from `aux_sign_label_per_step_kernel.cubin` on the collector's
/// stream so producer→consumer ordering inside the per-step body
/// is stream-implicit.
exp_aux_sign_label_per_step_kernel: cudarc::driver::CudaFunction,
/// B.2 Plan 3 Task 3: trade_attempt_rate_ema_update GPU kernel.
/// Single-block (1 thread) reduction + adaptive EMA of Flat→Positioned
/// transition rate into ISV[71]. Launched alongside reward_component_ema
@@ -2004,6 +2014,29 @@ impl GpuExperienceCollector {
"sp14-β: alloc exp_aux_dir_acc_buf (6 f32, collector): {e}"
)))?;
// SP14 β-migration step 3 (2026-05-07): per-rollout-step label
// producer kernel. Loaded on the collector's stream so its launch
// chains directly after `forward_online_f32` (which populates
// `exp_h_s2_f32`) without cross-stream sync. The kernel body lives
// in `aux_sign_label_per_step_kernel.cu`; the trajectory variant
// (`aux_sign_label_kernel.cu`) is unchanged — its end-of-rollout
// launch at line 3776 still fires for the trainer's replay-batch
// labels in `collect_experiences_gpu`. Per
// `pearl_no_host_branches_in_captured_graph` the load happens
// here in the constructor; per-step launches use the cached
// `CudaFunction` handle.
let exp_aux_sign_label_per_step_kernel = {
let m = stream.context()
.load_cubin(SP14_AUX_SIGN_LABEL_PER_STEP_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!(
"sp14-β: aux_sign_label_per_step cubin (collector): {e}"
)))?;
m.load_function("aux_sign_label_per_step_kernel")
.map_err(|e| MLError::ModelError(format!(
"sp14-β: aux_sign_label_per_step_kernel load (collector): {e}"
)))?
};
// SP15 Phase 1.3.b-followup (2026-05-07): per-env DD state tile
// sized `[alloc_episodes * 6]`. The redesigned `dd_state_kernel`
// writes 6 stats per env to this tile (Path A → Path B
@@ -2378,6 +2411,7 @@ impl GpuExperienceCollector {
exp_aux_nb_softmax_buf,
exp_aux_nb_label_buf,
exp_aux_dir_acc_buf,
exp_aux_sign_label_per_step_kernel,
trade_attempt_rate_ema_kernel,
plan_threshold_update_kernel,
state_kl_kernel,
@@ -4409,6 +4443,108 @@ impl GpuExperienceCollector {
}
}
// ── 2b. SP14 β-migration step 3 (Option B, 2026-05-07): aux
// next-bar forward + per-step label producer.
//
// Runs AFTER the captured forward graph (which populates
// `exp_h_s2_f32` via `forward_online_f32`) and BEFORE the
// expected_q kernel + the SP14 EGF chain (step 4) that consume
// `exp_aux_nb_softmax_buf`. Mirrors the trainer's
// `aux_heads_forward` pattern (gpu_dqn_trainer.rs:17159) — same
// orchestrator, same param tensors at indices [119..123), but
// bound to the collector's stream and operating on rollout-
// sized buffers (n_episodes rows) instead of the trainer's
// batch_size.
//
// ISV gating: the aux forward itself doesn't touch ISV; it
// produces the softmax tile that step 4's chain reads. Skip
// this whole block in the cold-start (test-scaffold) path
// where the trainer hasn't wired its params yet — without
// that wire `trainer_params_ptr == 0` and the aux forward
// would dereference NULL param pointers.
if self.trainer_params_ptr != 0 {
use crate::cuda_pipeline::batched_forward::f32_weight_ptrs_from_base;
use super::gpu_aux_heads::AUX_NEXT_BAR_K;
// Param tensor pointers via the same offset arithmetic the
// trainer's `aux_heads_forward` uses (gpu_dqn_trainer.rs:17166).
// The shared `f32_weight_ptrs_from_base` resolves all
// NUM_WEIGHT_TENSORS tensors against the trainer's flat
// params buffer; aux next-bar tensors are at indices
// [119..123) (nb_w1, nb_b1, nb_w2, nb_b2). Regime-head
// tensors at [123..127) are not used here — the rollout
// path doesn't compute a regime classification loss.
let aux_w_ptrs = f32_weight_ptrs_from_base(
self.trainer_params_ptr,
&self.param_sizes,
);
let nb_w1 = aux_w_ptrs[119];
let nb_b1 = aux_w_ptrs[120];
let nb_w2 = aux_w_ptrs[121];
let nb_b2 = aux_w_ptrs[122];
let sh2 = self.network_dims.1;
// Step A: per-step sign-label producer. Writes
// `exp_aux_nb_label_buf[ep] := sign(p_close[bar+lookahead] − p_close[bar])`
// for each ep in 0..n_episodes, where `bar = episode_starts[ep] + t`.
// Sentinel -1 fires when the lookahead window runs past
// `total_bars`; the CE loss / dir-acc / q_disagreement
// consumers all mask -1 rows out of denominators (see
// `aux_dir_acc_reduce_kernel.cu` for the canonical mask
// semantics). Lookahead matches the trajectory-wide
// `aux_sign_label_kernel` launch (line ~3776 below) so
// both encode the same K=2 "up vs not-up" boundary.
{
let lookahead_i32 = config.hindsight_lookahead.max(1);
let labels_dev = self.exp_aux_nb_label_buf.raw_ptr();
let targets_dev = targets_buf.dev_ptr;
let starts_dev = self.episode_starts_buf.raw_ptr();
let t_i32 = t as i32;
let blocks = ((n as u32) + 255) / 256;
unsafe {
self.stream
.launch_builder(&self.exp_aux_sign_label_per_step_kernel)
.arg(&targets_dev)
.arg(&starts_dev)
.arg(&t_i32)
.arg(&labels_dev)
.arg(&n_i32)
.arg(&total_bars)
.arg(&lookahead_i32)
.launch(LaunchConfig {
grid_dim: (blocks, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: 0,
})
.map_err(|e| MLError::ModelError(format!(
"sp14-β: aux_sign_label_per_step t={t}: {e}"
)))?;
}
}
// Step B: aux next-bar forward. Reads `exp_h_s2_f32`
// (n_episodes × shared_h2) populated by `forward_online_f32`
// above; writes hidden + logits + softmax tiles. The
// softmax tile (`exp_aux_nb_softmax_buf [n_episodes, 2]`)
// is the production output consumed by step 4's SP14 EGF
// chain. Per `pearl_no_host_branches_in_captured_graph`
// this launch is OUTSIDE the captured graph (post
// `cuGraphLaunch`) — same placement as the trainer's
// `aux_heads_forward` call relative to the trainer's
// captured graph.
self.exp_aux_heads_fwd.forward_next_bar(
&self.stream,
self.exp_h_s2_f32.raw_ptr(),
nb_w1, nb_b1, nb_w2, nb_b2,
n, // batch dim = n_episodes
sh2,
AUX_NEXT_BAR_K,
self.exp_aux_nb_hidden_buf.raw_ptr(),
self.exp_aux_nb_logits_buf.raw_ptr(),
self.exp_aux_nb_softmax_buf.raw_ptr(),
)?;
}
// ── 3. Convert logits → expected Q-values ───────────────────
// When num_atoms == 1, b_logits is already [N, 11] Q-values.
// When num_atoms > 1, we need softmax + expectation over z-support.
@@ -5879,6 +6015,14 @@ static REWARD_SHAPING_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/r
static SP13_AUX_SIGN_LABEL_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/aux_sign_label_kernel.cubin"));
/// SP14 β-migration step 3 (Option B, 2026-05-07): per-rollout-step
/// variant of `aux_sign_label_kernel`. The trajectory kernel writes
/// `[total]` after the step loop; this kernel writes `[n_episodes]`
/// every step using `bar = episode_starts[ep] + t`. See
/// `aux_sign_label_per_step_kernel.cu` for the kernel body.
static SP14_AUX_SIGN_LABEL_PER_STEP_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/aux_sign_label_per_step_kernel.cubin"));
fn compile_experience_kernels(
stream: &Arc<CudaStream>,
_state_dim: usize,

2
docs/dqn-wire-up-audit.md
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 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).
SP14 β-migration step 1 — pre-load EGF kernel handles in experience collector (2026-05-07): Step 1 of 5 of the Option B (collector-native) β-migration unwinding the `train-6fcml`-fix follow-up. Even with the empty-batch decay gate (commit `9d0c124ce`) the trainer-time `submit_aux_ops` path can only fire the EGF chain against batches whose Q-direction picks are dominated by Hold/Flat (the natural argmax distribution of the converged 4-way direction head); the genuine aux↔Q disagreement signal lives in the rollout path on `gpu_experience_collector`'s stream where Thompson sampling forces Short/Long picks every step. Cross-stream launches from a trainer-stream-captured graph would race the collector-stream forward producers; the cleanest fix is collector-native — let the collector load its own copies of the EGF kernel handles and launch them with `self.stream`. Mirror of the existing SP13 hold_rate observer pattern (`crates/ml/src/cuda_pipeline/gpu_experience_collector.rs:1820`) where the collector loads `SP13_HOLD_RATE_OBSERVER_CUBIN` + `SP13_APPLY_FIXED_ALPHA_EMA_CUBIN` once on its own stream because `batch_actions` lives on this stream and cross-stream sync would defeat the no-host-sync per-step contract. This commit allocates 4 `CudaFunction` fields on `GpuExperienceCollector` (`exp_sp13_aux_dir_acc_reduce_kernel`, `exp_sp13_aux_pred_to_isv_tanh_kernel`, `exp_sp14_q_disagreement_update_kernel`, `exp_sp14_alpha_grad_compute_kernel`) — same cubins as the trainer's, loaded once on `self.stream` so producer→consumer ordering inside the collector's per-step body becomes stream-implicit. Cubin static decls in `gpu_dqn_trainer.rs` flipped from `static` to `pub(crate) static` for `SP13_AUX_DIR_ACC_REDUCE_CUBIN`, `SP13_AUX_PRED_TO_ISV_TANH_CUBIN`, `SP14_Q_DISAGREEMENT_CUBIN`, `SP14_ALPHA_GRAD_CUBIN`. Per `feedback_no_partial_refactor.md` the β migration is split into 5 atomic commits — this commit is additive (new fields, no launches) so behavior is bit-identical; step 4 lands the launches and step 5 deletes the trainer-time path. Per `pearl_no_host_branches_in_captured_graph` the `cuModuleLoadData` calls all happen in the constructor before any per-step launch, identical pattern to SP13 hold_rate. Verification: `SQLX_OFFLINE=true cargo check -p ml --lib` clean (dev profile, warnings unrelated to this change); `cargo test -p ml --test sp14_oracle_tests` 2/2 non-GPU pass (7 GPU tests `ignored` on RTX 3050 Ti host). Files modified: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (4 visibility flips); `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` (+4 struct fields, +4 cubin loads in `new()`, +4 field placements in struct literal).