# CRT Phase A — forward_only Cost Investigation

**Date:** 2026-05-20
**Branch:** `ml-alpha-phase-a`
**Task:** A0 (read-only research)
**Author:** CRT-A0 agent

---

## 1. TL;DR

**This is Case 2: stateless K-window.** Every call to `forward_only` in
`crates/ml-alpha/src/trainer/perception.rs` processes the full K-window from
scratch. The Mamba2 SSM recurrent state (`h_s2`) is zero-initialised once at
scratch allocation time and is never updated between calls — the scan kernel
treats it as a fixed zero residual, not a carry-over from the previous call.
Concretely, `h_s2` is a `[B, hidden_dim]` buffer that never receives a write
after construction.

**However:** the spec's §3.2 claim "Trunk forward pass (already every event)"
is **WRONG** as of HEAD. The current harness at `harness.rs:242` places
`forward_only` INSIDE the `if self.event_count % stride == 0` gate — it is
called every 200 events, not every event. This directly contradicts §3.2 and is
the primary finding of this investigation. Moving `forward_only` outside the
stride gate would mean calling the full K=64 window scan 200× more often,
which is a ~200× GPU cost increase for that call alone.

**Recommendation for A0.5:** A0.5 MUST implement `forward_step` — an
incremental single-step SSM inference that carries persistent `h_s2` state
across calls instead of zeroing it per call. This is the only path to ≤ 2×
wall-time on Phase A. Task CRT.A.0.5 is not a no-op.

---

## 2. Evidence

### 2.1 forward_only is inside the stride gate

File: `crates/ml-backtesting/src/harness.rs`

```rust
// Line 242 — both forward_only AND step_decision are gated:
if self.event_count % stride == 0 && self.snapshot_window.len() == self.seq_len {
    let window: Vec<Mbp10RawInput> = self.snapshot_window.iter().cloned().collect();
    let probs_all = self.trainer.forward_only(&window)   // line 244
        .context("trainer.forward_only")?;
    // ...
    self.sim.broadcast_alpha(&last_probs)?;
    self.sim.step_decision_with_latency(raw.ts_ns, &self.sim_config)?;
}
```

The comment at line 239 says "At decision-stride boundaries: run forward inference + sim decision." Both operations are inside the same gate. The spec's §3.2 says "Trunk forward pass (already every event)" — this is incorrect.

### 2.2 forward_only signature requires the full K-window

File: `crates/ml-alpha/src/trainer/perception.rs`, line 2301:

```rust
pub fn forward_only(&mut self, snapshots: &[Mbp10RawInput]) -> Result<Vec<f32>> {
    let b_sz = self.cfg.n_batch;
    let k_seq = self.cfg.seq_len;
    anyhow::ensure!(
        snapshots.len() == b_sz * k_seq,   // line 2305 — requires B*K snapshots
        ...
    );
    ...
}
```

The function hard-enforces that exactly `n_batch * seq_len` snapshots are
provided. With the default `seq_len=64` and `n_batch=1`, every call requires
exactly 64 snapshots. There is no single-step variant.

### 2.3 Mamba2 h_s2 is never carried between calls

File: `crates/ml-alpha/src/mamba2_block.rs`, lines 151–165:

```rust
/// `h_s2` is zero-initialised once and never written (no residual carry
/// from a prior chunk in the supervised path). The scan kernel reads it
/// as a constant addition to h_enriched_seq.
pub struct Mamba2BlockForwardScratch {
    ...
    pub h_s2: GpuTensor,  // [n_batch, hidden_dim]  (zero residual)
    ...
}
```

And in `forward_train_seq_into` (`mamba2_block.rs:1116`), `h_s2` is passed as
the initial-state argument to the scan kernel:

```rust
.arg(scratch.h_s2.cuda_data())   // line 1116 — always zero, never updated
```

`h_s2` is allocated once in `Mamba2BlockForwardScratch::new` via
`GpuTensor::zeros` (line 191) and never written afterward. It is a constant
zero buffer. Each call to `forward_train_seq_into` starts the SSM scan from
h=0 regardless of what was computed on the previous call.

The same pattern holds for both Mamba2 layers:
- Layer 1: `self.trunk.mamba2_l1_mut().forward_train_seq_into(&self.vsn_out_d, &mut self.mamba2_fwd_scratch)` at line 2506
- Layer 2: `self.trunk.mamba2_l2_mut().forward_train_seq_into(&self.ln_a_out_d, &mut self.mamba2_l2_fwd_scratch)` at line 2531

Both scratch objects hold their own `h_s2: GpuTensor` initialised to zeros.

### 2.4 The CUDA Graph capture bakes in K=seq_len steps

File: `crates/ml-alpha/src/trainer/perception.rs`, lines 2358–2396:

