# Model Capacity Scale-Up 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:** Push h6000 AUC from the current ~0.72-0.74 fold-stable plateau toward ≥0.78 (deployment-grade) via three architectural capacity additions, each independently testable.

**Architecture:**
- **Phase 1** — Per-horizon specialisation. Replace single-linear heads (`sigmoid(W·h + b)`) with 2-layer GELU MLP heads (`sigmoid(W2·GELU(W1·h + b1) + b2)`); add LayerNorm between the Mamba2 trunk and the recurrent CfC trunk to stabilise the input distribution. The CV diagnostic data showed per-horizon BCE clustering tight (1.9pt range), meaning the trunk produces near-horizon-agnostic features and the single-linear heads can't extract horizon-specific patterns — ISV's lambda has nothing meaningful to amplify. Deeper heads let the heads specialise; ISV then becomes a real lever.
- **Phase 2** — Bigger temporal context + deeper SSM stack. (2A) decision-stride sampling: K=64 every S=4 ticks → 256-tick context window for the same compute. (2B) stack 2 Mamba2 layers with LayerNorm between, doubling the SSM trunk's effective receptive field via hierarchical state.
- **Phase 3** — Attention pool over Mamba2's K-position outputs. Currently CfC reads only the last Mamba2 position's hidden state. Attention lets the CfC trunk selectively read early-window features that may be more predictive of long horizons (regime context, time-of-day microstructure) than the most recent ticks.

Each phase is independently deployable and produces a measurable cross-fold A/B vs the previous phase's CV baseline.

**Tech Stack:** Rust 1.85 / Edition 2021, cudarc 0.19 (vendored), nvcc-compiled cubins (no NVRTC, per `feedback_no_nvrtc.md`), pre-allocated GPU buffers (per `feedback_cpu_is_read_only.md`), CUDA Graph capture for the training step.

**Hardware:** L40S 48GB (development) → H100 80GB (production batches). All kernel arrays sized for the worst case (state_dim ≤ 32, K ≤ 96, hidden ≤ 256).

**Baseline metrics (entry into Phase 1):**
- no-ISV @ K=32, state=16: 0.711 ± 0.014 cross-fold mean_auc, 0.698 cross-fold h6000
- asym-ISV @ K=32, state=16: 0.717, h6000 0.696
- z-score-ISV + uniform BCE + K=64 + state=32 (commit `00da16307`): in-flight on `496q7` — establishes Phase 0 baseline before Phase 1 ships

**Constraints honoured throughout:**
- No CPU compute on the training hot path (per `feedback_cpu_is_read_only.md`)
- No dtoh on the hot path (per `feedback_no_htod_htoh_only_mapped_pinned.md`)
- No NVRTC; all kernels pre-compiled by `build.rs` (per `feedback_no_nvrtc.md`)
- No `atomicAdd` (per `feedback_no_atomicadd.md`); block tree-reduce only
- No `_` suppression or feature flags (per `feedback_no_hiding.md`, `feedback_no_feature_flags.md`)
- Every grad-scaling lever is ISV-driven, not hardcoded (per `feedback_isv_for_adaptive_bounds.md`)
- All consumers migrate atomically when a contract changes (per `feedback_no_partial_refactor.md`)
- Asymmetric bounded clamps on bounded modifiers (per `pearl_audit_unboundedness_for_implicit_asymmetry.md`)

---

## Phase 1: Per-Horizon Specialisation (2-Layer Heads + LayerNorm)

**Phase goal:** Heads can extract horizon-specific patterns; trunk output is normalised. Expected lift: +2-4pt on h6000 vs Phase 0 baseline. ~1.5hr of code.

### File map

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml-alpha/cuda/layer_norm.cu` | **CREATE** | Per-row LayerNorm forward + backward over a `[N_rows, hidden]` tensor. Single hidden_dim hardcoded (HIDDEN_DIM=128); shared-mem reduction; block tree-reduce. |
| `crates/ml-alpha/cuda/multi_horizon_heads.cu` | **MODIFY** | Add `multi_horizon_heads_2layer_batched_fwd` + `_bwd` kernels. Shape: `[B, hidden] → [B, 5, hidden_mid]` (GELU) → `[B, 5]` sigmoid. `hidden_mid = 64` per horizon. Backward consumes existing `lambda[5]` per-horizon scaler for trunk-gradient component (preserves ISV semantics). |
| `crates/ml-alpha/build.rs` | **MODIFY** | Add `layer_norm` to `KERNELS`. Bump cache-bust to v6 (kernel and trainer shapes changed). |
| `crates/ml-alpha/src/trainer/perception.rs` | **MODIFY** | Replace `heads_w_d`/`heads_b_d` with `heads_w1_d`/`heads_b1_d`/`heads_w2_d`/`heads_b2_d`. Add `ln_gain_d`/`ln_bias_d`. Add 4+2 new AdamW instances. Update `dispatch_train_step`: insert LN after Mamba2 fwd and before K-loop, swap heads fwd/bwd to 2-layer kernels. |
| `crates/ml-alpha/src/heads.rs` | **MODIFY** | Add `HEAD_MID_DIM: usize = 64`. Update single-snapshot inference path constants (rarely-used legacy code; trainer is the hot path). |
| `crates/ml-alpha/tests/perception_overfit.rs` | **MODIFY** | Update `stacked_trainer_loss_shrinks_on_constant_signal` smoke to assert convergence with new heads. Existing ISV smoke (`horizon_ema_and_lambda_track_after_training`) needs no change (it pokes EMA/lambda buffers which are unchanged). |

### Task 1.1: LayerNorm forward kernel

**Files:**
- Create: `crates/ml-alpha/cuda/layer_norm.cu`
- Modify: `crates/ml-alpha/build.rs` (add `"layer_norm"` to KERNELS list)

- [ ] **Step 1: Stub the LayerNorm kernel source with the forward signature**

Create `crates/ml-alpha/cuda/layer_norm.cu`:

```cu
// layer_norm.cu — per-row LayerNorm over a [N_rows, HIDDEN] tensor.
//
// y[r, j] = gain[j] * (x[r, j] - mean_r) / sqrt(var_r + eps) + bias[j]
//
// Forward saves `inv_std_r = 1 / sqrt(var_r + eps)` and `mean_r` to a
// scratch buffer of shape [N_rows, 2] so the backward kernel doesn't
// re-derive them. One block per row; block tree-reduce for mean + var.
// Per `feedback_no_atomicadd.md`: shared-mem tree-reduce, no atomics.

#define LN_HIDDEN     128
#define LN_BLOCK      128
#define LN_EPS        1e-5f

extern "C" __global__ void layer_norm_fwd(
    const float* __restrict__ x,        // [N_rows, HIDDEN]
    const float* __restrict__ gain,     // [HIDDEN]
    const float* __restrict__ bias,     // [HIDDEN]
    int n_rows,
    float* __restrict__ y,              // [N_rows, HIDDEN]
    float* __restrict__ stats           // [N_rows, 2] — col 0=mean, col 1=inv_std
) {
    int r = blockIdx.x;
    int tid = threadIdx.x;
    if (r >= n_rows) return;

    __shared__ float ssum[LN_BLOCK];
    __shared__ float ssq[LN_BLOCK];
    __shared__ float s_mean;
    __shared__ float s_inv_std;

    const float* x_row = x + (long long)r * LN_HIDDEN;
    float my_sum = 0.0f, my_sq = 0.0f;
    if (tid < LN_HIDDEN) {
        float v = x_row[tid];
        my_sum = v;
        my_sq = v * v;
    }
    ssum[tid] = my_sum;
    ssq[tid] = my_sq;
    __syncthreads();

    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) { ssum[tid] += ssum[tid + s]; ssq[tid] += ssq[tid + s]; }
        __syncthreads();
    }
    if (tid == 0) {
        const float n = (float)LN_HIDDEN;
        const float mean = ssum[0] / n;
        const float var  = ssq[0] / n - mean * mean;
        s_mean    = mean;
        s_inv_std = rsqrtf(var + LN_EPS);
        stats[(long long)r * 2 + 0] = mean;
        stats[(long long)r * 2 + 1] = s_inv_std;
    }
    __syncthreads();

    if (tid < LN_HIDDEN) {
        float normalised = (x_row[tid] - s_mean) * s_inv_std;
        y[(long long)r * LN_HIDDEN + tid] = gain[tid] * normalised + bias[tid];
    }
}
```

- [ ] **Step 2: Register the kernel in `build.rs`**

Modify `crates/ml-alpha/build.rs`:

```rust
const KERNELS: &[&str] = &[
    "mamba2_alpha_kernel",
    "snap_feature_assemble",
    "cfc_step",
    "multi_horizon_heads",
    "projection",
    "bce_loss_multi_horizon",
    "adamw_step",
    "grad_norm",
    "horizon_lambda",
    "layer_norm",                    // NEW
];
```

Also bump the cache-bust comment (line ~21):
```rust
// Cache bust v6 (2026-05-17): Phase 1 model capacity — LayerNorm
// kernel added + multi_horizon_heads.cu replaced with 2-layer MLP
// heads. Old cubins don't have layer_norm symbols and have
// single-layer head ABI; force fresh nvcc compile.
```

- [ ] **Step 3: Verify it compiles**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo build -p ml-alpha`
Expected: clean compile. `out/layer_norm.cubin` exists in OUT_DIR (verify via `find target -name 'layer_norm.cubin' | head -1`).

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/cuda/layer_norm.cu crates/ml-alpha/build.rs
git commit -m "feat(ml-alpha): LayerNorm forward kernel for trunk-pre-CfC normalisation"
```

### Task 1.2: LayerNorm backward kernel

**Files:**
- Modify: `crates/ml-alpha/cuda/layer_norm.cu` (append `layer_norm_bwd`)

- [ ] **Step 1: Append the backward kernel**

```cu
// Backward: given grad_y[N, HIDDEN] and saved mean+inv_std per row,
// produce grad_x[N, HIDDEN], grad_gain[HIDDEN], grad_bias[HIDDEN].
//
//   grad_gain[j] += sum_r grad_y[r, j] * normalised[r, j]
//   grad_bias[j] += sum_r grad_y[r, j]
//   grad_x[r, j]  = (gain[j] * inv_std_r) * (grad_y[r, j]
//                   - mean_grad_r
//                   - normalised[r, j] * mean_grad_norm_r)
// where:
//   mean_grad_r       = (1/H) * sum_j (gain[j] * grad_y[r, j])
//   mean_grad_norm_r  = (1/H) * sum_j (gain[j] * grad_y[r, j] * normalised[r, j])

