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

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

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

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.