The three-state machine (warmup → capture → replay) captures the full kernel
chain — VSN → Mamba2×2 → LN×2 → transpose → attn-pool → CfC K-loop → heads —
for a fixed `(b_sz, k_seq, total_snaps)` shape. A `forward_step` variant would
need its own graph with `(b_sz, 1, b_sz)` shape and persistent state passed via
device pointers rather than the scratch zero buffer.

### 2.5 Production seq_len

From `crates/ml-alpha/examples/alpha_train.rs`, line 62:

```rust
#[arg(long, default_value_t = 64)]
seq_len: usize,
```

Default seq_len is 64. Each `forward_only` call runs 64 Mamba2 steps for both
layers.

---

## 3. Cost Analysis

### 3.1 Per-call kernel work

Each `forward_only` call executes:
1. Host-side staging fill: `B*K * ~60` float32 writes (< 1 µs, OUTSIDE graph)
2. CUDA Graph replay or first-pass dispatch:
   - `snap_feature_assemble_batched`: `B*K` rows
   - VSN forward: `B*K` rows
   - Mamba2 L1 `scan_fwd_seq`: K sequential steps over `[B, K, hidden_dim]`
   - LN_a forward: `B*K` rows
   - Mamba2 L2 `scan_fwd_seq`: K sequential steps
   - LN_b forward: `B*K` rows
   - Transpose: `[B, K, H]` → `[K, B, H]`
   - Attention pool: K steps
   - CfC K-loop: K steps
   - Multi-horizon heads: K positions × N_HORIZONS outputs
3. Stream sync + DtoH probs: `K * B * N_HORIZONS` floats (64 × 1 × 4 = 256 f32)

The Mamba2 scan is a sequential K-step recurrence — it cannot be parallelised
over K. Cost scales linearly with K.

### 3.2 Cost at stride=200 vs stride=1

At `decision_stride=200`, `forward_only` is called once per 200 events.
If the stride gate is removed without adding `forward_step`, `forward_only`
is called once per event: **200× more calls, each still running K=64 steps**.

With CUDA Graph replay the per-call overhead is low (mostly the DtoD staging
copy + graph launch), but the actual Mamba2 scan work is O(K) per call. Calling
K=64-step inference 200× more frequently is a ~200× increase in GPU kernel
work for the encoder path.

The controller kernels (`decision_policy_*` + `seed_inflight_limits_batched`)
are indeed light (block-per-backtest, O(1) work). The spec's §3.4 cost estimate
correctly identifies these as the dominant concern — but it implicitly assumes
the trunk forward pass stays at stride=200. It is not correct that the forward
pass is "already every event."

### 3.3 Correct cost model for `forward_step`

If `forward_step` is implemented with persistent h_s2 state:
- Per-call work: 1 Mamba2 step × 2 layers = O(hidden_dim × state_dim)
- No K-window scatter/gather: only 1 snapshot staged per call
- No CfC K-loop: only the final position's head output needed
- Estimated per-call GPU time: ~1/64 of current `forward_only`

At stride=1 with `forward_step`, total encoder GPU work ≈ 200/64 ≈ 3.1× current.
Combined with controller kernels (light, 200× more calls), total wall-time
should land in the 2–4× range, consistent with the §3.4 budget of ≤ 2×
(achievable if the forward step is sufficiently fast relative to controller work).

---

## 4. Recommendation for Task A0.5

**A0.5 is required and non-trivial.** The following kernel changes are needed:

### 4.1 New Rust struct: `Mamba2BlockStepScratch`

Mirrors `Mamba2BlockForwardScratch` but sized for `K=1`:
- `x: GpuTensor` — `[B, hidden_dim]`
- `a_proj: GpuTensor` — `[B, state_dim]`
- `b_proj: GpuTensor` — `[B, state_dim]`
- `h_s2: GpuTensor` — **`[B, hidden_dim]` — PERSISTENT, updated each call**
- `h_out: GpuTensor` — `[B, hidden_dim]` — single-step enriched output

Location: add to `crates/ml-alpha/src/mamba2_block.rs` alongside
`Mamba2BlockForwardScratch`.

### 4.2 New kernel call path: `Mamba2Block::step_into`

Mirrors `forward_train_seq_into` but:
- Input shape: `[B, 1, in_dim]` (single snapshot)
- Calls `kernel_fwd_seq` with `k_i32 = 1`
- After the scan, **copies `h_out` into `h_s2`** (carry-forward)
- Or: a dedicated `mamba2_alpha_scan_step` kernel that reads/writes `h_s2`
  in-place

The Mamba2 `scan_fwd_seq` kernel at
`crates/ml-alpha/src/kernels/mamba2_alpha_kernel.cu` (referenced at
`crates/ml-alpha/build.rs:11`) will need to be called with K=1. Check whether
the kernel handles K=1 without corner-case issues (e.g., shared-memory tile
sizing assumptions). If K=1 is not safe, a dedicated `scan_fwd_step` kernel
variant is required.