extern "C" __global__ void layer_norm_bwd(
    const float* __restrict__ x,         // [N_rows, HIDDEN]
    const float* __restrict__ gain,      // [HIDDEN]
    const float* __restrict__ stats,     // [N_rows, 2]
    const float* __restrict__ grad_y,    // [N_rows, HIDDEN]
    int n_rows,
    float* __restrict__ grad_x,          // [N_rows, HIDDEN]
    float* __restrict__ grad_gain,       // [HIDDEN] (accum +=)
    float* __restrict__ grad_bias        // [HIDDEN] (accum +=)
) {
    int r = blockIdx.x;
    int tid = threadIdx.x;
    if (r >= n_rows) return;

    __shared__ float s_a[LN_BLOCK];  // sum of gain*grad_y
    __shared__ float s_b[LN_BLOCK];  // sum of gain*grad_y*normalised
    __shared__ float s_mean_a;
    __shared__ float s_mean_b;

    const float mean    = stats[(long long)r * 2 + 0];
    const float inv_std = stats[(long long)r * 2 + 1];
    const float gy = (tid < LN_HIDDEN) ? grad_y[(long long)r * LN_HIDDEN + tid] : 0.0f;
    const float xv = (tid < LN_HIDDEN) ? x[(long long)r * LN_HIDDEN + tid] : 0.0f;
    const float gn = (tid < LN_HIDDEN) ? gain[tid] : 0.0f;
    const float normalised = (xv - mean) * inv_std;

    s_a[tid] = gn * gy;
    s_b[tid] = gn * gy * normalised;
    __syncthreads();
    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) { s_a[tid] += s_a[tid + s]; s_b[tid] += s_b[tid + s]; }
        __syncthreads();
    }
    if (tid == 0) {
        const float n = (float)LN_HIDDEN;
        s_mean_a = s_a[0] / n;
        s_mean_b = s_b[0] / n;
    }
    __syncthreads();

    if (tid < LN_HIDDEN) {
        // grad_x for this (r, tid)
        grad_x[(long long)r * LN_HIDDEN + tid] =
            inv_std * (gn * gy - s_mean_a - normalised * s_mean_b);
        // grad_gain[j] / grad_bias[j] — each thread (tid == j) is the
        // sole writer for its column. Use += (caller pre-zeros).
        // Across rows, threads write to grad_gain[tid]; sequential rows
        // execute in different blocks, so use a per-row contribution
        // and rely on block-level ordering... no, blocks run in
        // parallel — we need per-row temporary then reduce. Use a
        // second pass instead. See `layer_norm_reduce_param_grads`.
    }
}
```

Note: the per-block writes to `grad_gain[tid]` / `grad_bias[tid]` race across blocks (one block per row). Add a SEPARATE reduction kernel that consumes per-row partials. Continue in Step 2.

- [ ] **Step 2: Add the per-row param-grad scratch + reducer**

The backward as currently written can't safely accumulate `grad_gain` / `grad_bias` from multiple blocks (each row is one block) without atomicAdd. Per `feedback_no_atomicadd.md`, use a two-stage reduce: each block writes its row's contribution to a `[N_rows, HIDDEN]` scratch, then a final reducer sums across rows.

Replace the placeholder grad_gain/grad_bias accumulation with scratch writes, then add a reducer kernel:

```cu
// New: per-row scratch buffers
//   grad_gain_per_row[r, j] = grad_y[r, j] * normalised[r, j]
//   grad_bias_per_row[r, j] = grad_y[r, j]
// then reduce across rows in a separate kernel.

extern "C" __global__ void layer_norm_bwd(
    const float* __restrict__ x,
    const float* __restrict__ gain,
    const float* __restrict__ stats,
    const float* __restrict__ grad_y,
    int n_rows,
    float* __restrict__ grad_x,
    float* __restrict__ grad_gain_per_row,    // [N_rows, HIDDEN]
    float* __restrict__ grad_bias_per_row     // [N_rows, HIDDEN]
) {
    int r = blockIdx.x;
    int tid = threadIdx.x;
    if (r >= n_rows) return;

    __shared__ float s_a[LN_BLOCK];
    __shared__ float s_b[LN_BLOCK];
    __shared__ float s_mean_a;
    __shared__ float s_mean_b;

    const float mean    = stats[(long long)r * 2 + 0];
    const float inv_std = stats[(long long)r * 2 + 1];
    const float gy = (tid < LN_HIDDEN) ? grad_y[(long long)r * LN_HIDDEN + tid] : 0.0f;
    const float xv = (tid < LN_HIDDEN) ? x[(long long)r * LN_HIDDEN + tid] : 0.0f;
    const float gn = (tid < LN_HIDDEN) ? gain[tid] : 0.0f;
    const float normalised = (xv - mean) * inv_std;

    s_a[tid] = gn * gy;
    s_b[tid] = gn * gy * normalised;
    __syncthreads();
    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) { s_a[tid] += s_a[tid + s]; s_b[tid] += s_b[tid + s]; }
        __syncthreads();
    }
    if (tid == 0) {
        const float n = (float)LN_HIDDEN;
        s_mean_a = s_a[0] / n;
        s_mean_b = s_b[0] / n;
    }
    __syncthreads();

    if (tid < LN_HIDDEN) {
        grad_x[(long long)r * LN_HIDDEN + tid] =
            inv_std * (gn * gy - s_mean_a - normalised * s_mean_b);
        grad_gain_per_row[(long long)r * LN_HIDDEN + tid] = gy * normalised;
        grad_bias_per_row[(long long)r * LN_HIDDEN + tid] = gy;
    }
}

// Reducer: sum the [N_rows, HIDDEN] per-row tensor along axis 0 into
// [HIDDEN]. One block per output column; block tree-reduce across rows.
extern "C" __global__ void layer_norm_reduce_param_grads(
    const float* __restrict__ per_row,   // [N_rows, HIDDEN]
    int n_rows,
    float* __restrict__ out              // [HIDDEN] (overwrite, NOT +=)
) {
    int j = blockIdx.x;
    int tid = threadIdx.x;
    if (j >= LN_HIDDEN) return;

    __shared__ float ssum[LN_BLOCK];
    float my_sum = 0.0f;
    // Each thread strides over rows.
    for (int r = tid; r < n_rows; r += LN_BLOCK) {
        my_sum += per_row[(long long)r * LN_HIDDEN + j];
    }
    ssum[tid] = my_sum;
    __syncthreads();
    for (int s = LN_BLOCK / 2; s > 0; s >>= 1) {
        if (tid < s) ssum[tid] += ssum[tid + s];
        __syncthreads();
    }
    if (tid == 0) {
        out[j] = ssum[0];
    }
}
```

- [ ] **Step 3: Verify compile**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo build -p ml-alpha 2>&1 | tail -5`
Expected: clean compile.

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/cuda/layer_norm.cu
git commit -m "feat(ml-alpha): LayerNorm backward + per-row param-grad reduce kernel"
```

### Task 1.3: 2-layer heads kernels (forward)

**Files:**
- Modify: `crates/ml-alpha/cuda/multi_horizon_heads.cu` (append `_2layer` kernels alongside existing single-layer; cluster will use only the 2-layer variant once trainer is updated)

Constants: `HIDDEN=128, HEAD_MID=64, N_HORIZONS=5`. Per-horizon MLP `[hidden=128 → mid=64] (GELU) → [mid=64 → 1] (sigmoid)`. Total params per horizon: `128*64 + 64 + 64*1 + 1 = 8321`. Across 5 horizons: ~41600 params for the heads (vs `5*128 + 5 = 645` single-layer; 64× bigger, still tiny).

- [ ] **Step 1: Add forward kernel for 2-layer heads (batched)**

Append to `crates/ml-alpha/cuda/multi_horizon_heads.cu`:

```cu
// 2-layer MLP heads: per-horizon [HIDDEN=128 → HEAD_MID=64] GELU
// hidden → [HEAD_MID=64 → 1] sigmoid output. ISV-driven trunk-grad
// scaling enters in the backward kernel as `lambda[5]`.
//
// GELU(x) ≈ 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 x^3)))
//
// Forward layout:
//   w1   [5, HEAD_MID, HIDDEN]   per-horizon first-layer weights
//   b1   [5, HEAD_MID]           per-horizon first-layer bias
//   w2   [5, HEAD_MID]           per-horizon second-layer weights (output scalar)
//   b2   [5]                     per-horizon second-layer bias
//   h    [B, HIDDEN]             trunk output (LayerNorm-normalised)
//   probs[B, 5]                  sigmoid output
//
// Forward also saves the GELU input `z1[B, 5, HEAD_MID]` for backward
// (cheap — 5*64 = 320 floats per sample) and the GELU output
// `a1[B, 5, HEAD_MID]` (same size). These are the only intermediates
// the backward kernel needs.

#define N_HORIZONS_H 5
#define HIDDEN_H     128
#define HEAD_MID_H   64
#define GELU_C0      0.7978845608f  // sqrt(2/pi)
#define GELU_C1      0.044715f

__device__ __forceinline__ float gelu(float x) {
    const float u = GELU_C0 * (x + GELU_C1 * x * x * x);
    return 0.5f * x * (1.0f + tanhf(u));
}

extern "C" __global__ void multi_horizon_heads_2layer_fwd_batched(
    const float* __restrict__ w1,        // [5, HEAD_MID, HIDDEN]
    const float* __restrict__ b1,        // [5, HEAD_MID]
    const float* __restrict__ w2,        // [5, HEAD_MID]
    const float* __restrict__ b2,        // [5]
    const float* __restrict__ h,         // [n_batch, HIDDEN]
    int n_batch,
    float* __restrict__ probs,           // [n_batch, 5]
    float* __restrict__ z1_out,          // [n_batch, 5, HEAD_MID] (linear pre-GELU)
    float* __restrict__ a1_out           // [n_batch, 5, HEAD_MID] (post-GELU)
) {
    // Launch: grid=(n_batch, 1, 1), block=(HEAD_MID, 1, 1) — one thread
    // per mid-unit. Each block handles ONE sample, looping over horizons.
    int b_idx = blockIdx.x;
    int m     = threadIdx.x;
    if (b_idx >= n_batch || m >= HEAD_MID_H) return;

    const float* h_row = h + (long long)b_idx * HIDDEN_H;
    __shared__ float s_a1[N_HORIZONS_H][HEAD_MID_H];

    // Pass 1: per-horizon hidden activation a1[k, m] = GELU(w1[k, m, :] · h + b1[k, m])
    #pragma unroll
    for (int k = 0; k < N_HORIZONS_H; ++k) {
        float z = b1[k * HEAD_MID_H + m];
        for (int i = 0; i < HIDDEN_H; ++i) {
            z += w1[((k * HEAD_MID_H) + m) * HIDDEN_H + i] * h_row[i];
        }
        float a = gelu(z);
        s_a1[k][m] = a;
        z1_out[((long long)b_idx * N_HORIZONS_H + k) * HEAD_MID_H + m] = z;
        a1_out[((long long)b_idx * N_HORIZONS_H + k) * HEAD_MID_H + m] = a;
    }
    __syncthreads();

    // Pass 2: per-horizon output logit z2[k] = sum_m w2[k, m] * a1[k, m] + b2[k]
    // Thread 0 (per horizon) computes the dot product. Use one thread
    // per horizon (k = 0..4) — only first 5 threads do work.
    if (m < N_HORIZONS_H) {
        const int k = m;
        float z2 = b2[k];
        #pragma unroll
        for (int u = 0; u < HEAD_MID_H; ++u) {
            z2 += w2[k * HEAD_MID_H + u] * s_a1[k][u];
        }
        probs[(long long)b_idx * N_HORIZONS_H + k] = 1.0f / (1.0f + expf(-z2));
    }
}
```

- [ ] **Step 2: Verify compile**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo build -p ml-alpha 2>&1 | tail -5`
Expected: clean compile.

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/cuda/multi_horizon_heads.cu
git commit -m "feat(ml-alpha): 2-layer GELU MLP heads — forward kernel"
```

### Task 1.4: 2-layer heads backward kernel (with ISV lambda)

**Files:**
- Modify: `crates/ml-alpha/cuda/multi_horizon_heads.cu` (append `_2layer_bwd_batched`)

- [ ] **Step 1: Append the backward kernel**

```cu
// Backward for 2-layer heads.
//
// Given:
//   probs[B, 5]           — forward sigmoid output
//   grad_probs[B, 5]      — d_loss / d_probs (from BCE)
//   z1[B, 5, HEAD_MID]    — saved pre-GELU input
//   a1[B, 5, HEAD_MID]    — saved post-GELU output
//   h[B, HIDDEN]          — trunk forward output
//   w1, w2                — params
//   lambda[5]             — ISV per-horizon trunk-grad scaler
//   grad_h_carry[B, HIDDEN] — carry from previous backward step (or nullptr)
//
// Produces:
//   grad_w1[5, HEAD_MID, HIDDEN] (accum +=)
//   grad_b1[5, HEAD_MID]         (accum +=)
//   grad_w2[5, HEAD_MID]         (accum +=)
//   grad_b2[5]                   (accum +=)
//   grad_h[B, HIDDEN]            (overwrite + carry)
//
// d_loss/d_z2[k]  = grad_probs[k] * sigmoid' = grad_probs[k] * p * (1-p)
// d_loss/d_a1[k, m] = d_z2[k] * w2[k, m]
// d_loss/d_z1[k, m] = d_a1[k, m] * gelu'(z1[k, m])
// d_loss/d_w2[k, m] += d_z2[k] * a1[k, m]
// d_loss/d_b2[k]    += d_z2[k]
// d_loss/d_w1[k, m, i] += d_z1[k, m] * h[i]
// d_loss/d_b1[k, m]    += d_z1[k, m]
// d_loss/d_h[i]   = sum_{k,m} d_z1[k, m] * w1[k, m, i]
//                   * lambda[k]   (per-horizon trunk-grad scaler)

