# Mamba2 + GRN Integration Notes (Phase E.4.A Task 1)

**Status:** research output for the Phase E.4.A implementation plan.

## TL;DR

We don't need to port the production trainer's Mamba2 + GRN — both are
**deeply coupled to the production policy class** (SH2=256 trunk hidden
dim, ofi_embed[10], ISV-driven temporal routing, K=8 history, residual
to h_s2 from upstream trunk).

Instead, **use `ml_alpha::mamba2_block::Mamba2Block`** — a from-scratch
Phase-1d.1 Mamba2 implementation that's:
- Fully configurable (in_dim, hidden_dim, state_dim, seq_len)
- GPU-pure (`CudaSlice<f32>` for all weights/state)
- Has `forward_train` + `backward` + AdamW step
- Already used in `ml-alpha` for Phase 1d.2 (lifted AUC 0.50 → 0.66)
- **`ml-alpha` is already a workspace dep of `ml`** (`Cargo.toml` line:
  `ml-alpha.workspace = true # FxCacheReader for the alpha_dqn_h600_smoke fxcache loader`)

For GRN: similarly, the production GRN kernels are coupled to the
production trunk. Initial Phase E.4 will SKIP GRN — use `Mamba2Block`
output (`h_enriched [B, hidden_dim]`) directly as input to C51. GRN can
come back later as a Phase E.4.B add-on if needed.

## Mamba2 — production stack (NOT used)

Source: `crates/ml/src/cuda_pipeline/mamba2_temporal_kernel.cu`

**Kernels:**
- `mamba2_update_history(h_history, h_s2, ofi_embed, N, K, sh2, embed_dim)` — circular history shift + insert (sh2=256 typical)
- `isv_temporal_route(isv_embedding[8], w_temporal_route, b_temporal_route, temporal_weight[SH2], SH2)` — ISV-driven per-feature temporal weight
- `mamba2_copy_enriched(src, dst, N)` — graph-safe DtoD copy
- `mamba2_scan_projected_fwd(...)` — forward SSM scan; needs A_proj, B_proj from upstream cuBLAS, plus h_s2 for residual
- `mamba2_scan_projected_bwd(...)` — backward
- `mamba2_scale_d_enriched(...)` — gradient scaling
- `extract_ofi_embed_grad(...)` — OFI embed gradient extraction

**Why we can't use it directly for Phase E:**
- Assumes h_s2 (SH2=256) trunk encoding upstream
- Coupled to ofi_embed[10] from production state pipeline
- Uses ISV[8] for temporal routing (not Phase E ISV layout)
- Residual semantics (output = h_s2 + temporal_context) require trunk encoder upstream
- K=8 history hardcoded in production wiring

**Rust launcher**: private inside `gpu_dqn_trainer.rs` (35932 lines; methods
like `mamba2_scan_projected_fwd` at line 12309, `mamba2_update_history` at
12332). Not exposed for external use.

## Mamba2 — ml-alpha stack (USE THIS)

Source: `crates/ml-alpha/src/mamba2_block.rs` + `crates/ml-alpha/cuda/mamba2_alpha_kernel.cu`

### `Mamba2BlockConfig`

```rust
pub struct Mamba2BlockConfig {
    pub in_dim: usize,       // Phase E: 10 (STATE_DIM)
    pub hidden_dim: usize,   // Phase E: 32 (start; sweep 16/32/64)
    pub state_dim: usize,    // ≤ 16; start at 16 (kernel max)
    pub seq_len: usize,      // ≤ 32 (kernel max); start at 16
}
```

Kernel hard limits: `MAMBA2_KERNEL_STATE_MAX = 16`, `MAMBA2_KERNEL_SEQ_MAX = 32`.

### Architecture

```
input [B, K=seq_len, in_dim]
  │  cuBLAS GEMM (W_in)
  ▼
x [B, K, hidden_dim]
  ├─ cuBLAS GEMM (W_a) → a_proj [B, K, state_dim] (gate via sigmoid)
  ├─ cuBLAS GEMM (W_b) → b_proj [B, K, state_dim] (input)
  ▼
mamba2_alpha_scan_fwd kernel
  (selective SSM scan over K timesteps; sigmoid gating;
   W_c [hidden_dim, state_dim] mixes state into per-position hidden output)
  ▼
h_enriched [B, hidden_dim]   (zero residual — we don't add anything)
  │  cuBLAS GEMM (W_out)
  ▼
logit [B, 1]                  (only used by ml-alpha; Phase E will read h_enriched directly)
```

### Public API

- `Mamba2Block::new(config, stream)` — Xavier-init weights, load cubin
- `block.forward(input: &GpuTensor) -> Result<GpuTensor>` — inference; returns h_enriched
- `block.forward_train(input: &GpuTensor) -> Result<(GpuTensor, Mamba2ForwardCache)>` — training-mode; cache for backward
- `block.backward(d_h_enriched: &GpuTensor, cache: &Mamba2ForwardCache) -> Result<Mamba2BackwardGrads>` — analytical backward
- `block.param_count() -> usize`

