gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(sp14-c): aux trunk forward kernel + Rust wrapper + oracle test

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-08 00:56:23 +02:00
parent 4926fb7c65
commit cb6bca4629
7 changed files with 661 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
crates/ml/build.rs
View File

@@ -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,

172
crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu Normal file
View File

@@ -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 <cuda_runtime.h>
#include <math_constants.h>
#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;
}
}
}

164
crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs Normal file
View File

@@ -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<CudaStream>) -> Result<Self, MLError> {
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<CudaStream>,
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::<f32>() 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(())
}
}

29
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -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<f32>,
/// 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<f32>,
@@ -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::<f32>(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,

1
crates/ml/src/cuda_pipeline/mod.rs
View File

@@ -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;

250
crates/ml/tests/aux_trunk_oracle_tests.rs Normal file
View File

@@ -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<CudaStream> {
let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
ctx.default_stream()
}
fn load_aux_trunk_forward(stream: &Arc<CudaStream>) -> 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<f32> = (0..B * ENCODER_OUT_DIM)
.map(|_| next_f32(&mut rng_state) * 0.1)
.collect();
let w1_host: Vec<f32> = (0..ENCODER_OUT_DIM * H1)
.map(|_| next_f32(&mut rng_state) * 0.05)
.collect();
let b1_host: Vec<f32> = vec![0.0; H1];
let w2_host: Vec<f32> = (0..H1 * H2)
.map(|_| next_f32(&mut rng_state) * 0.05)
.collect();
let b2_host: Vec<f32> = vec![0.0; H2];
let w3_host: Vec<f32> = (0..H2 * AUX_HIDDEN_DIM)
.map(|_| next_f32(&mut rng_state) * 0.05)
.collect();
let b3_host: Vec<f32> = 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::<f32>() 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})"
);
}
}

35
docs/dqn-wire-up-audit.md
View File

@@ -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`).