diff --git a/crates/ml/src/cuda_pipeline/alpha_c51.cu b/crates/ml/src/cuda_pipeline/alpha_c51.cu index a1ca6adc8..81c70a702 100644 --- a/crates/ml/src/cuda_pipeline/alpha_c51.cu +++ b/crates/ml/src/cuda_pipeline/alpha_c51.cu @@ -247,6 +247,58 @@ extern "C" __global__ void alpha_c51_grad_kernel( } } +// ---------------------------------------------------------------------- +// (6) C51 cross-entropy gradient w.r.t. the LAYER INPUT +// ---------------------------------------------------------------------- +// +// Phase E.4.A.10: needed to chain the C51 head's gradient back into +// upstream encoders (Mamba2 in --temporal mode). The C51 layer +// computes `logits[a, k] = Σ_j W[a*K+k, j] · input[j] + b[a*K+k]`, +// followed by softmax and CE loss. The chain rule gives: +// +// dL/d_input[b, j] = Σ_{k} (p[b, a_taken, k] - m[b, k]) · W[a_taken*K+k, j] +// +// (Only the taken action's column contributes because dL/dlogit is +// nonzero only for the taken action's atom row.) +// +// One thread per (batch, input feature). Threads: batch × state_dim. +// For ep_len=600 × state_dim=32 → 19200 threads, single block grid is +// tiny; standard launch config. +extern "C" __global__ void alpha_c51_grad_input_kernel( + const float* __restrict__ probs, // [batch, n_actions, n_atoms] + const float* __restrict__ m, // [batch, n_atoms] + const int* __restrict__ actions, // [batch] + const float* __restrict__ W, // [n_actions * n_atoms, state_dim] + float* __restrict__ d_input, // [batch, state_dim] (overwritten) + int batch, + int state_dim, + int n_actions, + int n_atoms, + float scale // typ. 1/batch +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + const int total = batch * state_dim; + if (idx >= total) return; + + const int b_idx = idx / state_dim; + const int j = idx % state_dim; + const int a_taken = actions[b_idx]; + + float g = 0.0f; + // Loop over atoms; the partial gradient dlogit[a_taken, k]/d_input[j] + // is W[a_taken*n_atoms+k, j]. The atom-grad is (p[b,a_taken,k] - m[b,k]). + const int probs_base = (b_idx * n_actions + a_taken) * n_atoms; + const int m_base = b_idx * n_atoms; + const int w_base = a_taken * n_atoms * state_dim; + for (int k = 0; k < n_atoms; ++k) { + const float p = probs[probs_base + k]; + const float mv = m[m_base + k]; + const float w = W[w_base + k * state_dim + j]; + g += (p - mv) * w; + } + d_input[idx] = g * scale; +} + // ---------------------------------------------------------------------- // (4) Thompson-select with confidence gate // ---------------------------------------------------------------------- diff --git a/crates/ml/src/cuda_pipeline/alpha_kernels.rs b/crates/ml/src/cuda_pipeline/alpha_kernels.rs index 625cebada..f3ccf886b 100644 --- a/crates/ml/src/cuda_pipeline/alpha_kernels.rs +++ b/crates/ml/src/cuda_pipeline/alpha_kernels.rs @@ -646,6 +646,60 @@ pub unsafe fn launch_alpha_c51_expected_q( Ok(()) } +/// Launch `alpha_c51_grad_input_kernel`. Computes dL/d_input for the +/// C51 layer (gradient w.r.t. the upstream encoder's output). +/// Threads: one per (batch, input feature). Needed to chain C51's +/// gradient into Mamba2's backward pass in `--temporal`. +/// +/// # Safety +/// All buffer pointers MUST be valid device pointers. `d_input_dev` +/// must point at a writable `[batch * state_dim]`-float region. +pub unsafe fn launch_alpha_c51_grad_input( + stream: &cudarc::driver::CudaStream, + kernel: &cudarc::driver::CudaFunction, + probs_dev: u64, + m_dev: u64, + actions_dev: u64, + w_dev: u64, + d_input_dev: u64, + batch: i32, + state_dim: i32, + n_actions: i32, + n_atoms: i32, + scale: f32, +) -> Result<(), MLError> { + use cudarc::driver::{LaunchConfig, PushKernelArg}; + + debug_assert!(batch > 0, "batch must be positive"); + debug_assert!(state_dim > 0, "state_dim must be positive"); + debug_assert!(n_actions > 0, "n_actions must be positive"); + debug_assert!(n_atoms > 0 && n_atoms <= 64, "n_atoms must be in (0, 64]"); + + const BLOCK: u32 = 256; + let total = (batch as u32) * (state_dim as u32); + let grid_x = (total + BLOCK - 1) / BLOCK; + let cfg = LaunchConfig { + grid_dim: (grid_x.max(1), 1, 1), + block_dim: (BLOCK, 1, 1), + shared_mem_bytes: 0, + }; + stream + .launch_builder(kernel) + .arg(&probs_dev) + .arg(&m_dev) + .arg(&actions_dev) + .arg(&w_dev) + .arg(&d_input_dev) + .arg(&batch) + .arg(&state_dim) + .arg(&n_actions) + .arg(&n_atoms) + .arg(&scale) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("alpha_c51_grad_input launch: {e}")))?; + Ok(()) +} + /// Launch `alpha_c51_thompson_select_kernel`. GPU action selector with /// confidence gate. Replaces CPU `epsilon_greedy_gated` per /// `feedback_cpu_is_read_only`. One thread per batch sample; writes the diff --git a/docs/isv-slots.md b/docs/isv-slots.md index 104e86c15..71b7ec290 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -777,6 +777,26 @@ inference; effective policy adapts via ISV slot 543 (threshold), 545 allocated for the E.3 follow-up; E.4.B will add MoE-gate-entropy and Pearl-1 atom-headroom slots when those land. +### Phase E.4.A.10 — alpha_c51_grad_input (2026-05-15) + +`crates/ml/src/cuda_pipeline/alpha_c51.cu` adds a 6th kernel: +`alpha_c51_grad_input_kernel`. Computes `dL/d_input[b, j]` for the +C51 layer (gradient w.r.t. the upstream encoder's output). Required +to chain C51's gradient into Mamba2's backward in `--temporal`. + +NO ISV slot interaction — pure Q-layer gradient compute. + +T10 status: kernel landed but Mamba2 backward NOT yet wired in the +smoke binary because `ml_alpha::Mamba2Block::backward` takes +`d_logit [B, 1]` (post-W_out gradient) rather than +`d_h_enriched [B, hidden_dim]`. Two paths to complete T10: +1. Modify ml-alpha to expose `backward_from_h_enriched`. +2. Replicate the post-W_out backward logic inline. + +Deferred pending backtest validation — if frozen-Mamba2 already +lifts backtest Sharpe, T10 becomes optimisation rather than +prerequisite. + ### Phase E.4.A.6 — alpha_window_push (2026-05-15) `crates/ml/src/cuda_pipeline/alpha_window_push.cu`: shift+insert