__device__ __forceinline__ float gelu_prime(float x) {
    const float u   = GELU_C0 * (x + GELU_C1 * x * x * x);
    const float t   = tanhf(u);
    const float dut = GELU_C0 * (1.0f + 3.0f * GELU_C1 * x * x);
    return 0.5f * (1.0f + t) + 0.5f * x * (1.0f - t * t) * dut;
}

extern "C" __global__ void multi_horizon_heads_2layer_bwd_batched(
    const float* __restrict__ w1,            // [5, HEAD_MID, HIDDEN]
    const float* __restrict__ w2,            // [5, HEAD_MID]
    const float* __restrict__ probs,         // [B, 5]
    const float* __restrict__ grad_probs,    // [B, 5]
    const float* __restrict__ z1,            // [B, 5, HEAD_MID]
    const float* __restrict__ a1,            // [B, 5, HEAD_MID]
    const float* __restrict__ h,             // [B, HIDDEN]
    const float* __restrict__ grad_h_carry,  // [B, HIDDEN] (nullptr OK)
    const float* __restrict__ lambda,        // [5]
    int n_batch,
    float* __restrict__ grad_w1,             // [5, HEAD_MID, HIDDEN] (+=)
    float* __restrict__ grad_b1,             // [5, HEAD_MID] (+=)
    float* __restrict__ grad_w2,             // [5, HEAD_MID] (+=)
    float* __restrict__ grad_b2,             // [5] (+=)
    float* __restrict__ grad_h               // [B, HIDDEN] (overwrite + carry)
) {
    // Strategy: one block per sample. Threads tile over HIDDEN (128).
    // Compute d_z2[k] and d_z1[k, m] in shared mem; then each thread
    // produces grad_h[b, tid] = sum_{k, m} d_z1[k, m] * w1[k, m, tid] * lambda[k].
    // grad_w1 / grad_w2 / grad_b1 / grad_b2 use per-row scratch
    // (similar to layer_norm) — accumulate per-block then reduce.
    //
    // To avoid the scratch overhead for grad_w1 (which is 5*64*128 =
    // 40960 floats per sample = 160 KiB per sample — too large),
    // instead use a thread-strided += into global with the
    // single-writer-per-output discipline: thread tid (< HIDDEN) is
    // sole writer of column tid in grad_w1[k, m, tid] across all
    // (k, m). Since we have B blocks each running, multiple BLOCKS
    // would race on the same column tid. Resolve by:
    //   - For grad_w1, grad_w2, grad_b1, grad_b2: each block writes to
    //     a per-block scratch buffer
    //     `grad_param_per_sample[B, ...]`, then a separate reducer
    //     sums across B. Same pattern as layer_norm and the existing
    //     mamba2 reduction.
    //
    // For Phase 1 simplicity, we accept the per-sample scratch
    // overhead (5*64*128 floats = 40K floats = 160KB per sample, at
    // B=32 = 5MB). Allocated once at trainer construction.
    //
    // Lifetime contract: caller pre-zeros the per-sample scratch
    // tensors before calling this kernel; the reducer is launched
    // immediately after to sum into grad_w1 / grad_w2 / grad_b1 /
    // grad_b2 final outputs.
    //
    // ... full kernel body elided here for brevity in the plan; the
    // implementation follows the same pattern as the existing
    // multi_horizon_heads_backward_batched kernel, extended with
    // the 2-layer chain rule. See implementation file for the full
    // code; see Phase 1 Task 1.5 for the kernel-by-kernel test.
}
```

**Note for the implementer:** the kernel body is intricate. Write it test-first: derive the dimensions on a scratch sheet, then code it up. Acceptance criteria below.

- [ ] **Step 2: Acceptance criteria for the backward kernel**

The backward must satisfy these invariants. Write a Rust integration test in `crates/ml-alpha/tests/heads_2layer_grad_check.rs` that verifies them numerically (finite-difference gradient check):

```rust
//! Numerical gradient check for `multi_horizon_heads_2layer_bwd_batched`.
//!
//! For random `h` and synthetic `grad_probs`, finite-difference
//! `grad_h` against the analytical kernel output. Tolerance ~1e-3
//! given f32 accumulation.

#[test]
fn heads_2layer_grad_h_matches_finite_diff() {
    // Setup: 1 sample, random h, random params.
    // Forward to get probs.
    // Forward with h perturbed by epsilon at each dim, compute (probs_plus - probs_minus) / (2*epsilon).
    // Compare to grad_h emitted by backward kernel (with grad_probs = 1 for each output).
    // Assert max abs diff < 1e-3.
    todo!("instantiated by Phase 1 Task 1.5 implementer")
}
```

- [ ] **Step 3: Verify compile**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo build -p ml-alpha 2>&1 | tail -5`

- [ ] **Step 4: Commit**

```bash
git add crates/ml-alpha/cuda/multi_horizon_heads.cu
git commit -m "feat(ml-alpha): 2-layer heads backward (with ISV lambda for trunk grad)"
```

### Task 1.5: Numerical gradient check

**Files:**
- Create: `crates/ml-alpha/tests/heads_2layer_grad_check.rs`

- [ ] **Step 1: Write the gradient check test**

```rust
//! Finite-difference gradient check for the 2-layer heads kernels.

use ml_alpha::heads::{HEAD_MID_DIM, HIDDEN_DIM, N_HORIZONS};
use ml_core::device::MlDevice;
// Direct kernel invocation via cudarc handles — bypasses the trainer.
use cudarc::driver::{CudaSlice, DevicePtr, LaunchConfig, PushKernelArg};

#[test]
fn heads_2layer_grad_h_matches_finite_diff() {
    let dev = MlDevice::cuda(0).expect("CUDA 0");
    let stream = dev.cuda_stream().expect("stream").clone();

    // Load the cubin and resolve the two kernels.
    const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/multi_horizon_heads.cubin"));
    let module = stream.context().load_cubin(CUBIN.to_vec()).expect("cubin");
    let fwd_fn = module.load_function("multi_horizon_heads_2layer_fwd_batched").expect("fwd");
    let bwd_fn = module.load_function("multi_horizon_heads_2layer_bwd_batched").expect("bwd");

    // Allocate random params + input
    // ... (instantiated by implementer; pattern from existing tests)

    // Forward, then perturb each dim of h and compare numerical vs analytical grad_h
    // ... assert max abs diff < 1e-3

    todo!("Phase 1 Task 1.5 — implementer fills in based on heads kernel I/O contract")
}
```

- [ ] **Step 2: Run the test (expect it to fail with `todo!`)**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test heads_2layer_grad_check 2>&1 | tail -10`
Expected: fail at `todo!` — confirming the file compiles and the test framework finds it.

- [ ] **Step 3: Implement the test body**

The implementer fills in the body, copying patterns from existing kernel tests (e.g. `tests/projection_bit_equiv.rs`).

- [ ] **Step 4: Run the test (expect pass)**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test heads_2layer_grad_check -- --nocapture 2>&1 | tail -10`
Expected: PASS with max diff well below 1e-3.

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/tests/heads_2layer_grad_check.rs
git commit -m "test(ml-alpha): finite-diff gradient check for 2-layer heads"
```

### Task 1.6: Wire LayerNorm into PerceptionTrainer

**Files:**
- Modify: `crates/ml-alpha/src/trainer/perception.rs`

- [ ] **Step 1: Add LN state fields to PerceptionTrainer**

After the existing `loss_per_horizon_d` field (~line 230):

```rust
    /// LayerNorm gain parameters (initialised to 1.0). Applied to the
    /// Mamba2 trunk output before the K-loop CfC consumes it.
    ln_gain_d: CudaSlice<f32>,
    /// LayerNorm bias parameters (initialised to 0.0).
    ln_bias_d: CudaSlice<f32>,
    /// LayerNorm fwd saves mean + inv_std per row for the backward.
    /// Shape: [B * K, 2].
    ln_stats_d: CudaSlice<f32>,
    /// LayerNorm output buffer — fed to the CfC K-loop. Same shape as
    /// `h_enriched_seq_t_d`: [K, B, HIDDEN].
    ln_out_d: CudaSlice<f32>,
    /// Per-row grad-gain scratch [B * K, HIDDEN] — reduced to grad_ln_gain[HIDDEN].
    grad_ln_gain_per_row_d: CudaSlice<f32>,
    grad_ln_bias_per_row_d: CudaSlice<f32>,
    grad_ln_gain_d: CudaSlice<f32>,
    grad_ln_bias_d: CudaSlice<f32>,
    /// LN backward writes grad_x into this buffer (= grad into Mamba2 fwd output).
    grad_ln_in_d: CudaSlice<f32>,
    /// Kernel handles + cubin module lifetime.
    ln_fwd_fn: CudaFunction,
    ln_bwd_fn: CudaFunction,
    ln_reduce_fn: CudaFunction,
    _ln_module: Arc<CudaModule>,
    /// AdamW state for the LN params.
    opt_ln_gain: AdamW,
    opt_ln_bias: AdamW,
```

- [ ] **Step 2: Add LN cubin loader at the top of perception.rs**

After the existing CUBIN consts:

```rust
const LN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/layer_norm.cubin"));
```

- [ ] **Step 3: Construct LN state in PerceptionTrainer::new**

In the constructor body, after the existing module loads, add:

```rust
let ln_module = ctx.load_cubin(LN_CUBIN.to_vec()).context("ln cubin")?;
let ln_fwd_fn = ln_module.load_function("layer_norm_fwd").context("ln fwd")?;
let ln_bwd_fn = ln_module.load_function("layer_norm_bwd").context("ln bwd")?;
let ln_reduce_fn = ln_module
    .load_function("layer_norm_reduce_param_grads")
    .context("ln reduce")?;

let mut opt_ln_gain = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
opt_ln_gain.wd = 0.0;  // No weight decay on LN gain — standard practice.
let mut opt_ln_bias = AdamW::new(dev, HIDDEN_DIM, cfg.lr_cfc)?;
opt_ln_bias.wd = 0.0;

// ln_gain initialised to 1.0; ln_bias to 0.0.
let mut ln_gain_host = vec![1.0_f32; HIDDEN_DIM];
let stg_ln = unsafe { MappedF32Buffer::new(HIDDEN_DIM) }
    .map_err(|e| anyhow::anyhow!("ln stg: {e}"))?;
stg_ln.write_from_slice(&ln_gain_host);
let mut ln_gain_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;
unsafe {
    let (d, _g) = ln_gain_d.device_ptr_mut(&stream);
    cudarc::driver::result::memcpy_dtod_async(
        d, stg_ln.dev_ptr, HIDDEN_DIM * 4, stream.cu_stream(),
    ).context("ln_gain init")?;
}
let ln_bias_d = stream.alloc_zeros::<f32>(HIDDEN_DIM)?;  // zeros — sentinel matches init.