### 4.3 New method: `PerceptionTrainer::forward_step`

Signature:

```rust
pub fn forward_step(&mut self, snapshot: &Mbp10RawInput) -> Result<[f32; N_HORIZONS]>
```

Behaviour:
1. Stage 1 snapshot into pinned buffers (outside CUDA Graph)
2. Run VSN on 1 row
3. Run `Mamba2Block::step_into` (L1) — updates L1 h_s2 in-place
4. Run LN_a on 1 row
5. Run `Mamba2Block::step_into` (L2) — updates L2 h_s2 in-place
6. Run LN_b on 1 row
7. Run the final CfC step + head projection for 1 position
8. Sync + DtoH: `N_HORIZONS` floats (4 f32)
9. Return probs for the current (single) position

This method should have its own CUDA Graph (capture on second call, replay
thereafter). Because K=1, the graph is small and replay overhead dominates —
confirm with profiling that graph replay is faster than eager dispatch at K=1.

### 4.4 Harness changes (in A1)

Once `forward_step` exists, `harness.rs` changes are:

```rust
// Before (stride-gated):
if self.event_count % stride == 0 && self.snapshot_window.len() == self.seq_len {
    let window = self.snapshot_window.iter().cloned().collect::<Vec<_>>();
    let probs_all = self.trainer.forward_only(&window)?;
    ...
}

// After (every event, once window bootstrapped):
if self.snapshot_window.len() == self.seq_len {
    let probs = self.trainer.forward_step(&raw)?;
    self.sim.broadcast_alpha(&probs)?;
    self.sim.step_decision_with_latency(raw.ts_ns, &self.sim_config)?;
}
```

The sliding window (`snapshot_window: VecDeque<Mbp10RawInput>`) is still
needed for the warm-up period (until K events are seen) but thereafter serves
only as a `len() == self.seq_len` bootstrap check — the SSM state carries
forward in the L1/L2 scratch buffers.

### 4.5 Bootstrap concern

`forward_only` bootstraps Mamba2 with K=64 steps of context. `forward_step`
starts from h_s2=0. For the first `seq_len-1` events the model has no
accumulated context — the predictions will be from a cold-start state.

Mitigation options (implementer should choose):
a. Keep the first inference as `forward_only` (at `event_count == seq_len-1`),
   then transfer the final `h_s2` state from that call into the step scratch,
   and switch to `forward_step` thereafter. This requires a new method
   `PerceptionTrainer::transfer_step_state_from_window` or similar.
b. Accept cold-start predictions for the first K events; the Wiener-α EMA
   in conviction smoothing (spec §4.2) will dampen the cold-start noise.
c. Use option (b) but gate decisions on `event_count >= seq_len` — already
   present in the harness via the `self.snapshot_window.len() == self.seq_len`
   check.

Option (c) is the simplest: no changes to bootstrap semantics, the existing
warm-up guard stays, and the first K events remain decision-free. The h_s2
state simply accumulates naturally from zero.

---

## 5. Open Questions for A0.5 Implementer

1. **K=1 in scan_fwd_seq kernel**: Does the existing Mamba2 CUDA kernel at
   `mamba2_alpha_kernel.cu` handle `K=1` safely? Shared-memory block sizing and
   warp-tile assumptions may break at K=1. If so, a dedicated `scan_fwd_step`
   kernel variant is required. Verify with the K=1 case before implementing the
   full graph.

2. **h_s2 carry-forward mechanism**: The cleanest implementation writes h_out
   over h_s2 in-place at the end of each step. But `h_s2` in the current
   scratch is declared immutable (`cuda_data()` not `data_mut()`) and is passed
   read-only to the scan kernel. Either add a write-back pass after the scan, or
   restructure `step_into` to accept a mutable h_s2 that it writes in-place.

3. **CUDA Graph at K=1**: Graph replay overhead is typically 5–15 µs. At K=1,
   the actual kernel work may be shorter than the replay overhead. Profile the
   K=1 eager dispatch vs graph replay on the L40S to determine whether graphing
   is worth it.

4. **Session resets**: The SSM state in `forward_step` accumulates across
   session boundaries. The harness already has session-gap detection (via
   `ts_ns` gaps). On session-gap detection, `forward_step` should reset h_s2 to
   zero and require a new bootstrap period. Add a `reset_step_state()` method
   and call it from the harness on session-gap events.

5. **Spec §3.2 correction**: The spec states "Trunk forward pass (already every
   event)" — this is factually incorrect as of HEAD. The spec should be
   corrected (or a note added) to reflect that `forward_only` is currently
   stride-gated and that `forward_step` is the A0.5 deliverable. The A1
   implementer should know the spec was wrong on this point.
