diff --git a/docs/superpowers/plans/2026-04-14-direction-conditioned-magnitude.md b/docs/superpowers/plans/2026-04-14-direction-conditioned-magnitude.md new file mode 100644 index 000000000..c98236455 --- /dev/null +++ b/docs/superpowers/plans/2026-04-14-direction-conditioned-magnitude.md @@ -0,0 +1,740 @@ +# Direction-Conditioned Magnitude Head 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:** Feed direction E[Q] (3 scalars) as extra input to the magnitude branch FC, so the model can distinguish Long+Full from Flat+Full. + +**Architecture:** A small CUDA kernel computes E[Q] from direction logits and concatenates with h_s2 into a scratch buffer `[B, SH2+3]`. The magnitude branch FC GEMM uses this wider input (K=SH2+3). Backward splits dX back into d_h_s2 and d_Q_dir (discarded — detached, no backprop to direction). All other branches unchanged. + +**Tech Stack:** CUDA 12.4, Rust (cudarc), cuBLASLt SGEMM, nvcc via build.rs + +**Baseline:** Commits through b2eb1b774 — Layers 1+2+4 + asymmetric spread + per-branch scaling. + +--- + +### Task 1: CUDA Kernel — mag_concat_qdir + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/experience_kernels.cu` + +- [ ] **Step 1: Add the kernel to experience_kernels.cu** + +Append at the end of the file (before any trailing `#endif` or after the last kernel): + +```cuda +/* ================================================================== */ +/* Kernel: mag_concat_qdir */ +/* ================================================================== */ + +/** + * Concatenate h_s2 with direction E[Q] for magnitude branch conditioning. + * + * For each sample b: + * 1. Copy h_s2[b, 0..SH2] → concat_out[b, 0..SH2] + * 2. Compute E[Q_dir_a] for a=0,1,2 from dir_logits via softmax × support + * 3. Write E[Q_dir_0..2] → concat_out[b, SH2..SH2+3] + * + * Grid: ceil(B/256), Block: 256. One thread per sample. + */ +extern "C" __global__ void mag_concat_qdir( + const float* __restrict__ h_s2, /* [B, SH2] trunk activation */ + const float* __restrict__ dir_logits_v, /* [B, NA] value logits (shared head) */ + const float* __restrict__ dir_logits_b, /* [B*b0*NA] branch 0 logits (direction) */ + float* __restrict__ concat_out, /* [B, SH2+3] output */ + int B, + int SH2, + int NA, /* num_atoms */ + int b0_size, /* direction branch size (3) */ + const float* __restrict__ per_sample_support /* [B, 3]: v_min, v_max, delta_z */ +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= B) return; + + /* Step 1: copy h_s2 row */ + int out_stride = SH2 + 3; + for (int i = 0; i < SH2; i++) { + concat_out[b * out_stride + i] = h_s2[b * SH2 + i]; + } + + /* Step 2: compute E[Q_dir] for each direction action */ + float v_min = per_sample_support[b * 3 + 0]; + float dz = per_sample_support[b * 3 + 2]; + + for (int a = 0; a < b0_size && a < 3; a++) { + /* Combined value + advantage logits for direction action a */ + /* Branch logits layout: BRANCH-MAJOR [N × b0 × NA], so action a for sample b is: + * dir_logits_b[b * b0_size * NA + a * NA + z] */ + float max_logit = -1e30f; + for (int z = 0; z < NA; z++) { + float logit = dir_logits_v[b * NA + z] + dir_logits_b[b * b0_size * NA + a * NA + z]; + if (logit > max_logit) max_logit = logit; + } + float sum_exp = 0.0f; + for (int z = 0; z < NA; z++) { + float logit = dir_logits_v[b * NA + z] + dir_logits_b[b * b0_size * NA + a * NA + z]; + sum_exp += expf(logit - max_logit); + } + float eq = 0.0f; + for (int z = 0; z < NA; z++) { + float logit = dir_logits_v[b * NA + z] + dir_logits_b[b * b0_size * NA + a * NA + z]; + float prob = expf(logit - max_logit) / sum_exp; + float z_val = v_min + (float)z * dz; + eq += prob * z_val; + } + concat_out[b * out_stride + SH2 + a] = eq; + } +} +``` + +- [ ] **Step 2: Verify nvcc compilation** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +Expected: `Finished` with no nvcc errors (experience_kernels.cu compiles via build.rs with common header). + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/experience_kernels.cu +git commit -m "feat: mag_concat_qdir kernel — direction E[Q] conditioning for magnitude branch + +Computes softmax E[Q] for 3 direction actions from combined value+advantage +logits, concatenates with h_s2 trunk activation to produce [B, SH2+3] input +for the magnitude branch FC layer. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +### Task 2: Widen w_b1fc in param_sizes and allocate concat buffers + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Change param_sizes for w_b1fc** + +In `compute_param_sizes` (around line 392), change: + +```rust +cfg.adv_h * cfg.shared_h2, // [12] w_b1fc +``` + +To: + +```rust +cfg.adv_h * (cfg.shared_h2 + 3), // [12] w_b1fc (direction-conditioned: SH2+3 input) +``` + +- [ ] **Step 2: Add mag_concat_buf and d_mag_concat_buf fields to GpuDqnTrainer struct** + +Find the struct fields near `save_h_b1` (around line 587). Add after the branch hidden buffers: + +```rust +/// Magnitude branch concat input [B, SH2+3]: [h_s2; Q_dir] for direction conditioning. +mag_concat_buf: CudaSlice, +/// Backward scratch for magnitude concat dX [B, SH2+3]. +d_mag_concat_buf: CudaSlice, +``` + +- [ ] **Step 3: Add mag_concat pointers to CachedPtrs struct** + +Find the `CachedPtrs` struct (search for `pub(crate) struct CachedPtrs`). Add: + +```rust +pub(crate) mag_concat_buf: u64, +pub(crate) d_mag_concat_buf: u64, +``` + +- [ ] **Step 4: Allocate buffers in constructor** + +Find where `save_h_b1` is allocated (around line 2370). Add after: + +```rust +let mag_concat_dim = config.shared_h2 + 3; +let mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "mag_concat_buf")?; +let d_mag_concat_buf = alloc_f32(&stream, b * mag_concat_dim + kt, "d_mag_concat_buf")?; +``` + +(`kt` is the padding tail already used by surrounding allocations — find its value from the `save_h_b0` line.) + +- [ ] **Step 5: Initialize in CachedPtrs** + +In the CachedPtrs initialization block (around line 2907), add: + +```rust +mag_concat_buf: mag_concat_buf.raw_ptr(), +d_mag_concat_buf: d_mag_concat_buf.raw_ptr(), +``` + +- [ ] **Step 6: Add to struct literal** + +In the struct literal (around line 3059), add after `save_h_b1`: + +```rust +mag_concat_buf, +d_mag_concat_buf, +``` + +- [ ] **Step 7: Load mag_concat_qdir kernel** + +Add a field to the struct: + +```rust +mag_concat_kernel: CudaFunction, +``` + +In the constructor, after loading `expected_q_kernel` (search for `compile_expected_q_kernel`), load from the experience kernels cubin: + +```rust +let mag_concat_kernel = exp_cubin_module.load_function("mag_concat_qdir") + .map_err(|e| MLError::ModelError(format!("mag_concat_qdir load: {e}")))?; +``` + +NOTE: You'll need to get a reference to the experience kernels module. The pattern varies — the experience kernels cubin is loaded in `gpu_experience_collector.rs`, not `gpu_dqn_trainer.rs`. The trainer loads kernels from `DQN_UTILITY_CUBIN` and `experience_kernels.cubin`. Search for how `expected_q_kernel` is loaded — it uses `compile_expected_q_kernel` which loads `experience_kernels.cubin`: + +```rust +fn compile_expected_q_kernel(stream: &Arc) -> Result { +``` + +Add a similar function or extend `compile_expected_q_kernel` to also return the mag_concat function. + +Add `mag_concat_kernel` to the struct literal. + +- [ ] **Step 8: Add a method to launch the concat kernel** + +```rust +/// Build [h_s2; Q_dir] concat buffer for magnitude branch conditioning. +/// Must be called AFTER direction branch forward (branch 0 logits are ready). +fn launch_mag_concat(&self, v_logits_ptr: u64, b_logits_ptr: u64) -> Result<(), MLError> { + let b = self.config.batch_size; + let sh2 = self.config.shared_h2 as i32; + let na = self.config.num_atoms as i32; + let b0 = self.config.branch_0_size as i32; + let n = b as i32; + let blocks = ((b as u32 + 255) / 256).max(1); + let h_s2_ptr = self.ptrs.save_h_s2; + let concat_ptr = self.ptrs.mag_concat_buf; + let support_ptr = self.per_sample_support_ptr; + unsafe { + self.stream + .launch_builder(&self.mag_concat_kernel) + .arg(&h_s2_ptr) + .arg(&v_logits_ptr) + .arg(&b_logits_ptr) + .arg(&concat_ptr) + .arg(&n) + .arg(&sh2) + .arg(&na) + .arg(&b0) + .arg(&support_ptr) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("mag_concat_qdir: {e}")))?; + } + Ok(()) +} +``` + +- [ ] **Step 9: Verify build** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +Expected: clean build. + +- [ ] **Step 10: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "feat: widen w_b1fc param_sizes + allocate mag_concat buffers + +w_b1fc grows from [AH, SH2] to [AH, SH2+3] for direction conditioning. +mag_concat_buf [B, SH2+3] and d_mag_concat_buf allocated for forward/backward. +mag_concat_qdir kernel loaded and launch method added. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +### Task 3: Forward Pass — Magnitude Branch Uses Concat Input + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_forward.rs` +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` (callers) + +- [ ] **Step 1: Add mag_concat_dim field to CublasForward** + +In `crates/ml/src/cuda_pipeline/batched_forward.rs`, add a field to `CublasForward`: + +```rust +/// Magnitude branch FC input dimension: SH2+3 (direction conditioned). +mag_concat_dim: usize, +``` + +Initialize it in the constructor: + +```rust +mag_concat_dim: shared_h2 + 3, +``` + +- [ ] **Step 2: Add wider GEMM shape to cache** + +In the constructor of `CublasForward`, after the existing `unique_shapes.push((adv_h, ...))` for branch FC (around line 263), add: + +```rust +// Magnitude branch FC uses wider input: K=SH2+3 (direction-conditioned) +unique_shapes.push((adv_h, batch_size, shared_h2 + 3, shared_h2 + 3)); +``` + +Also add to `relu_bias_shapes`: + +```rust +(adv_h, batch_size, shared_h2 + 3, shared_h2 + 3), // h_b1 (mag, conditioned) +``` + +- [ ] **Step 3: Add mag_concat_ptr parameter to forward_online_raw** + +Change the signature of `forward_online_raw` to accept the concat pointer: + +```rust +pub fn forward_online_raw( + &self, + stream: &Arc, + states_ptr: u64, + w_ptrs: &[u64; super::gpu_dqn_trainer::NUM_WEIGHT_TENSORS], + h_s1_ptr: u64, h_s2_ptr: u64, h_v_ptr: u64, + h_b0_ptr: u64, h_b1_ptr: u64, h_b2_ptr: u64, h_b3_ptr: u64, + v_logits_ptr: u64, b_logits_ptr: u64, + mag_concat_ptr: u64, // [B, SH2+3] or 0 for no conditioning +) -> Result<(), MLError> { +``` + +- [ ] **Step 4: Branch 1 uses concat input in multi-stream path** + +In `forward_online_raw`, inside the multi-stream branch loop (`for d in 0..4`), replace the magnitude branch FC GEMM. Change: + +```rust +if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], + self.adv_h, b, self.shared_h2, self.shared_h2, bws, bwss, "h_bd").is_err() +{ + self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], h_s2_ptr, branch_h_ptrs[d], self.adv_h, b, self.shared_h2, d, "h_bd")?; + self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; +} +``` + +To use conditional input and K dimension: + +```rust +let (fc_input, fc_k) = if d == 1 && mag_concat_ptr != 0 { + (mag_concat_ptr, self.mag_concat_dim) +} else { + (h_s2_ptr, self.shared_h2) +}; +if self.sgemm_f32_fused_relu_bias(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], + self.adv_h, b, fc_k, fc_k, bws, bwss, "h_bd").is_err() +{ + self.sgemm_f32_branch(bs, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, d, "h_bd")?; + self.launch_add_bias_relu_f32_raw(bs, branch_h_ptrs[d], w_ptrs[w_fc_idx + 1], self.adv_h, b)?; +} +``` + +- [ ] **Step 5: Branch 1 uses concat input in sequential fallback** + +Apply the same change to the sequential fallback path (`else` branch): + +```rust +let (fc_input, fc_k) = if d == 1 && mag_concat_ptr != 0 { + (mag_concat_ptr, self.mag_concat_dim) +} else { + (h_s2_ptr, self.shared_h2) +}; +self.sgemm_f32(stream, w_ptrs[w_fc_idx], fc_input, branch_h_ptrs[d], self.adv_h, b, fc_k, "h_bd")?; +``` + +- [ ] **Step 6: Apply same changes to forward_online_f32** + +Add `mag_concat_ptr: u64` parameter. Same conditional logic for `d == 1`. + +- [ ] **Step 7: Apply same changes to forward_target_raw** + +Add `mag_concat_ptr: u64` parameter. Same conditional logic for `d == 1`. + +IMPORTANT: The target network also needs conditioning. The concat must be computed from the TARGET network's direction logits. The caller (gpu_dqn_trainer) must call `launch_mag_concat` with target v_logits and b_logits for this path too. Use a separate target concat buffer or reuse the same one (target forward runs after online forward). + +- [ ] **Step 8: Update forward_online and forward_target wrappers** + +The CudaSlice-based wrappers (`forward_online`, `forward_target`) call the raw versions. Add the `mag_concat_ptr` parameter (pass 0 to disable, or pass the actual ptr). + +- [ ] **Step 9: Update ALL callers in gpu_dqn_trainer.rs** + +Search for all calls to `forward_online_raw`, `forward_online_f32`, `forward_target_raw`. Each needs the new `mag_concat_ptr` argument. Pass `self.ptrs.mag_concat_buf` for online forward, and a target concat ptr for target forward. + +The training graph capture and the causal sensitivity forward also call these — they need the parameter too (can pass `0u64` for causal sensitivity since it's a perturbation analysis, or `self.ptrs.mag_concat_buf` for correctness). + +- [ ] **Step 10: Add mag_concat call before branch forward** + +In the training forward (the CUDA graph path), add `self.launch_mag_concat(v_logits_ptr, b_logits_ptr)?;` AFTER the value head forward completes and AFTER direction branch (d=0) forward completes, but BEFORE the magnitude branch (d=1) forward. + +The challenge: with multi-stream parallel branches, direction (d=0) and magnitude (d=1) run in parallel. For conditioning, magnitude must wait for direction. The simplest fix: run direction first (d=0), sync, launch concat, then run magnitude (d=1) + order (d=2) + urgency (d=3) in parallel. + +In the forward, change the multi-stream loop from `for d in 0..4` to: +1. Submit direction (d=0) on branch_streams[0] +2. Record branch_done_events[0] on branch_streams[0] +3. Main stream waits for branch_done_events[0] +4. Launch mag_concat on main stream +5. Submit magnitude (d=1), order (d=2), urgency (d=3) on branch_streams 1-3 + +- [ ] **Step 11: Verify build** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +Expected: clean build. + +- [ ] **Step 12: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/batched_forward.rs \ + crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "feat: magnitude branch forward uses [h_s2; Q_dir] concat input + +Direction branch runs first, then mag_concat_qdir produces [B, SH2+3]. +Magnitude FC GEMM uses K=SH2+3. Order and urgency branches unchanged. +All 3 forward paths updated: online_raw, online_f32, target_raw. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +### Task 4: Backward Pass — Magnitude Branch Uses Wider dX + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/batched_backward.rs` + +- [ ] **Step 1: Add mag_concat_dim to backward context** + +The backward struct needs `mag_concat_dim` (or it can reference the forward's). Add the field if needed (check if backward has its own `shared_h2` field). + +- [ ] **Step 2: Modify branch FC backward for d==1** + +In the backward pass's branch loop (around line 809), the branch FC backward for d==1 must use the wider input. Currently: + +```rust +// dW for branch FC: dW[AH, SH2] += dY^T @ h_s2 +self.launch_dw_only( + stream, + scratch_d_h_b[d], // dY [B, AH] + save_h_s2, // X [B, SH2] + grad_buf_base + goff_w_bfc[d], // dW [AH, SH2] + grad_buf_base + goff_b_bfc[d], // db [AH] + self.adv_h, // out_dim + self.shared_h2, // in_dim + b, +)?; +``` + +For d==1, change to: + +```rust +let (fc_input, fc_in_dim) = if d == 1 { + (mag_concat_ptr, self.shared_h2 + 3) // [B, SH2+3] +} else { + (save_h_s2, self.shared_h2) +}; + +self.launch_dw_only( + stream, + scratch_d_h_b[d], + fc_input, + grad_buf_base + goff_w_bfc[d], + grad_buf_base + goff_b_bfc[d], + self.adv_h, + fc_in_dim, + b, +)?; +``` + +The `mag_concat_ptr` is the saved forward concat buffer (not the backward one). + +- [ ] **Step 3: dX for d==1 uses wider W and outputs to d_mag_concat** + +Currently: + +```rust +let beta_s2 = if d == 0 { 0.0_f32 } else { 1.0_f32 }; +self.launch_dx_only( + stream, + scratch_d_h_b[d], // dY [B, AH] + w_fc, // W [AH, SH2] + scratch_d_h_s2, // dX [B, SH2] + self.adv_h, // out_dim + self.shared_h2, // in_dim + b, + beta_s2, +)?; +``` + +For d==1, output to d_mag_concat_buf then extract the first SH2 columns into scratch_d_h_s2: + +```rust +if d == 1 { + // Magnitude: dX = dY @ W_b1fc → [B, SH2+3] + self.launch_dx_only( + stream, + scratch_d_h_b[d], + w_fc, // W [AH, SH2+3] + d_mag_concat_ptr, // dX [B, SH2+3] + self.adv_h, + self.shared_h2 + 3, + b, + 0.0, // beta=0 (fresh write) + )?; + // Accumulate first SH2 columns of d_mag_concat into scratch_d_h_s2 + // Use a SAXPY or memcpy+add. Simplest: launch a small kernel or + // use cuBLAS SAXPY. The last 3 columns (d_Q_dir) are discarded (detached). + self.accumulate_d_h_s2_from_concat( + stream, d_mag_concat_ptr, scratch_d_h_s2, b, beta_s2, + )?; +} else { + self.launch_dx_only( + stream, + scratch_d_h_b[d], + w_fc, + scratch_d_h_s2, + self.adv_h, + self.shared_h2, + b, + beta_s2, + )?; +} +``` + +- [ ] **Step 4: Implement accumulate_d_h_s2_from_concat** + +Add a method that copies `d_mag_concat[:, 0:SH2]` into `scratch_d_h_s2` with beta accumulation. This can use `cuMemcpy2DAsync` (strided copy) or a simple CUDA kernel. The simplest approach: a SAXPY-like kernel that reads stride `SH2+3` and writes stride `SH2`: + +```rust +/// Accumulate the first SH2 columns of d_mag_concat [B, SH2+3] into d_h_s2 [B, SH2]. +/// beta=0: overwrite. beta=1: add. +fn accumulate_d_h_s2_from_concat( + &self, + stream: &Arc, + d_concat_ptr: u64, // [B, SH2+3] + d_h_s2_ptr: u64, // [B, SH2] + batch: usize, + beta: f32, +) -> Result<(), MLError> { + // Use a simple kernel: for each (b, i) where i < SH2: + // d_h_s2[b*SH2+i] = beta * d_h_s2[b*SH2+i] + d_concat[b*(SH2+3)+i] + // Total elements = B * SH2, launched as 1D grid. + let total = batch * self.shared_h2; + let blocks = ((total + 255) / 256) as u32; + let sh2 = self.shared_h2 as i32; + let sh2_plus_3 = (self.shared_h2 + 3) as i32; + let total_i32 = total as i32; + unsafe { + self.stream + .launch_builder(&self.strided_accumulate_kernel) + .arg(&d_concat_ptr) + .arg(&d_h_s2_ptr) + .arg(&sh2) + .arg(&sh2_plus_3) + .arg(&total_i32) + .arg(&beta) + .launch(LaunchConfig { + grid_dim: (blocks, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }) + .map_err(|e| MLError::ModelError(format!("strided_accumulate: {e}")))?; + } + Ok(()) +} +``` + +The `strided_accumulate_kernel` needs to be added to experience_kernels.cu: + +```cuda +extern "C" __global__ void strided_accumulate( + const float* __restrict__ src, /* [B, src_stride] */ + float* __restrict__ dst, /* [B, dst_stride] */ + int dst_stride, /* SH2 */ + int src_stride, /* SH2+3 */ + int total, /* B * dst_stride */ + float beta /* 0=overwrite, 1=add */ +) { + int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) return; + int b = idx / dst_stride; + int i = idx % dst_stride; + float val = src[b * src_stride + i]; + if (beta > 0.0f) { + dst[idx] += val; + } else { + dst[idx] = val; + } +} +``` + +- [ ] **Step 5: Add mag_concat_ptr and d_mag_concat_ptr parameters to backward** + +The backward function signature needs these two extra pointers. + +- [ ] **Step 6: Update all backward callers** + +Search for all calls to the backward function and add the new parameters. + +- [ ] **Step 7: Verify build** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +``` + +Expected: clean build. + +- [ ] **Step 8: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/batched_backward.rs \ + crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs \ + crates/ml/src/cuda_pipeline/experience_kernels.cu +git commit -m "feat: magnitude backward uses wider [SH2+3] dX with strided accumulate + +Branch 1 backward: dW uses [h_s2; Q_dir] as X, dX writes to d_mag_concat +[B, SH2+3], first SH2 columns accumulated into d_h_s2. d_Q_dir (last 3 +columns) discarded — detached, no backprop to direction head. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +### Task 5: Zero-Init New Weight Columns and Build Verification + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` + +- [ ] **Step 1: Zero the extra 3 columns of w_b1fc after xavier init** + +The xavier init fills the entire `w_b1fc[AH, SH2+3]` buffer with random values. The extra 3 columns (Q_dir conditioning) should start at zero so the model doesn't rely on random direction signals at epoch 0. + +Find where weights are initialized (search for `xavier_init` or `param_init`). After the full buffer is initialized, zero the last 3 columns of w_b1fc: + +```rust +// Zero the Q_dir conditioning columns of w_b1fc so training starts +// as if magnitude is unconditioned. The model learns to use Q_dir gradually. +let w_b1fc_offset = padded_byte_offset(¶m_sizes, 12); +let sh2_plus_3 = config.shared_h2 + 3; +for row in 0..config.adv_h { + let col_offset = w_b1fc_offset + (row * sh2_plus_3 + config.shared_h2) as u64 * 4; + stream.memset_zeros_at(params_ptr + col_offset, 3 * std::mem::size_of::()) + .map_err(|e| MLError::ModelError(format!("zero w_b1fc qdir cols row {row}: {e}")))?; +} +``` + +NOTE: The exact API for partial memset may differ. Alternative: use `cuMemsetD32Async` for 3 floats at a time, or write a small kernel. + +- [ ] **Step 2: Run smoke tests** + +```bash +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "avg_grad\|test result" | tail -5 +``` + +Expected: 19 tests pass, avg_grad non-zero. + +- [ ] **Step 3: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +git commit -m "feat: zero-init Q_dir conditioning columns of w_b1fc + +New 3 columns of magnitude FC weight start at zero so training begins +as if unconditioned. Model learns to use direction E[Q] gradually. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +### Task 6: Experience Collector Forward — Add Conditioning + +**Files:** +- Modify: `crates/ml/src/cuda_pipeline/gpu_experience_collector.rs` + +The experience collector has its own cuBLAS forward pass (`forward_online_f32`). It also needs the mag_concat conditioning. + +- [ ] **Step 1: Allocate mag_concat buffer in collector** + +Add `mag_concat_exp_buf: CudaSlice` field, allocate `[N, SH2+3]`, and load the `mag_concat_qdir` kernel. + +- [ ] **Step 2: Call mag_concat before action selection** + +In the experience collection timestep loop, after the forward pass produces v_logits and b_logits, call `mag_concat_qdir` to build the concat buffer, then the forward pass uses it for branch 1. + +Since the collector calls `forward_online_f32` which now takes `mag_concat_ptr`, pass the buffer pointer. + +- [ ] **Step 3: Verify build and smoke tests** + +```bash +SQLX_OFFLINE=true cargo check -p ml 2>&1 | tail -5 +SQLX_OFFLINE=true FOXHUNT_TEST_DATA=test_data/futures-baseline cargo test -p ml --lib -- smoke_tests --ignored --nocapture 2>&1 | grep "avg_grad\|test result" | tail -5 +``` + +Expected: clean build, 19 tests pass. + +- [ ] **Step 4: Commit** + +```bash +git add crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +git commit -m "feat: experience collector forward uses direction-conditioned magnitude + +mag_concat_qdir runs after forward pass produces direction logits. +Magnitude branch sees [h_s2; Q_dir] during experience collection. + +Co-Authored-By: Claude Opus 4.6 (1M context) " +``` + +--- + +### Task 7: Integration Test on H100 + +**Files:** +- No code changes — deployment test + +- [ ] **Step 1: Push and launch** + +```bash +git push origin main +./scripts/argo-train.sh dqn --epochs 50 --baseline --gpu-pool ci-training-h100 +``` + +- [ ] **Step 2: Monitor first 10 epochs** + +Check: +- `avg_grad` non-zero (~1.0-2.5) +- `val_Sharpe` positive (target: >15) +- Q-gap growing (target: >0.05 by epoch 5) +- No CUDA errors, no graph capture failures +- Magnitude Q-values should diverge for different direction actions + +```bash +argo logs -n foxhunt -c main | grep -E "val_Sharpe|Q-gap|avg_grad|EPOCH_DIAG" +``` + +- [ ] **Step 3: Monitor epoch 20-30** + +The key test: does the Q-gap freeze break? With reward rank normalization + Q-gap momentum + direction conditioning, the plateau at epoch 22 should not recur. + +- [ ] **Step 4: Document results** + +Record: peak val_Sharpe, sustained val_Sharpe range, Q-gap trajectory, any anomalies. Compare to train-w6qfd baseline (frozen at val_Sharpe=22.18, Q-gap=0.133 from epoch 22).