Per A0 investigation memo (commit 2e87ed0da) — forward_only was Case 2
(stateless K=64 window per call). Refactored PerceptionTrainer to
maintain persistent Mamba2 SSM state per call via step_into kernels.
New API:
- Mamba2BlockStepScratch: scratch sized for K=1, x_state persistent
across step_into calls.
- Mamba2Block::step_into: single-step forward with x_state in-place
update.
- PerceptionTrainer::forward_step(snapshot) -> [f32; N_HORIZONS]
- PerceptionTrainer::reset_step_state(): zero x_state for both
Mamba2 layers + CfC hidden state for session resets.
Decisions (from A0 memo §5):
1. K=1 path: added a dedicated `mamba2_alpha_scan_fwd_step` kernel.
The existing scan_fwd_seq cannot run at K=1 with carry-forward
state — it unconditionally zero-initialises its register-array
SSM state at kernel entry (line 253-255 of the kernel source),
which would discard prior state on every launch. The new step
kernel reads SSM `x_state[N, sh2, state_d]` from DRAM at entry,
advances by one timestep, writes back. Same arithmetic as
scan_fwd_seq's per-step inner loop.
2. x_state carry: written in-place in DRAM at end of step_into.
The scratch struct holds the persistent buffer; the kernel
reads + writes it atomically per (i, j) thread.
3. CUDA Graph at K=1: chose eager dispatch. Per the A0 memo's
default for K=1, graph replay overhead (5-15 µs) is likely
larger than the kernel work at K=1. Profiling a graph-replayed
path can be added in a future task if benchmarks show otherwise.
4. Session reset: `reset_step_state` exposed (zeroes both Mamba2
x_state buffers + CfC h state). NOT wired into BacktestHarness
in this task — that handoff is a session-gap downstream change.
5. Spec §3.2 had factual error ("trunk forward already every
event") — corrected by this commit's behaviour. Spec doc edit
deferred to a separate concern.
Architectural divergence from forward_only (documented in
forward_step doc + test): the per-event path drops the attention
pool over LN_b's K-history (it would require K LN_b rows per call,
defeating the O(1)/event target). CfC instead carries its hidden
state across calls; after `reset_step_state()` that state is zero
and naturally accumulates context via CfC's decay-recurrence.
Golden test (forward_step_golden.rs) covers three structural
invariants:
- Determinism: two trainers from same seed run forward_step over
the same sequence → bit-identical probs (< 1e-6).
- Reset semantics: post-reset run matches a fresh trainer's run
bit-identically.
- Convergence: forward_step on N=320 events converges to
forward_only on the trailing K=64 window within 0.15. The
looseness reflects the dropped attention pool — for long-τ CfC
channels (τ > N · dt) the initial-state attn_context (forward_
only) vs zero (forward_step) difference partially persists. Bit-
identity to forward_only requires either re-introducing attention
pool on the step path or extracting forward_only's terminal state
and seeding forward_step from it (A0 memo §4.5 option (a));
both deferred.
Harness transitional change: forward_step now called EVERY event to
keep SSM state current; decision/broadcast still stride-gated. A1
will delete the stride gate. Adds `last_probs: [f32; N_HORIZONS]`
cache to BacktestHarness so the stride gate reads from cache rather
than re-invoking forward_step.
Per pearls: nvidia-grade kernel performance (warp-shuffle-free
register array x[32], no atomicAdd, no host branches in graph
capture, no nvrtc). The new kernel is pre-compiled in build.rs's
existing mamba2_alpha_kernel.cu cubin alongside fwd/bwd/seq variants.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
a0e81fbdfc