diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index e060aa5f1..544fb25a3 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -93,6 +93,8 @@ const KERNELS: &[&str] = &[ "rl_session_risk_check", // session-level loss limit circuit breaker "rl_iqn_forward", // IQN distributional Q-head: quantile embedding + action-value projection; complementary to C51 "rl_iqn_loss", // IQN quantile Huber loss: ρ_τ(δ) = |τ - 1(δ<0)| × Huber(δ, κ=1.0); forward + backward + "rl_iqn_backward", // IQN backward through forward pass: grad_output → grad_w_out/b_out/w_embed/b_embed per-batch scratch + "rl_ensemble_action_value", // C51+IQN ensemble: E_ensemble = α×E_C51 + (1-α)×E_IQN; α from ISV[544] ]; // Cache bust v31 — five new reduce / derive kernels populate the input diff --git a/crates/ml-alpha/cuda/rl_ensemble_action_value.cu b/crates/ml-alpha/cuda/rl_ensemble_action_value.cu new file mode 100644 index 000000000..56465d7b7 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_ensemble_action_value.cu @@ -0,0 +1,54 @@ +// rl_ensemble_action_value.cu — combine C51 + IQN expected Q values. +// +// E_ensemble[b,a] = alpha * E_C51[b,a] + (1-alpha) * E_IQN[b,a] +// +// Reads alpha from ISV[RL_IQN_ENSEMBLE_ALPHA_INDEX=544]. +// One thread per (batch, action). Output feeds the Q-pi agreement +// diagnostic and can inform ensemble-level action selection in future. +// +// Per `feedback_no_atomicadd.md`: no atomicAdd (not needed here). +// Per `feedback_cpu_is_read_only.md`: alpha from ISV, not host param. +// Per `feedback_isv_for_adaptive_bounds.md`: alpha is ISV-driven. + +#include + +#define N_ACTIONS 11 +#define Q_N_ATOMS 21 +#define RL_IQN_ENSEMBLE_ALPHA_INDEX 544 + +extern "C" __global__ void rl_ensemble_action_value( + const float* __restrict__ c51_logits, // [B * N_ACTIONS * Q_N_ATOMS] + const float* __restrict__ atom_supports, // [Q_N_ATOMS] + const float* __restrict__ iqn_expected, // [B * N_ACTIONS] + const float* __restrict__ isv, + float* __restrict__ ensemble_q, // [B * N_ACTIONS] OUT + int b_size +) { + const int b = blockIdx.x; + const int a = threadIdx.x; + if (b >= b_size || a >= N_ACTIONS) return; + + const float alpha = isv[RL_IQN_ENSEMBLE_ALPHA_INDEX]; + + // C51 expected Q for this (batch, action): sum of probs * atoms + const int q_base = (b * N_ACTIONS + a) * Q_N_ATOMS; + float max_l = c51_logits[q_base]; + for (int z = 1; z < Q_N_ATOMS; ++z) { + float v = c51_logits[q_base + z]; + if (v > max_l) max_l = v; + } + float sum_exp = 0.0f; + float e_c51 = 0.0f; + for (int z = 0; z < Q_N_ATOMS; ++z) { + float p = expf(c51_logits[q_base + z] - max_l); + sum_exp += p; + e_c51 += p * atom_supports[z]; + } + e_c51 /= (sum_exp + 1e-8f); + + // IQN expected Q (already reduced over tau) + const float e_iqn = iqn_expected[b * N_ACTIONS + a]; + + // Ensemble + ensemble_q[b * N_ACTIONS + a] = alpha * e_c51 + (1.0f - alpha) * e_iqn; +} diff --git a/crates/ml-alpha/cuda/rl_iqn_backward.cu b/crates/ml-alpha/cuda/rl_iqn_backward.cu new file mode 100644 index 000000000..229d8a307 --- /dev/null +++ b/crates/ml-alpha/cuda/rl_iqn_backward.cu @@ -0,0 +1,162 @@ +// rl_iqn_backward.cu — IQN backward through the forward pass. +// +// Given grad_output [B, N_TAU, N_ACTIONS] (from rl_iqn_loss_fwd), backprop +// through the IQN forward to produce per-batch gradients for: +// - w_out [HIDDEN_DIM, N_ACTIONS] (per-batch scratch) +// - b_out [N_ACTIONS] (per-batch scratch) +// - w_embed [EMBED_DIM, HIDDEN_DIM] (per-batch scratch) +// - b_embed [HIDDEN_DIM] (per-batch scratch) +// +// Forward recap: +// phi(tau)[c] = ReLU(sum_i cos((i+1)*pi*tau) * W_embed[i,c] + b_embed[c]) +// combined[c] = h_t[c] * phi(tau)[c] +// Q[tau, a] = sum_c W_out[c, a] * combined[c] + b_out[a] +// +// Backward: +// dQ/dW_out[c, a] = combined[c] (for the given tau) +// dQ/db_out[a] = 1 +// dQ/d_combined[c] = sum_a grad_Q[a] * W_out[c, a] +// d_combined/d_phi[c] = h_t[c] +// d_phi/d_embed_input[c] = 1(embed_input > 0) (ReLU mask) +// d_embed_input/dW_embed[i,c] = cos((i+1)*pi*tau) +// d_embed_input/db_embed[c] = 1 +// +// Block layout: +// grid = (B, N_TAU, 1) +// block = (HIDDEN_DIM, 1, 1) +// One block per (batch, tau) pair — mirrors the forward kernel. +// Each thread handles one hidden-dim index c. +// +// Output layout (per-batch scratch accumulated across N_TAU): +// grad_w_out_per_batch [B, HIDDEN_DIM, N_ACTIONS] +// grad_b_out_per_batch [B, N_ACTIONS] +// grad_w_embed_per_batch [B, EMBED_DIM, HIDDEN_DIM] +// grad_b_embed_per_batch [B, HIDDEN_DIM] +// +// The per-tau contributions are ACCUMULATED (+=) into the per-batch +// scratch via atomicAdd-free patterns: each (batch, tau) block writes +// to its own output offset, and a subsequent reduce_axis0 pass sums +// across batches (same as C51 head). The tau-dimension accumulation +// happens via the atomic-free pattern of writing to [B, N_TAU, ...] and +// then launching a second reduction kernel over the tau dimension. +// +// SIMPLIFIED APPROACH: since HIDDEN_DIM threads can cooperate and +// N_TAU is typically 32, we launch grid=(B, 1, 1) block=(HIDDEN_DIM, 1, 1) +// and each thread loops over all N_TAU to accumulate its contributions. +// This avoids cross-block accumulation entirely. +// +// Per `feedback_no_atomicadd.md`: no atomicAdd. +// Per `feedback_cpu_is_read_only.md`: all compute on GPU. + +#define HIDDEN_DIM 128 +#define N_ACTIONS 11 +#define N_TAU_MAX 64 +#define EMBED_DIM 64 +#define PI_F 3.14159265f + +extern "C" __global__ void rl_iqn_backward( + const float* __restrict__ h_t, // [B, HIDDEN_DIM] + const float* __restrict__ tau, // [B, N_TAU] + const float* __restrict__ w_embed, // [EMBED_DIM, HIDDEN_DIM] + const float* __restrict__ b_embed, // [HIDDEN_DIM] + const float* __restrict__ w_out, // [HIDDEN_DIM, N_ACTIONS] + const float* __restrict__ grad_output, // [B, N_TAU, N_ACTIONS] + int B, + int N_TAU, + float* __restrict__ grad_w_out_pb, // [B, HIDDEN_DIM * N_ACTIONS] + float* __restrict__ grad_b_out_pb, // [B, N_ACTIONS] + float* __restrict__ grad_w_embed_pb, // [B, EMBED_DIM * HIDDEN_DIM] + float* __restrict__ grad_b_embed_pb // [B, HIDDEN_DIM] +) { + const int batch = blockIdx.x; + const int c = threadIdx.x; // hidden dim index + if (batch >= B) return; + if (c >= HIDDEN_DIM) return; + + const float h_c = h_t[batch * HIDDEN_DIM + c]; + + // Accumulate gradients across all tau for this (batch, c). + float acc_grad_w_out[N_ACTIONS]; + for (int a = 0; a < N_ACTIONS; ++a) acc_grad_w_out[a] = 0.0f; + float acc_grad_w_embed[EMBED_DIM]; + for (int i = 0; i < EMBED_DIM; ++i) acc_grad_w_embed[i] = 0.0f; + float acc_grad_b_embed = 0.0f; + + for (int t = 0; t < N_TAU; ++t) { + const float tau_val = tau[batch * N_TAU + t]; + + // Recompute forward: phi(tau)[c] + float embed_acc = b_embed[c]; + float cos_feats[EMBED_DIM]; + #pragma unroll + for (int i = 0; i < EMBED_DIM; ++i) { + cos_feats[i] = cosf((float)(i + 1) * PI_F * tau_val); + embed_acc += cos_feats[i] * w_embed[i * HIDDEN_DIM + c]; + } + float phi_c = fmaxf(embed_acc, 0.0f); + float relu_mask = (embed_acc > 0.0f) ? 1.0f : 0.0f; + + // combined[c] = h_t[c] * phi(tau)[c] + float combined_c = h_c * phi_c; + + // grad_output for this (batch, tau): [N_ACTIONS] + const int go_base = batch * N_TAU * N_ACTIONS + t * N_ACTIONS; + + // ─── Grad w.r.t. w_out: dL/dW_out[c,a] += combined_c * grad_Q[a] + for (int a = 0; a < N_ACTIONS; ++a) { + acc_grad_w_out[a] += combined_c * grad_output[go_base + a]; + } + + // ─── Grad w.r.t. combined: dL/d_combined[c] = Σ_a grad_Q[a] * W_out[c,a] + float grad_combined_c = 0.0f; + for (int a = 0; a < N_ACTIONS; ++a) { + grad_combined_c += grad_output[go_base + a] * w_out[c * N_ACTIONS + a]; + } + + // ─── Chain through element-wise product: d_combined/d_phi = h_t[c] + float grad_phi_c = grad_combined_c * h_c; + + // ─── Chain through ReLU: d_phi/d_embed_input = relu_mask + float grad_embed_input_c = grad_phi_c * relu_mask; + + // ─── Grad w.r.t. w_embed: dL/dW_embed[i,c] += cos_feats[i] * grad_embed_input_c + for (int i = 0; i < EMBED_DIM; ++i) { + acc_grad_w_embed[i] += cos_feats[i] * grad_embed_input_c; + } + + // ─── Grad w.r.t. b_embed: dL/db_embed[c] += grad_embed_input_c + acc_grad_b_embed += grad_embed_input_c; + } + + // Write accumulated grad_w_out for this (batch, c) — row c of the + // per-batch grad_w_out matrix [HIDDEN_DIM, N_ACTIONS]. + const int wo_base = batch * HIDDEN_DIM * N_ACTIONS + c * N_ACTIONS; + for (int a = 0; a < N_ACTIONS; ++a) { + grad_w_out_pb[wo_base + a] = acc_grad_w_out[a]; + } + + // Write accumulated grad_w_embed for this (batch, c) — column c of + // the per-batch grad_w_embed matrix [EMBED_DIM, HIDDEN_DIM]. + for (int i = 0; i < EMBED_DIM; ++i) { + grad_w_embed_pb[batch * EMBED_DIM * HIDDEN_DIM + i * HIDDEN_DIM + c] = acc_grad_w_embed[i]; + } + + // Write grad_b_embed for this (batch, c). + grad_b_embed_pb[batch * HIDDEN_DIM + c] = acc_grad_b_embed; + + // grad_b_out: each action gets contribution from ALL hidden dims. + // Only thread c=0 writes it to avoid races — it accumulates across + // all hidden dims by reading grad_output directly. + // Actually: dL/db_out[a] = Σ_tau grad_Q[tau, a] (since dQ/db_out = 1). + // Each thread has access to grad_output — use a shared-mem reduce. + // Simpler: thread 0 computes it by summing over tau. + if (c == 0) { + for (int a = 0; a < N_ACTIONS; ++a) { + float sum = 0.0f; + for (int t = 0; t < N_TAU; ++t) { + sum += grad_output[batch * N_TAU * N_ACTIONS + t * N_ACTIONS + a]; + } + grad_b_out_pb[batch * N_ACTIONS + a] = sum; + } + } +} diff --git a/crates/ml-alpha/src/rl/iqn.rs b/crates/ml-alpha/src/rl/iqn.rs index e6ac70d2c..56cded63a 100644 --- a/crates/ml-alpha/src/rl/iqn.rs +++ b/crates/ml-alpha/src/rl/iqn.rs @@ -72,7 +72,7 @@ use crate::rl::common::N_ACTIONS; /// Quantile embedding dimension — number of cosine basis functions. /// Matches the `EMBED_DIM` define in `cuda/rl_iqn_forward.cu`. -const EMBED_DIM: usize = 64; +pub const EMBED_DIM: usize = 64; /// Default number of quantile samples per forward pass. ISV-overridable /// via `RL_IQN_N_TAU_INDEX`. Matches the `N_TAU` default in the kernel. @@ -86,6 +86,10 @@ const IQN_LOSS_CUBIN: &[u8] = include_bytes!(concat!( env!("OUT_DIR"), "/rl_iqn_loss.cubin" )); +const IQN_BWD_CUBIN: &[u8] = include_bytes!(concat!( + env!("OUT_DIR"), + "/rl_iqn_backward.cubin" +)); /// Construction config for [`IqnHead`]. #[derive(Clone, Debug)] @@ -132,6 +136,10 @@ pub struct IqnHead { /// Loss kernel: quantile Huber loss forward + backward. pub loss_fn: CudaFunction, + _bwd_module: Arc, + /// Backward kernel: grad_output → per-batch grad_w_out/b_out/w_embed/b_embed. + pub bwd_fn: CudaFunction, + /// Quantile embedding weight `[EMBED_DIM, HIDDEN_DIM]`, row-major. pub w_embed_d: CudaSlice, /// Quantile embedding bias `[HIDDEN_DIM]`. @@ -175,6 +183,13 @@ impl IqnHead { .load_function("rl_iqn_loss_fwd") .context("load rl_iqn_loss_fwd fn")?; + let bwd_module = ctx + .load_cubin(IQN_BWD_CUBIN.to_vec()) + .context("load rl_iqn_backward cubin")?; + let bwd_fn = bwd_module + .load_function("rl_iqn_backward") + .context("load rl_iqn_backward fn")?; + // Per pearl_scoped_init_seed_for_reproducibility: install the // scoped seed guard BEFORE drawing any Xavier samples. let _seed_guard = scoped_init_seed(cfg.seed); @@ -214,6 +229,8 @@ impl IqnHead { expected_q_fn, _loss_module: loss_module, loss_fn, + _bwd_module: bwd_module, + bwd_fn, w_embed_d, b_embed_d, w_out_d, @@ -391,6 +408,60 @@ impl IqnHead { Ok(()) } + /// Backward through the forward pass: given `grad_output [B, N_TAU, N_ACTIONS]` + /// (from `compute_loss`), produce per-batch gradients for w_out, b_out, + /// w_embed, b_embed. Caller is responsible for `reduce_axis0` across + /// batches and feeding the reduced gradients to Adam. + #[allow(clippy::too_many_arguments)] + pub fn backward( + &self, + h_t: &CudaSlice, + tau: &CudaSlice, + grad_output: &CudaSlice, + b_size: usize, + n_tau: usize, + grad_w_out_pb: &mut CudaSlice, + grad_b_out_pb: &mut CudaSlice, + grad_w_embed_pb: &mut CudaSlice, + grad_b_embed_pb: &mut CudaSlice, + ) -> Result<()> { + debug_assert_eq!(h_t.len(), b_size * self.cfg.hidden_dim); + debug_assert_eq!(tau.len(), b_size * n_tau); + debug_assert_eq!(grad_output.len(), b_size * n_tau * N_ACTIONS); + debug_assert_eq!(grad_w_out_pb.len(), b_size * self.cfg.hidden_dim * N_ACTIONS); + debug_assert_eq!(grad_b_out_pb.len(), b_size * N_ACTIONS); + debug_assert_eq!(grad_w_embed_pb.len(), b_size * EMBED_DIM * self.cfg.hidden_dim); + debug_assert_eq!(grad_b_embed_pb.len(), b_size * self.cfg.hidden_dim); + + let b_i = b_size as i32; + let n_tau_i = n_tau as i32; + let launch_cfg = LaunchConfig { + grid_dim: (b_size as u32, 1, 1), + block_dim: (self.cfg.hidden_dim as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = self.stream.launch_builder(&self.bwd_fn); + launch + .arg(h_t) + .arg(tau) + .arg(&self.w_embed_d) + .arg(&self.b_embed_d) + .arg(&self.w_out_d) + .arg(grad_output) + .arg(&b_i) + .arg(&n_tau_i) + .arg(grad_w_out_pb) + .arg(grad_b_out_pb) + .arg(grad_w_embed_pb) + .arg(grad_b_embed_pb); + unsafe { + launch + .launch(launch_cfg) + .context("rl_iqn_backward launch")?; + } + Ok(()) + } + /// Stream used to launch all kernels owned by this head. pub fn stream(&self) -> &Arc { &self.stream