### AdamW step

```rust
pub struct Mamba2AdamW { ... }
impl Mamba2AdamW {
    pub fn new(block: &Mamba2Block, config: Mamba2AdamWConfig) -> Result<Self>
    pub fn step(&mut self, block: &mut Mamba2Block, grads: &Mamba2BackwardGrads) -> Result<()>
    pub fn set_learning_rate(&mut self, lr: f32)
}
```

### Phase E config

For initial wiring:
- `in_dim = 10` (Phase E state)
- `hidden_dim = 32` (sweep candidate)
- `state_dim = 16` (kernel max)
- `seq_len = 16` (window length)

Memory footprint: small — `W_in[32, 10]`, `W_a[16, 32]`, `W_b[16, 32]`, `W_c[32, 16]`, `W_out[1, 32]` = 320 + 512 + 512 + 512 + 32 = ~1900 params. Trivial.

### Output for C51

`h_enriched` is `[B, hidden_dim] = [B, 32]`. This **replaces the 10-dim state input** to the C51 head. Smoke binary changes:
- C51 W reshape: `[N_ACTIONS * n_atoms, hidden_dim]` instead of `[N_ACTIONS * n_atoms, STATE_DIM]`
- Pass `hidden_dim_i = 32 as i32` to `launch_alpha_c51_forward` in place of `state_dim_i`

### What we're NOT taking from production GRN

Production GRN (`grn_kernel.cu`) is 8 kernels: `grn_elu_inplace`, `grn_glu_forward`, `grn_residual_layernorm_forward`, plus 5 backward kernels.

For Phase E.4.A we **skip GRN entirely**. Mamba2's output goes straight into C51. If E.4.A passes the gates, we add GRN in E.4.B alongside MoE.

This is a deliberate scope reduction — adopting Mamba2's expressive trunk first, then adding GRN as a separate experiment if Sharpe doesn't lift enough. Easier to debug, easier to A/B.

## Controller-at-inference (Task 2 also captured here)

Source: `crates/ml/src/cuda_pipeline/stacker_threshold_controller.cu`

The controller kernel does NOT have any `if (training)` branches. It's a
pure GPU controller — reads rollout scalars, updates ISV slots, runs
Pearl A + Pearl D Wiener-α math.

**Wiener state lifecycle:** allocated once at smoke startup, lives
across all episodes. The state (`[sample_var, diff_var, x_lag]` for
slot 545) is supposed to persist — that's how the Wiener-α adapts.

**For backtest eval:** the same lifecycle works. Allocate
`ctl_wiener_dev` at the start of eval, fire the controller after each
eval episode. Wiener state accumulates across episodes (across cost
cells too — that's fine; the controller is observing the policy under
shifting cost regimes which IS the point of ISV-continual learning).

If we want a fresh controller per cost cell, we'd zero `ctl_wiener_dev`
and reset `isv_host[STACKER_THRESHOLD_INDEX]` between cells. Decision:
**preserve across cells** for the first cut — this gives the most
natural "live deployment" simulation.

## Revisions to Phase E.4.A plan

Original Tasks 8-10 sketched custom Mamba2 + GRN kernel wiring. The
revision is much simpler:

- **Task 8** (new): Wire `Mamba2Block` into smoke. Add `use ml_alpha::mamba2_block::{Mamba2Block, Mamba2BlockConfig};` to smoke binary. Construct a `Mamba2Block` at startup. Per-step: build a `GpuTensor` from the window buffer; call `forward_train`; pass `h_enriched` to C51 forward.
- **Task 9** (new): Wire `Mamba2AdamW` step. Replace plain SGD with AdamW for Mamba2 params; keep plain SGD for C51 W/b. (Or use ml-alpha's AdamW for all params — simpler.)
- **Task 10** (revised): Backward pass — call `block.backward(d_h_enriched, &cache)` to get `Mamba2BackwardGrads`, then `adamw.step(&mut block, &grads)`. The d_h_enriched flowing into Mamba2 is the gradient from C51's `dW_input` (i.e., gradient w.r.t. its input features).
- **Skipping GRN** for E.4.A.

This simplifies the plan substantially. Three tasks (8/9/10) become "use existing crate", not "write custom kernels".

Risk surfaces:
- Type interop: `Mamba2Block` uses `GpuTensor` (`ml_core::cuda_autograd::gpu_tensor`); smoke uses raw `CudaSlice<f32>` + device pointers. Need to bridge.
- `forward_train` builds a cache that must persist until `backward`. Lifetimes: the cache holds device tensors; we hold it in a `Option<Mamba2ForwardCache>` between forward and backward.
- C51 grad currently doesn't compute `d_input` (gradient w.r.t. the input state). We need to add that to provide the input gradient for Mamba2 backward. **This is a NEW kernel: `alpha_c51_grad_input_kernel`.**

The C51 grad-input kernel is the one piece of new CUDA we genuinely need. ~50 lines.
