diff --git a/crates/ml-alpha/build.rs b/crates/ml-alpha/build.rs index 6117b3dca..7d38b55c2 100644 --- a/crates/ml-alpha/build.rs +++ b/crates/ml-alpha/build.rs @@ -27,6 +27,7 @@ const KERNELS: &[&str] = &[ "bucket_transition_kernels", // Per-horizon CfC Phase 1→2 transition: tau_sort, bucket_assign, bucket_iqr, channels_in_bucket, heads_compact, zero_off_bucket (ALPHA fix 2026-05-21) "cfc_step_per_branch", // Per-horizon CfC Phase 2: fused per-(batch, branch) fwd + bwd over [25,25,25,25,28] buckets "heads_block_diagonal_fwd", // Per-horizon CfC Phase 2: heads w_skip projection with compact ragged storage (640→128 floats) + "aux_trunk", // SDD-3 Layer B3: smaller single-bucket CfC trunk (AUX_HIDDEN=64) for outcome-supervision (D-labels) ]; // Cache bust v16 (2026-05-21): ALPHA fix — channels_in_bucket_kernel + zero_off_bucket_kernel replace tau_reorder_kernel. diff --git a/crates/ml-alpha/cuda/aux_trunk.cu b/crates/ml-alpha/cuda/aux_trunk.cu new file mode 100644 index 000000000..34e798dcc --- /dev/null +++ b/crates/ml-alpha/cuda/aux_trunk.cu @@ -0,0 +1,241 @@ +// aux_trunk.cu — single-bucket CfC step (forward + backward) at AUX_HIDDEN=64. +// +// Per `pearl_separate_aux_trunk_when_shared_starves`: this is the kernel +// behind a SECOND, smaller CfC trunk used for outcome-supervision (D-labels). +// It runs in parallel with the main per-branch CfC trunk on the same encoder +// output, with INDEPENDENT parameters (no shared w_in/w_rec/b/tau) so that +// BCE gradient flow on the main trunk does not starve the aux supervision. +// +// Design contract: +// • Single bucket — no per-horizon channel splitting, no +// channels_in_bucket lookup, no bucket_dim_k. Aux supervision is +// per-K at the head (B4 wiring), not at the τ-state. +// • Hidden dim is AUX_HIDDEN=64 (half of main trunk's HIDDEN_DIM=128). +// • Algorithm matches `cfc_step_per_branch.cu`: continuous-time CfC +// decay update `h_new = h_old·decay + (1-decay)·tanh(pre)`, with +// `decay = exp(-dt / max(tau, 1e-6))`. +// • Input feature dim is parameterised (`feat_dim` runtime argument) +// so a single cubin handles both raw-snap input (FEATURE_DIM=40) and +// post-encoder input (HIDDEN_DIM=128). The host wrapper passes the +// correct dim; the kernel iterates `feat_dim` in the input MVM. +// +// Per `feedback_no_atomicadd.md`: no atomicAdd anywhere. Each thread +// writes to its own (batch, c) slot in forward, and to its own per-batch +// grad slice in backward. Cross-batch reduction is the caller's +// responsibility via the existing `reduce_axis0` infrastructure. +// +// Per `pearl_cooperative_staging_eliminates_redundant_reads`: the per- +// batch `x` and `h_old` rows are SHARED across all output channels in a +// block. Stage them into shared memory once at block entry so the K-loop +// reads from smem instead of issuing AUX_HIDDEN × redundant DRAM reads. +// +// Per `pearl_no_host_branches_in_captured_graph`: no host branching; +// the kernel takes feat_dim as a runtime arg but does not consult any +// host-side state during execution. +// +// Per `feedback_nvidia_grade_perf_for_kernels.md`: warp-uniform branches +// (the only conditional is `batch < B`, identical for every thread in a +// block since batch == blockIdx.x), no atomicAdd, coalesced loads via +// cooperative staging. + +#define AUX_HIDDEN 64 +#define MAX_FEAT_DIM 128 // upper bound for shared-memory staging; covers + // both FEATURE_DIM=40 and HIDDEN_DIM=128 callers + +// ───────────────────────────────────────────────────────────────────── +// aux_trunk_fwd: forward pass. +// +// Launch: +// grid = (B, 1, 1) +// block = (AUX_HIDDEN = 64, 1, 1) +// shared_mem_bytes = (feat_dim + AUX_HIDDEN) * sizeof(float) +// +// Per (batch, c) thread computes: +// pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k] +// decay = exp(-dt / max(tau[c], 1e-6)) +// h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre) +// +// W_in is [AUX_HIDDEN × feat_dim], W_rec is [AUX_HIDDEN × AUX_HIDDEN]. +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void aux_trunk_fwd( + const float* __restrict__ w_in, // [AUX_HIDDEN × feat_dim] + const float* __restrict__ w_rec, // [AUX_HIDDEN × AUX_HIDDEN] + const float* __restrict__ b, // [AUX_HIDDEN] + const float* __restrict__ tau, // [AUX_HIDDEN] + const float* __restrict__ x, // [B × feat_dim] + const float* __restrict__ h_old, // [B × AUX_HIDDEN] + float dt, + int B, + int feat_dim, + float* __restrict__ h_new // [B × AUX_HIDDEN] +) { + extern __shared__ float smem[]; + float* x_local = smem; // [feat_dim] + float* h_old_local = smem + feat_dim; // [AUX_HIDDEN] + + int batch = blockIdx.x; + int c = threadIdx.x; + + if (batch >= B) return; + + // Cooperative staging of per-batch x[batch, *] and h_old[batch, *]. + // Each thread loads ceil(feat_dim / AUX_HIDDEN) = up to ~2 x slots + // and exactly one h_old slot (since blockDim.x == AUX_HIDDEN). + for (int i = c; i < feat_dim; i += AUX_HIDDEN) { + x_local[i] = x[batch * feat_dim + i]; + } + h_old_local[c] = h_old[batch * AUX_HIDDEN + c]; + __syncthreads(); + + // pre = b[c] + Σ_k W_in[c, k] * x_local[k] + Σ_k W_rec[c, k] * h_old_local[k] + float pre = b[c]; + for (int k = 0; k < feat_dim; ++k) { + pre += w_in[c * feat_dim + k] * x_local[k]; + } + for (int k = 0; k < AUX_HIDDEN; ++k) { + pre += w_rec[c * AUX_HIDDEN + k] * h_old_local[k]; + } + + float decay = expf(-dt / fmaxf(tau[c], 1e-6f)); + h_new[batch * AUX_HIDDEN + c] = + h_old_local[c] * decay + (1.0f - decay) * tanhf(pre); +} + +// ───────────────────────────────────────────────────────────────────── +// aux_trunk_bwd: backward pass. +// +// Launch: +// grid = (B, 1, 1) +// block = (AUX_HIDDEN = 64, 1, 1) +// shared_mem_bytes = (feat_dim + AUX_HIDDEN) * sizeof(float) +// +// Per-batch grad slices are written by this kernel (OVERWRITE semantics); +// cross-batch reduction is the caller's responsibility via existing +// `reduce_axis0` infrastructure. +// +// grad shapes: +// grad_w_in : [B × AUX_HIDDEN × feat_dim] per-batch scratch +// grad_w_rec : [B × AUX_HIDDEN × AUX_HIDDEN] per-batch scratch +// grad_b : [B × AUX_HIDDEN] per-batch scratch +// grad_tau : [B × AUX_HIDDEN] per-batch scratch +// grad_h_old : [B × AUX_HIDDEN] per-batch (decay +// contribution only; +// cross-channel reduction +// deferred to caller if +// unrolled K>1) +// grad_x : [B × feat_dim] per-batch (reduced +// across c via shared- +// memory tree-reduce) +// +// grad_x cross-channel reduction is done inside this kernel via a +// block-resident shared-memory tree-reduce (no atomicAdd per +// `feedback_no_atomicadd`). +// ───────────────────────────────────────────────────────────────────── +extern "C" __global__ void aux_trunk_bwd( + const float* __restrict__ w_in, // [AUX_HIDDEN × feat_dim] + const float* __restrict__ w_rec, // [AUX_HIDDEN × AUX_HIDDEN] + const float* __restrict__ b, // [AUX_HIDDEN] + const float* __restrict__ tau, // [AUX_HIDDEN] + const float* __restrict__ x, // [B × feat_dim] + const float* __restrict__ h_old, // [B × AUX_HIDDEN] + const float* __restrict__ grad_h_new, // [B × AUX_HIDDEN] + float dt, + int B, + int feat_dim, + float* __restrict__ grad_w_in, // [B × AUX_HIDDEN × feat_dim] + float* __restrict__ grad_w_rec, // [B × AUX_HIDDEN × AUX_HIDDEN] + float* __restrict__ grad_b, // [B × AUX_HIDDEN] + float* __restrict__ grad_tau, // [B × AUX_HIDDEN] + float* __restrict__ grad_h_old, // [B × AUX_HIDDEN] + float* __restrict__ grad_x // [B × feat_dim] +) { + // Shared layout: + // [0 .. feat_dim) : x_local + // [feat_dim .. feat_dim+AH) : h_old_local + // [feat_dim+AH .. feat_dim+AH+AH) : d_pre_smem (per-channel d_pre + // for cross-channel grad_x reduce) + extern __shared__ float smem[]; + float* x_local = smem; + float* h_old_local = smem + feat_dim; + float* d_pre_smem = smem + feat_dim + AUX_HIDDEN; + + int batch = blockIdx.x; + int c = threadIdx.x; + + if (batch >= B) return; + + // Cooperative staging — identical pattern to forward. + for (int i = c; i < feat_dim; i += AUX_HIDDEN) { + x_local[i] = x[batch * feat_dim + i]; + } + h_old_local[c] = h_old[batch * AUX_HIDDEN + c]; + __syncthreads(); + + // Recompute forward pre + tanh to derive d_pre. + float pre = b[c]; + for (int k = 0; k < feat_dim; ++k) { + pre += w_in[c * feat_dim + k] * x_local[k]; + } + for (int k = 0; k < AUX_HIDDEN; ++k) { + pre += w_rec[c * AUX_HIDDEN + k] * h_old_local[k]; + } + const float tau_eps = 1e-6f; + const float tau_raw = tau[c]; + const float tau_eff = fmaxf(tau_raw, tau_eps); + const float decay = expf(-dt / tau_eff); + const float s = tanhf(pre); + const float dh = grad_h_new[batch * AUX_HIDDEN + c]; + const float d_pre = dh * (1.0f - decay) * (1.0f - s * s); + const float d_decay = dh * (h_old_local[c] - s); + + // Per-batch grad-scratch writes. Each thread (batch, c) is the sole + // writer of grad_*[batch, c, ...] — no race, no atomicAdd. + grad_b[batch * AUX_HIDDEN + c] = d_pre; + + // grad_tau receives 0 when tau hits the clamp floor + // (strict d(max)/d(tau) = 0 below tau_eps). + const float gate = (tau_raw > tau_eps) ? 1.0f : 0.0f; + grad_tau[batch * AUX_HIDDEN + c] = + gate * d_decay * decay * dt / (tau_eff * tau_eff); + + // grad_w_in[batch, c, k] = d_pre * x_local[k] + for (int k = 0; k < feat_dim; ++k) { + const long long off = + (long long)batch * AUX_HIDDEN * feat_dim + + (long long)c * feat_dim + k; + grad_w_in[off] = d_pre * x_local[k]; + } + // grad_w_rec[batch, c, k] = d_pre * h_old_local[k] + for (int k = 0; k < AUX_HIDDEN; ++k) { + const long long off = + (long long)batch * AUX_HIDDEN * AUX_HIDDEN + + (long long)c * AUX_HIDDEN + k; + grad_w_rec[off] = d_pre * h_old_local[k]; + } + + // grad_h_old[batch, c] direct decay contribution. + // (Cross-channel term Σ_j d_pre[j] * W_rec[j, c] is left for the + // caller to add if BPTT through h_old is needed — for K=1 per-step + // aux supervision this direct term is sufficient.) + grad_h_old[batch * AUX_HIDDEN + c] = dh * decay; + + // grad_x[batch, k] = Σ_c d_pre[c] * W_in[c, k] + // Cross-channel reduction via shared memory: stash d_pre to smem, + // then have each thread compute its assigned k-slice by summing + // across c. This avoids atomicAdd while keeping the reduction + // block-local. (Per pearl_no_atomicadd + the block tree-reduce + // pattern used in reduce_axis0.cu.) + d_pre_smem[c] = d_pre; + __syncthreads(); + + // Each of the AUX_HIDDEN threads handles a strided subset of feat_dim + // output positions. With feat_dim ≤ MAX_FEAT_DIM = 128 and + // AUX_HIDDEN = 64, each thread covers at most 2 k positions. + for (int k = c; k < feat_dim; k += AUX_HIDDEN) { + float acc = 0.0f; + for (int j = 0; j < AUX_HIDDEN; ++j) { + acc += d_pre_smem[j] * w_in[j * feat_dim + k]; + } + grad_x[batch * feat_dim + k] = acc; + } +} diff --git a/crates/ml-alpha/src/cfc/aux_trunk.rs b/crates/ml-alpha/src/cfc/aux_trunk.rs new file mode 100644 index 000000000..349fd2831 --- /dev/null +++ b/crates/ml-alpha/src/cfc/aux_trunk.rs @@ -0,0 +1,385 @@ +//! SDD-3 Layer B3 — auxiliary CfC trunk for outcome-supervision (D-labels). +//! +//! A SECOND, smaller CfC trunk that operates on the same encoder output +//! as the main BCE trunk, but with INDEPENDENT parameters and a smaller +//! hidden dim. Lives next to the main per-branch trunk in `cfc::trunk`. +//! +//! ## Why separate trunk (not shared) +//! +//! Per `pearl_separate_aux_trunk_when_shared_starves`: aux supervision +//! plateauing at ln(K) on a shared trunk signals starvation. Separate +//! aux trunk + own optimiser group prevents the BCE direction signal +//! from dominating gradient flow toward the aux head. +//! +//! Per the E3 design memo: asymmetric stop-grad at encoder boundary — +//! BCE drives encoder; aux gets gradient only into its own params at +//! first; lifted conditionally once aux CE < 0.4 + dir_acc > 0.85 +//! within 200 steps. +//! +//! ## Algorithm +//! +//! Identical to `cfc_step_per_branch` but stripped of the per-bucket +//! channel split (single bucket, single launch grid). +//! +//! pre = b[c] + Σ_k W_in[c, k] * x[batch, k] + Σ_k W_rec[c, k] * h_old[batch, k] +//! decay = exp(-dt / max(tau[c], 1e-6)) +//! h_new[batch, c] = h_old[batch, c] * decay + (1 - decay) * tanh(pre) +//! +//! ## Initialisation +//! +//! Per `pearl_scoped_init_seed_for_reproducibility`: `AuxTrunk::new` +//! installs a `scoped_init_seed(seed)` guard so the Xavier draws are +//! deterministic given `seed`. +//! +//! * `w_in`: Xavier uniform `[-sqrt(6/(feat_dim+AUX_HIDDEN)), +]` +//! * `w_rec`: Xavier uniform `[-sqrt(6/(2*AUX_HIDDEN)), +]` +//! * `b`: zeros +//! * `tau`: log-uniform over `[2 s, 200 s]` (single-bucket so a wide +//! span across all channels is fine; no per-horizon assignment). +//! +//! ## Constraints honoured +//! +//! * `feedback_no_atomicadd.md` — no atomicAdd in fwd or bwd kernels; +//! bwd writes per-batch grad scratch and the caller reduces axis 0. +//! * `feedback_no_htod_htoh_only_mapped_pinned.md` — uploads go through +//! mapped-pinned staging like the rest of `cfc/step.rs`. +//! * `feedback_no_nvrtc.md` — pre-compiled cubin via build.rs. +//! * `pearl_no_host_branches_in_captured_graph` — kernel uses no host +//! branching; the `feat_dim` runtime arg is read once and used as +//! loop bound. + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use cudarc::driver::{ + CudaFunction, CudaModule, CudaSlice, CudaStream, DevicePtrMut, LaunchConfig, PushKernelArg, +}; +use ml_core::cuda_autograd::init::scoped_init_seed; +use ml_core::device::MlDevice; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +use crate::pinned_mem::MappedF32Buffer; + +const KERNEL_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_trunk.cubin")); + +/// Aux trunk hidden dimension. Half the main trunk's HIDDEN_DIM=128. +/// Kept constant (not configurable) so cubin compile-time `#define +/// AUX_HIDDEN` stays in lockstep with the host wrapper. +pub const AUX_HIDDEN: usize = 64; + +/// Upper bound on the input feature dim. Matches the cubin's +/// `MAX_FEAT_DIM = 128`. Covers both raw-snap input (FEATURE_DIM=40) +/// and post-encoder input (HIDDEN_DIM=128). +pub const MAX_AUX_FEAT_DIM: usize = 128; + +/// Tau initialisation range (seconds). Single-bucket trunk gets a wide +/// log-uniform span so the aux head sees diverse memory horizons within +/// one trunk. Matches the spirit of the main trunk's [10 ms, 1000 s] +/// span but narrowed since aux supervision targets ~1k-step D-label +/// outcomes — too-long τ wastes capacity on horizons the head doesn't +/// supervise. +const AUX_TAU_LO_S: f32 = 2.0; +const AUX_TAU_HI_S: f32 = 200.0; + +/// Host-side weight storage for `AuxTrunk` (used by checkpoint round-trip). +#[derive(Clone, Debug)] +pub struct AuxTrunkWeights { + pub w_in: Vec, // [AUX_HIDDEN × feat_dim] + pub w_rec: Vec, // [AUX_HIDDEN × AUX_HIDDEN] + pub b: Vec, // [AUX_HIDDEN] + pub tau: Vec, // [AUX_HIDDEN] + pub feat_dim: usize, +} + +/// Construction config for `AuxTrunk`. +#[derive(Clone, Debug)] +pub struct AuxTrunkConfig { + /// Input feature dimensionality. Must be ≤ `MAX_AUX_FEAT_DIM` so the + /// cubin's shared-memory staging buffer fits. + pub feat_dim: usize, + /// Random seed for Xavier + tau init. Reproducibility-critical per + /// `pearl_scoped_init_seed_for_reproducibility`. + pub seed: u64, +} + +/// Auxiliary CfC trunk. Owns device weights + a cached cubin module +/// with `aux_trunk_fwd` / `aux_trunk_bwd` function handles. +pub struct AuxTrunk { + cfg: AuxTrunkConfig, + stream: Arc, + + // Cubin module kept alive so the cached `CudaFunction` handles stay + // valid for the trunk's lifetime. + _module: Arc, + pub fwd_fn: CudaFunction, + pub bwd_fn: CudaFunction, + + // Device weights. Public so the trainer (B5) can mutate via AdamW + // and the checkpoint envelope (B5) can memcpy_htod from a loaded + // file. + pub w_in_d: CudaSlice, // [AUX_HIDDEN × feat_dim] + pub w_rec_d: CudaSlice, // [AUX_HIDDEN × AUX_HIDDEN] + pub b_d: CudaSlice, // [AUX_HIDDEN] + pub tau_d: CudaSlice, // [AUX_HIDDEN] +} + +impl AuxTrunk { + /// Construct a fresh `AuxTrunk` with seeded Xavier + log-uniform τ + /// initialisation. Loads the `aux_trunk.cubin` and resolves the + /// `aux_trunk_fwd` / `aux_trunk_bwd` function handles once. + pub fn new(dev: &MlDevice, cfg: AuxTrunkConfig) -> Result { + anyhow::ensure!( + cfg.feat_dim > 0 && cfg.feat_dim <= MAX_AUX_FEAT_DIM, + "AuxTrunk feat_dim={} out of (0, {}]", + cfg.feat_dim, + MAX_AUX_FEAT_DIM, + ); + + let stream: Arc = dev.cuda_stream().context("aux_trunk stream")?.clone(); + let ctx = dev.cuda_context().context("aux_trunk ctx")?; + let module = ctx + .load_cubin(KERNEL_CUBIN.to_vec()) + .context("load aux_trunk cubin")?; + let fwd_fn = module + .load_function("aux_trunk_fwd") + .context("load aux_trunk_fwd")?; + let bwd_fn = module + .load_function("aux_trunk_bwd") + .context("load aux_trunk_bwd")?; + + // Deterministic init guard — every random draw below funnels + // through `scoped_init_seed`'s thread-local RNG so the same + // cfg.seed produces the same weights bit-for-bit. + let _seed_guard = scoped_init_seed(cfg.seed); + let mut r = ChaCha8Rng::seed_from_u64(cfg.seed); + let scale_in = (6.0_f32 / (cfg.feat_dim + AUX_HIDDEN) as f32).sqrt(); + let scale_rec = (6.0_f32 / (2 * AUX_HIDDEN) as f32).sqrt(); + let w_in: Vec = (0..AUX_HIDDEN * cfg.feat_dim) + .map(|_| r.gen_range(-scale_in..scale_in)) + .collect(); + let w_rec: Vec = (0..AUX_HIDDEN * AUX_HIDDEN) + .map(|_| r.gen_range(-scale_rec..scale_rec)) + .collect(); + let b: Vec = vec![0.0; AUX_HIDDEN]; + // tau: log-uniform over [AUX_TAU_LO_S, AUX_TAU_HI_S]. + let lo = AUX_TAU_LO_S.ln(); + let hi = AUX_TAU_HI_S.ln(); + let tau: Vec = (0..AUX_HIDDEN) + .map(|_| { + let u: f32 = r.gen_range(0.0..1.0); + (lo + u * (hi - lo)).exp() + }) + .collect(); + + let w_in_d = upload(&stream, &w_in)?; + let w_rec_d = upload(&stream, &w_rec)?; + let b_d = upload(&stream, &b)?; + let tau_d = upload(&stream, &tau)?; + + Ok(Self { + cfg, + stream, + _module: module, + fwd_fn, + bwd_fn, + w_in_d, + w_rec_d, + b_d, + tau_d, + }) + } + + pub fn config(&self) -> &AuxTrunkConfig { + &self.cfg + } + + pub fn stream(&self) -> &Arc { + &self.stream + } + + /// Download the trunk's current weights into host vectors. Used by + /// checkpoint save and unit tests. + pub fn download_weights(&self) -> Result { + let mut w_in = vec![0.0_f32; self.w_in_d.len()]; + let mut w_rec = vec![0.0_f32; self.w_rec_d.len()]; + let mut b = vec![0.0_f32; self.b_d.len()]; + let mut tau = vec![0.0_f32; self.tau_d.len()]; + self.stream + .memcpy_dtoh(&self.w_in_d, w_in.as_mut_slice()) + .context("aux_trunk dtoh w_in")?; + self.stream + .memcpy_dtoh(&self.w_rec_d, w_rec.as_mut_slice()) + .context("aux_trunk dtoh w_rec")?; + self.stream + .memcpy_dtoh(&self.b_d, b.as_mut_slice()) + .context("aux_trunk dtoh b")?; + self.stream + .memcpy_dtoh(&self.tau_d, tau.as_mut_slice()) + .context("aux_trunk dtoh tau")?; + Ok(AuxTrunkWeights { + w_in, + w_rec, + b, + tau, + feat_dim: self.cfg.feat_dim, + }) + } +} + +/// Launch the fused aux-trunk forward kernel. +/// +/// Buffer contract: +/// * `x_d` : `[B × feat_dim]` per-batch encoder output +/// * `h_old_d` : `[B × AUX_HIDDEN]` per-batch prior aux-trunk state +/// * `h_new_d` : `[B × AUX_HIDDEN]` overwrite — new aux-trunk state +/// +/// Shared memory layout: `(feat_dim + AUX_HIDDEN) * sizeof(float)` +/// covering `x_local` + `h_old_local` cooperative staging. +#[allow(clippy::too_many_arguments)] +pub fn aux_trunk_fwd_gpu( + stream: &Arc, + func: &CudaFunction, + feat_dim: usize, + w_in_d: &CudaSlice, + w_rec_d: &CudaSlice, + b_d: &CudaSlice, + tau_d: &CudaSlice, + x_d: &CudaSlice, + h_old_d: &CudaSlice, + dt_s: f32, + b_sz: i32, + h_new_d: &mut CudaSlice, +) -> Result<()> { + debug_assert_eq!(w_in_d.len(), AUX_HIDDEN * feat_dim); + debug_assert_eq!(w_rec_d.len(), AUX_HIDDEN * AUX_HIDDEN); + debug_assert_eq!(b_d.len(), AUX_HIDDEN); + debug_assert_eq!(tau_d.len(), AUX_HIDDEN); + debug_assert_eq!(x_d.len(), (b_sz as usize) * feat_dim); + debug_assert_eq!(h_old_d.len(), (b_sz as usize) * AUX_HIDDEN); + debug_assert_eq!(h_new_d.len(), (b_sz as usize) * AUX_HIDDEN); + + let feat_dim_i: i32 = feat_dim as i32; + let cfg = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (AUX_HIDDEN as u32, 1, 1), + // x_local (feat_dim floats) + h_old_local (AUX_HIDDEN floats) + shared_mem_bytes: ((feat_dim + AUX_HIDDEN) * std::mem::size_of::()) as u32, + }; + let mut launch = stream.launch_builder(func); + launch + .arg(w_in_d) + .arg(w_rec_d) + .arg(b_d) + .arg(tau_d) + .arg(x_d) + .arg(h_old_d) + .arg(&dt_s) + .arg(&b_sz) + .arg(&feat_dim_i) + .arg(h_new_d); + unsafe { + launch.launch(cfg).context("aux_trunk_fwd launch")?; + } + Ok(()) +} + +/// Launch the fused aux-trunk backward kernel. +/// +/// Per-batch grad scratch is OVERWRITTEN by this kernel; the caller is +/// responsible for the cross-batch reduction (e.g. `reduce_axis0`). +/// +/// Shared memory layout: `(feat_dim + 2 * AUX_HIDDEN) * sizeof(float)` +/// covering `x_local` + `h_old_local` + `d_pre_smem` (the per-channel +/// `d_pre` staging used for the in-block `grad_x` tree-reduce). +#[allow(clippy::too_many_arguments)] +pub fn aux_trunk_bwd_gpu( + stream: &Arc, + func: &CudaFunction, + feat_dim: usize, + w_in_d: &CudaSlice, + w_rec_d: &CudaSlice, + b_d: &CudaSlice, + tau_d: &CudaSlice, + x_d: &CudaSlice, + h_old_d: &CudaSlice, + grad_h_new_d: &CudaSlice, + dt_s: f32, + b_sz: i32, + grad_w_in_d: &mut CudaSlice, + grad_w_rec_d: &mut CudaSlice, + grad_b_d: &mut CudaSlice, + grad_tau_d: &mut CudaSlice, + grad_h_old_d: &mut CudaSlice, + grad_x_d: &mut CudaSlice, +) -> Result<()> { + let b_sz_u = b_sz as usize; + debug_assert_eq!(w_in_d.len(), AUX_HIDDEN * feat_dim); + debug_assert_eq!(w_rec_d.len(), AUX_HIDDEN * AUX_HIDDEN); + debug_assert_eq!(b_d.len(), AUX_HIDDEN); + debug_assert_eq!(tau_d.len(), AUX_HIDDEN); + debug_assert_eq!(x_d.len(), b_sz_u * feat_dim); + debug_assert_eq!(h_old_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(grad_h_new_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(grad_w_in_d.len(), b_sz_u * AUX_HIDDEN * feat_dim); + debug_assert_eq!(grad_w_rec_d.len(), b_sz_u * AUX_HIDDEN * AUX_HIDDEN); + debug_assert_eq!(grad_b_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(grad_tau_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(grad_h_old_d.len(), b_sz_u * AUX_HIDDEN); + debug_assert_eq!(grad_x_d.len(), b_sz_u * feat_dim); + + let feat_dim_i: i32 = feat_dim as i32; + let cfg = LaunchConfig { + grid_dim: (b_sz as u32, 1, 1), + block_dim: (AUX_HIDDEN as u32, 1, 1), + shared_mem_bytes: ((feat_dim + 2 * AUX_HIDDEN) * std::mem::size_of::()) as u32, + }; + let mut launch = stream.launch_builder(func); + launch + .arg(w_in_d) + .arg(w_rec_d) + .arg(b_d) + .arg(tau_d) + .arg(x_d) + .arg(h_old_d) + .arg(grad_h_new_d) + .arg(&dt_s) + .arg(&b_sz) + .arg(&feat_dim_i) + .arg(grad_w_in_d) + .arg(grad_w_rec_d) + .arg(grad_b_d) + .arg(grad_tau_d) + .arg(grad_h_old_d) + .arg(grad_x_d); + unsafe { + launch.launch(cfg).context("aux_trunk_bwd launch")?; + } + Ok(()) +} + +// ── pinned-staging upload helper (mirrors cfc/step.rs::upload) ──────── + +fn upload(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("aux_trunk upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream + .alloc_zeros::(n) + .context("aux_trunk upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + staging.dev_ptr, + nbytes, + stream.cu_stream(), + ) + .context("aux_trunk upload DtoD")?; + } + } + Ok(dst) +} diff --git a/crates/ml-alpha/src/cfc/mod.rs b/crates/ml-alpha/src/cfc/mod.rs index 1ed75c03a..a996b6b95 100644 --- a/crates/ml-alpha/src/cfc/mod.rs +++ b/crates/ml-alpha/src/cfc/mod.rs @@ -3,9 +3,11 @@ //! amendment: this module owns the CfC layer that sits on top of a //! frozen Mamba2 sequence encoder. +pub mod aux_trunk; pub mod bucket_routing; pub mod snap_features; pub mod step; pub mod trunk; +pub use aux_trunk::{AuxTrunk, AuxTrunkConfig, AuxTrunkWeights, AUX_HIDDEN, MAX_AUX_FEAT_DIM}; pub use trunk::{CfcConfig, CfcTrunk}; diff --git a/crates/ml-alpha/tests/aux_trunk.rs b/crates/ml-alpha/tests/aux_trunk.rs new file mode 100644 index 000000000..8b7c76fe2 --- /dev/null +++ b/crates/ml-alpha/tests/aux_trunk.rs @@ -0,0 +1,326 @@ +//! GPU oracle tests for `aux_trunk.cu` (single-bucket CfC at AUX_HIDDEN=64). +//! +//! Run with: +//! SQLX_OFFLINE=true cargo test -p ml-alpha --test aux_trunk \ +//! -- --ignored --nocapture +//! +//! These tests are `#[ignore]` because they require a CUDA device. +//! They drive the cubin directly through the host wrappers in +//! `crates/ml-alpha/src/cfc/aux_trunk.rs` and compare against a +//! deterministic CPU naive reference (per +//! `feedback_no_cpu_test_fallbacks.md` the production path stays GPU- +//! only; the CPU implementation here exists ONLY as a unit-test oracle). + +use anyhow::Result; +use ml_alpha::cfc::aux_trunk::{ + aux_trunk_bwd_gpu, aux_trunk_fwd_gpu, AuxTrunk, AuxTrunkConfig, AUX_HIDDEN, +}; +use ml_core::device::MlDevice; + +// Pick a representative feat_dim. Aux trunk's production caller (B5) +// will pass HIDDEN_DIM=128 (encoder output). Using a non-128 dim here +// also exercises the runtime-`feat_dim` parameterisation. +const TEST_FEAT_DIM: usize = 96; + +fn upload( + stream: &std::sync::Arc, + host: &[f32], +) -> Result> { + let mut d = stream.alloc_zeros::(host.len())?; + stream.memcpy_htod(host, &mut d)?; + Ok(d) +} + +fn download( + stream: &std::sync::Arc, + d: &cudarc::driver::CudaSlice, +) -> Result> { + let mut h = vec![0.0_f32; d.len()]; + stream.memcpy_dtoh(d, &mut h)?; + Ok(h) +} + +/// Forward shape correctness + value match against a naive CPU oracle. +#[test] +#[ignore = "requires CUDA"] +fn fwd_matches_naive_reference() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let trunk = AuxTrunk::new( + &dev, + AuxTrunkConfig { + feat_dim: TEST_FEAT_DIM, + seed: 0xAEB1_BEEF_u64, + }, + )?; + + let b_sz: usize = 3; + let x: Vec = (0..b_sz * TEST_FEAT_DIM) + .map(|i| ((i as f32) * 0.0011).sin() * 0.5) + .collect(); + let h_old: Vec = (0..b_sz * AUX_HIDDEN) + .map(|i| ((i as f32) * 0.0037).cos() * 0.3) + .collect(); + + let x_d = upload(&stream, &x)?; + let h_old_d = upload(&stream, &h_old)?; + let mut h_new_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN)?; + + let dt: f32 = 1.0; + aux_trunk_fwd_gpu( + &stream, + &trunk.fwd_fn, + TEST_FEAT_DIM, + &trunk.w_in_d, + &trunk.w_rec_d, + &trunk.b_d, + &trunk.tau_d, + &x_d, + &h_old_d, + dt, + b_sz as i32, + &mut h_new_d, + )?; + stream.synchronize()?; + let h_new = download(&stream, &h_new_d)?; + + // Pull weights down for the CPU oracle. + let weights = trunk.download_weights()?; + let w_in = &weights.w_in; + let w_rec = &weights.w_rec; + let b_vec = &weights.b; + let tau = &weights.tau; + + for batch in 0..b_sz { + for c in 0..AUX_HIDDEN { + let mut pre = b_vec[c]; + for k in 0..TEST_FEAT_DIM { + pre += w_in[c * TEST_FEAT_DIM + k] * x[batch * TEST_FEAT_DIM + k]; + } + for k in 0..AUX_HIDDEN { + pre += w_rec[c * AUX_HIDDEN + k] * h_old[batch * AUX_HIDDEN + k]; + } + let decay = (-dt / tau[c].max(1e-6)).exp(); + let expected = + h_old[batch * AUX_HIDDEN + c] * decay + (1.0 - decay) * pre.tanh(); + let got = h_new[batch * AUX_HIDDEN + c]; + assert!( + (got - expected).abs() < 5e-4, + "h_new[{},{}] = {got} vs expected {expected} (diff {})", + batch, + c, + (got - expected).abs(), + ); + } + } + Ok(()) +} + +/// Backward grad shapes finite, non-zero, and bias-grad finite-difference +/// matches within tolerance. Full per-weight FD validation is expensive +/// (O(AUX_HIDDEN²·feat_dim) reference evaluations); we sample a handful +/// of bias entries which is sufficient to catch wiring breakage. +#[test] +#[ignore = "requires CUDA"] +fn bwd_finite_diff_matches_bias_sample() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let trunk = AuxTrunk::new( + &dev, + AuxTrunkConfig { + feat_dim: TEST_FEAT_DIM, + seed: 0x12_34_56_78, + }, + )?; + + let b_sz: usize = 2; + let x: Vec = (0..b_sz * TEST_FEAT_DIM) + .map(|i| ((i as f32) * 0.0021).sin() * 0.4) + .collect(); + let h_old: Vec = (0..b_sz * AUX_HIDDEN) + .map(|i| ((i as f32) * 0.0049).cos() * 0.25) + .collect(); + // Pick a structured grad_h_new (1.0 at one slot per batch, 0 else) + // so that grad_b finite-diff has a clean closed-form expectation. + let grad_h_new: Vec = (0..b_sz * AUX_HIDDEN) + .map(|i| ((i % 7) as f32) * 0.03) + .collect(); + + let weights = trunk.download_weights()?; + let w_in = weights.w_in.clone(); + let w_rec = weights.w_rec.clone(); + let b_vec = weights.b.clone(); + let tau = weights.tau.clone(); + + let dt: f32 = 1.0; + let x_d = upload(&stream, &x)?; + let h_old_d = upload(&stream, &h_old)?; + let grad_h_new_d = upload(&stream, &grad_h_new)?; + + let mut grad_w_in_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN * TEST_FEAT_DIM)?; + let mut grad_w_rec_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN * AUX_HIDDEN)?; + let mut grad_b_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN)?; + let mut grad_tau_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN)?; + let mut grad_h_old_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN)?; + let mut grad_x_d = stream.alloc_zeros::(b_sz * TEST_FEAT_DIM)?; + + aux_trunk_bwd_gpu( + &stream, + &trunk.bwd_fn, + TEST_FEAT_DIM, + &trunk.w_in_d, + &trunk.w_rec_d, + &trunk.b_d, + &trunk.tau_d, + &x_d, + &h_old_d, + &grad_h_new_d, + dt, + b_sz as i32, + &mut grad_w_in_d, + &mut grad_w_rec_d, + &mut grad_b_d, + &mut grad_tau_d, + &mut grad_h_old_d, + &mut grad_x_d, + )?; + stream.synchronize()?; + + // All grads finite, no NaN/Inf. + let g_w_in = download(&stream, &grad_w_in_d)?; + let g_w_rec = download(&stream, &grad_w_rec_d)?; + let g_b = download(&stream, &grad_b_d)?; + let g_tau = download(&stream, &grad_tau_d)?; + let g_h_old = download(&stream, &grad_h_old_d)?; + let g_x = download(&stream, &grad_x_d)?; + for (name, v) in [ + ("grad_w_in", &g_w_in), + ("grad_w_rec", &g_w_rec), + ("grad_b", &g_b), + ("grad_tau", &g_tau), + ("grad_h_old", &g_h_old), + ("grad_x", &g_x), + ] { + assert!( + v.iter().all(|x| x.is_finite()), + "{name} contains NaN/Inf", + ); + } + // Kernel must have written something non-zero for at least the bias + // (the only grad guaranteed to be non-trivial given non-zero + // grad_h_new entries). + assert!( + g_b.iter().any(|x| x.abs() > 0.0), + "grad_b is all zero — kernel didn't write?", + ); + + // Finite-difference oracle for grad_b on the first batch, channel 0. + // Hand-coded forward (CPU) to derive d/d b[c] of the implicit loss + // L = Σ_{batch,c} grad_h_new[batch, c] * h_new[batch, c]. + let h_new_fwd = |b_local: &[f32]| -> Vec { + let mut out = vec![0.0_f32; b_sz * AUX_HIDDEN]; + for batch in 0..b_sz { + for c in 0..AUX_HIDDEN { + let mut pre = b_local[c]; + for k in 0..TEST_FEAT_DIM { + pre += w_in[c * TEST_FEAT_DIM + k] * x[batch * TEST_FEAT_DIM + k]; + } + for k in 0..AUX_HIDDEN { + pre += w_rec[c * AUX_HIDDEN + k] * h_old[batch * AUX_HIDDEN + k]; + } + let decay = (-dt / tau[c].max(1e-6)).exp(); + out[batch * AUX_HIDDEN + c] = + h_old[batch * AUX_HIDDEN + c] * decay + (1.0 - decay) * pre.tanh(); + } + } + out + }; + let loss_for = |b_local: &[f32]| -> f32 { + let h_new = h_new_fwd(b_local); + let mut acc = 0.0_f32; + for i in 0..b_sz * AUX_HIDDEN { + acc += grad_h_new[i] * h_new[i]; + } + acc + }; + + let eps = 1e-3_f32; + // Sample 4 channels at modest stride; full sweep is overkill for a + // smoke. The kernel writes per-batch grad slices — sum across batch + // to match the loss gradient w.r.t. the shared bias param. + for &c in &[0_usize, 17, 33, 50] { + let mut b_plus = b_vec.clone(); + b_plus[c] += eps; + let mut b_minus = b_vec.clone(); + b_minus[c] -= eps; + let fd = (loss_for(&b_plus) - loss_for(&b_minus)) / (2.0 * eps); + + // Sum kernel's per-batch grad_b across batch to match dL/db[c]. + let mut analytical = 0.0_f32; + for batch in 0..b_sz { + analytical += g_b[batch * AUX_HIDDEN + c]; + } + let abs_err = (fd - analytical).abs(); + let rel_tol = 5e-3 * fd.abs().max(1e-4); + assert!( + abs_err < rel_tol.max(2e-3), + "grad_b[{c}] mismatch: fd={fd}, analytical={analytical}, abs_err={abs_err}", + ); + } + Ok(()) +} + +/// Smoke that AUX_HIDDEN=64 path runs without OOM or NaN at a batch +/// size representative of the trainer (which uses B=32-128 per batch). +/// Designed to be cheap enough for RTX 3050 sm_86 local validation. +#[test] +#[ignore = "requires CUDA"] +fn fwd_smoke_large_batch_no_nan_no_oom() -> Result<()> { + let dev = MlDevice::cuda(0)?; + let stream = dev.cuda_stream()?.clone(); + let trunk = AuxTrunk::new( + &dev, + AuxTrunkConfig { + feat_dim: 128, // matches HIDDEN_DIM — actual B5 wiring config + seed: 0xFEED_FACE_u64, + }, + )?; + + let b_sz: usize = 64; + let feat_dim = 128usize; + let x: Vec = (0..b_sz * feat_dim) + .map(|i| ((i as f32) * 0.0007).sin()) + .collect(); + let h_old: Vec = vec![0.0; b_sz * AUX_HIDDEN]; // cold start + let x_d = upload(&stream, &x)?; + let h_old_d = upload(&stream, &h_old)?; + let mut h_new_d = stream.alloc_zeros::(b_sz * AUX_HIDDEN)?; + + aux_trunk_fwd_gpu( + &stream, + &trunk.fwd_fn, + feat_dim, + &trunk.w_in_d, + &trunk.w_rec_d, + &trunk.b_d, + &trunk.tau_d, + &x_d, + &h_old_d, + 1.0, + b_sz as i32, + &mut h_new_d, + )?; + stream.synchronize()?; + let h_new = download(&stream, &h_new_d)?; + assert!( + h_new.iter().all(|v| v.is_finite()), + "h_new contains NaN/Inf after cold-start forward", + ); + // Output is bounded by tanh ∈ (-1, 1); with h_old = 0 and small init, + // |h_new| should sit well inside that range. + assert!( + h_new.iter().all(|v| v.abs() < 1.5), + "h_new produced out-of-range values (|v| ≥ 1.5)", + ); + Ok(()) +}