# CRT.train ISV-Driven λ Controller — Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development. Each task is a checkbox list of steps.

**Goal:** Replace static `SMOOTHNESS_LAMBDA_RATIO` with a signal-driven λ controller anchored on observed h30 jitter, with permanent floor.

**Architecture:** New `smoothness_lambda_controller.cu` kernel reads `raw_per_h` from `output_smoothness` (already emitted), maintains a Wiener-α-floor-0.5 EMA of jitter per horizon, derives per-horizon target as `jitter_ema[h30] × (K_h30 / K_h)`, and emits new `λ[h]` for next step's `output_smoothness` launch. λ floor = 1e-4 (permanent floor pattern). Launch order: ... → BCE → output_smoothness → smoothness_lambda_controller → backward K-loop.

**Tech Stack:** Same as the predecessor (Rust 1.85+, cudarc 0.19, CUDA 12.4, pre-compiled cubin).

**Predecessor commits (already shipped):** `bf5ecd109..70ecfc0fe` (the 10-commit static-λ chain).

---

## File structure (locked decisions)

| File | Action |
|------|--------|
| `crates/ml-alpha/cuda/smoothness_lambda_controller.cu` | CREATE |
| `crates/ml-alpha/build.rs` | Add `"smoothness_lambda_controller"` to `KERNELS`, bump v12→v13 |
| `crates/ml-alpha/src/heads.rs` | DELETE `SMOOTHNESS_LAMBDA_RATIO` |
| `crates/ml-alpha/src/trainer/perception.rs` | Add `jitter_ema_d`, `jitter_first_obs_d`, controller fn handle + module Arc; change λ init from `base × ratio[h]` to `[LAMBDA_FLOOR; 5]`; launch controller after `output_smoothness` in `step_batched`; remove `SMOOTHNESS_LAMBDA_RATIO` import |
| `crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs` | CREATE — 4 invariant tests |

---

## Task 1: Create kernel `smoothness_lambda_controller.cu`

- [ ] **1.1** Write `crates/ml-alpha/cuda/smoothness_lambda_controller.cu`:

```cuda
// smoothness_lambda_controller.cu — ISV-driven per-horizon λ for the
// output_smoothness regularizer.
//
// Reads `raw_per_h[5]` (emitted by output_smoothness_loss_and_grad),
// maintains a per-horizon Wiener-α-floor EMA of observed jitter,
// derives per-horizon target by anchoring on observed h30 jitter
// scaled by HORIZONS[0]/HORIZONS[h], and emits next-step λ[h] with a
// permanent floor.
//
// Per `pearl_controller_anchors_isv_driven`: target is signal-derived,
// not a constant. Per `pearl_first_observation_bootstrap`: first
// observation replaces EMA directly. Per
// `pearl_blend_formulas_must_have_permanent_floor`: λ has a permanent
// floor so the controller never fully self-disables. Per
// `pearl_wiener_alpha_floor_for_nonstationary`: α floored at 0.5 since
// the controller's target drifts as the policy co-adapts.
//
// Per `feedback_no_atomicadd`: 5 threads, single block, single writer
// per (h) slot; no atomics. Per `pearl_no_host_branches_in_captured_graph`:
// no host branching; thread-id gating only.

#define SLC_N_HORIZONS  5
#define SLC_LAMBDA_FLOOR 1.0e-4f
#define SLC_TARGET_EPS   1.0e-9f
#define SLC_ALPHA_FLOOR  0.5f

// HORIZONS = {30, 100, 300, 1000, 6000}.
// Ratios HORIZONS[0]/HORIZONS[h] = {1, 0.3, 0.1, 0.03, 0.005}.
// Constant array known at compile time.
__device__ __constant__ float TARGET_K_RATIO[SLC_N_HORIZONS] = {
    1.0f,
    30.0f / 100.0f,
    30.0f / 300.0f,
    30.0f / 1000.0f,
    30.0f / 6000.0f,
};

extern "C" __global__ void smoothness_lambda_controller(
    const float* __restrict__ raw_per_h,   // [5]  emitted by output_smoothness
    float* __restrict__ jitter_ema,         // [5]  in/out — EMA state
    int*   __restrict__ first_obs,          // [1]  in/out — sentinel
    float  base_lambda,                     // scalar — amplitude knob
    float* __restrict__ lambda_out          // [5]  output — λ for next step
) {
    const int h = threadIdx.x;
    if (h >= SLC_N_HORIZONS) return;

    __shared__ float s_jitter_after[SLC_N_HORIZONS];

    // Pass 1: EMA update with sentinel bootstrap.
    const float raw_h = raw_per_h[h];
    const int sentinel = first_obs[0];
    float jitter_h;
    if (sentinel == 0) {
        jitter_h = raw_h;
    } else {
        jitter_h = (1.0f - SLC_ALPHA_FLOOR) * jitter_ema[h] + SLC_ALPHA_FLOOR * raw_h;
    }
    jitter_ema[h] = jitter_h;
    s_jitter_after[h] = jitter_h;

    // Single-writer of sentinel — thread h=0 only.
    if (h == 0 && sentinel == 0) {
        first_obs[0] = 1;
    }
    __syncthreads();

    // Pass 2: derive target and update λ.
    // target[h] = jitter_ema[0] * TARGET_K_RATIO[h]
    //           = jitter_ema[0] for h=0 (self-target)
    //           < jitter_ema[0] for h>0
    const float jitter_h0 = s_jitter_after[0];
    const float target_h = jitter_h0 * TARGET_K_RATIO[h];

    // Excess controller: ratio - 1, clamped at zero (only push UP).
    // When observed > target: excess > 0 → λ grows
    // When observed ≤ target: excess = 0 → λ relaxes toward base_lambda × 1 = base_lambda
    // Floor: λ ≥ LAMBDA_FLOOR.
    const float safe_target = fmaxf(target_h, SLC_TARGET_EPS);
    const float excess_ratio = fmaxf(0.0f, jitter_h / safe_target - 1.0f);
    const float lambda_new = base_lambda * (1.0f + excess_ratio);
    lambda_out[h] = fmaxf(SLC_LAMBDA_FLOOR, lambda_new);
}
```

