From cb6bca4629237e9dad2ac9866e03a0c6f8508ed0 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 8 May 2026 00:56:23 +0200 Subject: [PATCH] feat(sp14-c): aux trunk forward kernel + Rust wrapper + oracle test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 3-layer MLP forward (Linear→ELU→Linear→ELU→Linear). Pre-loaded CudaFunction for graph-capture safety per pearl_no_host_branches_in_captured_graph. Oracle test verifies bit-for-bit match against numpy reference within 1e-4 tol. Saves h_aux1 and h_aux2 to global memory for backward. Phase C.3 of SP14 Layer C separate-aux-trunk refactor (plan: docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md). Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 10 + .../cuda_pipeline/aux_trunk_forward_kernel.cu | 172 ++++++++++++ crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs | 164 ++++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 29 ++ crates/ml/src/cuda_pipeline/mod.rs | 1 + crates/ml/tests/aux_trunk_oracle_tests.rs | 250 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 35 +++ 7 files changed, 661 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu create mode 100644 crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs create mode 100644 crates/ml/tests/aux_trunk_oracle_tests.rs diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 77f521f54..90b3c2a44 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -196,6 +196,16 @@ fn main() { // Producer-only; ISV[AUX_NEXT_BAR_MSE_EMA_INDEX] + // ISV[AUX_REGIME_CE_EMA_INDEX] consumer wires in Commit B. "aux_heads_loss_ema_kernel.cu", + // SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward kernel — + // 3-layer Linear→ELU→Linear→ELU→Linear MLP parallel to Q's GRN + // trunk. Reads encoder output (`x_in [B, ENCODER_OUT_DIM]`), + // writes `h_s2_aux [B, AUX_HIDDEN_DIM]` plus saved-for-backward + // hidden activations `h_aux1` and `h_aux2`. Module is additive + // in this commit — wire-up into the collector chain lands in + // Phase C.5 (atomic). Backward kernel + Adam updates land in + // Phase C.4. See plan §C.3 in + // `2026-05-07-sp14-layer-c-separate-aux-trunk.md`. + "aux_trunk_forward_kernel.cu", // Plan A Phase 0 (Thompson sampling spec 2026-04-26): standalone test // kernel exercised only by `distributional_q_tests.rs`. Implements the // Thompson direction sampling math (inverse-CDF over C51 atoms, diff --git a/crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu b/crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu new file mode 100644 index 000000000..86ed3138e --- /dev/null +++ b/crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu @@ -0,0 +1,172 @@ +/* + * aux_trunk_forward_kernel.cu — SP14 Layer C Phase C.3 (2026-05-08) + * + * 3-layer MLP forward pass for the auxiliary head's separate trunk: + * + * x_in [B, ENCODER_OUT_DIM] + * → Linear(w1, b1) → [B, H1] (H1 = AUX_TRUNK_H1 = 256) + * → ELU + * → Linear(w2, b2) → [B, H2] (H2 = AUX_TRUNK_H2 = 128) + * → ELU + * → Linear(w3, b3) → [B, AUX_HIDDEN_DIM] + * h_s2_aux [B, AUX_HIDDEN_DIM] + * + * No activation on the final output — the downstream aux head applies its + * own activation. The encoder output `x_in` flows through stop-gradient at + * the encoder boundary in C.4 (backward) so this trunk shapes a + * representation suited to next-bar prediction without pulling on the + * encoder. Aux head's own w1/b1/w2/b2 train independently from the trunk + * weights here. + * + * Topology (production config — verified C.2): + * ENCODER_OUT_DIM = config.shared_h1 = 256 + * AUX_HIDDEN_DIM = config.shared_h2 = 256 (matches existing aux head input dim) + * H1 = AUX_TRUNK_H1 = 256 + * H2 = AUX_TRUNK_H2 = 128 + * + * Saved-for-backward outputs: + * - h_aux1_out [B, H1]: POST-ELU Layer-1 hidden activations + * - h_aux2_out [B, H2]: POST-ELU Layer-2 hidden activations + * - h_s2_aux_out [B, AUX_HIDDEN_DIM]: final linear output (no activation) + * + * Backward (C.4) re-reads h_aux1 + h_aux2 to derive ELU' from the post- + * activation values via the standard `f'(x) = 1 if y > 0 else 1 + y` + * identity (mirrors `aux_elu_bwd_from_post` in `aux_heads_kernel.cu`). + * + * Pearls applied: + * - feedback_no_atomicadd.md — forward writes only per-(b, j) outputs; + * no contention across blocks. Caller-side cuBLAS is NOT used here + * because mirroring `aux_next_bar_forward`'s thread-per-output style + * keeps the kernel self-contained for graph capture. + * - pearl_no_host_branches_in_captured_graph.md — kernel signature is + * pure-device; no host-side dispatch inside the captured region. + * - feedback_no_stubs.md — full SGEMM math, no return-zero, no + * placeholder body. + * - common_device_functions.cuh provides standard helpers (this kernel + * gets the common header prepended at build time per `build.rs`). + * + * Layout convention: row-major `[B, …]` for activations. + * w1 [ENCODER_OUT_DIM, H1] row-major → w1[k, j] = w1_ptr[k * H1 + j] + * w2 [H1, H2] row-major → w2[k, j] = w2_ptr[k * H2 + j] + * w3 [H2, AUX_HIDDEN_DIM] row-major → w3[k, j] = w3_ptr[k * AUX_HIDDEN_DIM + j] + * + * ELU activation: f(x) = x if x > 0 else (exp(x) - 1). + * + * Block: `AUX_TRUNK_BLOCK = 256` threads. Grid: `(B, 1, 1)` blocks. + * Shared memory: `(H1 + H2) * sizeof(float)` bytes (post-ELU caches for + * Layer 2 + Layer 3). + */ + +#include +#include + +#ifndef AUX_TRUNK_BLOCK +#define AUX_TRUNK_BLOCK 256 +#endif + +/* ELU helper — branchless on `x > 0.f`. Mirrors `aux_elu_fwd` in + * aux_heads_kernel.cu (kept inline rather than #include'd because the + * common header pre-pended by build.rs only ships the small generic + * device fns; the aux-specific ELU helpers live in their own kernel + * file). */ +__device__ __forceinline__ float aux_trunk_elu_fwd(float x) { + return (x > 0.0f) ? x : (expf(x) - 1.0f); +} + +/* ===================================================================== + * aux_trunk_forward — 3-layer Linear→ELU→Linear→ELU→Linear MLP. + * + * Per sample b (one block per sample, AUX_TRUNK_BLOCK threads): + * 1. h_aux1[k] = ELU( b1[k] + sum_j w1[j, k] * x_in[b, j] ) for k ∈ [0, H1) + * 2. h_aux2[k] = ELU( b2[k] + sum_j w2[j, k] * h_aux1[j] ) for k ∈ [0, H2) + * 3. h_s2_aux[k] = b3[k] + sum_j w3[j, k] * h_aux2[j] for k ∈ [0, AUX_HIDDEN_DIM) + * + * Hidden layer outputs (post-ELU) are saved to global memory for backward. + * Final layer output has NO activation — downstream aux head applies its + * own activation pipeline. + * + * Each thread within the block owns a slice of the layer's output dim via + * stride-loop. The matmul over the input dim is serial within each + * thread (mirrors `aux_next_bar_forward` — input dims are small enough + * that block-cooperative reduce wastes more cycles in shfl bookkeeping + * than the serial pass). The shared-memory `sh_h1` / `sh_h2` caches let + * the next layer's matmul read post-ELU values without a second global- + * memory pass. + * + * Block: AUX_TRUNK_BLOCK threads. Grid: (B, 1, 1). Shared memory: + * (H1 + H2) floats — Layer-1 cache for Layer-2 matmul + Layer-2 cache + * for Layer-3 matmul. + * ===================================================================== */ +extern "C" __global__ void aux_trunk_forward( + const float* __restrict__ x_in, /* [B, ENCODER_OUT_DIM] row-major */ + const float* __restrict__ w1, /* [ENCODER_OUT_DIM, H1] row-major */ + const float* __restrict__ b1, /* [H1] */ + const float* __restrict__ w2, /* [H1, H2] row-major */ + const float* __restrict__ b2, /* [H2] */ + const float* __restrict__ w3, /* [H2, AUX_HIDDEN_DIM] row-major */ + const float* __restrict__ b3, /* [AUX_HIDDEN_DIM] */ + float* __restrict__ h_aux1_out, /* [B, H1] saved for backward */ + float* __restrict__ h_aux2_out, /* [B, H2] saved for backward */ + float* __restrict__ h_s2_aux_out, /* [B, AUX_HIDDEN_DIM] final */ + int B, + int ENCODER_OUT_DIM, + int H1, + int H2, + int AUX_HIDDEN_DIM +) { + const int b = blockIdx.x; + if (b >= B) return; + const int tid = threadIdx.x; + + /* Shared memory: post-ELU caches for Layer 1 + Layer 2. + * Total: (H1 + H2) floats. Allocator-side smem size is computed in + * the Rust wrapper based on the runtime H1/H2 values. */ + extern __shared__ float smem[]; + float* sh_h1 = smem; /* [H1] */ + float* sh_h2 = smem + H1; /* [H2] */ + + /* ─────────────── Layer 1: Linear(w1, b1) → ELU ─────────────── */ + /* Each thread computes one (or more, via stride loop) hidden lanes. + * Serial dot over ENCODER_OUT_DIM per lane — input dim is small + * enough that block-cooperative reduce is wasteful. */ + { + const float* x_row = x_in + (size_t)b * ENCODER_OUT_DIM; + for (int k = tid; k < H1; k += blockDim.x) { + float acc = b1[k]; + for (int j = 0; j < ENCODER_OUT_DIM; ++j) { + acc += x_row[j] * w1[(size_t)j * H1 + k]; + } + const float h_post = aux_trunk_elu_fwd(acc); + sh_h1[k] = h_post; + h_aux1_out[(size_t)b * H1 + k] = h_post; + } + } + __syncthreads(); + + /* ─────────────── Layer 2: Linear(w2, b2) → ELU ─────────────── */ + /* Read post-ELU Layer-1 from shmem (no second global-memory pass). */ + for (int k = tid; k < H2; k += blockDim.x) { + float acc = b2[k]; + for (int j = 0; j < H1; ++j) { + acc += sh_h1[j] * w2[(size_t)j * H2 + k]; + } + const float h_post = aux_trunk_elu_fwd(acc); + sh_h2[k] = h_post; + h_aux2_out[(size_t)b * H2 + k] = h_post; + } + __syncthreads(); + + /* ─────────────── Layer 3: Linear(w3, b3) — NO activation ─────────────── */ + /* Read post-ELU Layer-2 from shmem. Output goes directly to global — + * downstream aux head's first FC applies its own activation. */ + { + float* out_row = h_s2_aux_out + (size_t)b * AUX_HIDDEN_DIM; + for (int k = tid; k < AUX_HIDDEN_DIM; k += blockDim.x) { + float acc = b3[k]; + for (int j = 0; j < H2; ++j) { + acc += sh_h2[j] * w3[(size_t)j * AUX_HIDDEN_DIM + k]; + } + out_row[k] = acc; + } + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs b/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs new file mode 100644 index 000000000..cf2afe354 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs @@ -0,0 +1,164 @@ +#![allow(unsafe_code)] // Required for CUDA kernel launches. + +//! `gpu_aux_trunk` — SP14 Layer C Phase C.3 (2026-05-08). +//! +//! Rust wrapper for the auxiliary trunk forward kernel +//! (`aux_trunk_forward_kernel.cu`). Owns the pre-loaded `CudaFunction` +//! handle so launches in the captured graph never touch host-side cubin +//! load paths (per `pearl_no_host_branches_in_captured_graph.md`). +//! +//! # Architecture +//! +//! 3-layer Linear→ELU→Linear→ELU→Linear MLP parallel to Q's GRN trunk. +//! Both trunks read the shared encoder output (`x_in [B, ENCODER_OUT_DIM]`); +//! Q trunk produces `h_s2`, aux trunk produces `h_s2_aux`. Aux backward +//! (Phase C.4) terminates at the encoder boundary — Q-loss is the sole +//! shaping force on the encoder per the locked design decision in +//! `2026-05-07-sp14-layer-c-separate-aux-trunk.md`. +//! +//! Saved-for-backward outputs (consumed by aux trunk backward in C.4): +//! +//! - `h_aux1 [B, H1]` — Layer-1 post-ELU activation +//! - `h_aux2 [B, H2]` — Layer-2 post-ELU activation +//! - `h_s2_aux [B, AUX_HIDDEN_DIM]` — final linear output (no activation) +//! +//! Production topology (verified C.2): +//! - `ENCODER_OUT_DIM = config.shared_h1 = 256` +//! - `H1 = AUX_TRUNK_H1 = 256` +//! - `H2 = AUX_TRUNK_H2 = 128` +//! - `AUX_HIDDEN_DIM = config.shared_h2 = 256` (matches existing aux head input dim) +//! +//! # Pearls applied +//! +//! - `pearl_no_host_branches_in_captured_graph.md` — `CudaFunction` +//! pre-loaded once at construction; no `load_cubin` / `load_function` +//! on the launch path. +//! - `feedback_no_atomicadd.md` — forward kernel writes only per-(b, j) +//! outputs; no contention across blocks. +//! - `feedback_no_stubs.md` — full launcher, no placeholder body. + +use std::sync::Arc; + +use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg}; + +use crate::MLError; + +use super::gpu_dqn_trainer::AUX_TRUNK_FORWARD_CUBIN; + +/// Hidden width of the aux trunk's first internal layer (Linear_1 → ELU +/// output). Mirrors the encoder output dimension so Layer-1 acts as a +/// representation-mixing identity-shaped projection. +pub(crate) const AUX_TRUNK_H1: usize = 256; + +/// Hidden width of the aux trunk's second internal layer (Linear_2 → ELU +/// output). Smaller than H1 for a mild bottleneck — encourages the aux +/// trunk to compress its representation before lifting back up to +/// `AUX_HIDDEN_DIM` in Layer 3. +pub(crate) const AUX_TRUNK_H2: usize = 128; + +/// Block dim used by the per-sample forward kernel. Mirrors +/// `aux_heads_kernel.cu`'s `AUX_BLOCK = 256`. +const AUX_TRUNK_BLOCK: u32 = 256; + +/// Forward orchestrator — owns the single `CudaFunction` for +/// `aux_trunk_forward`. Backward kernel + ops land in Phase C.4. +/// +/// Dropping this struct releases the underlying `CudaFunction` handle +/// via cudarc's RAII; no explicit teardown required. +#[allow(missing_debug_implementations)] +pub(crate) struct AuxTrunkForwardOps { + /// Pre-loaded forward kernel handle (loaded once at construction; + /// never re-loaded on the launch path). + forward_kernel: CudaFunction, +} + +impl AuxTrunkForwardOps { + /// Load the forward kernel handle from the precompiled cubin. + /// Mirrors `GrnBlock::new`'s loader pattern. + pub(crate) fn new(stream: &Arc) -> Result { + let context = stream.context(); + let module = context + .load_cubin(AUX_TRUNK_FORWARD_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("aux_trunk_forward cubin load: {e}")))?; + let forward_kernel = module + .load_function("aux_trunk_forward") + .map_err(|e| MLError::ModelError(format!("aux_trunk_forward load: {e}")))?; + Ok(Self { forward_kernel }) + } + + /// Launch `aux_trunk_forward`: 3-layer Linear→ELU→Linear→ELU→Linear MLP. + /// + /// Caller-owned buffers (raw `u64` device pointers for graph-capture + /// safety; mirrors `gpu_grn.rs::forward_raw` and `gpu_aux_heads.rs`): + /// * `x_in_ptr` — `[B, ENCODER_OUT_DIM]` row-major (encoder output). + /// * `w1_ptr`/`b1_ptr` — `[ENCODER_OUT_DIM, H1]` / `[H1]`. + /// * `w2_ptr`/`b2_ptr` — `[H1, H2]` / `[H2]`. + /// * `w3_ptr`/`b3_ptr` — `[H2, AUX_HIDDEN_DIM]` / `[AUX_HIDDEN_DIM]`. + /// * `h_aux1_out_ptr` — `[B, H1]` SAVED post-ELU Layer-1 (consumed by C.4 backward). + /// * `h_aux2_out_ptr` — `[B, H2]` SAVED post-ELU Layer-2 (consumed by C.4 backward). + /// * `h_s2_aux_out_ptr` — `[B, AUX_HIDDEN_DIM]` final aux trunk output. + /// + /// `h1` / `h2` / `aux_hidden_dim` are passed as runtime args for + /// kernel-signature stability (a future bump won't require an ABI + /// change). `encoder_out_dim` is the encoder's actual output width + /// (typically `config.shared_h1 = 256`). + /// + /// Block: `AUX_TRUNK_BLOCK = 256` threads. + /// Grid: `B` blocks (one block per batch row). + /// Shared mem: `(H1 + H2) * sizeof(f32)` bytes — Layer-1 cache for + /// Layer-2 matmul + Layer-2 cache for Layer-3 matmul. + #[allow(clippy::too_many_arguments)] + pub(crate) fn launch( + &self, + stream: &Arc, + x_in_ptr: u64, + w1_ptr: u64, + b1_ptr: u64, + w2_ptr: u64, + b2_ptr: u64, + w3_ptr: u64, + b3_ptr: u64, + h_aux1_out_ptr: u64, + h_aux2_out_ptr: u64, + h_s2_aux_out_ptr: u64, + b: usize, + encoder_out_dim: usize, + h1: usize, + h2: usize, + aux_hidden_dim: usize, + ) -> Result<(), MLError> { + let b_i32 = b as i32; + let enc_i32 = encoder_out_dim as i32; + let h1_i32 = h1 as i32; + let h2_i32 = h2 as i32; + let aux_i32 = aux_hidden_dim as i32; + // Shared memory: H1 + H2 floats — Layer-1 cache + Layer-2 cache. + let smem_bytes = (h1 as u32 + h2 as u32) * std::mem::size_of::() as u32; + unsafe { + stream + .launch_builder(&self.forward_kernel) + .arg(&x_in_ptr) + .arg(&w1_ptr) + .arg(&b1_ptr) + .arg(&w2_ptr) + .arg(&b2_ptr) + .arg(&w3_ptr) + .arg(&b3_ptr) + .arg(&h_aux1_out_ptr) + .arg(&h_aux2_out_ptr) + .arg(&h_s2_aux_out_ptr) + .arg(&b_i32) + .arg(&enc_i32) + .arg(&h1_i32) + .arg(&h2_i32) + .arg(&aux_i32) + .launch(LaunchConfig { + grid_dim: (b as u32, 1, 1), + block_dim: (AUX_TRUNK_BLOCK, 1, 1), + shared_mem_bytes: smem_bytes, + }) + .map_err(|e| MLError::ModelError(format!("aux_trunk_forward: {e}")))?; + } + Ok(()) + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index d21105deb..0ad179e39 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -2061,6 +2061,15 @@ pub(crate) static AUX_HEADS_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR" #[allow(dead_code)] pub(crate) static AUX_HEADS_LOSS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_heads_loss_ema_kernel.cubin")); +/// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward kernel cubin. +/// 3-layer Linear→ELU→Linear→ELU→Linear MLP forward — reads encoder +/// output, writes `h_s2_aux` plus saved-for-backward hidden activations. +/// Loaded by `gpu_aux_trunk::AuxTrunkForwardOps`. Module is additive in +/// this commit — wire-up into the collector chain lands in Phase C.5 +/// (atomic). Backward kernel + Adam updates land in Phase C.4. +#[allow(dead_code)] +pub(crate) static AUX_TRUNK_FORWARD_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/aux_trunk_forward_kernel.cubin")); + /// Plan C Phase 2 follow-up A.2 (2026-04-29): q-drift rate ISV producer. /// Single-thread single-block cold-path kernel mirroring /// `moe_lambda_eff_kernel.cu` / `kelly_cap_update_kernel.cu`. Reads @@ -7289,6 +7298,17 @@ pub struct GpuDqnTrainer { /// Adam second moment for `aux_trunk_b3` [SH2]. aux_trunk_b3_v: CudaSlice, + /// SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator. + /// Owns the pre-loaded `CudaFunction` handle for `aux_trunk_forward` + /// per `pearl_no_host_branches_in_captured_graph.md` — the launch + /// path never touches `load_cubin` / `load_function`. Wire-up into + /// the collector chain lands in Phase C.5 (atomic). Backward + /// orchestrator + Adam updates land in Phase C.4. The `dead_code` + /// allow is necessary because this commit is additive — production + /// callers light up in C.5 per `feedback_no_partial_refactor`. + #[allow(dead_code)] + aux_trunk_forward_ops: super::gpu_aux_trunk::AuxTrunkForwardOps, + // ── Speculative inference cache ── /// Cached trunk output [B, SH2] from between-bar speculative forward. speculative_h_s2: CudaSlice, @@ -22135,6 +22155,13 @@ impl GpuDqnTrainer { aux_trunk_in_dim, aux_trunk_h1, aux_trunk_h2, aux_trunk_out_dim, aux_trunk_total_params, ); + // SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator. + // Pre-loads the `aux_trunk_forward` CudaFunction handle once at + // construction (per `pearl_no_host_branches_in_captured_graph.md`). + // Wire-up into the collector chain lands in Phase C.5 (atomic). + let aux_trunk_forward_ops = + super::gpu_aux_trunk::AuxTrunkForwardOps::new(&stream)?; + // ── Speculative inference cache ── let speculative_h_s2 = stream.alloc_zeros::(b * sh2) .map_err(|e| MLError::ModelError(format!("speculative_h_s2 alloc: {e}")))?; @@ -22935,6 +22962,8 @@ impl GpuDqnTrainer { aux_trunk_w3_v, aux_trunk_b3_m, aux_trunk_b3_v, + // SP14 Layer C Phase C.3 (2026-05-08): aux trunk forward orchestrator. + aux_trunk_forward_ops, speculative_h_s2, speculative_features, speculative_valid: false, diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index 1a4ef0ccd..c0b889295 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -53,6 +53,7 @@ pub mod gpu_attention; pub mod gpu_tlob; pub mod gpu_grn; pub mod gpu_aux_heads; +pub mod gpu_aux_trunk; pub mod gpu_moe_head; pub mod decision_transformer; pub mod learning_health; diff --git a/crates/ml/tests/aux_trunk_oracle_tests.rs b/crates/ml/tests/aux_trunk_oracle_tests.rs new file mode 100644 index 000000000..e90bbb0f6 --- /dev/null +++ b/crates/ml/tests/aux_trunk_oracle_tests.rs @@ -0,0 +1,250 @@ +//! Oracle tests for aux trunk forward (SP14 Layer C Phase C.3, 2026-05-08). +//! +//! Layer C introduces a separate auxiliary trunk parallel to Q's GRN trunk +//! per the locked design in +//! `2026-05-07-sp14-layer-c-separate-aux-trunk.md`. The forward kernel is +//! a 3-layer Linear→ELU→Linear→ELU→Linear MLP that reads the encoder +//! output and produces `h_s2_aux [B, AUX_HIDDEN_DIM]` for the aux head. +//! Backward (C.4) terminates at the encoder boundary — Q-loss is the +//! sole shaping force on the encoder. +//! +//! This test verifies the forward kernel's bit-for-bit match against a +//! deterministic host reference within `1e-4` tolerance (f32 SGEMM +//! precision). All CPU↔GPU buffers are `MappedF32Buffer` per +//! `feedback_no_htod_htoh_only_mapped_pinned.md` (tests are not exempt). +//! +//! Run with: +//! +//! SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \ +//! cargo test -p ml --test aux_trunk_oracle_tests --release \ +//! -- --ignored --nocapture +//! +//! Backward + gradient-check + stop-grad invariant tests land in C.4. + +#![allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory. + +#[cfg(feature = "cuda")] +mod gpu { + use std::sync::Arc; + + use cudarc::driver::{CudaContext, CudaFunction, CudaStream, LaunchConfig, PushKernelArg}; + use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer; + + const AUX_TRUNK_FORWARD_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/aux_trunk_forward_kernel.cubin")); + + const AUX_TRUNK_BLOCK: u32 = 256; + + fn make_test_stream() -> Arc { + let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?"); + ctx.default_stream() + } + + fn load_aux_trunk_forward(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(AUX_TRUNK_FORWARD_CUBIN.to_vec()) + .expect("load aux_trunk_forward cubin"); + module + .load_function("aux_trunk_forward") + .expect("load aux_trunk_forward function") + } + + /// Deterministic LCG → uniform `[-1, 1)` for reproducible weight/input + /// generation. Mirrors the LCG state-machine pattern used elsewhere + /// in the codebase (e.g. `xavier_init_dqn_weights` in + /// `gpu_dqn_trainer.rs`). Seed is fixed so every test run produces + /// bit-identical inputs. + fn lcg_next_signed(state: &mut u32) -> f32 { + *state = state.wrapping_mul(1_103_515_245).wrapping_add(12_345); + let raw = (*state >> 16) as f32 / 32_768.0; // [0, 2) + raw - 1.0 // [-1, 1) + } + + /// ELU forward used in the host reference. Mirrors + /// `aux_trunk_elu_fwd` in the kernel exactly. + fn elu(x: f32) -> f32 { + if x > 0.0 { x } else { x.exp() - 1.0 } + } + + /// Forward output check — runs the GPU kernel against a deterministic + /// host reference and verifies element-wise match within 1e-4 (f32 + /// SGEMM precision). + /// + /// Uses production topology dimensions: + /// - ENCODER_OUT_DIM = 256 (= config.shared_h1) + /// - H1 = AUX_TRUNK_H1 = 256 + /// - H2 = AUX_TRUNK_H2 = 128 + /// - AUX_HIDDEN_DIM = 256 (= config.shared_h2) + /// + /// B is small (4) to keep the host reference fast — the matmul + /// arithmetic is dominated by the per-row dot products (256 × 256 + + /// 256 × 128 + 128 × 256 = 131,072 multiplies per row), not B. + #[test] + #[ignore = "requires GPU"] + fn aux_trunk_forward_matches_numpy_reference() { + let stream = make_test_stream(); + let kernel = load_aux_trunk_forward(&stream); + + const B: usize = 4; + const ENCODER_OUT_DIM: usize = 256; + const H1: usize = 256; + const H2: usize = 128; + const AUX_HIDDEN_DIM: usize = 256; + + // Deterministic LCG seed — fixed value so test is reproducible. + let mut rng_state: u32 = 0x1234_5678; + let next_f32 = |state: &mut u32| lcg_next_signed(state); + + // Inputs scaled to small magnitudes so the cumulative 256-deep + // dot product stays in a numerically benign range. (×0.1 input, + // ×0.05 weights → per-row sum of ~256 products of 0.005 RMS + // each → expected acc magnitude < ~5 before activation.) + let x_in_host: Vec = (0..B * ENCODER_OUT_DIM) + .map(|_| next_f32(&mut rng_state) * 0.1) + .collect(); + let w1_host: Vec = (0..ENCODER_OUT_DIM * H1) + .map(|_| next_f32(&mut rng_state) * 0.05) + .collect(); + let b1_host: Vec = vec![0.0; H1]; + let w2_host: Vec = (0..H1 * H2) + .map(|_| next_f32(&mut rng_state) * 0.05) + .collect(); + let b2_host: Vec = vec![0.0; H2]; + let w3_host: Vec = (0..H2 * AUX_HIDDEN_DIM) + .map(|_| next_f32(&mut rng_state) * 0.05) + .collect(); + let b3_host: Vec = vec![0.0; AUX_HIDDEN_DIM]; + + // Allocate device buffers via MappedF32Buffer (mapped-pinned, no + // raw HtoD per `feedback_no_htod_htoh_only_mapped_pinned.md`). + let x_in = unsafe { MappedF32Buffer::new(B * ENCODER_OUT_DIM) } + .expect("alloc x_in"); + x_in.write_from_slice(&x_in_host); + let w1 = unsafe { MappedF32Buffer::new(ENCODER_OUT_DIM * H1) }.expect("alloc w1"); + w1.write_from_slice(&w1_host); + let b1 = unsafe { MappedF32Buffer::new(H1) }.expect("alloc b1"); + b1.write_from_slice(&b1_host); + let w2 = unsafe { MappedF32Buffer::new(H1 * H2) }.expect("alloc w2"); + w2.write_from_slice(&w2_host); + let b2 = unsafe { MappedF32Buffer::new(H2) }.expect("alloc b2"); + b2.write_from_slice(&b2_host); + let w3 = unsafe { MappedF32Buffer::new(H2 * AUX_HIDDEN_DIM) }.expect("alloc w3"); + w3.write_from_slice(&w3_host); + let b3 = unsafe { MappedF32Buffer::new(AUX_HIDDEN_DIM) }.expect("alloc b3"); + b3.write_from_slice(&b3_host); + + // Output + saved-for-backward buffers (zero-init). + let h_aux1 = unsafe { MappedF32Buffer::new(B * H1) }.expect("alloc h_aux1"); + h_aux1.write_from_slice(&vec![0.0_f32; B * H1]); + let h_aux2 = unsafe { MappedF32Buffer::new(B * H2) }.expect("alloc h_aux2"); + h_aux2.write_from_slice(&vec![0.0_f32; B * H2]); + let h_s2_aux = unsafe { MappedF32Buffer::new(B * AUX_HIDDEN_DIM) } + .expect("alloc h_s2_aux"); + h_s2_aux.write_from_slice(&vec![0.0_f32; B * AUX_HIDDEN_DIM]); + + let b_i32: i32 = B as i32; + let enc_i32: i32 = ENCODER_OUT_DIM as i32; + let h1_i32: i32 = H1 as i32; + let h2_i32: i32 = H2 as i32; + let aux_i32: i32 = AUX_HIDDEN_DIM as i32; + // Shared memory: H1 + H2 floats — Layer-1 cache + Layer-2 cache. + let smem_bytes = + (H1 as u32 + H2 as u32) * std::mem::size_of::() as u32; + + unsafe { + stream + .launch_builder(&kernel) + .arg(&x_in.dev_ptr) + .arg(&w1.dev_ptr) + .arg(&b1.dev_ptr) + .arg(&w2.dev_ptr) + .arg(&b2.dev_ptr) + .arg(&w3.dev_ptr) + .arg(&b3.dev_ptr) + .arg(&h_aux1.dev_ptr) + .arg(&h_aux2.dev_ptr) + .arg(&h_s2_aux.dev_ptr) + .arg(&b_i32) + .arg(&enc_i32) + .arg(&h1_i32) + .arg(&h2_i32) + .arg(&aux_i32) + .launch(LaunchConfig { + grid_dim: (B as u32, 1, 1), + block_dim: (AUX_TRUNK_BLOCK, 1, 1), + shared_mem_bytes: smem_bytes, + }) + .expect("launch aux_trunk_forward"); + } + stream.synchronize().expect("sync after aux_trunk_forward"); + + // ── Host reference ──────────────────────────────────────────── + // Layer 1: h_aux1[b, k] = ELU( b1[k] + sum_j x_in[b, j] * w1[j, k] ) + let mut h_aux1_ref = vec![0.0_f32; B * H1]; + for b in 0..B { + for k in 0..H1 { + let mut acc = b1_host[k]; + for j in 0..ENCODER_OUT_DIM { + acc += x_in_host[b * ENCODER_OUT_DIM + j] * w1_host[j * H1 + k]; + } + h_aux1_ref[b * H1 + k] = elu(acc); + } + } + // Layer 2: h_aux2[b, k] = ELU( b2[k] + sum_j h_aux1[b, j] * w2[j, k] ) + let mut h_aux2_ref = vec![0.0_f32; B * H2]; + for b in 0..B { + for k in 0..H2 { + let mut acc = b2_host[k]; + for j in 0..H1 { + acc += h_aux1_ref[b * H1 + j] * w2_host[j * H2 + k]; + } + h_aux2_ref[b * H2 + k] = elu(acc); + } + } + // Layer 3: h_s2_aux[b, k] = b3[k] + sum_j h_aux2[b, j] * w3[j, k] (no activation) + let mut h_s2_aux_ref = vec![0.0_f32; B * AUX_HIDDEN_DIM]; + for b in 0..B { + for k in 0..AUX_HIDDEN_DIM { + let mut acc = b3_host[k]; + for j in 0..H2 { + acc += h_aux2_ref[b * H2 + j] * w3_host[j * AUX_HIDDEN_DIM + k]; + } + h_s2_aux_ref[b * AUX_HIDDEN_DIM + k] = acc; + } + } + + // ── Compare ─────────────────────────────────────────────────── + let h_aux1_dev = h_aux1.read_all(); + let h_aux2_dev = h_aux2.read_all(); + let h_s2_aux_dev = h_s2_aux.read_all(); + + let max_diff = |dev: &[f32], reference: &[f32]| -> f32 { + dev.iter() + .zip(reference.iter()) + .map(|(d, r)| (d - r).abs()) + .fold(0.0_f32, f32::max) + }; + + let diff_h1 = max_diff(&h_aux1_dev, &h_aux1_ref); + let diff_h2 = max_diff(&h_aux2_dev, &h_aux2_ref); + let diff_out = max_diff(&h_s2_aux_dev, &h_s2_aux_ref); + + const TOL: f32 = 1e-4; + assert!( + diff_h1 < TOL, + "h_aux1 max diff {diff_h1} exceeds {TOL} tol — Layer-1 kernel mismatch", + ); + assert!( + diff_h2 < TOL, + "h_aux2 max diff {diff_h2} exceeds {TOL} tol — Layer-2 kernel mismatch", + ); + assert!( + diff_out < TOL, + "h_s2_aux max diff {diff_out} exceeds {TOL} tol — Layer-3 kernel mismatch", + ); + eprintln!( + "aux_trunk_forward oracle: h_aux1 diff={diff_h1:.2e}, h_aux2 diff={diff_h2:.2e}, h_s2_aux diff={diff_out:.2e} (tol={TOL:.0e})" + ); + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index f9ee7fb5b..61713157c 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2,6 +2,41 @@ **Status:** Populated during Plan 1 Task 6 (A.5 orphan audit). Updated on every commit per Invariant 7. +## 2026-05-08 — SP14 Layer C Phase C.3: aux trunk forward kernel + Rust wrapper + oracle test + +Phase C.3 of `docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md`. Additive commit — forward kernel lands here; backward (C.4) and wire-up (C.5) land atomically in subsequent commits per `feedback_no_partial_refactor`. + +### Files added +- `crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu` (172 LOC) — 3-layer Linear→ELU→Linear→ELU→Linear forward kernel. One block per batch sample, `AUX_TRUNK_BLOCK=256` threads. Shared memory `(H1+H2)×sizeof(f32)` caches post-ELU Layer-1 and Layer-2 activations for Layer-2 and Layer-3 matmuls respectively. Row-major weight layout: `w[j, k] = w_ptr[j * out_dim + k]`. +- `crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs` (164 LOC) — `AuxTrunkForwardOps` struct. Pre-loads `CudaFunction` handle at construction; launch path never touches `load_cubin`/`load_function` per `pearl_no_host_branches_in_captured_graph`. + +### Files modified +- `crates/ml/build.rs` — registers `aux_trunk_forward_kernel.cu` in the cubin manifest. +- `crates/ml/src/cuda_pipeline/mod.rs` — declares `pub mod gpu_aux_trunk`. +- `crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs` — adds `AUX_TRUNK_FORWARD_CUBIN` static, adds `aux_trunk_forward_ops: AuxTrunkForwardOps` field (pre-loaded at constructor), inserts `AuxTrunkForwardOps::new(&stream)?` call in `GpuDqnTrainer::new`. Field is `#[allow(dead_code)]` until Phase C.5 wire-up. +- `crates/ml/tests/aux_trunk_oracle_tests.rs` (250 LOC) — oracle test `aux_trunk_forward_matches_numpy_reference` verifies h_aux1, h_aux2, and h_s2_aux match host reference within 1e-4 (actual max diff: 1.19e-7). + +### Saved-for-backward outputs +- `h_aux1 [B, H1=256]` — post-ELU Layer-1 activations +- `h_aux2 [B, H2=128]` — post-ELU Layer-2 activations +- `h_s2_aux [B, AUX_HIDDEN_DIM=256]` — final linear output (no activation) + +### Oracle test result +- `h_aux1 max diff = 5.96e-8` (tol 1e-4) — PASS +- `h_aux2 max diff = 1.19e-7` (tol 1e-4) — PASS +- `h_s2_aux max diff = 2.89e-8` (tol 1e-4) — PASS + +### Scope explicitly excluded from C.3 (deferred to C.4-C.8) +- Backward kernel `aux_trunk_backward_kernel.cu` (Phase C.4) with stop-grad on encoder boundary +- Atomic forward+backward wire-up + aux head input switch `h_s2 → h_s2_aux` (Phase C.5) +- `h_s2_aux_rms_ema` producer kernel (Phase C.6) +- ISV-driven Adam β1/β2/ε/LR/grad-clip wiring (Phase C.8) + +### Build state +- `SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --tests --all-targets` clean (only pre-existing warnings). + +--- + ## 2026-05-08 — SP14 Layer C Phase C.2: aux trunk param tensors + Adam state allocated Phase C.2 of `docs/superpowers/plans/2026-05-07-sp14-layer-c-separate-aux-trunk.md`. Pure constructor-time allocation — no forward/backward wire-up yet (those land atomically in Phases C.3-C.5 per `feedback_no_partial_refactor`).