# 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;
        }
        /* Normalize Q_dir to unit-ish scale: divide by delta_z.
         * Raw E[Q] ~ 0.1 while h_s2 ~ 1.0 post-ReLU. Without normalization,
         * the 3 Q_dir columns of W_b1fc get 10× smaller gradients → learn 10× slower.
         * dz ~ 0.01-0.1, so E[Q]/dz ~ 1-10 — similar scale to h_s2 activations. */
        concat_out[b * out_stride + SH2 + a] = eq / fmaxf(dz, 1e-6f);
    }
}
```

- [ ] **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) <noreply@anthropic.com>"
```

---

### 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<f32>,
/// Backward scratch for magnitude concat dX [B, SH2+3].
d_mag_concat_buf: CudaSlice<f32>,
```

- [ ] **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<CudaStream>) -> Result<CudaFunction, MLError> {
```

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) <noreply@anthropic.com>"
```

---

### 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<CudaStream>,
    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) <noreply@anthropic.com>"
```

---

### 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<CudaStream>,
    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) <noreply@anthropic.com>"
```

---

### 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(&param_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::<f32>())
        .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) <noreply@anthropic.com>"
```

---

### 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<f32>` 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) <noreply@anthropic.com>"
```

---

### 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 <workflow-name> -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).