- [ ] **1.2** Commit:
```bash
git add crates/ml-alpha/cuda/smoothness_lambda_controller.cu
git commit -m "feat(crt-train): add ISV-driven smoothness_lambda_controller kernel"
```

---

## Task 2: Register kernel in `build.rs`

- [ ] **2.1** Edit `crates/ml-alpha/build.rs`. In the `KERNELS` array, after `"output_smoothness"`, add `"smoothness_lambda_controller"` with comment. Bump cache-bust v12→v13:
```rust
    "output_smoothness",             // CRT.train: per-horizon adjacent-position prob-jitter penalty
    "smoothness_lambda_controller",  // CRT.train: ISV-driven λ controller anchored on h30 jitter
];
```
And replace the v12 comment with `// Cache bust v13 (2026-05-21): smoothness_lambda_controller.cu added — ISV-driven λ.`

- [ ] **2.2** Verify:
```bash
SQLX_OFFLINE=true CARGO_FEATURE_CUDA=1 cargo build -p ml-alpha 2>&1 | grep -E "smoothness_lambda_controller|^error" | head
```
Expected: `compiled cuda/smoothness_lambda_controller.cu -> .../smoothness_lambda_controller.cubin`. No errors.

- [ ] **2.3** Commit:
```bash
git add crates/ml-alpha/build.rs
git commit -m "build(crt-train): register smoothness_lambda_controller.cu in build.rs"
```

---

## Task 3: Delete `SMOOTHNESS_LAMBDA_RATIO` from `heads.rs`

- [ ] **3.1** Edit `crates/ml-alpha/src/heads.rs`. DELETE the entire `pub const SMOOTHNESS_LAMBDA_RATIO: [f32; N_HORIZONS] = [ ... ];` declaration (added in Task 3 of the predecessor plan, around lines 27-39). Also delete its doc-comment block.

- [ ] **3.2** This will break `perception.rs` import — that's expected; Task 4 below replaces the import.

- [ ] **3.3** Do NOT commit yet — Task 4 ships in the same commit (atomic refactor per `feedback_no_partial_refactor`).

---

## Task 4: Wire controller into `PerceptionTrainer` (perception.rs)

This is the largest task. It includes: new fields, new cubin/handle, new λ init pattern, controller launch in `step_batched`, removal of the static-ratio code.

