docs(phase-e-4-a): mamba2 + grn integration research — use ml-alpha Mamba2Block

Findings:
- Production Mamba2 (gpu_dqn_trainer) is coupled to SH2=256 trunk +
  ofi_embed + ISV[8] temporal routing — not portable to Phase E.
- ml-alpha::mamba2_block::Mamba2Block is from-scratch, fully
  configurable (in_dim/hidden_dim/state_dim/seq_len), GPU-pure with
  forward_train/backward/AdamW. Used in Phase 1d.2 to lift AUC 0.50
  to 0.66. ml-alpha is already a workspace dep of ml.
- GRN skipped for E.4.A — Mamba2 output goes straight to C51 head.
  Reintroduce GRN in E.4.B if Sharpe gates don't pass.
- Controller-at-inference: kernel has no training-mode branches;
  Wiener state preserved across episodes/cost cells for natural
  live-deployment simulation.

Revises Tasks 8-10 of the plan: use Mamba2Block API instead of
writing custom kernels. Only new CUDA needed: alpha_c51_grad_input
(C51 gradient w.r.t. input features, for Mamba2 backward chain).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

docs/superpowers/specs/2026-05-15-mamba2-grn-integration-notes.md

@@ -0,0 +1,170 @@
+# 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.