# T10 — `Mamba2Block::backward_from_h_enriched` patch sketch

**Status:** Ready-to-apply patch (Phase E.4.A T10 prerequisite).
**Date:** 2026-05-15
**Target file:** `crates/ml-alpha/src/mamba2_block.rs`
**Why needed:** see `pearl_c51_thompson_closed_phase_e3_gap.md` and
`docs/superpowers/specs/2026-05-15-mamba2-grn-integration-notes.md`.

## Problem

`Mamba2Block::backward(cache, d_logit: [B, 1])` takes the gradient
w.r.t. the SCALAR logit output (the post-W_out projection). For
Phase E we want to backprop a `[B, hidden_dim]` gradient computed
by the new `alpha_c51_grad_input` kernel — that's the C51 head's
gradient w.r.t. its input, which IS `h_enriched`.

The existing backward chains:

    d_logit [B, 1]
       │  W_out backward
       ▼
    d_h_enriched [B, hidden_dim]
       │  (steps 5–1 as before)
       ▼
    dw_in, db_in, dw_a, db_a, dw_b, db_b, dw_c

For Phase E we want to ENTER at `d_h_enriched` directly and skip
the W_out step (since C51 replaces W_out's role).

## Patch

Add a new public method to `impl Mamba2Block` in
`crates/ml-alpha/src/mamba2_block.rs`. Same signature shape as
`backward` but takes `d_h_enriched` instead of `d_logit`. Returns
`Mamba2BackwardGrads` with `dw_out` and `db_out` set to zero tensors
(matching production AdamW's expectations — Adam steps on zero grad
are no-ops with the correct moment decay).

```rust
/// Backward variant that BYPASSES the W_out projection. Phase E
/// uses Mamba2's `h_enriched` directly as input to a downstream
/// classifier (C51 head), not Mamba2's own scalar logit. The
/// `d_logit`-based `backward` is unsuitable because Phase E never
/// computes a `d_logit` — there is no logit in its loss path.
///
/// `d_h_enriched` MUST have shape `[B, hidden_dim]` matching
/// `cache.h_enriched`. Returns `Mamba2BackwardGrads` where
/// `dw_out` and `db_out` are zero-initialised tensors of the
/// correct shapes (AdamW step on these is a no-op with the
/// correct moment-decay; the W_out parameters effectively freeze
/// — which is the correct semantics since Phase E never uses them).
pub fn backward_from_h_enriched(
    &self,
    cache: &Mamba2ForwardCache,
    d_h_enriched: &GpuTensor,
) -> Result<Mamba2BackwardGrads> {
    let c = &self.config;
    let n_batch = cache.h_enriched.shape()[0];
    let n_rows = n_batch * c.seq_len;

    if d_h_enriched.shape() != [n_batch, c.hidden_dim] {
        return Err(anyhow!(
            "Mamba2Block::backward_from_h_enriched: \
             d_h_enriched shape {:?} != [{}, {}]",
            d_h_enriched.shape(), n_batch, c.hidden_dim
        ));
    }

    // ── 6′ SKIPPED. Phase E doesn't go through W_out. ────────────
    // (Zero w_out grads are constructed at return.)

    // ── 5′. Allocate scan-backward scratch buffers ────────────────
    let per_chan_n = n_batch * c.hidden_dim * c.seq_len * c.state_dim;
    let per_sample_n = n_batch * c.hidden_dim * c.state_dim;
    let mut d_a_per_channel = self.stream
        .alloc_zeros::<f32>(per_chan_n)
        .map_err(|e| anyhow!("alloc d_a_per_channel: {e}"))?;
    let mut d_b_per_channel = self.stream
        .alloc_zeros::<f32>(per_chan_n)
        .map_err(|e| anyhow!("alloc d_b_per_channel: {e}"))?;
    let mut d_w_c_per_sample = self.stream
        .alloc_zeros::<f32>(per_sample_n)
        .map_err(|e| anyhow!("alloc d_w_c_per_sample: {e}"))?;
    let mut d_h_s2 = self.stream
        .alloc_zeros::<f32>(n_batch * c.hidden_dim)
        .map_err(|e| anyhow!("alloc d_h_s2: {e}"))?;

    let block_threads: u32 = 32;
    let grid_y_h: u32 =
        ((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
    let bwd_cfg = LaunchConfig {
        grid_dim: (n_batch as u32, grid_y_h, 1),
        block_dim: (block_threads, 1, 1),
        shared_mem_bytes: 0,
    };
    let n_i32 = n_batch as i32;
    let k_i32 = c.seq_len as i32;
    let sh2_i32 = c.hidden_dim as i32;
    let st_i32 = c.state_dim as i32;
    unsafe {
        self.stream
            .launch_builder(&self.kernel_bwd)
            .arg(cache.a_proj.cuda_data())
            .arg(cache.b_proj.cuda_data())
            .arg(d_h_enriched.cuda_data())   // <-- the caller's gradient
            .arg(&self.w_c)
            .arg(&mut d_a_per_channel)
            .arg(&mut d_b_per_channel)
            .arg(&mut d_w_c_per_sample)
            .arg(&mut d_h_s2)
            .arg(&n_i32)
            .arg(&k_i32)
            .arg(&sh2_i32)
            .arg(&st_i32)
            .launch(bwd_cfg)
            .map_err(|e| anyhow!("mamba2_alpha_scan_bwd launch: {e}"))?;
    }

    // ── Reduce d_a_per_channel and d_b_per_channel ───────────────
    let red_grid_z: u32 =
        ((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
    let red_cfg = LaunchConfig {
        grid_dim: (n_batch as u32, c.seq_len as u32, red_grid_z),
        block_dim: (block_threads, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut d_a_proj_flat: CudaSlice<f32> = self.stream
        .alloc_zeros::<f32>(n_rows * c.state_dim)
        .map_err(|e| anyhow!("alloc d_a_proj_flat: {e}"))?;
    let mut d_b_proj_flat: CudaSlice<f32> = self.stream
        .alloc_zeros::<f32>(n_rows * c.state_dim)
        .map_err(|e| anyhow!("alloc d_b_proj_flat: {e}"))?;
    unsafe {
        self.stream
            .launch_builder(&self.kernel_reduce_d_proj)
            .arg(&d_a_per_channel)
            .arg(&mut d_a_proj_flat)
            .arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32)
            .launch(red_cfg)
            .map_err(|e| anyhow!("reduce d_a_proj: {e}"))?;
        self.stream
            .launch_builder(&self.kernel_reduce_d_proj)
            .arg(&d_b_per_channel)
            .arg(&mut d_b_proj_flat)
            .arg(&n_i32).arg(&k_i32).arg(&sh2_i32).arg(&st_i32)
            .launch(red_cfg)
            .map_err(|e| anyhow!("reduce d_b_proj: {e}"))?;
    }

    // ── Reduce d_w_c_per_sample → dw_c ──────────────────────────
    let red_w_c_cfg = LaunchConfig {
        grid_dim: (c.hidden_dim as u32, red_grid_z, 1),
        block_dim: (block_threads, 1, 1),
        shared_mem_bytes: 0,
    };
    let mut dw_c: CudaSlice<f32> = self.stream
        .alloc_zeros::<f32>(c.hidden_dim * c.state_dim)
        .map_err(|e| anyhow!("alloc dw_c: {e}"))?;
    unsafe {
        self.stream
            .launch_builder(&self.kernel_reduce_d_w_c)
            .arg(&d_w_c_per_sample)
            .arg(&mut dw_c)
            .arg(&n_i32).arg(&sh2_i32).arg(&st_i32)
            .launch(red_w_c_cfg)
            .map_err(|e| anyhow!("reduce dw_c: {e}"))?;
    }

    // ── 3′. W_b backward ─────────────────────────────────────────
    let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim])
        .map_err(|e| anyhow!("reshape d_b_proj: {e}"))?;
    let x_act = LinearActivations { input: cache.x.clone() };
    let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self
        .w_b.inner
        .backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight,
                              &self.cublas, &self.stream)
        .map_err(|e| anyhow!("w_b backward: {e}"))?;

    // ── 2′. W_a backward ─────────────────────────────────────────
    let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim])
        .map_err(|e| anyhow!("reshape d_a_proj: {e}"))?;
    let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self
        .w_a.inner
        .backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight,
                              &self.cublas, &self.stream)
        .map_err(|e| anyhow!("w_a backward: {e}"))?;

    let d_x = d_x_from_a.add(&d_x_from_b, &self.stream)
        .map_err(|e| anyhow!("sum d_x branches: {e}"))?;

    // ── 1′. W_in backward ────────────────────────────────────────
    let input_act = LinearActivations { input: cache.input_2d.clone() };
    let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self
        .w_in.inner
        .backward_with_slices(&d_x, &input_act, &self.w_in.weight,
                              &self.cublas, &self.stream)
        .map_err(|e| anyhow!("w_in backward: {e}"))?;

    // ── 6′ replacement: zero W_out grads ────────────────────────
    // AdamW will see zero grad → effectively freezes w_out
    // parameters. Correct semantics — Phase E doesn't use them.
    let dw_out = GpuTensor::zeros(&[1, c.hidden_dim], &self.stream)
        .map_err(|e| anyhow!("alloc zero dw_out: {e}"))?;
    let db_out = GpuTensor::zeros(&[1], &self.stream)
        .map_err(|e| anyhow!("alloc zero db_out: {e}"))?;

    Ok(Mamba2BackwardGrads {
        dw_in, db_in, dw_a, db_a, dw_b, db_b, dw_c, dw_out, db_out,
    })
}
```

## Application steps

1. Open `crates/ml-alpha/src/mamba2_block.rs`.
2. Inside `impl Mamba2Block`, append the method above immediately
   after the existing `backward` method (around line ~549).
3. Run ml-alpha's existing tests to confirm no regression:
   ```bash
   SQLX_OFFLINE=true cargo test -p ml-alpha --release
   ```
   Existing `backward` is untouched, so all existing tests should
   pass without modification.
4. (Optional) Add a unit test that verifies the new method
   produces gradients consistent with manual `backward` chained
   through the W_out backward — i.e., feeding
   `d_logit = h_enriched @ W_out.T + b_out` partial gradient
   should produce the same scan/projection grads as
   `backward_from_h_enriched(d_h_enriched_recovered_via_W_out)`.

## Wiring in the smoke binary (after applying the patch)

The smoke binary's batched training would:

```rust
// (existing) C51 forward on h_enriched → probs_current
// (existing) C51 grad → dW_c51, db_c51

// NEW: compute d_h_enriched (gradient into Mamba2 output)
let mut d_h_enriched_dev = stream.alloc_zeros::<f32>(
    (ep_len as usize) * cli.mamba2_hidden_dim
)?;
{
    let (p_ptr, _) = probs_current_dev.device_ptr(&stream);
    let (m_ptr, _) = m_dev.device_ptr(&stream);
    let (a_ptr, _) = actions_dev.device_ptr(&stream);
    let (w_ptr, _) = w_dev.device_ptr(&stream);
    let (d_ptr, _) = d_h_enriched_dev.device_ptr_mut(&stream);
    unsafe {
        launch_alpha_c51_grad_input(
            &stream, &c51_grad_input_kernel,
            p_ptr, m_ptr, a_ptr, w_ptr, d_ptr,
            ep_len, c51_input_dim as i32, n_act_i, n_atoms_i,
            1.0 / ep_len as f32,
        )?;
    }
}

// NEW: wrap d_h_enriched as GpuTensor, call backward_from_h_enriched
let d_h_tensor = GpuTensor::new(
    d_h_enriched_dev,
    vec![ep_len as usize, cli.mamba2_hidden_dim],
)?;
let block = mamba2_block.as_ref().expect("Mamba2Block missing");
let grads = block.backward_from_h_enriched(&cache, &d_h_tensor)?;

// NEW: AdamW step
adamw.step(&mut block_mut, &grads)?;
```

Two changes needed in the smoke wiring beyond this:
1. `mamba2_block` needs to be `Option<Mut<Mamba2Block>>` since
   AdamW takes `&mut`. Simplest: wrap in `Option<RefCell<...>>`
   or just borrow mutably at the call site.
2. The batched training currently re-uses the inference-time
   `cache.h_enriched` from a per-step forward, but `backward`
   needs the FULL cache (input_2d, x, a_proj, b_proj). We'd need
   to do ONE batched forward at training time over a
   `[ep_len, K, in_dim]` tensor to get the cache that matches
   the gradient.

## Estimated effort
- Patch application + ml-alpha test verification: 30-45 min
- Smoke binary integration: 1-2 hr
- Backtest binary integration: 30-45 min
- Smoke + backtest validation runs: 1-2 hr
- Total: ~4-6 hr focused session

Conditional on E.4.A T14 result: if temporal already lifts
backtest Sharpe substantially with frozen Mamba2, T10 is upside.
If T14 is flat or worse, T10 may not be the lever and the
optimization budget should go elsewhere.