- [ ] **4.1** In `crates/ml-alpha/src/trainer/perception.rs`:

  a. **Remove** the `SMOOTHNESS_LAMBDA_RATIO` from the `use crate::heads::{...}` line (the import added in predecessor Task 5).

  b. **Add** the controller cubin constant adjacent to `SMOOTHNESS_CUBIN`:
  ```rust
  const SMOOTHNESS_CONTROLLER_CUBIN: &[u8] = include_bytes!(
      concat!(env!("OUT_DIR"), "/smoothness_lambda_controller.cubin")
  );
  ```

  c. **Add** new struct fields adjacent to the existing `smoothness_*` fields (after `_smoothness_module`):
  ```rust
      /// Per-horizon EMA of `raw_per_h` from output_smoothness. Updated
      /// by `smoothness_lambda_controller` each step. Drives the
      /// controller's target derivation.
      smoothness_jitter_ema_d: CudaSlice<f32>,
      /// First-observation sentinel (per pearl_first_observation_bootstrap).
      /// 0 → next step bootstraps EMA from raw_per_h; 1 → Wiener-α update.
      smoothness_jitter_first_obs_d: CudaSlice<i32>,
      /// Cached handle for `smoothness_lambda_controller`.
      smoothness_controller_fn: CudaFunction,
      _smoothness_controller_module: Arc<CudaModule>,
  ```

  d. **In the constructor**, after the existing smoothness module load, add the controller module load:
  ```rust
          let smoothness_controller_module = ctx
              .load_cubin(SMOOTHNESS_CONTROLLER_CUBIN.to_vec())
              .context("load smoothness_lambda_controller cubin")?;
          let smoothness_controller_fn = smoothness_controller_module
              .load_function("smoothness_lambda_controller")
              .context("load smoothness_lambda_controller")?;
  ```

  e. **In the constructor**, **REPLACE** the existing λ init block (which computes `base × ratio[h]` per the predecessor Task 5) with a uniform-floor init:
  ```rust
          // ── CRT.train: output-smoothness regularizer state ──
          // λ is now ISV-driven by smoothness_lambda_controller — initialised
          // to LAMBDA_FLOOR per horizon and updated each step inside the
          // captured graph. The base_lambda amplitude knob enters the
          // kernel as a scalar arg at launch time.
          const LAMBDA_FLOOR_INIT: f32 = 1.0e-4;
          let smoothness_lambda_host: Vec<f32> = vec![LAMBDA_FLOOR_INIT; N_HORIZONS];
          let smoothness_lambda_d = {
              let staging = unsafe { MappedF32Buffer::new(N_HORIZONS) }
                  .map_err(|e| anyhow::anyhow!("smoothness λ staging: {e}"))?;
              staging.write_from_slice(&smoothness_lambda_host);
              let mut dst = stream.alloc_zeros::<f32>(N_HORIZONS)
                  .context("smoothness_lambda_d alloc")?;
              let nbytes = N_HORIZONS * std::mem::size_of::<f32>();
              unsafe {
                  let (dst_ptr, _g) = dst.device_ptr_mut(&stream);
                  cudarc::driver::result::memcpy_dtod_async(
                      dst_ptr, staging.dev_ptr, nbytes, stream.cu_stream(),
                  ).context("smoothness λ DtoD")?;
              }
              dst
          };
          let smoothness_loss_d = stream.alloc_zeros::<f32>(1)
              .context("smoothness_loss_d alloc")?;
          let smoothness_loss_host_d = unsafe { MappedF32Buffer::new(1) }
              .map_err(|e| anyhow::anyhow!("smoothness_loss_host_d: {e}"))?;
          let smoothness_loss_per_horizon_d = stream.alloc_zeros::<f32>(N_HORIZONS)
              .context("smoothness_loss_per_horizon_d alloc")?;
          let smoothness_loss_per_horizon_host_d = unsafe { MappedF32Buffer::new(N_HORIZONS) }
              .map_err(|e| anyhow::anyhow!("smoothness_loss_per_horizon_host_d: {e}"))?;
          // Controller state buffers.
          let smoothness_jitter_ema_d = stream.alloc_zeros::<f32>(N_HORIZONS)
              .context("smoothness_jitter_ema_d alloc")?;
          let smoothness_jitter_first_obs_d = stream.alloc_zeros::<i32>(1)
              .context("smoothness_jitter_first_obs_d alloc")?;
  ```

  f. **In the `Ok(Self { ... })` block**, add the new controller fields (alphabetically near the other `smoothness_*` fields):
  ```rust
              smoothness_jitter_ema_d,
              smoothness_jitter_first_obs_d,
              smoothness_controller_fn,
              _smoothness_controller_module: smoothness_controller_module,
  ```

  g. **In `step_batched`** (inside the captured-graph region), find the smoothness launch added in predecessor Task 6 (around line 1991-2008). IMMEDIATELY AFTER the smoothness launch block's closing brace, insert the controller launch:
  ```rust

          // ── 5d. ISV-driven λ controller ──
          //       Reads the raw per-horizon mean-sq-diff just emitted by
          //       the output_smoothness kernel, updates jitter_ema, derives
          //       per-horizon target anchored on h30, and produces λ for
          //       the NEXT step's output_smoothness launch. Same captured
          //       graph; one-step delay between observation and effect
          //       (standard closed-loop pattern).
          let base_lambda = self.cfg.smoothness_base_lambda;
          let smooth_ctrl_cfg = LaunchConfig {
              grid_dim: (1, 1, 1),
              block_dim: (N_HORIZONS as u32, 1, 1),
              shared_mem_bytes: 0,
          };
          {
              let mut launch = self.stream.launch_builder(&self.smoothness_controller_fn);
              launch
                  .arg(&self.smoothness_loss_per_horizon_d)
                  .arg(&mut self.smoothness_jitter_ema_d)
                  .arg(&mut self.smoothness_jitter_first_obs_d)
                  .arg(&base_lambda)
                  .arg(&mut self.smoothness_lambda_d);
              unsafe { launch.launch(smooth_ctrl_cfg).context("smoothness_lambda_controller launch")?; }
          }
  ```

