diff --git a/docs/superpowers/specs/2026-05-15-mamba2-grn-integration-notes.md b/docs/superpowers/specs/2026-05-15-mamba2-grn-integration-notes.md new file mode 100644 index 000000000..fcddb7c56 --- /dev/null +++ b/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` 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` — 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` — analytical backward +- `block.param_count() -> usize` + +### AdamW step + +```rust +pub struct Mamba2AdamW { ... } +impl Mamba2AdamW { + pub fn new(block: &Mamba2Block, config: Mamba2AdamWConfig) -> Result + 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` + 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` 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.