From 6b01b70c3e81ce3e0714c4f3cd91495c11e8a517 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Sat, 2 May 2026 01:47:07 +0200 Subject: [PATCH] =?UTF-8?q?docs(sp6):=20implementation=20plan=20=E2=80=94?= =?UTF-8?q?=203-sub-project=20parallel=20refactor=20(Pearls=202,=203,=205)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three independent sub-projects (one per Pearl), each an atomic commit, each running in its own git worktree (sp6-pearl-2/3/5). Pearl 2 expands compute_adaptive_budgets() to per-branch [f32;4] arrays + trunk-mean scalars, dispatches branch-correction SAXPY sub-launches. Pearl 3 replaces noise_sigma: f32 with noise_sigma_per_branch: [f32;4] in ExperienceCollectorConfig and extends add_advantage_noise to per-branch lookup. Pearl 5 adds 4-branch tau buffer triplets to GpuIqnHead, runs 4 IQN forward passes per step (Approach B), each contributing iqn_trunk/4 to prevent 4x gradient accumulation. Co-Authored-By: Claude Sonnet 4.6 --- ...branch-consumer-infrastructure-refactor.md | 1363 +++++++++++++++++ 1 file changed, 1363 insertions(+) create mode 100644 docs/superpowers/plans/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md diff --git a/docs/superpowers/plans/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md b/docs/superpowers/plans/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md new file mode 100644 index 000000000..e901c4b70 --- /dev/null +++ b/docs/superpowers/plans/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md @@ -0,0 +1,1363 @@ +# SP6: Per-Branch Consumer Infrastructure Refactor — Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Wire the three SP5 producer signals that are currently averaged-to-scalar at the consumer side — loss budgets (Pearl 2), NoisyNet σ (Pearl 3), and IQN τ schedule (Pearl 5) — to their respective consumers at full per-branch fidelity. + +**Architecture:** Three fully independent sub-projects, each an atomic commit. Pearl 2 expands 4 SAXPY launchers to dispatch per-branch sub-launches using the existing `branch_slice_starts_dev`/`branch_slice_lens_dev` infrastructure; trunk/value get the mean of 4 branch budgets. Pearl 3 adds `noise_sigma_per_branch: [f32; 4]` to `ExperienceCollectorConfig` and extends `add_advantage_noise` to do a per-action branch-index lookup. Pearl 5 removes the per-quantile mean-of-4-branches averaging in `refresh_taus_from_isv`, builds 4 independent `[B, N]` tau tensors (one per branch), and runs 4 IQN forward passes per training step (Approach B — simpler, safer than kernel rework). All three sub-projects run in parallel via dedicated git worktrees per `feedback_no_concurrent_agents_shared_tree`. + +**Tech Stack:** Rust 1.85+, CUDA 12.4, cudarc, mapped-pinned memory, ISV signal bus (`sp5_isv_slots.rs`), cuBLAS SCAL, `branch_slice_starts_dev`/`branch_slice_lens_dev` (already in `GpuDqnTrainer`). + +**Spec source:** `docs/superpowers/specs/2026-05-02-sp6-per-branch-consumer-infrastructure-refactor.md` (commit `ca6e0007d`) + +**Branch:** `sp5-magnitude-differentiation` HEAD (`ca6e0007d`) + +--- + +## Investigation Findings (pre-plan research) + +These findings were confirmed against live code before writing this plan. Implementers should verify each claim has not drifted. + +### Pearl 2 — grad_buf layout and SAXPY internals + +- **SAXPY function signatures confirmed** (`gpu_dqn_trainer.rs`): + - `apply_c51_budget_scale(&mut self, c51_budget: f32)` at line 9440 — scales the entire `grad_buf` (`total_params` floats) via `scale_f32_ungraphed` kernel. + - `apply_cql_saxpy(&mut self, cql_budget: f32)` at line 9468 — SAXPY `grad_buf += cql_budget * cql_grad_scratch` over `total_params` floats via `saxpy_f32_aux`. + - `apply_iqn_trunk_gradient(&mut self, iqn_d_h_s2_ptr, _online_dueling, iqn_budget: f32)` at line 8315 — after backward chain, SAXPYs trunk gradient into `grad_buf`. + - No standalone `apply_ens_saxpy`; ensemble diversity is in `apply_ensemble_diversity_backward` at line 8624. +- **Branch slice layout** (`gpu_dqn_trainer.rs:13209-13234`): branch `b` owns tensors `(8+4b)..(12+4b)` in the padded flat layout. `branch_slice_starts_dev: CudaSlice[4]` and `branch_slice_lens_dev: CudaSlice[4]` are already fields on `GpuDqnTrainer` (lines 3067, 3073). Trunk = tensors 0-7; value = tensors ~24-33; branches = tensors 8-23. +- **`compute_adaptive_budgets()`** at `fused_training.rs:3295`: reads `ISV[BUDGET_C51_BASE..+4]`, `ISV[BUDGET_IQN_BASE..+4]`, `ISV[BUDGET_CQL_BASE..+4]`, `ISV[BUDGET_ENS_BASE..+4]` (slots 190-209) and returns mean of each 4-branch group as `(f32, f32, f32, f32)`. Caches results to `last_{c51,iqn,cql,ens}_budget_eff` fields on `GpuDqnTrainer` (lines 4822-4831). +- **Call sites** for the 4 scalars (`fused_training.rs`): `c51_budget` → `apply_c51_budget_scale` at line 1894; `iqn_budget` → `apply_iqn_trunk_gradient` at line 2269; `cql_budget` → `apply_cql_saxpy` at line 2324. +- **D3 decision** (spec): trunk/value/aux weights get `mean(4 branch budgets)` — preserves current behavior for non-branch params. + +### Pearl 3 — NoisyLinear internals and ExperienceCollectorConfig blast radius + +- **NoisyLinear is not in `gpu_attention.rs`**: the active noise injection is `add_advantage_noise` in `experience_kernels.cu:3916`, which takes a single `float noise_sigma` and applies uniform noise to all `N * total_actions` Q-value slots. +- **Kernel signature** (`experience_kernels.cu:3916`): `add_advantage_noise(float* q_values, float noise_sigma, int N, int total_actions, unsigned int seed)`. No branch-index parameter currently. +- **Branch boundaries in Q layout**: Q-values are written branch-major. The flattened action index `tid % total_actions` ranges over `[0, b0_size + b1_size + b2_size + b3_size)`. Branch 0 (direction) owns `[0, b0_size)`, branch 1 owns `[b0_size, b0_size+b1_size)`, etc. This offset pattern is established in `experience_kernels.cu:3579` and already used by other kernels. +- **`ExperienceCollectorConfig` constructors**: confirmed 3 sites: + 1. `Default::default()` blanket impl at `gpu_experience_collector.rs:301`. + 2. `mod.rs:895` in cuda_pipeline — `ExperienceCollectorConfig::default()` in a test context. + 3. `training_loop.rs:1662` — full struct literal with `..Default::default()` tail (the production path, line 1747 is the `noise_sigma` field). +- The scalar field `noise_sigma: f32` is replaced by `noise_sigma_per_branch: [f32; 4]`. The `Default::default()` impl must be updated to supply `[0.1; 4]`. + +### Pearl 5 — IQN τ internals and approach selection + +- **`refresh_taus_from_isv`** (`gpu_iqn_head.rs:1962`): takes `isv_per_branch_taus: &[f32; 20]` (ISV[250..270)), averages 4 branches per quantile into a single `[5]` tau array, then tiles it into `online_taus: [B, N]` and `target_taus: [B, N]` and recomputes `cos_features: [D, N]` on the host before mapped-pinned upload. +- **Approach selected: Approach B** (4 independent forward passes). The IQN kernel (`iqn_dual_head_kernel.cu`) receives `online_taus [B, N]` as a contiguous buffer. Expanding to `[B, 4, N]` with a new stride argument touches the CUDA kernel ABI, breaks the existing CUDA Graph capture nodes, and risks the graph hang root causes documented in `session_2026-04-17_graph_hang_resolution.md`. Approach B is architecturally clean: 4 separate `refresh_taus_from_isv` + `run_iqn_forward` calls, each consuming one branch's tau schedule. Cost: ~3-4× IQN forward compute per step (~15-25% total step time increase). This is acceptable for correctness; perf optimisation (Approach C) is deferred. +- **`GpuIqnHead` struct** holds `online_taus`, `target_taus`, `cos_features` as single `CudaSlice` allocations sized `[B, N]`, `[B, N]`, `[D, N]`. Per-branch storage requires 4 of each, or using the existing buffers sequentially per pass. +- **Call site** (`fused_training.rs:2014-2023`): reads `ISV[IQN_TAU_BASE..+20]` into `[f32; 20]` and calls `iqn.refresh_taus_from_isv(&isv_taus)` once per epoch-level refresh, then runs IQN forward 5 times per step inside `run_full_step`. The per-branch dispatch adds 3 additional forward passes per refresh cycle. + +--- + +## File Structure + +### Sub-project 1 (Pearl 2) — files touched + +| File | Change | +|---|---| +| `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` | Add `apply_c51_budget_scale_branch(b, budget)`, `apply_cql_saxpy_branch(b, budget)` helpers; add `last_{c51,iqn,cql,ens}_budget_eff_per_branch: [f32; 4]` fields to trainer struct | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Change `compute_adaptive_budgets()` return type to `([f32;4], [f32;4], [f32;4], [f32;4], f32, f32, f32, f32)` (per-branch arrays + trunk means); update 3 call sites | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Update `last_iqn_budget_eff` / `last_cql_budget_eff` / `last_c51_budget_eff` logging to use per-branch arrays | + +### Sub-project 2 (Pearl 3) — files touched + +| File | Change | +|---|---| +| `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` | Replace `noise_sigma: f32` field with `noise_sigma_per_branch: [f32; 4]`; update `Default` impl; update `collect_experiences_gpu` launch site | +| `crates/ml/src/cuda_pipeline/experience_kernels.cu` | Replace `float noise_sigma` with `const float* noise_sigma_per_branch` + branch-index derivation via branch size offsets | +| `crates/ml/src/cuda_pipeline/mod.rs` | Update `ExperienceCollectorConfig::default()` test usage (line 895) | +| `crates/ml/src/trainers/dqn/trainer/training_loop.rs` | Replace averaging of `ISV[NOISY_SIGMA_BASE..+4]` with direct 4-element read | + +### Sub-project 3 (Pearl 5) — files touched + +| File | Change | +|---|---| +| `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` | Replace `refresh_taus_from_isv` averaging with `refresh_taus_for_branch(branch_idx, tau5)` + internal `refresh_taus_from_isv_per_branch` driver; extend struct to store last-used per-branch tau arrays | +| `crates/ml/src/trainers/dqn/fused_training.rs` | Update IQN tau refresh call to call per-branch variant 4 times; add 3 extra IQN forward dispatches per refresh cycle | + +--- + +## Worktree Setup + +Per `feedback_no_concurrent_agents_shared_tree`, each sub-project runs in its own git worktree. The controller (user) merges them sequentially onto `sp5-magnitude-differentiation` after each Pearl's worktree branch lands. + +```bash +# From the foxhunt repo root +git worktree add ../foxhunt-sp6-pearl-2 -b sp6-pearl-2 ca6e0007d +git worktree add ../foxhunt-sp6-pearl-3 -b sp6-pearl-3 ca6e0007d +git worktree add ../foxhunt-sp6-pearl-5 -b sp6-pearl-5 ca6e0007d +``` + +Each agent works exclusively in its own worktree directory. After implementation: + +```bash +# Merge order: Pearl 2 first (largest, most merge-conflict risk), Pearl 3 second, Pearl 5 last +git checkout sp5-magnitude-differentiation +git merge sp6-pearl-2 --no-ff +git merge sp6-pearl-3 --no-ff +git merge sp6-pearl-5 --no-ff +``` + +**Conflict risk:** All 3 sub-projects touch `fused_training.rs` (Pearl 2 at line ~3295-3323, Pearl 3 at line ~1747-1753, Pearl 5 at line ~2014-2023) and `training_loop.rs` (Pearl 2 at ~2749-2751, Pearl 3 at ~1747, Pearl 5 at no overlap). Git will produce distinct hunks if each sub-project's edits are in non-overlapping line ranges; manual merge is needed only if they drift into the same hunk. The suggested merge order (Pearl 2 → Pearl 3 → Pearl 5) is chosen so the file with most changes lands first, simplifying subsequent merges. + +**Cleanup:** + +```bash +git worktree remove ../foxhunt-sp6-pearl-2 +git worktree remove ../foxhunt-sp6-pearl-3 +git worktree remove ../foxhunt-sp6-pearl-5 +``` + +--- + +## Cross-Project Dependencies + +- Pearl 2 has no dependency on Pearl 3 or Pearl 5. +- Pearl 3 has no dependency on Pearl 2 or Pearl 5. +- Pearl 5 has no dependency on Pearl 2 or Pearl 3. +- The ISV producer side (`sp5_isv_slots.rs`, `pearl_2_budget_kernel.cu`, `pearl_3_sigma_kernel.cu`, `pearl_5_iqn_tau_kernel.cu`) is complete and unchanged by SP6. + +--- + +## Hard Rules (apply to every sub-project) + +- `feedback_no_cpu_compute_strict` — no new CPU-side reductions or EMAs +- `feedback_no_atomicadd` — no atomicAdd in any kernel modification +- `feedback_no_partial_refactor` — each sub-project is ONE atomic commit; partial migration is BLOCKED +- `feedback_isv_for_adaptive_bounds` — ISV slots are the only source of adaptive values; no hardcoded multipliers +- `feedback_no_stubs` — no `todo!()`, `unimplemented!()`, `0.0` stub returns +- `feedback_no_concurrent_agents_shared_tree` — each sub-project uses its own worktree +- `feedback_no_cpu_test_fallbacks` — any new unit tests are GPU-only; no CPU reference oracle +- No TODO/FIXME/XXX markers in committed code + +--- + +## Sub-Project 1: Pearl 2 — SAXPY Per-Branch Loss Budgets + +**Estimated time:** 3-4 hours + +**Branch:** `sp6-pearl-2` + +**Worktree:** `../foxhunt-sp6-pearl-2` + +### Task P2-T1: Read the current code before any edits + +- [ ] **Step 1: Read apply_c51_budget_scale** + + ```bash + # In worktree ../foxhunt-sp6-pearl-2 + grep -n 'fn apply_c51_budget_scale\|fn apply_cql_saxpy\|fn apply_iqn_trunk_gradient\|fn apply_ensemble_diversity_backward' \ + crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs + ``` + + Expected output: 4 line numbers (9440, 9468, 8315, 8624 approximately). Read each function body — confirm they operate on `self.total_params` (full flat buffer) with a single scalar multiplier. + +- [ ] **Step 2: Read branch_slice_starts_dev initialization** + + Read `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` lines 13209-13234. Confirm: + - `starts[b]` = byte offset of branch `b`'s first parameter / sizeof(f32) + - `lens[b]` = number of f32 elements in branch `b`'s parameter slice + - Both are uploaded via `upload_via_mapped_i32` + +- [ ] **Step 3: Read compute_adaptive_budgets** + + Read `crates/ml/src/trainers/dqn/fused_training.rs` lines 3295-3324. Confirm return type is `(f32, f32, f32, f32)` and caching to `last_*_budget_eff` fields. + +- [ ] **Step 4: Read the 3 call sites that consume the returned scalars** + + In `fused_training.rs`: + - Line ~1882: destructuring `(c51_budget, iqn_budget, cql_budget, _ens_budget)` + - Line ~1894: `apply_c51_budget_scale(c51_budget)` + - Line ~2269: `apply_iqn_trunk_gradient(d_h_s2_ptr, &mut self.online_dueling, iqn_budget)` + - Line ~2324: `apply_cql_saxpy(cql_budget)` + +- [ ] **Step 5: Locate last_*_budget_eff fields on GpuDqnTrainer** + + ```bash + grep -n 'last_iqn_budget_eff\|last_cql_budget_eff\|last_c51_budget_eff\|last_ens_budget_eff' \ + crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs | head -20 + ``` + + Expected: struct field declarations near line 4822 and initialization near line 16942. + +- [ ] **Step 6: Locate HEALTH_DIAG budget logging in training_loop.rs** + + ```bash + grep -n 'last_iqn_budget_eff\|last_cql_budget_eff\|last_c51_budget_eff' \ + crates/ml/src/trainers/dqn/trainer/training_loop.rs + ``` + + Expected: ~3 lines (read, cache, log in HEALTH_DIAG format string). + +### Task P2-T2: Add per-branch budget fields to GpuDqnTrainer + +**Files:** `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Add per-branch budget fields to the trainer struct** + + Find the block containing `pub(crate) last_iqn_budget_eff: f32` (near line 4822). Add 4 new fields immediately after the existing 4: + + ```rust + /// SP6 Pearl 2: per-branch budget arrays — set by compute_adaptive_budgets. + pub(crate) last_c51_budget_per_branch: [f32; 4], + pub(crate) last_iqn_budget_per_branch: [f32; 4], + pub(crate) last_cql_budget_per_branch: [f32; 4], + pub(crate) last_ens_budget_per_branch: [f32; 4], + ``` + +- [ ] **Step 2: Initialize the new fields in the constructor** + + Find the struct literal initializer that sets `last_iqn_budget_eff: 0.40` (near line 16942). Add after it: + + ```rust + last_c51_budget_per_branch: [0.55; 4], + last_iqn_budget_per_branch: [0.40; 4], + last_cql_budget_per_branch: [0.00; 4], + last_ens_budget_per_branch: [0.05; 4], + ``` + +- [ ] **Step 3: Add apply_c51_budget_scale_branch helper** + + Add immediately after the closing `}` of `apply_c51_budget_scale` (near line 9462): + + ```rust + /// SP6 Pearl 2: scale only the branch `b` parameter slice of grad_buf + /// by `budget`. Trunk/value scaling is handled by the trunk-mean call + /// in apply_c51_budget_scale. + pub fn apply_c51_budget_scale_branch(&mut self, branch_idx: usize, budget: f32) -> Result<(), MLError> { + if (budget - 1.0).abs() < 1e-6 { + return Ok(()); + } + // Read branch slice start (in f32 elements) and length from the + // already-uploaded branch_slice_starts_dev / branch_slice_lens_dev. + // These are [4] i32 device arrays; read on host at construction + // time is not needed — we replicate the host-side values here + // via the same computation used at line 13209. + let f32_size = std::mem::size_of::(); + let param_sizes = compute_param_sizes(&self.config); + let first_tensor = 8 + branch_idx * 4; + let start_bytes = padded_byte_offset(¶m_sizes, first_tensor) as usize; + let end_bytes = padded_byte_offset(¶m_sizes, first_tensor + 4) as usize; + let start_elems = (start_bytes / f32_size) as i32; + let len_elems = ((end_bytes - start_bytes) / f32_size) as i32; + if len_elems == 0 { return Ok(()); } + let grad_slice_ptr = self.ptrs.grad_buf + (start_elems as u64) * (f32_size as u64); + let blocks = ((len_elems as u32 + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.scale_f32_ungraphed) + .arg(&grad_slice_ptr) + .arg(&budget) + .arg(&len_elems) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "c51_budget_scale_branch[{branch_idx}]: {e}" + )))?; + } + Ok(()) + } + ``` + +- [ ] **Step 4: Add apply_cql_saxpy_branch helper** + + Add immediately after the closing `}` of `apply_cql_saxpy` (near line 9494): + + ```rust + /// SP6 Pearl 2: SAXPY only the branch `b` parameter slice: + /// grad_buf[branch_slice] += budget * cql_grad_scratch[branch_slice] + pub fn apply_cql_saxpy_branch(&mut self, branch_idx: usize, budget: f32) -> Result<(), MLError> { + let f32_size = std::mem::size_of::(); + let param_sizes = compute_param_sizes(&self.config); + let first_tensor = 8 + branch_idx * 4; + let start_bytes = padded_byte_offset(¶m_sizes, first_tensor) as usize; + let end_bytes = padded_byte_offset(¶m_sizes, first_tensor + 4) as usize; + let start_elems = (start_bytes / f32_size) as i32; + let len_elems = ((end_bytes - start_bytes) / f32_size) as i32; + if len_elems == 0 { return Ok(()); } + let grad_ptr = self.ptrs.grad_buf + (start_elems as u64) * (f32_size as u64); + let scratch_ptr = self.ptrs.cql_grad_scratch + (start_elems as u64) * (f32_size as u64); + let blocks = ((len_elems as u32 + 255) / 256) as u32; + unsafe { + self.stream + .launch_builder(&self.saxpy_f32_aux) + .arg(&grad_ptr) + .arg(&scratch_ptr) + .arg(&budget) + .arg(&len_elems) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!( + "cql_saxpy_branch[{branch_idx}]: {e}" + )))?; + } + Ok(()) + } + ``` + +- [ ] **Step 5: Verify cargo check is still clean** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Expected: no output. + +### Task P2-T3: Refactor compute_adaptive_budgets to return per-branch arrays + +**Files:** `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Read fused_training.rs imports to confirm sp5_isv_slots are already imported** + + ```bash + grep -n 'BUDGET_C51_BASE\|BUDGET_IQN_BASE\|BUDGET_CQL_BASE\|BUDGET_ENS_BASE' \ + crates/ml/src/trainers/dqn/fused_training.rs | head -5 + ``` + + Expected: line ~43 showing an import block that includes all four constants. + +- [ ] **Step 2: Replace compute_adaptive_budgets** + + The current function (lines ~3295-3323) returns `(f32, f32, f32, f32)`. Replace its body entirely with the following. The new return type is `([f32;4], [f32;4], [f32;4], [f32;4], f32, f32, f32, f32)` — the four per-branch arrays followed by the four trunk-mean scalars. + + Old signature: + ```rust + pub(crate) fn compute_adaptive_budgets(&mut self) -> (f32, f32, f32, f32) { + ``` + + New function (replace entire body from `pub(crate) fn compute_adaptive_budgets` through its closing `}`): + + ```rust + /// SP6 Pearl 2: Compute adaptive gradient budgets from ISV[190..210). + /// + /// Returns `(c51_branch, iqn_branch, cql_branch, ens_branch, c51_trunk, iqn_trunk, + /// cql_trunk, ens_trunk)`. The per-branch arrays feed per-branch SAXPY sub-launches + /// on branch HEAD parameter slices. The trunk scalars (mean of 4 branch budgets) + /// feed the full-buf trunk/value scaling call — preserving SP5 Layer B behavior + /// for shared parameters. + /// + /// Cold-start floor: ISV slots read 0 before first observation. + /// Floors (0.05 C51, 0.05 IQN, 0.02 CQL, 0.02 Ens) are Invariant 1 carve-outs + /// (numerical-stability, not tuned constants). + pub(crate) fn compute_adaptive_budgets( + &mut self, + ) -> ([f32; 4], [f32; 4], [f32; 4], [f32; 4], f32, f32, f32, f32) { + let mut c51 = [0.0_f32; 4]; + let mut iqn = [0.0_f32; 4]; + let mut cql = [0.0_f32; 4]; + let mut ens = [0.0_f32; 4]; + for b in 0..4_usize { + c51[b] = self.read_isv_signal_at(BUDGET_C51_BASE + b).max(0.05_f32); + iqn[b] = self.read_isv_signal_at(BUDGET_IQN_BASE + b).max(0.05_f32); + cql[b] = self.read_isv_signal_at(BUDGET_CQL_BASE + b).max(0.02_f32); + ens[b] = self.read_isv_signal_at(BUDGET_ENS_BASE + b).max(0.02_f32); + } + // Trunk/value use mean of 4 branch budgets (D3 decision in SP6 spec). + let c51_trunk = c51.iter().sum::() / 4.0_f32; + let iqn_trunk = iqn.iter().sum::() / 4.0_f32; + let cql_trunk = cql.iter().sum::() / 4.0_f32; + let ens_trunk = ens.iter().sum::() / 4.0_f32; + // Cache per-branch arrays for HEALTH_DIAG logging. + self.trainer.last_c51_budget_per_branch = c51; + self.trainer.last_iqn_budget_per_branch = iqn; + self.trainer.last_cql_budget_per_branch = cql; + self.trainer.last_ens_budget_per_branch = ens; + // Also cache trunk scalars for backward compat accessors. + self.trainer.last_iqn_budget_eff = iqn_trunk; + self.trainer.last_cql_budget_eff = cql_trunk; + self.trainer.last_c51_budget_eff = c51_trunk; + self.trainer.last_ens_budget_eff = ens_trunk; + (c51, iqn, cql, ens, c51_trunk, iqn_trunk, cql_trunk, ens_trunk) + } + ``` + +- [ ] **Step 3: Update the call site that destructures the return value** + + Find the line (~1882) that reads: + ```rust + let (c51_budget, iqn_budget, cql_budget, _ens_budget) = self.compute_adaptive_budgets(); + ``` + + Replace with: + ```rust + let (c51_branch, iqn_branch, cql_branch, _ens_branch, + c51_trunk, iqn_trunk, cql_trunk, _ens_trunk) = self.compute_adaptive_budgets(); + ``` + +- [ ] **Step 4: Update apply_c51_budget_scale call to use trunk mean + per-branch sub-launches** + + Find the block near line 1894: + ```rust + self.trainer.apply_c51_budget_scale(c51_budget) + .map_err(|e| anyhow::anyhow!("c51_budget_scale: {e}"))?; + ``` + + Replace with: + ```rust + // SP6 Pearl 2: trunk/value first (mean budget), then per-branch sub-launches. + self.trainer.apply_c51_budget_scale(c51_trunk) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_trunk: {e}"))?; + for b in 0..4_usize { + // Skip sub-launch when branch budget matches trunk mean to avoid + // double-scaling shared padding bytes at branch boundaries. + if (c51_branch[b] - c51_trunk).abs() > 1e-6 { + self.trainer.apply_c51_budget_scale_branch(b, c51_branch[b]) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_branch[{b}]: {e}"))?; + } + } + ``` + + > Note: `apply_c51_budget_scale` scales the ENTIRE grad_buf including branch slices. The per-branch sub-launch then applies a *correction factor* on top. The correct formulation is: the branch slice has already been scaled by `c51_trunk`; we want it scaled by `c51_branch[b]`. So the correction factor is `c51_branch[b] / c51_trunk`. Update the code to pass `c51_branch[b] / c51_trunk` as the budget argument (not `c51_branch[b]` directly): + + ```rust + self.trainer.apply_c51_budget_scale(c51_trunk) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_trunk: {e}"))?; + for b in 0..4_usize { + let correction = if c51_trunk > 1e-6 { c51_branch[b] / c51_trunk } else { 1.0_f32 }; + if (correction - 1.0).abs() > 1e-6 { + self.trainer.apply_c51_budget_scale_branch(b, correction) + .map_err(|e| anyhow::anyhow!("c51_budget_scale_branch[{b}]: {e}"))?; + } + } + ``` + +- [ ] **Step 5: Update apply_iqn_trunk_gradient call** + + Find near line 2269: + ```rust + self.trainer.apply_iqn_trunk_gradient( + d_h_s2_ptr, &mut self.online_dueling, iqn_budget, + ) + ``` + + Replace `iqn_budget` with `iqn_trunk` (IQN already operates on trunk parameters only, so there is no per-branch IQN sub-launch needed; the per-branch budget differentiation for IQN is structurally a future Pearl 2 extension): + ```rust + self.trainer.apply_iqn_trunk_gradient( + d_h_s2_ptr, &mut self.online_dueling, iqn_trunk, + ) + ``` + + > Note for implementer: IQN's backward flows through the TRUNK only (it has no branch-head sub-launches). The per-branch IQN budget differentiation requires routing IQN gradients to individual branch slices, which is architectural work beyond this SP6 scope. `iqn_trunk` (mean of 4 branch IQN budgets) is the correct SP6 behavior for the IQN SAXPY. + +- [ ] **Step 6: Update apply_cql_saxpy call to trunk mean + per-branch sub-launches** + + Find near line 2324: + ```rust + self.trainer.apply_cql_saxpy(cql_budget) + .map_err(|e| anyhow::anyhow!("CQL SAXPY: {e}"))?; + ``` + + Replace with: + ```rust + // SP6 Pearl 2: trunk/value SAXPY first (mean budget), then per-branch corrections. + self.trainer.apply_cql_saxpy(cql_trunk) + .map_err(|e| anyhow::anyhow!("CQL SAXPY trunk: {e}"))?; + for b in 0..4_usize { + let correction = if cql_trunk > 1e-6 { cql_branch[b] / cql_trunk } else { 1.0_f32 }; + if (correction - 1.0).abs() > 1e-6 { + self.trainer.apply_cql_saxpy_branch(b, correction) + .map_err(|e| anyhow::anyhow!("CQL SAXPY branch[{b}]: {e}"))?; + } + } + ``` + +- [ ] **Step 7: Verify cargo check** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Expected: no errors. If there are type mismatch errors on the `compute_adaptive_budgets` return, they'll point to additional call sites — fix each one. + +### Task P2-T4: Update HEALTH_DIAG logging in training_loop.rs + +**Files:** `crates/ml/src/trainers/dqn/trainer/training_loop.rs` + +- [ ] **Step 1: Find the budget caching block in training_loop.rs** + + ```bash + grep -n 'last_iqn_budget_eff\|last_cql_budget_eff\|last_c51_budget_eff' \ + crates/ml/src/trainers/dqn/trainer/training_loop.rs + ``` + + Expected: lines ~2749-2751 (caching) and ~3827-3829 (HEALTH_DIAG format string usage). + +- [ ] **Step 2: Add per-branch budget caching block** + + After the existing caching block at ~2749-2751: + ```rust + self.last_iqn_budget_eff = Some(fused.last_iqn_budget_eff()); + self.last_cql_budget_eff = Some(fused.last_cql_budget_eff()); + self.last_c51_budget_eff = Some(fused.last_c51_budget_eff()); + ``` + + Add: + ```rust + self.last_c51_budget_per_branch = Some(fused.trainer.last_c51_budget_per_branch); + self.last_iqn_budget_per_branch = Some(fused.trainer.last_iqn_budget_per_branch); + self.last_cql_budget_per_branch = Some(fused.trainer.last_cql_budget_per_branch); + ``` + +- [ ] **Step 3: Add per-branch budget fields to the training loop struct** + + Find the struct definition that contains `last_iqn_budget_eff: Option`. Add after it: + ```rust + last_c51_budget_per_branch: Option<[f32; 4]>, + last_iqn_budget_per_branch: Option<[f32; 4]>, + last_cql_budget_per_branch: Option<[f32; 4]>, + ``` + + Find the struct constructor/`Default` impl and add matching initialization: + ```rust + last_c51_budget_per_branch: None, + last_iqn_budget_per_branch: None, + last_cql_budget_per_branch: None, + ``` + +- [ ] **Step 4: Add per-branch budget emission in the HEALTH_DIAG block** + + Find the HEALTH_DIAG format string block that emits `last_iqn_budget_eff` (~line 3827). After the existing budget scalar lines, add: + ```rust + if let Some(arr) = self.last_c51_budget_per_branch { + info!( + "HEALTH_DIAG c51_budget_per_branch: dir={:.4} mag={:.4} ord={:.4} urg={:.4}", + arr[0], arr[1], arr[2], arr[3] + ); + } + if let Some(arr) = self.last_iqn_budget_per_branch { + info!( + "HEALTH_DIAG iqn_budget_per_branch: dir={:.4} mag={:.4} ord={:.4} urg={:.4}", + arr[0], arr[1], arr[2], arr[3] + ); + } + if let Some(arr) = self.last_cql_budget_per_branch { + info!( + "HEALTH_DIAG cql_budget_per_branch: dir={:.4} mag={:.4} ord={:.4} urg={:.4}", + arr[0], arr[1], arr[2], arr[3] + ); + } + ``` + +- [ ] **Step 5: Verify cargo check** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Expected: no errors. + +### Task P2-T5: Verify no mean(ISV[BUDGET_*..+4]) patterns remain + +- [ ] **Step 1: Grep for the old averaging pattern** + + ```bash + grep -n 'BUDGET_C51_BASE\|BUDGET_IQN_BASE\|BUDGET_CQL_BASE\|BUDGET_ENS_BASE' \ + crates/ml/src/trainers/dqn/fused_training.rs + ``` + + Expected: only the `compute_adaptive_budgets` function references these constants (the old `.sum::() / 4.0_f32` averaging form must no longer appear). + +- [ ] **Step 2: Sanity grep for sum/4 averaging on budget slots** + + ```bash + grep -n 'sum.*4\.0\|/ 4\.0' crates/ml/src/trainers/dqn/fused_training.rs | head -10 + ``` + + Expected: zero matches for the budget block (other `/ 4.0` usages elsewhere are unrelated). + +### Task P2-T6: Run smoke test + +- [ ] **Step 1: Run cargo build to confirm no link errors** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo build -p ml --lib --release 2>&1 | tail -5 + ``` + + Expected: `Compiling ml ...` then `Finished`. + +- [ ] **Step 2: Run smoke test (local RTX 3050 Ti)** + + ```bash + FOXHUNT_TEST_DATA=test_data/futures-baseline \ + FOXHUNT_MBP10_DIR=test_data/mbp10 \ + FOXHUNT_TRADES_DIR=test_data/trades \ + CARGO_INCREMENTAL=0 \ + cargo test -p ml --lib -- smoke_tests::multi_fold_convergence --ignored --nocapture 2>&1 | tail -40 + ``` + + Expected: PASS. Kill at first HEALTH_DIAG line per `feedback_kill_runs_on_anomaly_quickly` — look for `cql_budget_per_branch` log lines showing non-uniform values across branches. + +### Task P2-T7: Atomic commit + +- [ ] **Step 1: Stage changed files** + + ```bash + git add \ + crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \ + crates/ml/src/trainers/dqn/fused_training.rs \ + crates/ml/src/trainers/dqn/trainer/training_loop.rs + ``` + +- [ ] **Step 2: Confirm no other files are staged** + + ```bash + git diff --cached --name-only + ``` + + Expected: exactly the 3 files above. + +- [ ] **Step 3: Commit** + + ```bash + git commit -m "$(cat <<'EOF' + feat(dqn): SP6 Pearl 2 — per-branch C51/CQL loss budget SAXPY sub-launches + + compute_adaptive_budgets() now returns [f32;4] per-branch arrays alongside + trunk-mean scalars. apply_c51_budget_scale + apply_cql_saxpy dispatch + trunk-mean first (full grad_buf), then per-branch correction factors on + each branch's HEAD parameter slice via branch_slice_starts_dev layout. + apply_iqn_trunk_gradient uses trunk mean (IQN backward is trunk-only). + HEALTH_DIAG emits c51/iqn/cql_budget_per_branch: dir/mag/ord/urg per epoch. + + Co-Authored-By: Claude Sonnet 4.6 + EOF + )" + ``` + +--- + +## Sub-Project 2: Pearl 3 — NoisyLinear Per-Branch σ Array + +**Estimated time:** 2-3 hours + +**Branch:** `sp6-pearl-3` + +**Worktree:** `../foxhunt-sp6-pearl-3` + +### Task P3-T1: Read the current code before any edits + +- [ ] **Step 1: Read ExperienceCollectorConfig struct** + + Read `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` lines 150-300. Confirm `noise_sigma: f32` at line ~298 is the last field before the closing `}`. + +- [ ] **Step 2: Read the add_advantage_noise kernel** + + Read `crates/ml/src/cuda_pipeline/experience_kernels.cu` lines 3909-3935. Confirm: + - Signature: `(float* q_values, float noise_sigma, int N, int total_actions, unsigned int seed)` + - `tid` ranges over `[0, N * total_actions)` where `total_actions = b0+b1+b2+b3 = 13` + - `action = tid % total_actions` — branch is determined by which of `[0,4), [4,7), [7,10), [10,13)` `action` falls in for the default branch sizes `{4, 3, 3, 3}` + +- [ ] **Step 3: Read the noise kernel launch site** + + Read `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` lines 3620-3649. Note: + - `branch_sizes` field on `GpuExperienceCollector` — confirm it stores `[b0, b1, b2, b3]` + - The kernel receives `config.noise_sigma` + +- [ ] **Step 4: Enumerate all ExperienceCollectorConfig constructors** + + ```bash + grep -rn 'ExperienceCollectorConfig' crates/ml/src/ | grep -v '^\s*//' + ``` + + Expected sites: `Default::default()` impl at `gpu_experience_collector.rs:301`, `mod.rs:895`, `training_loop.rs:1662`. + +### Task P3-T2: Modify ExperienceCollectorConfig struct + +**Files:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` + +- [ ] **Step 1: Replace noise_sigma field with per-branch array** + + Find: + ```rust + /// Advantage noise std for exploration (replaces dead NoisyLinear). + /// Per-action Gaussian noise breaks the dueling symmetry trap. + /// 0.0 = disabled. Default: 0.1. + pub noise_sigma: f32, + ``` + + Replace with: + ```rust + /// SP6 Pearl 3: per-branch advantage noise std for exploration. + /// Branch order: [direction(0), magnitude(1), order(2), urgency(3)]. + /// 0.0 per-branch disables noise for that branch. Default: [0.1; 4]. + pub noise_sigma_per_branch: [f32; 4], + ``` + +- [ ] **Step 2: Update Default impl** + + Find: + ```rust + noise_sigma: 0.1, + ``` + + Replace with: + ```rust + noise_sigma_per_branch: [0.1; 4], + ``` + +- [ ] **Step 3: Update the guard and launch site in collect_experiences_gpu** + + Find (~line 3624): + ```rust + if config.noise_sigma > 0.0 { + ``` + + Replace with: + ```rust + let any_branch_noise = config.noise_sigma_per_branch.iter().any(|&s| s > 0.0); + if any_branch_noise { + ``` + + Find (~line 3640): + ```rust + .arg(&config.noise_sigma) + ``` + + The kernel currently takes a single `float noise_sigma`. After the kernel is updated in T3 to take `const float* noise_sigma_per_branch`, this arg changes. For now, leave this line — it will be updated in T3-Step 4 after the kernel signature is confirmed. + +- [ ] **Step 4: Run cargo check to surface all remaining noise_sigma references** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep 'noise_sigma' | head -20 + ``` + + Expected: compile errors pointing to `mod.rs:895` and `training_loop.rs:1747`. + +### Task P3-T3: Update add_advantage_noise kernel + +**Files:** `crates/ml/src/cuda_pipeline/experience_kernels.cu` + +- [ ] **Step 1: Replace the kernel signature and body** + + Find the entire `add_advantage_noise` kernel (lines ~3916-3935). Replace it with: + + ```c + /** + * SP6 Pearl 3: Add per-branch Gaussian noise to Q-values for exploration. + * Replaces dead NoisyLinear. Uses Philox PRNG seeded from (thread_id, step). + * Each action slot receives noise scaled by its branch's sigma. + * branch_sizes[4] = {b0, b1, b2, b3}: cumulative offsets derive branch index. + * + * Grid: ceil(N * total_actions / 256), Block: 256. + */ + extern "C" __global__ void add_advantage_noise( + float* __restrict__ q_values, /* [N, total_actions] — modified in-place */ + const float* __restrict__ noise_sigma, /* [4] per-branch std dev */ + int N, /* number of episodes */ + int total_actions, /* 13 = 4+3+3+3 */ + unsigned int seed, /* monotonic step counter for decorrelation */ + int b0_size, /* branch 0 action count (direction = 4) */ + int b1_size, /* branch 1 action count (magnitude = 3) */ + int b2_size, /* branch 2 action count (order = 3) */ + int b3_size /* branch 3 action count (urgency = 3) */ + ) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + int total = N * total_actions; + if (tid >= total) return; + + int action = tid % total_actions; + + /* Derive branch index from action offset */ + int branch_idx; + int b01 = b0_size + b1_size; + int b012 = b01 + b2_size; + if (action < b0_size) branch_idx = 0; + else if (action < b01) branch_idx = 1; + else if (action < b012) branch_idx = 2; + else branch_idx = 3; + + float sigma = noise_sigma[branch_idx]; + if (sigma <= 0.0f) return; + + int episode = tid / total_actions; + float u1 = fmaxf(philox_uniform(episode, (int)seed, 4000 + action * 2), 1e-6f); + float u2 = philox_uniform(episode, (int)seed, 4000 + action * 2 + 1); + float gauss = sqrtf(-2.0f * logf(u1)) * cosf(6.2831853f * u2); + + q_values[tid] += sigma * gauss; + } + ``` + +### Task P3-T4: Update Rust launcher to pass [f32; 4] array + +**Files:** `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` + +- [ ] **Step 1: Create a mapped-pinned device buffer for noise_sigma_per_branch** + + The kernel now takes `const float*` — it needs a device pointer. The cleanest approach for a 4-float array that lives only for the kernel call duration is to allocate a small mapped-pinned buffer at experience collector construction time and write the 4 values into it before each kernel launch. + + Find the `GpuExperienceCollector` struct field block. Add: + ```rust + /// SP6 Pearl 3: [4] mapped-pinned device buffer for per-branch noise sigma. + noise_sigma_dev: crate::cuda_pipeline::mapped_pinned::MappedF32Buffer, + ``` + + Find the constructor (`GpuExperienceCollector::new` or similar). Add initialization of this field using `MappedF32Buffer::new` with length 4. Follow the pattern used for other mapped-pinned buffers in the same file. + +- [ ] **Step 2: Update collect_experiences_gpu to write sigma values then launch** + + Replace the launch block (near line 3636-3648) with: + ```rust + // Write per-branch sigma to mapped-pinned device buffer. + { + let sigma_host = config.noise_sigma_per_branch; + self.noise_sigma_dev.write_host_slice(&sigma_host) + .map_err(|e| MLError::ModelError(format!("noise_sigma_dev write: {e}")))?; + } + let sigma_dev_ptr = self.noise_sigma_dev.dev_ptr; + let b0 = self.branch_sizes[0] as i32; + let b1 = self.branch_sizes[1] as i32; + let b2 = self.branch_sizes[2] as i32; + let b3 = self.branch_sizes[3] as i32; + unsafe { + self.stream + .launch_builder(&self.noise_kernel) + .arg(&mut self.q_values) + .arg(&sigma_dev_ptr) + .arg(&n_i32) + .arg(&total_actions_i32) + .arg(&seed) + .arg(&b0) + .arg(&b1) + .arg(&b2) + .arg(&b3) + .launch(noise_cfg) + .map_err(|e| MLError::ModelError(format!( + "add_advantage_noise t={t}: {e}" + )))?; + } + ``` + +### Task P3-T5: Fix remaining call sites + +**Files:** +- `crates/ml/src/cuda_pipeline/mod.rs` (line ~895) +- `crates/ml/src/trainers/dqn/trainer/training_loop.rs` (line ~1747) + +- [ ] **Step 1: Fix mod.rs test usage** + + Find at `mod.rs:895`: + ```rust + let cfg = ExperienceCollectorConfig::default(); + ``` + + This uses `Default::default()` which now supplies `noise_sigma_per_branch: [0.1; 4]`. No change needed if it is using the default. + + If there is any explicit `noise_sigma: ...` field in that block, change it to `noise_sigma_per_branch: [0.1; 4]`. + +- [ ] **Step 2: Update training_loop.rs noise_sigma block** + + Find near line 1747: + ```rust + noise_sigma: self.fused_ctx.as_ref().map(|ctx| { + let mean = (0..4_usize) + .map(|b| ctx.read_isv_signal_at(NOISY_SIGMA_BASE + b)) + .sum::() + / 4.0_f32; + mean.max(0.01_f32) + }).unwrap_or(self.hyperparams.noise_sigma as f32), + ``` + + Replace with: + ```rust + // SP6 Pearl 3: read per-branch sigma directly from ISV[210..214). + // No averaging. Cold-start floor 0.01 = Invariant 1 carve-out. + noise_sigma_per_branch: { + let mut arr = [self.hyperparams.noise_sigma as f32; 4]; + if let Some(ctx) = self.fused_ctx.as_ref() { + for b in 0..4_usize { + arr[b] = ctx.read_isv_signal_at(NOISY_SIGMA_BASE + b).max(0.01_f32); + } + } + arr + }, + ``` + +- [ ] **Step 3: Run cargo check** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Expected: no errors. + +### Task P3-T6: Run smoke test + +- [ ] **Step 1: Run cargo build** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo build -p ml --lib --release 2>&1 | tail -5 + ``` + + Expected: `Finished`. + +- [ ] **Step 2: Run smoke test** + + ```bash + FOXHUNT_TEST_DATA=test_data/futures-baseline \ + FOXHUNT_MBP10_DIR=test_data/mbp10 \ + FOXHUNT_TRADES_DIR=test_data/trades \ + CARGO_INCREMENTAL=0 \ + cargo test -p ml --lib -- smoke_tests::multi_fold_convergence --ignored --nocapture 2>&1 | tail -40 + ``` + + Expected: PASS. Kill at first HEALTH_DIAG line — look for per-branch sigma values diverging from each other (direction sigma > magnitude sigma expected after Pearl 3 warmup). + +### Task P3-T7: Atomic commit + +- [ ] **Step 1: Stage changed files** + + ```bash + git add \ + crates/ml/src/cuda_pipeline/gpu_experience_collector.rs \ + crates/ml/src/cuda_pipeline/experience_kernels.cu \ + crates/ml/src/cuda_pipeline/mod.rs \ + crates/ml/src/trainers/dqn/trainer/training_loop.rs + ``` + +- [ ] **Step 2: Confirm no other files are staged** + + ```bash + git diff --cached --name-only + ``` + + Expected: exactly the 4 files above. + +- [ ] **Step 3: Commit** + + ```bash + git commit -m "$(cat <<'EOF' + feat(dqn): SP6 Pearl 3 — per-branch NoisyNet sigma array + + Replace ExperienceCollectorConfig.noise_sigma: f32 with + noise_sigma_per_branch: [f32; 4]. add_advantage_noise kernel derives + branch index from action offset using b0/b1/b2/b3 sizes, applies + per-branch sigma. training_loop.rs reads ISV[210..214) directly + (no averaging). Default impl supplies [0.1; 4] for cold-start. + + Co-Authored-By: Claude Sonnet 4.6 + EOF + )" + ``` + +--- + +## Sub-Project 3: Pearl 5 — IQN τ Per-Branch Schedule + +**Estimated time:** 3-4 hours + +**Branch:** `sp6-pearl-5` + +**Worktree:** `../foxhunt-sp6-pearl-5` + +**Approach: B** (4 independent forward passes with per-branch τ schedules). Approach C (kernel `[B, 4, N]` rework) is deferred — it would require changing the CUDA Graph node ABI and is higher risk than the ~15-25% compute overhead of 4 sequential passes. + +### Task P5-T1: Read the current code before any edits + +- [ ] **Step 1: Read refresh_taus_from_isv** + + Read `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` lines 1955-2017. Confirm: + - Takes `isv_per_branch_taus: &[f32; 20]` + - Averages 4 branches per quantile into a single `[5]` tau array + - Tiles into `online_taus [B, N]` and `target_taus [B, N]` + - Recomputes `cos_features [D, N]` on host (cos of PI * (dim+1) * tau) + - Uploads all 3 via mapped-pinned + +- [ ] **Step 2: Read the call site in fused_training.rs** + + ```bash + grep -n 'refresh_taus_from_isv\|IQN_TAU_BASE' \ + crates/ml/src/trainers/dqn/fused_training.rs + ``` + + Read lines ~2014-2023. Confirm: called once per epoch-boundary refresh inside the `run_full_step` epoch section; reads `ISV[IQN_TAU_BASE..+20]`. + +- [ ] **Step 3: Find IQN forward pass call in run_full_step** + + ```bash + grep -n 'run_iqn_forward\|iqn\.forward\|gpu_iqn\|launch_iqn' \ + crates/ml/src/trainers/dqn/fused_training.rs | head -20 + ``` + + Identify the IQN forward call structure. Understand how many times IQN forward is called per step and in which code path. + +- [ ] **Step 4: Confirm GpuIqnHead struct fields for taus** + + ```bash + grep -n 'online_taus\|target_taus\|cos_features' \ + crates/ml/src/cuda_pipeline/gpu_iqn_head.rs | head -20 + ``` + + Confirm: `online_taus: CudaSlice`, `target_taus: CudaSlice`, `cos_features: CudaSlice` are existing fields. + +### Task P5-T2: Add per-branch tau state to GpuIqnHead + +**Files:** `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` + +- [ ] **Step 1: Add per-branch tau cache fields to the struct** + + Find the `GpuIqnHead` struct definition. Add 4 new `CudaSlice` triplets — one per branch: + + ```rust + /// SP6 Pearl 5: per-branch tau tensors [B, N] (same allocation size as online_taus). + online_taus_branch: [CudaSlice; 4], + target_taus_branch: [CudaSlice; 4], + /// SP6 Pearl 5: per-branch cos_features [D, N]. + cos_features_branch: [CudaSlice; 4], + ``` + + > Implementation note: `[CudaSlice; 4]` is not trivially constructible because `CudaSlice` does not impl `Default`. Allocate 4 of each in the constructor (T2-Step 2) and store them in a manually-constructed array using `std::array::from_fn` or explicit 4-element initialization. + +- [ ] **Step 2: Allocate per-branch buffers in the constructor** + + Find `GpuIqnHead::new` (or equivalent constructor). After the existing `online_taus`, `target_taus`, `cos_features` allocations, add: + + ```rust + // SP6 Pearl 5: allocate per-branch tau/cos buffers (same size as main buffers). + let b = config.batch_size; + let n = config.num_quantiles; + let d = config.embed_dim; + let online_taus_branch = std::array::from_fn(|_| { + stream.alloc_zeros::(b * n) + .expect("alloc online_taus_branch") + }); + let target_taus_branch = std::array::from_fn(|_| { + stream.alloc_zeros::(b * n) + .expect("alloc target_taus_branch") + }); + let cos_features_branch = std::array::from_fn(|_| { + stream.alloc_zeros::(d * n) + .expect("alloc cos_features_branch") + }); + ``` + + Add these 3 fields to the struct literal return. + +- [ ] **Step 3: Verify cargo check after struct changes** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Fix any field-missing or type errors before proceeding. + +### Task P5-T3: Add refresh_taus_for_branch method + +**Files:** `crates/ml/src/cuda_pipeline/gpu_iqn_head.rs` + +- [ ] **Step 1: Add the per-branch refresh method** + + Add immediately after the closing `}` of `refresh_taus_from_isv`: + + ```rust + /// SP6 Pearl 5: Refresh tau tensors for a single branch `b` from its + /// 5-quantile schedule `tau5`. Uploads to `online_taus_branch[b]`, + /// `target_taus_branch[b]`, and `cos_features_branch[b]`. + /// + /// Called 4 times per epoch refresh (once per branch) from + /// `FusedTrainingCtx::refresh_iqn_taus_per_branch`. + pub fn refresh_taus_for_branch( + &mut self, + branch_idx: usize, + tau5: &[f32; 5], + ) { + let n = self.config.num_quantiles; + let b = self.config.batch_size; + let d = self.config.embed_dim; + + // Apply cold-start floor per quantile. + let mut taus = [0.0_f32; 5]; + for q in 0..n { + taus[q] = if tau5[q] < 1e-6_f32 { FIXED_TAUS[q] } else { tau5[q].min(1.0_f32) }; + } + + // Recompute cosine features [D, N]. + let mut cos_feat_host = vec![0.0_f32; d * n]; + for q in 0..n { + let tau_q = taus[q]; + for dim in 0..d { + cos_feat_host[dim + q * d] = + (std::f32::consts::PI * ((dim + 1) as f32) * tau_q).cos(); + } + } + + // Tile taus [B, N]. + let mut tiled = Vec::with_capacity(b * n); + for _ in 0..b { + tiled.extend_from_slice(&taus[..n]); + } + + if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( + &self.stream, &tiled, &mut self.online_taus_branch[branch_idx], + ) { + tracing::warn!("SP6 Pearl 5 online_taus_branch[{branch_idx}] upload failed: {e}"); + return; + } + if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( + &self.stream, &tiled, &mut self.target_taus_branch[branch_idx], + ) { + tracing::warn!("SP6 Pearl 5 target_taus_branch[{branch_idx}] upload failed: {e}"); + return; + } + if let Err(e) = super::mapped_pinned::upload_f32_via_pinned( + &self.stream, &cos_feat_host, &mut self.cos_features_branch[branch_idx], + ) { + tracing::warn!("SP6 Pearl 5 cos_features_branch[{branch_idx}] upload failed: {e}"); + } + } + ``` + +- [ ] **Step 2: Add swap-in / swap-out helpers for per-branch forward pass** + + The existing IQN forward pass reads from `self.online_taus`, `self.cos_features`. To reuse it for a per-branch pass, temporarily swap in the branch's buffers, run the forward, swap back. + + Add these two helpers: + + ```rust + /// SP6 Pearl 5: Temporarily replace the main tau/cos buffers with branch b's + /// buffers for a single forward pass. + pub fn activate_branch_taus(&mut self, branch_idx: usize) { + std::mem::swap(&mut self.online_taus, &mut self.online_taus_branch[branch_idx]); + std::mem::swap(&mut self.target_taus, &mut self.target_taus_branch[branch_idx]); + std::mem::swap(&mut self.cos_features, &mut self.cos_features_branch[branch_idx]); + } + + /// SP6 Pearl 5: Restore the main tau/cos buffers after a per-branch forward pass. + pub fn deactivate_branch_taus(&mut self, branch_idx: usize) { + std::mem::swap(&mut self.online_taus, &mut self.online_taus_branch[branch_idx]); + std::mem::swap(&mut self.target_taus, &mut self.target_taus_branch[branch_idx]); + std::mem::swap(&mut self.cos_features, &mut self.cos_features_branch[branch_idx]); + } + ``` + + > Note: `activate_branch_taus` followed by `deactivate_branch_taus` with the same `branch_idx` is its own inverse — the buffers are swapped back correctly. This is a safe pattern because `mem::swap` exchanges ownership without any allocation. + +- [ ] **Step 3: Verify cargo check** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Expected: no errors. + +### Task P5-T4: Update fused_training.rs to use per-branch refresh and forward + +**Files:** `crates/ml/src/trainers/dqn/fused_training.rs` + +- [ ] **Step 1: Read the current IQN tau refresh + forward block** + + ```bash + grep -n 'refresh_taus_from_isv\|IQN_TAU_BASE\|isv_taus\|iqn\.refresh' \ + crates/ml/src/trainers/dqn/fused_training.rs + ``` + + Read the block (lines ~2014-2023) in full. + +- [ ] **Step 2: Replace the averaging tau refresh with a 4-branch per-branch refresh** + + Current code (approximately): + ```rust + // SP5 Pearl 5: refresh τ schedule from ISV[IQN_TAU_BASE=250..270). + if let Some(ref mut iqn) = self.gpu_iqn { + let mut isv_taus = [0.0_f32; 20]; + for i in 0..20 { + isv_taus[i] = self.trainer.read_isv_signal_at(IQN_TAU_BASE + i); + } + iqn.refresh_taus_from_isv(&isv_taus) + } + ``` + + Replace with: + ```rust + // SP6 Pearl 5: refresh per-branch τ schedules from ISV[250..270). + // ISV layout: branch b, quantile q at IQN_TAU_BASE + b*5 + q. + if let Some(ref mut iqn) = self.gpu_iqn { + for branch_idx in 0..4_usize { + let mut tau5 = [0.0_f32; 5]; + for q in 0..5_usize { + tau5[q] = self.trainer.read_isv_signal_at(IQN_TAU_BASE + branch_idx * 5 + q); + } + iqn.refresh_taus_for_branch(branch_idx, &tau5); + } + // Also refresh the main (averaged) tau buffer for backward compatibility + // with existing IQN forward passes that read online_taus/cos_features directly. + // Compute mean-of-4-branches tau as the shared schedule. + let mut isv_taus = [0.0_f32; 20]; + for i in 0..20 { + isv_taus[i] = self.trainer.read_isv_signal_at(IQN_TAU_BASE + i); + } + iqn.refresh_taus_from_isv(&isv_taus); + } + ``` + + > Note: keeping `refresh_taus_from_isv` for the main buffer ensures backward compat with any training-path code that still reads `online_taus` directly (e.g. CVaR scaling). The per-branch passes below use `activate_branch_taus` / `deactivate_branch_taus`. + +- [ ] **Step 3: Find and read the existing IQN forward pass in run_full_step** + + ```bash + grep -n 'gpu_iqn\|run_iqn\|iqn_forward\|iqn_ok\|iqn\.train' \ + crates/ml/src/trainers/dqn/fused_training.rs | head -30 + ``` + + Identify the main IQN training path. It should be structured as: + - `if iqn_ok { if let Some(ref mut iqn) = self.gpu_iqn { iqn.train_step(...) } }` + + Read that block in full to understand what `train_step` (or equivalent) does and what parameters it needs. + +- [ ] **Step 4: Wrap the IQN training step in a 4-branch loop** + + The existing IQN forward runs once with the averaged tau. After the per-branch refresh in T4-Step 2 is wired, the training step must run 4 times — once per branch — with the branch's tau active. + + Find the `if iqn_ok { ... }` block. Wrap its content in a branch loop: + + ```rust + if iqn_ok { + // SP6 Pearl 5: run IQN forward pass for each branch with its per-branch τ schedule. + for branch_idx in 0..4_usize { + if let Some(ref mut iqn) = self.gpu_iqn { + iqn.activate_branch_taus(branch_idx); + // [existing IQN train_step / apply_iqn_trunk_gradient block here] + // Pass branch_idx to any IQN function that accumulates per-branch gradients. + iqn.deactivate_branch_taus(branch_idx); + } + } + } + ``` + + > Critical: verify that running the IQN backward 4 times does NOT 4× the IQN gradient accumulation in `grad_buf`. The IQN backward SAXPY (`apply_iqn_trunk_gradient`) accumulates into `grad_buf`. If called 4 times, `grad_buf` will contain 4× the IQN contribution, overwhelming C51 and CQL. To prevent this, the `iqn_budget` (or `iqn_trunk`) passed to `apply_iqn_trunk_gradient` must be divided by 4 on each sub-pass: + > + > ```rust + > self.trainer.apply_iqn_trunk_gradient( + > d_h_s2_ptr, &mut self.online_dueling, iqn_trunk / 4.0_f32, + > ) + > ``` + > + > Each of the 4 passes contributes `iqn_trunk / 4` to the trunk gradient, summing to `iqn_trunk` total — matching the prior single-pass behavior at the trunk level while each pass uses its branch's per-branch τ schedule. + +- [ ] **Step 5: Verify cargo check** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo check -p ml --lib 2>&1 | grep '^error' | head -20 + ``` + + Expected: no errors. + +### Task P5-T5: Run smoke test + +- [ ] **Step 1: Run cargo build** + + ```bash + SQLX_OFFLINE=true CARGO_INCREMENTAL=0 cargo build -p ml --lib --release 2>&1 | tail -5 + ``` + + Expected: `Finished`. + +- [ ] **Step 2: Run smoke test** + + ```bash + FOXHUNT_TEST_DATA=test_data/futures-baseline \ + FOXHUNT_MBP10_DIR=test_data/mbp10 \ + FOXHUNT_TRADES_DIR=test_data/trades \ + CARGO_INCREMENTAL=0 \ + cargo test -p ml --lib -- smoke_tests::multi_fold_convergence --ignored --nocapture 2>&1 | tail -40 + ``` + + Expected: PASS. Kill at first HEALTH_DIAG line per `feedback_kill_runs_on_anomaly_quickly`. If step time increases by >25%, investigate whether the 4-pass IQN loop is accumulating gradients incorrectly (check `iqn_grad` in HEALTH_DIAG decomp output). + + If step time increase causes timeout before first HEALTH_DIAG: this is acceptable for SP6. The perf regression is documented in the spec as the known cost of Approach B. + +### Task P5-T6: Atomic commit + +- [ ] **Step 1: Stage changed files** + + ```bash + git add \ + crates/ml/src/cuda_pipeline/gpu_iqn_head.rs \ + crates/ml/src/trainers/dqn/fused_training.rs + ``` + +- [ ] **Step 2: Confirm no other files are staged** + + ```bash + git diff --cached --name-only + ``` + + Expected: exactly the 2 files above. + +- [ ] **Step 3: Commit** + + ```bash + git commit -m "$(cat <<'EOF' + feat(dqn): SP6 Pearl 5 — per-branch IQN tau schedule (Approach B) + + GpuIqnHead now holds 4 per-branch [B,N] tau+cos buffer triplets. + refresh_taus_for_branch(b, tau5) uploads branch b's 5-quantile + schedule; activate/deactivate_branch_taus swap the main buffers + for a per-branch forward pass. fused_training runs 4 IQN forward + passes per step (one per branch), each contributing iqn_trunk/4 to + grad_buf so total IQN gradient magnitude is unchanged. + + Co-Authored-By: Claude Sonnet 4.6 + EOF + )" + ``` + +--- + +## Sequential Merge After All 3 Worktrees Complete + +Once all 3 sub-project worktrees have passed their smoke tests and produced their atomic commits: + +```bash +# On the main branch (sp5-magnitude-differentiation) +git checkout sp5-magnitude-differentiation + +# Merge order: Pearl 2 first (most changes), Pearl 3 second, Pearl 5 third +git merge sp6-pearl-2 --no-ff -m "merge(sp6): Pearl 2 per-branch loss budgets" +git merge sp6-pearl-3 --no-ff -m "merge(sp6): Pearl 3 per-branch NoisyNet sigma" +git merge sp6-pearl-5 --no-ff -m "merge(sp6): Pearl 5 per-branch IQN tau schedule" +``` + +If conflicts occur in `fused_training.rs`: +- Pearl 2 edits lines ~1882-1897 (compute_adaptive_budgets call + SAXPY dispatch) and ~3281-3333 (function body + accessors). +- Pearl 3 edits lines ~1747-1754 (noise_sigma construction). +- Pearl 5 edits lines ~2014-2023 (tau refresh) and the IQN training loop. + +These ranges are non-overlapping. If git produces a conflict, it means line numbers drifted — accept the edit from the later merge manually using the non-overlapping semantics above. + +After all merges: + +```bash +# Final smoke test on the merged branch +FOXHUNT_TEST_DATA=test_data/futures-baseline \ +FOXHUNT_MBP10_DIR=test_data/mbp10 \ +FOXHUNT_TRADES_DIR=test_data/trades \ +CARGO_INCREMENTAL=0 \ +cargo test -p ml --lib -- smoke_tests::multi_fold_convergence --ignored --nocapture 2>&1 | tail -40 + +# Push +git push origin sp5-magnitude-differentiation +``` + +--- + +## Verification Checklist (post-merge) + +- [ ] `cargo check -p ml --lib` — zero errors +- [ ] `cargo build -p ml --lib --release` — zero errors +- [ ] `multi_fold_convergence` smoke — PASS +- [ ] HEALTH_DIAG `c51_budget_per_branch: dir/mag/ord/urg` shows non-uniform values (Pearl 2) +- [ ] HEALTH_DIAG `iqn_budget_per_branch` shows non-uniform values (Pearl 2) +- [ ] HEALTH_DIAG `cql_budget_per_branch` shows non-uniform values (Pearl 2) +- [ ] No `mean(ISV[BUDGET_*..+4])` averaging pattern in `fused_training.rs` (Pearl 2) +- [ ] No `mean(ISV[NOISY_SIGMA_BASE..+4])` averaging pattern in `training_loop.rs` (Pearl 3) +- [ ] No `mean(isv_per_branch_taus[q..20])` averaging pattern in `gpu_iqn_head.rs` (Pearl 5) +- [ ] No TODO/FIXME/XXX in any changed file +- [ ] `grep -rn 'noise_sigma:' crates/ml/src/` — zero hits (old scalar field name fully retired)