- [ ] **4.2** Verify:
```bash
SQLX_OFFLINE=true cargo check -p ml-alpha --all-targets 2>&1 | grep -E "^error" | head
```
Expected: clean.

- [ ] **4.3** If you've removed `SMOOTHNESS_LAMBDA_RATIO` import successfully and the build still compiles, commit BOTH file changes atomically:
```bash
git add crates/ml-alpha/src/heads.rs crates/ml-alpha/src/trainer/perception.rs
git commit -m "feat(crt-train): wire ISV-driven λ controller; remove SMOOTHNESS_LAMBDA_RATIO"
```

---

## Task 5: Write controller invariant tests

- [ ] **5.1** Create `crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs`. Use the same upload/download helpers as `output_smoothness_grad_finite_diff.rs` (mapped-pinned + memcpy_dtod_async).

Tests to implement (4 #[test] functions):

1. `first_observation_bootstraps_ema_and_sets_floor`:
   - Init `first_obs=0`, `jitter_ema=[0;5]`, raw_per_h=[0.5, 0.3, 0.1, 0.05, 0.01], base_lambda=0.001
   - After kernel: `jitter_ema == raw_per_h`, `first_obs==1`, `λ[h] == LAMBDA_FLOOR` for all h (because excess_ratio = 0 on the bootstrap step: target = raw_per_h[0]*ratio[h], and jitter_ema[h] = raw_per_h[h] which is at-target when raw_per_h IS scaled by ratio; assert λ near LAMBDA_FLOOR or base_lambda, NOT runaway)
   - (Actually on first observation, jitter_ema[h] = raw_per_h[h] for all h. target[h] = raw_per_h[0]*ratio[h]. excess_ratio[h] = (raw_per_h[h] / (raw_per_h[0]*ratio[h])) - 1. If the test setup has raw_per_h[h] = raw_per_h[0]*ratio[h], excess_ratio = 0 → λ[h] = max(LAMBDA_FLOOR, base_lambda). Use a test setup that satisfies this for clean assertion.)

2. `steady_state_at_target_yields_base_lambda`:
   - Pre-set `first_obs=1`, `jitter_ema = [0.5, 0.15, 0.05, 0.015, 0.0025]` (= 0.5 × [1, 0.3, 0.1, 0.03, 0.005])
   - raw_per_h same as above
   - base_lambda = 0.01
   - Expected: jitter_ema stays at target (Wiener-α update at α=0.5: same value), excess_ratio = 0, λ[h] = max(LAMBDA_FLOOR, 0.01) = 0.01 for all h

3. `excess_at_h6000_lifts_lambda_proportionally`:
   - Pre-set `first_obs=1`, `jitter_ema = [0.5, 0.15, 0.05, 0.015, 0.025]` (h6000 at 10× target)
   - raw_per_h same
   - base_lambda = 0.01
   - After kernel: jitter_ema[h6000] EMA update → 0.5*0.025 + 0.5*0.025 = 0.025 (raw matches old EMA, no change)
   - Wait: if raw == old_ema, jitter_ema doesn't change. Use raw_per_h[h6000] = 0.025 (matches the pre-set ema)
   - target[h6000] = jitter_ema_new[0]*0.005 = 0.5*0.005 = 0.0025
   - excess_ratio = 0.025/0.0025 - 1 = 9.0
   - λ[h6000] = 0.01 * (1 + 9) = 0.10
   - Assert λ[h6000] is within 0.1 ± 0.005 (allowing for ema update on h0 if it shifts)

4. `lambda_floor_when_base_lambda_zero`:
   - Pre-set `first_obs=1`, `jitter_ema = [0.5; 5]`, raw matching
   - base_lambda = 0.0
   - Expected: λ[h] = LAMBDA_FLOOR = 1e-4 for all h (max(1e-4, 0 * (1 + excess)) = 1e-4)

- [ ] **5.2** Build:
```bash
SQLX_OFFLINE=true cargo test -p ml-alpha --test smoothness_lambda_controller_invariants --no-run 2>&1 | tail
```
Expected: clean.

If local CUDA available:
```bash
SQLX_OFFLINE=true cargo test -p ml-alpha --test smoothness_lambda_controller_invariants -- --nocapture 2>&1 | tail -30
```
Expected: 4 tests pass.

- [ ] **5.3** Commit:
```bash
git add crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs
git commit -m "test(crt-train): invariant tests for smoothness_lambda_controller"
```

---

## Task 6: Push + validate

- [ ] **6.1** Run full check + perception_overfit:
```bash
cd /home/jgrusewski/Work/foxhunt
SQLX_OFFLINE=true cargo check -p ml-alpha --all-targets 2>&1 | tail
SQLX_OFFLINE=true cargo test -p ml-alpha --test perception_overfit -- --nocapture 2>&1 | tail -30  # if local CUDA
```
Expected: clean check; perception_overfit 9/9 (smoothness=0 base_lambda + LAMBDA_FLOOR = essentially-no-op).

- [ ] **6.2** Push:
```bash
git push origin ml-alpha-phase-a
```

- [ ] **6.3** Submit validation retrain at `smoothness-base-lambda=0.001` (the amplitude knob; the controller scales adaptively from there):
```bash
argo submit -n foxhunt --from=wftmpl/alpha-perception \
  -p commit-sha=$(git rev-parse HEAD) \
  -p git-branch=ml-alpha-phase-a \
  -p gpu-pool=ci-training-l40s \
  -p smoothness-base-lambda=0.001 \
  -p n-train-seqs=8000 \
  -p n-val-seqs=1000 \
  -p epochs=5
```

Watch logs for `final_smooth_loss_per_horizon` in the summary JSON. Expected: h30 ≈ h6000_target × 200, h6000 ≈ h30 / 200 (controller pulls h6000 to its 200× lower target).

---

## Self-Review

1. Spec extension §3 mapped to Task 1 (kernel), Task 4 (wiring). ✓
2. Spec extension §3.4 removals mapped to Task 3 (heads.rs) + Task 4 (perception.rs constructor). ✓
3. Spec extension §5 invariants mapped to Task 5 (4 tests). ✓
4. No placeholders. No TODO markers. ✓
5. Atomic refactor: Task 4 ships heads.rs deletion + perception.rs wiring in ONE commit (per `feedback_no_partial_refactor`). ✓
6. Memory rules cited:
   - `feedback_no_atomicadd` ✓ (5 threads, single writer per slot)
   - `feedback_no_htod_htoh_only_mapped_pinned` ✓ (λ init still uses MappedF32Buffer)
   - `feedback_no_nvrtc` ✓ (build.rs registration)
   - `pearl_no_host_branches_in_captured_graph` ✓ (controller launch has no host branches)
   - `pearl_first_observation_bootstrap` ✓ (sentinel pattern in kernel)
   - `pearl_blend_formulas_must_have_permanent_floor` ✓ (LAMBDA_FLOOR is permanent)
   - `pearl_wiener_alpha_floor_for_nonstationary` ✓ (α=0.5 floor)
   - `pearl_controller_anchors_isv_driven` ✓ (target anchored on observed h30 jitter)
   - `feedback_no_partial_refactor` ✓ (atomic delete+wire in Task 4)

Ready to execute.