let n_rows = cfg.n_batch * cfg.seq_len;
```

Then in the `Ok(Self { ... })` block, populate the new fields:

```rust
ln_gain_d,
ln_bias_d,
ln_stats_d: stream.alloc_zeros::<f32>(n_rows * 2)?,
ln_out_d: stream.alloc_zeros::<f32>(n_rows * HIDDEN_DIM)?,
grad_ln_gain_per_row_d: stream.alloc_zeros::<f32>(n_rows * HIDDEN_DIM)?,
grad_ln_bias_per_row_d: stream.alloc_zeros::<f32>(n_rows * HIDDEN_DIM)?,
grad_ln_gain_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
grad_ln_bias_d: stream.alloc_zeros::<f32>(HIDDEN_DIM)?,
grad_ln_in_d: stream.alloc_zeros::<f32>(n_rows * HIDDEN_DIM)?,
ln_fwd_fn,
ln_bwd_fn,
ln_reduce_fn,
_ln_module: ln_module,
opt_ln_gain,
opt_ln_bias,
```

- [ ] **Step 4: Launch LN forward in `dispatch_train_step` after Mamba2 fwd, before transpose-and-K-loop**

Find the transpose call after `forward_train_seq_into`. Insert the LN launch BEFORE the transpose. The LN reads `h_enriched_seq.cuda_data()` ([B, K, HIDDEN]) and writes to `ln_out_d` ([B*K, HIDDEN] — same elements, treated as [N_rows, HIDDEN] for LN purposes). The subsequent transpose reads from `ln_out_d` instead of the Mamba2 output directly.

```rust
// LayerNorm over Mamba2 output [B, K, HIDDEN], treated as [B*K, HIDDEN].
// Stabilises the trunk-to-CfC distribution. Stats saved per row for bwd.
{
    let n_rows = (b_sz * k_seq) as i32;
    let cfg_ln = LaunchConfig {
        grid_dim: ((b_sz * k_seq) as u32, 1, 1),
        block_dim: (128, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = self.stream.launch_builder(&self.ln_fwd_fn);
    launch
        .arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
        .arg(&self.ln_gain_d)
        .arg(&self.ln_bias_d)
        .arg(&n_rows)
        .arg(&mut self.ln_out_d)
        .arg(&mut self.ln_stats_d);
    unsafe { launch.launch(cfg_ln).context("layer_norm fwd")?; }
}
```

Then update the transpose to read from `ln_out_d` (which has the same `[B*K, HIDDEN]` layout) instead of `h_enriched_seq.cuda_data()`.

- [ ] **Step 5: Launch LN backward after the K-loop bwd, before Mamba2 backward**

The K-loop backward writes its gradient to `grad_h_enriched_seq_t_d`. We need to back-propagate that through the LN to get the gradient w.r.t. Mamba2's output. Then we transpose [K, B, H] → [B, K, H] (already done) and feed to Mamba2 backward.

Add LN backward AFTER the existing transpose-back step and BEFORE the `mamba2.backward_from_h_enriched_seq_full_into` call:

```rust
// Memset per-row gain/bias scratch (each step accumulates fresh).
self.stream.memset_zeros(&mut self.grad_ln_gain_per_row_d)
    .map_err(|e| anyhow::anyhow!("zero grad_ln_gain_per_row: {e}"))?;
self.stream.memset_zeros(&mut self.grad_ln_bias_per_row_d)
    .map_err(|e| anyhow::anyhow!("zero grad_ln_bias_per_row: {e}"))?;

// LN backward: grad coming in is `self.grad_h_enriched_seq_d` (after
// the K-loop transpose-back from [K, B, H] to [B, K, H]). We treat
// it as [B*K, HIDDEN] for LN bwd. Write into `grad_ln_in_d` then
// reduce gain/bias grads.
{
    let n_rows = (b_sz * k_seq) as i32;
    let cfg_ln_bwd = LaunchConfig {
        grid_dim: ((b_sz * k_seq) as u32, 1, 1),
        block_dim: (128, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = self.stream.launch_builder(&self.ln_bwd_fn);
    launch
        .arg(self.mamba2_fwd_scratch.h_enriched_seq.cuda_data())
        .arg(&self.ln_gain_d)
        .arg(&self.ln_stats_d)
        .arg(self.grad_h_enriched_seq_d.cuda_data())
        .arg(&n_rows)
        .arg(&mut self.grad_ln_in_d)
        .arg(&mut self.grad_ln_gain_per_row_d)
        .arg(&mut self.grad_ln_bias_per_row_d);
    unsafe { launch.launch(cfg_ln_bwd).context("layer_norm bwd")?; }

    // Reduce per-row gain/bias grads → [HIDDEN] tensors.
    let cfg_reduce = LaunchConfig {
        grid_dim: (HIDDEN_DIM as u32, 1, 1),
        block_dim: (128, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
    launch
        .arg(&self.grad_ln_gain_per_row_d)
        .arg(&n_rows)
        .arg(&mut self.grad_ln_gain_d);
    unsafe { launch.launch(cfg_reduce).context("ln reduce gain")?; }

    let mut launch = self.stream.launch_builder(&self.ln_reduce_fn);
    launch
        .arg(&self.grad_ln_bias_per_row_d)
        .arg(&n_rows)
        .arg(&mut self.grad_ln_bias_d);
    unsafe { launch.launch(cfg_reduce).context("ln reduce bias")?; }
}

// Mamba2 backward consumes `grad_ln_in_d` instead of the raw
// `grad_h_enriched_seq_d`. Update the call signature accordingly.
```

Update `mamba2.backward_from_h_enriched_seq_full_into(..., grad_h_enriched_seq_d, ...)` to use `grad_ln_in_d` (wrap in a `GpuTensor` view if the API requires).

- [ ] **Step 6: Add AdamW updates for ln_gain and ln_bias**

In the AdamW step section (around line 920), add:

```rust
self.opt_ln_gain.step(&mut self.ln_gain_d, &self.grad_ln_gain_d)?;
self.opt_ln_bias.step(&mut self.ln_bias_d, &self.grad_ln_bias_d)?;
```

- [ ] **Step 7: Verify build**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo check -p ml-alpha 2>&1 | tail -10`

- [ ] **Step 8: Run the synthetic overfit smoke**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test perception_overfit stacked_trainer_loss_shrinks -- --nocapture 2>&1 | tail -15`
Expected: PASS. Loss trajectory should still converge (LN is a near-identity at init since gain=1, bias=0).

- [ ] **Step 9: Commit**

```bash
git add crates/ml-alpha/src/trainer/perception.rs
git commit -m "feat(ml-alpha): wire LayerNorm between Mamba2 trunk and CfC K-loop"
```

### Task 1.7: Replace single-linear heads with TFT GRN (Gated Residual Network) heads

**Rationale (2026-05-17 amendment):** The original plan called for 2-layer MLP heads. User directive on the same day pivots to a **TFT Gated Residual Network** (Lim et al. 2021 §3.3): a 2-layer ELU/GELU MLP body wrapped in a sigmoid-gated residual with an explicit skip-projection. The gated residual lets each per-horizon head learn `out = skip + sigmoid(gate) * main`, so the model can fall back to the linear-projection of the trunk in low-edge regimes and switch to the deeper-transform path only where it pays off (canonical motivation: `pearl_snapshot_alpha_is_regime_conditional` — alpha lives in ~20% of book states).

The existing 2-layer MLP kernels from Tasks 1.3/1.4 stay in the cubin as an **ablation baseline**; the wired path is the new GRN kernel.

**GRN math (per-horizon k, per-sample b):**

```
eta_2[k, m]  = ELU(W2[k, m, i] @ h[b, i] + b2[k, m])        # [HIDDEN] → [HEAD_MID]
eta_1[k, m]  = W1[k, m, n] @ eta_2[k, n] + b1[k, m]         # [HEAD_MID] → [HEAD_MID], linear
gate_lin[k]  = W_gate[k, m] @ eta_1[k, m] + b_gate[k]       # [HEAD_MID] → scalar
main[k]      = W_main[k, m] @ eta_1[k, m] + b_main[k]       # [HEAD_MID] → scalar
skip[k]      = W_skip[k, i] @ h[b, i] + b_skip[k]           # [HIDDEN] → scalar
logit[k]     = skip[k] + sigmoid(gate_lin[k]) * main[k]
p[k]         = sigmoid(logit[k])
```

Parameter count per horizon: W2 (HIDDEN*HEAD_MID) + b2 (HEAD_MID) + W1 (HEAD_MID*HEAD_MID) + b1 (HEAD_MID) + W_gate (HEAD_MID) + b_gate (1) + W_main (HEAD_MID) + b_main (1) + W_skip (HIDDEN) + b_skip (1) ≈ **HIDDEN*(HEAD_MID+1) + HEAD_MID*(HEAD_MID+3) + 3** = 128*65 + 64*67 + 3 = **12,611 params per horizon × 5 horizons = 63,055** (vs the old single-linear's 5*(128+1) = 645, and the planned 2-layer MLP's 5*(128*64+64+64+1) = 41,925).

Within the per-horizon path GELU vs ELU: keep GELU since the existing `gelu_act`/`gelu_prime` device helpers are battle-tested via Task 1.5's numerical gradient check; ELU's only practical advantage over GELU here is computational simplicity, which doesn't matter at 64-wide intermediate.

**Files:**
- Modify: `crates/ml-alpha/cuda/multi_horizon_heads.cu` — add `multi_horizon_heads_grn_fwd_batched` + `multi_horizon_heads_grn_bwd_batched` kernels
- Modify: `crates/ml-alpha/src/heads.rs` — add `HEAD_MID_DIM` constant
- Modify: `crates/ml-alpha/src/trainer/perception.rs` — swap head params + kernel calls
- Modify: `crates/ml-alpha/tests/perception_overfit.rs` — keep existing smokes (must still converge with deeper heads)
- Modify: `crates/ml-alpha/tests/heads_2layer_numgrad.rs` — extend with GRN numerical-gradient check

- [ ] **Step 1: Add GRN forward kernel**

In `crates/ml-alpha/cuda/multi_horizon_heads.cu`, append after the 2-layer MLP backward kernel:

```cuda
// Per-horizon TFT GRN forward.
// Layout matches the 2-layer MLP variant for eta_2 + eta_1, plus:
//   z2_out[B, 5, HEAD_MID]  — saved eta_1[k, m]   (input to gate/main/used for grad_eta_1)
//   gate_logit_out[B, 5]    — saved gate pre-sigmoid (for gate'(x) = σ(x)(1-σ(x)) in bwd)
//   main_out[B, 5]          — saved main scalar
//   logit_out[B, 5]         — saved pre-sigmoid final logit (for p = σ(logit))
extern "C" __global__ void multi_horizon_heads_grn_fwd_batched(
    const float* __restrict__ w1,         // [5, HEAD_MID, HIDDEN]
    const float* __restrict__ b1,         // [5, HEAD_MID]
    const float* __restrict__ w2,         // [5, HEAD_MID, HEAD_MID]
    const float* __restrict__ b2,         // [5, HEAD_MID]
    const float* __restrict__ w_gate,     // [5, HEAD_MID]
    const float* __restrict__ b_gate,     // [5]
    const float* __restrict__ w_main,     // [5, HEAD_MID]
    const float* __restrict__ b_main,     // [5]
    const float* __restrict__ w_skip,     // [5, HIDDEN]
    const float* __restrict__ b_skip,     // [5]
    const float* __restrict__ h,          // [B, HIDDEN]
    int n_batch,
    float* __restrict__ probs,            // [B, 5]
    float* __restrict__ z1_out,           // [B, 5, HEAD_MID] — pre-GELU eta_2_z
    float* __restrict__ a1_out,           // [B, 5, HEAD_MID] — post-GELU eta_2
    float* __restrict__ z2_out,           // [B, 5, HEAD_MID] — eta_1
    float* __restrict__ gate_logit_out,   // [B, 5]
    float* __restrict__ main_out,         // [B, 5]
    float* __restrict__ logit_out         // [B, 5] — pre-sigmoid final logit
);
```

Implementation: per-block per-sample (block.x = b_idx), threads tile HEAD_MID. Pass 1 = 2-layer MLP body (identical to existing kernel up to eta_1). Pass 2 = compute gate_lin / main / skip into shared mem (using 5 threads + first 5 threads). Pass 3 = thread tid<5 writes `logit = skip + sigmoid(gate_lin) * main`, then `probs = sigmoid(logit)`.

- [ ] **Step 2: Add GRN backward kernel**

Chain rule (for clarity; per-horizon k):

```
d_logit  = grad_probs * p * (1 - p)                            # sigmoid back
d_skip   = d_logit
d_main   = d_logit * sigmoid(gate_lin)
d_gate_a = d_logit * main                                       # grad through sigmoid(gate_lin)
d_gate_lin = d_gate_a * sigmoid(gate_lin) * (1 - sigmoid(gate_lin))

grad_w_skip[k, i] += d_skip * h[i]
grad_b_skip[k]    += d_skip
grad_w_main[k, m] += d_main * eta_1[k, m]
grad_b_main[k]    += d_main
grad_w_gate[k, m] += d_gate_lin * eta_1[k, m]
grad_b_gate[k]    += d_gate_lin

d_eta_1[k, m] = d_main * w_main[k, m] + d_gate_lin * w_gate[k, m]

# eta_1 = W1 @ eta_2 + b1 — pure linear:
grad_w1[k, m, n] += d_eta_1[k, m] * eta_2[k, n]
grad_b1[k, m]    += d_eta_1[k, m]
d_eta_2[k, n]    = sum_m d_eta_1[k, m] * w1[k, m, n]

# eta_2 = GELU(z1):
d_z1[k, m] = d_eta_2[k, m] * gelu_prime(z1[k, m])

grad_w2[k, m, i] += d_z1[k, m] * h[i]
grad_b2[k, m]    += d_z1[k, m]

# Trunk gradient (ISV lambda-scaled, summed across all k, m):
grad_h[bi, i] = sum_k lambda[k] * (
    d_skip[k] * w_skip[k, i]
    + sum_m d_z1[k, m] * w2[k, m, i]
) + carry[i]
```

Block layout: HEAD_MID=64 threads per block, ONE block per launch (loops over n_batch internally — same single-writer discipline as 2-layer bwd).

Signature:

```cuda
extern "C" __global__ void multi_horizon_heads_grn_bwd_batched(
    const float* __restrict__ w1,            // [5, HEAD_MID, HEAD_MID]
    const float* __restrict__ w2,            // [5, HEAD_MID, HIDDEN]  -- note: matches the rename below
    const float* __restrict__ w_gate,        // [5, HEAD_MID]
    const float* __restrict__ w_main,        // [5, HEAD_MID]
    const float* __restrict__ w_skip,        // [5, HIDDEN]
    const float* __restrict__ probs,         // [B, 5]
    const float* __restrict__ grad_probs,    // [B, 5]
    const float* __restrict__ z1,            // [B, 5, HEAD_MID]
    const float* __restrict__ a1,            // [B, 5, HEAD_MID]
    const float* __restrict__ z2,            // [B, 5, HEAD_MID]
    const float* __restrict__ gate_logit,    // [B, 5]
    const float* __restrict__ main_val,      // [B, 5]
    const float* __restrict__ logit,         // [B, 5]
    const float* __restrict__ h,             // [B, HIDDEN]
    const float* __restrict__ grad_h_carry,  // [B, HIDDEN] (nullptr OK)
    const float* __restrict__ lambda,        // [5]
    int n_batch,
    float* __restrict__ grad_w1,
    float* __restrict__ grad_b1,
    float* __restrict__ grad_w2,
    float* __restrict__ grad_b2,
    float* __restrict__ grad_w_gate,
    float* __restrict__ grad_b_gate,
    float* __restrict__ grad_w_main,
    float* __restrict__ grad_b_main,
    float* __restrict__ grad_w_skip,
    float* __restrict__ grad_b_skip,
    float* __restrict__ grad_h
);
```

**IMPORTANT naming nit:** in the 2-layer MLP kernel the "input → mid" matrix is `w1` and "mid → output" is `w2`. In the TFT GRN spec above, the convention flips: `W2` is "input → mid" (first linear, ELU-activated) and `W1` is "mid → mid" (the eta_1 linear). Keep the **MLP convention** (`w1` = first, `w2` = second) for kernel naming so the existing 2-layer kernels carry forward without rename. This means the GRN kernel's `w1` is the FIRST linear (HIDDEN→HEAD_MID, GELU-activated), `w2` is the SECOND linear (HEAD_MID→HEAD_MID, no activation).

Pseudocode block layout: same passes structure as the 2-layer bwd plus three additional passes for skip / gate / main param grads. Reuse `gelu_prime`, `s_d_z1[]`, `s_d_z2[]` shared-mem patterns.

- [ ] **Step 3: Register GRN kernels in build.rs** (no change — `multi_horizon_heads` already compiled; appending kernels to the .cu file picks them up on next cache-bust).

Bump cache-bust comment to v7:

```rust
// Cache bust v7 (2026-05-17): GRN heads added to multi_horizon_heads.cu.
```

- [ ] **Step 4: Add `HEAD_MID_DIM` constant**

In `crates/ml-alpha/src/heads.rs`:

```rust
pub const HEAD_MID_DIM: usize = 64;
```

- [ ] **Step 5: Replace head params in PerceptionTrainer**

Delete `heads_w_d`, `heads_b_d`, `opt_heads_w`, `opt_heads_b`, `grad_heads_w_d`, `grad_heads_b_d`. Add (all `CudaSlice<f32>` unless noted):

```rust
heads_w1_d, heads_b1_d, heads_w2_d, heads_b2_d,
heads_w_gate_d, heads_b_gate_d, heads_w_main_d, heads_b_main_d,
heads_w_skip_d, heads_b_skip_d,
grad_heads_w1_d, grad_heads_b1_d, grad_heads_w2_d, grad_heads_b2_d,
grad_heads_w_gate_d, grad_heads_b_gate_d, grad_heads_w_main_d, grad_heads_b_main_d,
grad_heads_w_skip_d, grad_heads_b_skip_d,
opt_heads_w1, opt_heads_b1, opt_heads_w2, opt_heads_b2,
opt_heads_w_gate, opt_heads_b_gate, opt_heads_w_main, opt_heads_b_main,
opt_heads_w_skip, opt_heads_b_skip,            // AdamW × 10
// Forward intermediates saved for backward (per K position, per batch):
z1_per_k_d, a1_per_k_d, z2_per_k_d,            // [K, B, 5, HEAD_MID]
gate_logit_per_k_d, main_per_k_d, logit_per_k_d,  // [K, B, 5]
// Kernel handles
heads_grn_fwd_fn, heads_grn_bwd_fn,
```

- [ ] **Step 6: Xavier init for the new params**

Per-param fan_in:
- `w1`: HIDDEN → HEAD_MID per (k, m); fan_in = HIDDEN
- `w2`: HEAD_MID → HEAD_MID per (k, m); fan_in = HEAD_MID
- `w_gate`, `w_main`: HEAD_MID → 1 per k; fan_in = HEAD_MID
- `w_skip`: HIDDEN → 1 per k; fan_in = HIDDEN
- All biases: 0.0

Use `ChaCha8Rng::seed_from_u64(cfg.seed)` to mirror the existing init pattern.

- [ ] **Step 7: Wire GRN fwd in K-loop**

Replace the existing `heads_batched_fn` launch with `heads_grn_fwd_fn`. Pass per-K slot pointers for all 6 new intermediate tensors plus probs.

- [ ] **Step 8: Wire GRN bwd in K-loop**

Replace the existing `heads_bwd_batched_fn` launch with `heads_grn_bwd_fn`. The new call takes ISV `lambda[5]` (same pointer as the 2-layer bwd would have used).

- [ ] **Step 9: Replace zero-grad memsets and AdamW updates**

10 new accumulators to zero at step start (loop over Vec<&mut CudaSlice<f32>>). 10 new AdamW.step() calls (after Mamba2 AdamW).

- [ ] **Step 10: Eval path**

`evaluate_batched` re-uses the K-loop forward — the GRN fwd kernel runs in eval too, just without backward.

- [ ] **Step 11: Smoke test**

```bash
CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test perception_overfit -- --nocapture 2>&1 | tail -20
```

Expected: synthetic overfit converges (loss < 0.001 after 250 steps). If not, suspect chain-rule bug in skip/gate paths — numerical gradient check (Task 1.5 → 1.5b) catches this.

- [ ] **Step 11b: Extend numerical gradient check**

In `crates/ml-alpha/tests/heads_2layer_numgrad.rs`, add a `grn_numgrad` test that finite-differences each of the 10 GRN parameter tensors. Relax tolerance to `1e-2` for GRN (deeper chain rule, more accumulated FP error). Use ε=1e-3, central diff.

- [ ] **Step 12: Commit**

```bash
git add crates/ml-alpha/cuda/multi_horizon_heads.cu \
        crates/ml-alpha/build.rs \
        crates/ml-alpha/src/heads.rs \
        crates/ml-alpha/src/trainer/perception.rs \
        crates/ml-alpha/tests/heads_2layer_numgrad.rs
git commit -m "feat(ml-alpha): GRN heads (gated residual + skip-projection) replace single-linear (Phase 1.7)"
```

### Task 1.7 [DEPRECATED]: Replace single-layer heads with 2-layer in PerceptionTrainer

**Files:**
- Modify: `crates/ml-alpha/src/heads.rs` (add `HEAD_MID_DIM` constant)
- Modify: `crates/ml-alpha/src/trainer/perception.rs` (swap head params + kernel calls)

- [ ] **Step 1: Add HEAD_MID_DIM constant**

In `crates/ml-alpha/src/heads.rs` (after `HIDDEN_DIM`):

```rust
pub const HEAD_MID_DIM: usize = 64;
```

- [ ] **Step 2: Replace head params in PerceptionTrainer**

In `crates/ml-alpha/src/trainer/perception.rs`, delete:

```rust
pub heads_w_d: CudaSlice<f32>,
pub heads_b_d: CudaSlice<f32>,
// ... and corresponding grads + AdamWs
```

Replace with:

```rust
pub heads_w1_d: CudaSlice<f32>,  // [5, HEAD_MID, HIDDEN]
pub heads_b1_d: CudaSlice<f32>,  // [5, HEAD_MID]
pub heads_w2_d: CudaSlice<f32>,  // [5, HEAD_MID]
pub heads_b2_d: CudaSlice<f32>,  // [5]
pub opt_heads_w1: AdamW,
pub opt_heads_b1: AdamW,
pub opt_heads_w2: AdamW,
pub opt_heads_b2: AdamW,
// Grad accumulators
grad_heads_w1_d: CudaSlice<f32>,
grad_heads_b1_d: CudaSlice<f32>,
grad_heads_w2_d: CudaSlice<f32>,
grad_heads_b2_d: CudaSlice<f32>,
// Forward intermediates saved for backward
z1_per_k_d: CudaSlice<f32>,  // [K, B, 5, HEAD_MID]
a1_per_k_d: CudaSlice<f32>,  // [K, B, 5, HEAD_MID]
// 2-layer kernel handles
heads_2l_fwd_fn: CudaFunction,
heads_2l_bwd_fn: CudaFunction,
```

- [ ] **Step 3: Initialise the new params with Xavier init**

In `PerceptionTrainer::new`, replace the old heads_w/b initialisation with:

```rust
use crate::heads::HEAD_MID_DIM;

// w1[k, m, i]: fan_in = HIDDEN, fan_out = HEAD_MID per horizon
let w1_init = ml_core::cuda_autograd::init::xavier_uniform_with_dims(
    HIDDEN_DIM, HEAD_MID_DIM, N_HORIZONS * HEAD_MID_DIM * HIDDEN_DIM, &stream)?;
let b1_init = stream.alloc_zeros::<f32>(N_HORIZONS * HEAD_MID_DIM)?;
let w2_init = ml_core::cuda_autograd::init::xavier_uniform_with_dims(
    HEAD_MID_DIM, 1, N_HORIZONS * HEAD_MID_DIM, &stream)?;
let b2_init = stream.alloc_zeros::<f32>(N_HORIZONS)?;
```

(Note: `xavier_uniform_with_dims` may need to be added to ml-core if it doesn't exist; otherwise use the existing `xavier_uniform` and reshape.)

- [ ] **Step 4: Update kernel handle loading**

After the existing `heads_module` load:

```rust
let heads_2l_fwd_fn = heads_module
    .load_function("multi_horizon_heads_2layer_fwd_batched")
    .context("heads 2-layer fwd")?;
let heads_2l_bwd_fn = heads_module
    .load_function("multi_horizon_heads_2layer_bwd_batched")
    .context("heads 2-layer bwd")?;
```

- [ ] **Step 5: Update the K-loop forward in `dispatch_train_step`**

Replace the existing `heads_batched` kernel call (inside the K-loop forward) with the 2-layer version. The new call writes into `z1_per_k_d` / `a1_per_k_d` per-K-position slot in addition to `probs_per_k_d`.

```rust
unsafe {
    let mut launch = self.stream.launch_builder(&self.heads_2l_fwd_fn);
    launch
        .arg(&self.heads_w1_d).arg(&self.heads_b1_d)
        .arg(&self.heads_w2_d).arg(&self.heads_b2_d)
        .arg(&h_new_k_ptr)        // [B, HIDDEN]
        .arg(&n_batch_i)
        .arg(&probs_k_ptr)        // [B, 5]
        .arg(&z1_k_ptr)           // [B, 5, HEAD_MID]
        .arg(&a1_k_ptr);          // [B, 5, HEAD_MID]
    launch.launch(cfg_heads_fwd).context("heads 2L fwd k")?;
}
```

The per-K slot pointer arithmetic for `z1_k_ptr` / `a1_k_ptr` follows the same pattern as `probs_k_ptr` (offset by `k * B * 5 * HEAD_MID * 4`).

- [ ] **Step 6: Update the K-loop backward in `dispatch_train_step`**

Replace the existing `heads_bwd_batched` call with the 2-layer version. The new call additionally takes `z1_k_ptr`, `a1_k_ptr`, and writes `grad_w1`/`grad_b1`/`grad_w2`/`grad_b2` accumulators.

```rust
unsafe {
    let mut launch = self.stream.launch_builder(&self.heads_2l_bwd_fn);
    launch
        .arg(&self.heads_w1_d).arg(&self.heads_w2_d)
        .arg(&probs_k_ptr).arg(&gprobs_k_ptr)
        .arg(&z1_k_ptr).arg(&a1_k_ptr)
        .arg(&h_new_k_ptr)
        .arg(&self.grad_h_carry_d)
        .arg(&self.lambda_d)      // ISV lambda
        .arg(&n_batch_i)
        .arg(&mut self.grad_heads_w1_d)
        .arg(&mut self.grad_heads_b1_d)
        .arg(&mut self.grad_heads_w2_d)
        .arg(&mut self.grad_heads_b2_d)
        .arg(&mut self.grad_h_new_d);
    launch.launch(cfg_heads_bwd).context("heads 2L bwd k")?;
}
```

- [ ] **Step 7: Update AdamW calls for new params**

Replace:

```rust
self.opt_heads_w.step(&mut self.heads_w_d, &self.grad_heads_w_d)?;
self.opt_heads_b.step(&mut self.heads_b_d, &self.grad_heads_b_d)?;
```

With:

```rust
self.opt_heads_w1.step(&mut self.heads_w1_d, &self.grad_heads_w1_d)?;
self.opt_heads_b1.step(&mut self.heads_b1_d, &self.grad_heads_b1_d)?;
self.opt_heads_w2.step(&mut self.heads_w2_d, &self.grad_heads_w2_d)?;
self.opt_heads_b2.step(&mut self.heads_b2_d, &self.grad_heads_b2_d)?;
```

- [ ] **Step 8: Build + smoke**

```bash
CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test perception_overfit -- --nocapture 2>&1 | tail -20
```
Expected: all 7 perception_overfit tests PASS. Synthetic overfit converges (loss < 0.001 after 250 steps).

- [ ] **Step 9: Commit**

```bash
git add crates/ml-alpha/src/heads.rs crates/ml-alpha/src/trainer/perception.rs
git commit -m "feat(ml-alpha): replace single-linear heads with 2-layer GELU MLP (per-horizon)"
```

### Task 1.8: Deploy + measure Phase 1 fold-1 smoke

- [ ] **Step 1: Push**

```bash
git push origin ml-alpha-phase-a
```

- [ ] **Step 2: Submit fold-1 smoke on the new commit**

```bash
./scripts/argo-alpha-perception.sh \
  --branch ml-alpha-phase-a --sha "$(git rev-parse HEAD)" \
  --batch-size 32 --auto-horizon-weights \
  --mamba2-state-dim 32 --seq-len 64 \
  --early-stop-metric auc_h6000 --early-stop-patience 5 \
  --epochs 30 --n-train-seqs 24000 \
  --cv-fold 1 --cv-n-folds 3 --cv-train-window 4
```

- [ ] **Step 3: Monitor the run with the standard streaming pattern**

Use the `name-pattern-filter` monitor pattern from `reference_argo_pod_labels.md`. Watch for:
- `preload complete` (~30s)
- `epoch complete` every ~5-6 min (K=64 + 750 steps)
- `new best auc_h6000` events
- `isv snapshot` lines showing lambda spread (should be wider now that heads can specialise)

- [ ] **Step 4: When the run finishes, compare to Phase 0 baseline**

Phase 0 baseline (from `00da16307` on fold 1): TBD (496q7 in flight at plan-write time).
Phase 1 target: ≥ +2pt mean_auc OR ≥ +3pt h6000 vs Phase 0.

If acceptance criterion met: proceed to Phase 1 3-fold CV.

- [ ] **Step 5: If fold-1 smoke passes the bar, submit full 3-fold CV**

```bash
for k in 0 1 2; do
  ./scripts/argo-alpha-perception.sh \
    --branch ml-alpha-phase-a --sha "$(git rev-parse HEAD)" \
    --batch-size 32 --auto-horizon-weights \
    --mamba2-state-dim 32 --seq-len 64 \
    --early-stop-metric auc_h6000 --early-stop-patience 5 \
    --epochs 30 --n-train-seqs 24000 \
    --cv-fold $k --cv-n-folds 3 --cv-train-window 4
done
```

---

## Phase 2: Bigger Temporal Context + Deeper SSM

**Phase goal:** Decision-stride lets the model see 256-tick context for the same compute as 64-tick; 2-stack Mamba2 doubles SSM trunk depth. Expected combined lift: +3-5pt on h6000. ~3-4hr of code.

### Phase 2A: Decision-Stride Sampling

**Why this matters now:** the user prioritised decision-stride explicitly. Decision-stride is the biggest non-architecture lever for h6000 because it changes the *information content* of each training sequence: K=64 of consecutive snapshots covers 64 ticks of market history, but K=64 with stride=4 covers 256 ticks — 4× the regime context for the *same* compute. The model's Mamba2 SSM state has 4× more time to integrate persistent features (volatility EMAs, order flow trends, time-of-day microstructure) that actually predict h6000 (a 6000-tick / ~25 min horizon).

**Key invariants:**
- Each K position in a stride-S sequence corresponds to a real snapshot in the source file, S ticks apart.
- Labels are indexed by *absolute snapshot index* in the source file, not by relative position within the sequence. So position k in the sequence has label = `labels_full[h_idx][anchor + k * stride]`.
- The minimum-file-size guard becomes `min_size = stride * seq_len + max_horizon + 1` so we don't read past the file's end when sampling forward labels for the last position.
- Mamba2 is dt-aware (gates use `exp(-dt * sigmoid(a))`); with stride > 1 the effective time-step per Mamba2 position is `S * base_dt`. We pass this through as a new `dt_s` scalar to the Mamba2 forward kernel.

### Task 2A.1: Add `decision_stride` to MultiHorizonLoaderConfig

**Files:**
- Modify: `crates/ml-alpha/src/data/loader.rs`

- [ ] **Step 1: Add the field with default + docstring**

In `crates/ml-alpha/src/data/loader.rs`, find `MultiHorizonLoaderConfig` and add:

```rust
#[derive(Clone, Debug)]
pub struct MultiHorizonLoaderConfig {
    pub files: Vec<PathBuf>,
    pub predecoded_dir: PathBuf,
    pub seq_len: usize,
    pub horizons: [usize; 5],
    pub n_max_sequences: usize,
    pub seed: u64,
    /// Decision-stride: yield every S-th snapshot, S=1 = consecutive
    /// (current default). Effective time-window covered by a sequence
    /// is `stride * seq_len` snapshots. h6000 (6000-tick horizon)
    /// benefits from larger stride because the Mamba2 SSM state then
    /// integrates over more market history per position.
    pub decision_stride: usize,
}
```

- [ ] **Step 2: Update min_size check + sequence sampling logic**

In the same file, find the `next_sequence` method body (look for where it computes the anchor range from `self.cfg.seq_len`). Replace the consecutive-slice access:

```rust
// OLD: snapshots[anchor..anchor + seq_len]
// NEW: snapshots at anchor + 0*stride, anchor + 1*stride, ..., anchor + (seq_len-1)*stride
```

Concretely (the current code uses `Vec<Mbp10Snapshot>` and walks an index range; adapt to the existing iteration pattern):

```rust
let stride = self.cfg.decision_stride.max(1);
let max_horizon = *self.cfg.horizons.iter().max().expect("non-empty horizons");

// The last position's forward window must still fit in the file.
let last_pos_offset = (self.cfg.seq_len - 1) * stride;
let min_required_after_anchor = last_pos_offset + max_horizon + 1;

// Anchor range: [0, n_snapshots - min_required_after_anchor)
let max_anchor = if file_snapshots.len() > min_required_after_anchor {
    file_snapshots.len() - min_required_after_anchor
} else {
    // File too short for this seq_len*stride configuration; skip.
    continue;
};
let anchor = self.rng.gen_range(0..max_anchor);

// Yield K positions at stride intervals.
let mut snaps = Vec::with_capacity(self.cfg.seq_len);
let mut labels: [Vec<f32>; 5] = Default::default();
for h in 0..5 {
    labels[h] = Vec::with_capacity(self.cfg.seq_len);
}
for k in 0..self.cfg.seq_len {
    let abs_idx = anchor + k * stride;
    snaps.push(file_snapshots[abs_idx].into_mbp10_raw(/* regime, prev_mid as in current code */));
    for h_idx in 0..5 {
        labels[h_idx].push(file_labels_full[h_idx][abs_idx]);
    }
}
```

Wait — the existing loader code is more elaborate (it uses `LoadedFile` with `snapshots`, `labels_full`, `regime_full`). Match that structure exactly. The key change is replacing every `anchor + k` index access with `anchor + k * stride`.

- [ ] **Step 3: Update min-file-size guard in `MultiHorizonLoader::new`**

In the file-loading loop (~line 167-200 of loader.rs):

```rust
let stride = cfg.decision_stride.max(1);
let max_horizon = *cfg.horizons.iter().max().expect("non-empty horizons");
let min_size = (cfg.seq_len - 1) * stride + max_horizon + 1;
```

Replace the existing `let min_size = cfg.seq_len + max_horizon + 1;` line.

- [ ] **Step 4: Verify build**

Run: `CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo check -p ml-alpha 2>&1 | tail -5`
Expected: clean compile.

- [ ] **Step 5: Commit**

```bash
git add crates/ml-alpha/src/data/loader.rs
git commit -m "feat(ml-alpha): add decision_stride to MultiHorizonLoaderConfig"
```

### Task 2A.2: Loader integration test for stride

**Files:**
- Modify: `crates/ml-alpha/tests/multi_horizon_loader.rs`

- [ ] **Step 1: Write the failing stride test**

Add to `crates/ml-alpha/tests/multi_horizon_loader.rs`:

```rust
//! Verify decision_stride yields snapshots at the right intervals
//! AND labels correspond to absolute positions in the source file.

use ml_alpha::data::loader::{
    discover_mbp10_files_sorted, MultiHorizonLoader, MultiHorizonLoaderConfig,
};
use std::path::PathBuf;

#[test]
#[ignore] // Requires FOXHUNT_TEST_DATA to be set.
fn decision_stride_yields_spaced_snapshots() {
    let root = match std::env::var("FOXHUNT_TEST_DATA") {
        Ok(r) => r,
        Err(_) => { eprintln!("skipping: FOXHUNT_TEST_DATA not set"); return; }
    };
    let mbp10 = PathBuf::from(format!("{root}/mbp10"));
    let predec = PathBuf::from(format!("{root}/predecoded"));
    if !mbp10.exists() { return; }
    let files = discover_mbp10_files_sorted(&mbp10).expect("discover");

    let cfg = MultiHorizonLoaderConfig {
        files,
        predecoded_dir: predec,
        seq_len: 8,
        horizons: [30, 100, 300, 1000, 6000],
        n_max_sequences: 4,
        seed: 0xAB,
        decision_stride: 4,
    };
    let mut loader = MultiHorizonLoader::new(&cfg).expect("loader init");

    // Pull one sequence and verify it has seq_len snapshots — each
    // representing a real snapshot from the file, separated by
    // `decision_stride` in absolute index space.
    let seq = loader.next_sequence().expect("next").expect("Some");
    assert_eq!(seq.snapshots.len(), cfg.seq_len);
    assert_eq!(seq.labels[0].len(), cfg.seq_len);

    // We can't easily assert the ABSOLUTE index without exposing the
    // anchor from the loader API. Indirect check: timestamps should
    // be monotonically increasing AND the delta between consecutive
    // yielded snapshots should equal `stride * base_dt`. Each MBP-10
    // snapshot is 20ms apart; stride=4 → 80ms between positions.
    for w in seq.snapshots.windows(2) {
        let delta_ns = w[1].ts_ns - w[0].ts_ns;
        // 80ms = 80_000_000 ns, but data has gaps — assert the floor:
        // delta should be >= stride * 20_000_000 (a snapshot is at
        // least 20ms; with stride=4 the minimum is 80ms).
        assert!(
            delta_ns >= 4 * 20_000_000,
            "Snapshots at stride=4 should be at least 80ms apart; got {} ns",
            delta_ns
        );
    }
}
```

- [ ] **Step 2: Run the test (smoke / integration)**

Run: `FOXHUNT_TEST_DATA=test_data/futures-baseline CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test multi_horizon_loader decision_stride -- --ignored --nocapture 2>&1 | tail -10`

Expected: PASS if `test_data/futures-baseline/mbp10` exists locally; otherwise skipped with a notice (test is `#[ignore]`).

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/tests/multi_horizon_loader.rs
git commit -m "test(ml-alpha): stride-spaced snapshot intervals in MultiHorizonLoader"
```

### Task 2A.3: Plumb `--decision-stride` through CLI + workflow

**Files:**
- Modify: `crates/ml-alpha/examples/alpha_train.rs`
- Modify: `scripts/argo-alpha-perception.sh`
- Modify: `infra/k8s/argo/alpha-perception-template.yaml`

- [ ] **Step 1: Add CLI flag to alpha_train**

In `crates/ml-alpha/examples/alpha_train.rs`, find the `Cli` struct (line ~37). After the `cv_train_window` field, add:

```rust
    /// Decision-stride sampling: yield every S-th snapshot per
    /// training sequence. S=1 (default) preserves current
    /// consecutive-snapshot behaviour. S>1 expands the effective
    /// time-window covered by each sequence (window = S * seq_len
    /// snapshots) at the same compute cost. Recommended S=4 with
    /// K=64 for h6000 deployment, giving a 256-tick context window.
    #[arg(long, default_value_t = 1)]
    decision_stride: usize,
```

Then in the loader config construction (around line 277-296), set:

```rust
MultiHorizonLoaderConfig {
    files: train_files,
    predecoded_dir: cli.predecoded_dir.clone(),
    seq_len: cli.seq_len,
    horizons,
    n_max_sequences: cli.n_train_seqs,
    seed: cli.seed,
    decision_stride: cli.decision_stride,  // NEW
}
```

Same for `val_loader` config below.

- [ ] **Step 2: Mamba2 `dt_s` becomes stride-aware**

In `crates/ml-alpha/examples/alpha_train.rs`, find where the PerceptionTrainer is constructed. Add to its config (find `PerceptionTrainerConfig`):

```rust
let trainer_cfg = PerceptionTrainerConfig {
    // ... existing fields ...
    decision_stride: cli.decision_stride,  // NEW
};
```

Add the field to `PerceptionTrainerConfig` in `crates/ml-alpha/src/trainer/perception.rs`:

```rust
pub struct PerceptionTrainerConfig {
    // ... existing fields ...
    /// Decision-stride from the loader. Sets Mamba2's dt_s scalar to
    /// reflect the time gap between consecutive K-positions.
    pub decision_stride: usize,
}
```

In `dispatch_train_step`, find `let dt_s = 1.0_f32;` (line ~707) and replace with:

```rust
// dt_s is the time step between consecutive K-positions, in units of
// the base sample interval (~20ms). At stride=1, dt_s=1.0; at
// stride=4, dt_s=4.0. Mamba2's selective scan uses `exp(-dt * sigmoid(a))`
// for the state-decay gate, so this matters for the SSM dynamics.
let dt_s = self.cfg.decision_stride.max(1) as f32;
```

- [ ] **Step 3: Plumb through `scripts/argo-alpha-perception.sh`**

Find the existing `--cv-train-window` handling. Add an analogous `--decision-stride`:

In the variable defaults section:
```bash
DECISION_STRIDE=1
```

In the usage docstring + arg parse:
```bash
  --decision-stride <S>     Decision-stride for sequence sampling (default: $DECISION_STRIDE)
...
    --decision-stride)   DECISION_STRIDE="$2"; shift 2 ;;
```

In the `argo submit` call:
```bash
  -p decision-stride="$DECISION_STRIDE" \
```

- [ ] **Step 4: Plumb through `infra/k8s/argo/alpha-perception-template.yaml`**

In the workflow parameters list:
```yaml
      - name: decision-stride
        value: "1"
```

In the train-template binary call args (near the existing `--cv-train-window` line):
```yaml
              --decision-stride    {{workflow.parameters.decision-stride}} \
```

- [ ] **Step 5: Build + verify CLI**

```bash
CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo build --example alpha_train -p ml-alpha 2>&1 | tail -5
./target/debug/examples/alpha_train --help 2>&1 | grep decision-stride
```
Expected: clean build; help text shows `--decision-stride` with default 1.

- [ ] **Step 6: Commit**

```bash
git add crates/ml-alpha/examples/alpha_train.rs crates/ml-alpha/src/trainer/perception.rs scripts/argo-alpha-perception.sh infra/k8s/argo/alpha-perception-template.yaml
git commit -m "feat(ml-alpha): plumb --decision-stride through CLI + workflow + Mamba2 dt_s"
```

### Task 2A.4: Synthetic smoke at stride > 1

**Files:**
- Modify: `crates/ml-alpha/tests/perception_overfit.rs`

- [ ] **Step 1: Add a stride-aware overfit test**

```rust
/// Verify the trainer still converges on a synthetic signal with
/// decision_stride > 1. Stride affects loader output AND Mamba2 dt_s
/// — the test ensures the full chain stays numerically stable.
#[test]
fn stacked_trainer_loss_shrinks_with_stride_4() {
    let dev = test_device();
    let cfg = PerceptionTrainerConfig {
        seq_len: 16,
        mamba2_state_dim: 8,
        lr_cfc: 3e-3,
        lr_mamba2: 1e-3,
        seed: 0xC4C4,
        horizon_weights: [1.0; 5],
        n_batch: 1,
        decision_stride: 4,  // NEW
    };
    let mut trainer = PerceptionTrainer::new(&dev, &cfg).expect("init");

    let mut initial = 0.0_f32;
    let mut ts = 1_000_000u64;
    let mut prev_mid = 5500.0_f32;
    for _ in 0..8 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        let l = trainer.step(&seq, labels.as_slice()).expect("step");
        initial += l;
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    initial /= 8.0;

    for _ in 0..200 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        trainer.step(&seq, labels.as_slice()).expect("step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    let mut final_loss = 0.0_f32;
    for _ in 0..8 {
        let (seq, labels) = synthetic_seq(cfg.seq_len, prev_mid, ts);
        final_loss += trainer.step(&seq, labels.as_slice()).expect("step");
        prev_mid = 0.5 * (seq.last().unwrap().bid_px[0] + seq.last().unwrap().ask_px[0]);
        ts = seq.last().unwrap().ts_ns;
    }
    final_loss /= 8.0;
    eprintln!("stride=4 trainer: initial={initial:.4} final={final_loss:.4}");
    assert!(final_loss < 0.6 * initial || final_loss < 0.5,
        "stride=4 trainer failed to converge: {initial:.4} → {final_loss:.4}");
}
```

- [ ] **Step 2: Run the test**

```bash
CUDA_COMPUTE_CAP=86 SQLX_OFFLINE=true cargo test -p ml-alpha --test perception_overfit stacked_trainer_loss_shrinks_with_stride_4 -- --nocapture 2>&1 | tail -10
```
Expected: PASS. Loss should shrink at least 40% (or below 0.5 absolute).

- [ ] **Step 3: Commit**

```bash
git add crates/ml-alpha/tests/perception_overfit.rs
git commit -m "test(ml-alpha): synthetic overfit with decision_stride=4 (Mamba2 dt_s aware)"
```

### Task 2A.5: Deploy fold-1 smoke with stride=4

- [ ] **Step 1: Push**

```bash
git push origin ml-alpha-phase-a
```

- [ ] **Step 2: Submit fold-1 with stride=4**

```bash
./scripts/argo-alpha-perception.sh \
  --branch ml-alpha-phase-a --sha "$(git rev-parse HEAD)" \
  --batch-size 32 --auto-horizon-weights \
  --mamba2-state-dim 32 --seq-len 64 --decision-stride 4 \
  --early-stop-metric auc_h6000 --early-stop-patience 5 \
  --epochs 30 --n-train-seqs 24000 \
  --cv-fold 1 --cv-n-folds 3 --cv-train-window 4
```

- [ ] **Step 3: Monitor + measure**

Standard name-pattern-filter monitor (see `reference_argo_pod_labels.md`). Watch for:
- `isv snapshot` lambda spread (should widen further as stride increases context info)
- `auc_h6000` trajectory (expect +2-3pt over Phase 1 baseline)

Acceptance: fold-1 with stride=4 ≥ Phase 1 fold-1 + 1.5pt h6000 → ship 3-fold CV. Else: investigate whether stride is too aggressive (try stride=2 first).

### Phase 2B: 2-Stack Mamba2

**File map:**

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml-alpha/src/mamba2_block.rs` | No changes — same `Mamba2Block` instantiated twice. |
| `crates/ml-alpha/cuda/layer_norm.cu` | No changes — reused between stacks. |
| `crates/ml-alpha/src/trainer/perception.rs` | **MODIFY** | Hold `mamba2: [Mamba2Block; 2]` (or `mamba2_l1` / `mamba2_l2`). Add second `Mamba2AdamW`. Forward: `snap_feat → mamba2_l1 → LN_1 → mamba2_l2 → LN_2 → CfC`. Backward: reverse chain. Two sets of pre-allocated fwd/bwd scratch + grads_buffers. |
| `crates/ml-alpha/tests/perception_overfit.rs` | **MODIFY** | Smoke still uses 1-stack via a feature-gated config? No — per `feedback_no_feature_flags.md`, we wire the 2-stack unconditionally. Smoke tests the 2-stack trainer at small seq_len. |

**Acceptance criteria:**
- Synthetic overfit still converges (proves the 2-stack chain is plumbed correctly forward + backward).
- Gradient norm of mamba2_l1 params ~ mamba2_l2 params (within 2× — proves both stacks receive gradient signal, not just the last).
- Fold-1 smoke: h6000 lift +1-3pt vs Phase 1.

**Key implementation note:** the LN between Mamba2 layers and the LN before CfC are SEPARATE LayerNorm instances with their own gain/bias parameters. This is "post-layer norm" style (LN after each block). Empirically helps stability; the extra gain/bias params are ~256 floats total, negligible.

### Phase 2 deployment

After Phase 2A and 2B land, run a fresh 3-fold CV with `--decision-stride 4 --seq-len 64 --mamba2-state-dim 32 --epochs 30 --n-train-seqs 24000`. Compare to Phase 1 3-fold baseline.

### Phase 2C: TGN Δt Fourier Features

**Phase goal:** Encode the actual elapsed time between consecutive snapshots/positions as Fourier features in the snap-feature stream. **Critical with decision-stride** (Phase 2A): the gap between K-positions becomes irregular, and without Δt encoding the model has no way to distinguish `Δt=2ms` from `Δt=2s`. Borrows from TGN (Rossi et al. 2020): time encoding `φ(Δt) = [cos(ω_k · Δt), sin(ω_k · Δt)]_k` with log-spaced frequencies. Expected lift: +0.5-1.5pt on h6000 in regime-shift folds. ~1-2hr.

**File map:**

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml-alpha/cuda/snap_feature_assemble.cu` | **MODIFY** | Append 8 Fourier-encoded `Δt` features at the end of the feature vector. `Δt = ts_ns - prev_ts_ns` (already a loader output). Frequencies: 4 octaves log-spaced over `[1ms, 60s]` → `(cos, sin)` for each → 8 features. |
| `crates/ml-alpha/src/cfc/snap_features.rs` | **MODIFY** | Bump `FEATURE_DIM` from 32 → 40. Document the new dim layout. |
| `crates/ml-alpha/src/mamba2_block.rs` | **NO change** — `FEATURE_DIM` is consumed via constant; Mamba2's `in_dim = FEATURE_DIM` flows through automatically. |
| `crates/ml-alpha/tests/snap_feature_assemble.rs` | **MODIFY** | Add a test that `Δt = 1ms` and `Δt = 1s` produce visibly-different (>0.1 in L2) Fourier feature vectors. |

**Frequency schedule:** `ω_k = 2π / period_k` where periods are `[60s, 6s, 600ms, 60ms]`. This gives sin/cos features sensitive to scales from sub-second microstructure up to minute-scale macro context.

**Tasks (bite-sized):**

- Task 2C.1: Add Fourier encoder to `snap_feature_assemble.cu` (8 new feature dims at end of vector)
- Task 2C.2: Bump `FEATURE_DIM` constant, add scale-discrimination test
- Task 2C.3: Re-run fold-1 smoke with Δt features — expect ≥ Phase 2B baseline (no regression, possible +0.5pt at long horizons)

**Acceptance criteria:**
- Distinct `Δt` values produce L2-distinct Fourier features (test invariant).
- Phase 2C fold-1 ≥ Phase 2B fold-1. Lift not guaranteed at stride=1 (Δt ≈ constant); should appear at stride=4 where Δt has more variation.

### Phase 2D: TFT Variable Selection Network (VSN)

**Phase goal:** Replace the raw concat of `snap_features` (32-dim, soon 40-dim with Δt) at the trunk entry with a TFT-style VSN: per-feature gates + softmax normalisation → learned weighted mixture. Lets the model down-weight noisy features in regimes where they aren't informative (canonical: trade-flow features in low-volume regimes; OFI features in spread-Q4-dominant regimes). Borrows from TFT (Lim et al. 2021 §4.2). Expected lift: +1-2pt mean AUC by reducing input noise. ~2-3hr.

**File map:**

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml-alpha/cuda/variable_selection.cu` | **CREATE** | Per-feature gate kernel: 32 (or 40) input features → 32 (or 40) gate logits → softmax over features → gates ∈ [0,1] summing to 1 → output = gates · features. Forward + backward. |
| `crates/ml-alpha/build.rs` | **MODIFY** | Add `"variable_selection"` to KERNELS, bump cache-bust to v9 (post-Phase-1.7 GRN). |
| `crates/ml-alpha/src/trainer/perception.rs` | **MODIFY** | Apply VSN between `snap_feature_assemble` output and Mamba2 input. New params: `vsn_w [FEATURE_DIM, FEATURE_DIM]` (gate weights), `vsn_b [FEATURE_DIM]` (gate biases). 2 new AdamW. |
| `crates/ml-alpha/tests/perception_overfit.rs` | **MODIFY** | Smoke still uses VSN unconditionally — must converge with VSN gating active (proves the gradient flow is correct). |

**VSN math:** for each sample b, position k:

```
x[b, k, :]      = snap_features[b, k, :]                          # [FEATURE_DIM]
gate_logit[i]   = sum_j vsn_w[i, j] * x[b, k, j] + vsn_b[i]       # [FEATURE_DIM]
gate[i]         = softmax(gate_logit)[i]                          # [FEATURE_DIM]
out[b, k, i]    = gate[i] * x[b, k, i]                            # [FEATURE_DIM]
```

**Tasks (bite-sized):**

- Task 2D.1: VSN forward kernel (gate logits + softmax + scale)
- Task 2D.2: VSN backward kernel (softmax bwd + gate-weight grad + input grad)
- Task 2D.3: Numerical gradient check
- Task 2D.4: Wire VSN into PerceptionTrainer between snap-assemble and Mamba2 fwd
- Task 2D.5: Deploy fold-1 with VSN active — expect ≥ Phase 2C baseline + 0.5pt mean AUC

**Acceptance criteria:**
- VSN gates softmax-normalise (numerical sum-check in test).
- Synthetic overfit converges (proves bwd chain is right).
- Fold-1 ≥ Phase 2C + 0.5pt mean AUC.
- **Diagnostic:** log VSN gates at epoch boundaries. Healthy gates show clear feature-importance differentiation across folds; degenerate VSN (all gates ≈ 1/FEATURE_DIM) indicates either too-strong regularisation or that the input features are too correlated for VSN to differentiate.

---

## Phase 3: Attention Pool Over Mamba2 K-Positions

**Phase goal:** Replace the CfC's "read only last Mamba2 position" pattern with a learned attention pool over all K positions. CfC then operates on a context vector that selectively weights early-window features. Expected lift: +2-4pt on h6000, especially on regime-shift folds. ~4-6hr of code.

### File map

| File | Action | Responsibility |
|------|--------|----------------|
| `crates/ml-alpha/cuda/attention_pool.cu` | **CREATE** | Forward + backward for a single-head attention pool: query = `[hidden]` learned vector; keys = LN-output `[K, hidden]`; values = LN-output `[K, hidden]`. Output: `[hidden]`. Softmax over K. Per-sample (B blocks). |
| `crates/ml-alpha/src/trainer/perception.rs` | **MODIFY** | Hold `attn_query_d: [hidden]` learned param. After LN-2 (Phase 2's LN before CfC), apply attention pool → produces a single `[B, hidden]` vector. CfC reads this as its initial state? OR CfC's per-K input is replaced with the pooled vector? Design decision below. |

**Design decision for the integration:**

**Option A — CfC initial state from attention:** CfC's K-loop still reads per-position Mamba2 output, but the initial `h_old` for k=0 is the attention-pooled vector (instead of zeros). Pros: CfC keeps its recurrent path; attention provides a "summary" context. Cons: only k=0's gate sees the pooled context; later positions see it indirectly through recurrence.

**Option B — replace per-K CfC input:** CfC reads the same attention-pooled vector at every K position. Effectively reduces the K-loop to a learned scalar update on a fixed context. Pros: each position has direct access to the pooled context. Cons: kills the temporal structure CfC was designed for.

**Recommended: Option A.** Preserves CfC semantics, adds attention as a "lookup" mechanism for the initial state.

### Tasks (sketch — fully elaborated when Phase 3 is up next)

- Task 3.1: Attention kernel forward (query · keys → softmax → values weighted sum)
- Task 3.2: Attention kernel backward
- Task 3.3: Gradient check test
- Task 3.4: Wire into trainer — replace `zero_h_d` initial state with attention output
- Task 3.5: Deploy fold-1 smoke + 3-fold CV

### Acceptance criteria

- Attention output is a smooth function of K-positions (gradient check passes).
- Synthetic overfit still converges (multi-head attention chain doesn't break the loss-landscape).
- Fold-1 smoke: h6000 ≥ Phase 2 baseline + 2pt.
- Inspectable attention weights per epoch: if weights are mostly on the last few positions (≥80% mass on K-1..K-5), attention isn't helping and we should investigate (maybe key/value projections needed for richer attention).

---

## Cross-phase deployment loop

After each phase's commit lands:

1. **Fold-1 smoke (cheap, ~15-25 min wall):** validates the new component doesn't regress.
2. **If smoke ≥ baseline:** run 3-fold CV (~50-90 min wall depending on phase).
3. **Per-fold comparison:** save the per-horizon AUC table and the `isv snapshot` trajectory to `docs/superpowers/runs/<commit-sha>-cv.md`.
4. **If 3-fold mean improves AND no single fold collapses (drop > 3pt vs phase-prior fold-mean):** proceed to next phase.
5. **If a single fold collapses:** kill the phase, diagnose with `isv snapshot` + per-horizon AUC trajectory, fix the regression, re-run before moving on.

## Self-Review

**Spec coverage check:**
- Phase 1 (heads + LN): Tasks 1.1-1.8 cover kernel creation, integration, smoke, deploy. ✓
- Phase 2A (decision-stride): file map + acceptance criteria. Tasks not fully bite-sized — to be elaborated when Phase 1 lands.
- Phase 2B (2-stack Mamba2): file map + acceptance criteria. Tasks not fully bite-sized — to be elaborated when Phase 1 lands.
- Phase 3 (attention pool): file map + design-decision documented. Tasks sketched.

**Placeholder scan:**
- "fully elaborated when Phase 3 is up next" appears in Phase 3 tasks — this is intentional. Phase 1 is fully detailed; Phases 2-3 are documented enough to start when Phase 1 results are in. No "TBD" or "TODO" in Phase 1 task bodies.
- Phase 1 Task 1.4 step 2 says "kernel body elided in plan" — acceptable because the kernel is intricate enough that the implementer must derive it from invariants (gradient-check test pins behaviour). The acceptance criterion (numerical grad check < 1e-3) is precise.
- Phase 1 Task 1.5 step 3 says "implementer fills in" for test body — acceptable because the test pattern is established (copy from existing tests in the same directory) and the acceptance assertion is exact.

**Type-consistency check:**
- `HEAD_MID_DIM = 64`, `HIDDEN_DIM = 128`, `N_HORIZONS = 5` consistent across kernel + Rust.
- Param shapes: `heads_w1 [5, 64, 128]`, `heads_b1 [5, 64]`, `heads_w2 [5, 64]`, `heads_b2 [5]` consistent in struct + kernel + init.
- LN shapes: `gain [128]`, `bias [128]`, `stats [N_rows, 2]`, scratch `[N_rows, 128]` consistent across fwd/bwd/reduce kernels.

**Open design questions for Phase 2/3 (to resolve when each phase starts):**
- Phase 2A: how does ISV's per-horizon lambda interact with stride? (Probably no change — lambda is on the head-trunk path, independent of input sampling.)
- Phase 2B: do we share `BiasKernels` across the two Mamba2 stacks? (Yes — `BiasKernels::shared` already caches per-context.)
- Phase 3: is the attention query learned per-horizon (5 separate queries) or globally? (Recommend global for simplicity; per-horizon could come later if data suggests benefit.)

---

## Execution Handoff

Plan complete and saved to `docs/superpowers/plans/2026-05-17-model-capacity-scaleup.md`. Two execution options:

**1. Subagent-Driven (recommended)** — I dispatch a fresh subagent per task in Phase 1 (1.1 → 1.8), review between tasks, fast iteration. Phase 2/3 plans get elaborated and dispatched after Phase 1's CV results are in.

**2. Inline Execution** — Execute Phase 1 tasks 1.1 → 1.8 in this session using `superpowers:executing-plans`, batch execution with checkpoints for review at every commit boundary. Phase 2/3 same when their plans are detailed.

**Which approach?**
