feat(sp14-c): aux trunk forward kernel + Rust wrapper + oracle test
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
@@ -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
172
crates/ml/src/cuda_pipeline/aux_trunk_forward_kernel.cu
Normal 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
164
crates/ml/src/cuda_pipeline/gpu_aux_trunk.rs
Normal 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(())
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
|
||||
@@ -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
250
crates/ml/tests/aux_trunk_oracle_tests.rs
Normal 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})"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user