diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 1090db709..7be987c37 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -1359,6 +1359,26 @@ fn main() { // seeds + 1 scratch buffer + 4 registry entries + 4 dispatch // arms only. "dd_trajectory_kernel.cu", + // SP20 Phase 1.1 (2026-05-09): aux confidence p50/std stats + // producer. Single-block, BLOCK=256 kernel reading + // `aux_logits [B, 3]` (the SP14-C aux head's 3-class direction + // logits) and emitting `[p50, std]` of the per-row signal + // `aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/3` into a + // `MappedF32Buffer<2>`. p50 uses the inlined `sp4_histogram_p99` + // pattern with cumulative-from-bottom (target = ⌈B/2⌉) + // substituted for cumulative-from-top, per + // `pearl_fused_per_group_statistics_oracle` (one fused stream + // for both stats) and `feedback_no_atomicadd` (per-warp tile + // binning + block tree-reduce sums; no atomicAdd anywhere). + // std uses two block tree-reductions sharing one shmem tile + // sequentially. Component 5 / Kernel 3 of the SP20 fused- + // producer chain (`sp20_emas_compute` + `sp20_controllers_ + // compute` land in Phase 1.2 + 1.3). Phase 1.4 wires the + // production launch site; this commit adds kernel + Rust + // launcher (`sp20_stats_compute.rs`) + GPU oracle tests + // (`tests/sp20_stats_compute_test.rs`) atomically per + // `feedback_no_partial_refactor`. + "sp20_stats_compute_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/mod.rs b/crates/ml/src/cuda_pipeline/mod.rs index c0b889295..16f8a107e 100644 --- a/crates/ml/src/cuda_pipeline/mod.rs +++ b/crates/ml/src/cuda_pipeline/mod.rs @@ -85,6 +85,13 @@ pub use sp4_wiener_ema::{ pearls_ad_update, WienerState, ALPHA_META, EPS_DIV, EPS_CLAMP_FLOOR, }; +// SP20 Phase 1.1 (2026-05-09): aux_conf p50 + std producer kernel +// launcher. Component 5 / Kernel 3 of the SP20 fused-producer chain +// (the other two — `sp20_emas_compute` and `sp20_controllers_compute` — +// land in Phase 1.2 and 1.3). Streams `aux_logits [B, 3]` once and +// emits `[p50, std]` into a `MappedF32Buffer<2>` that the EMA producer +// (Phase 1.2) blends into ISV slots [510..520). +pub mod sp20_stats_compute; // `launch_apply_pearls` is `pub(crate)` and consumed only inside this // crate via `use crate::cuda_pipeline::sp4_wiener_ema::launch_apply_pearls`. // `pearls_ad_update` stays `pub` for the test-oracle path in diff --git a/crates/ml/src/cuda_pipeline/sp20_stats_compute.rs b/crates/ml/src/cuda_pipeline/sp20_stats_compute.rs new file mode 100644 index 000000000..f194f387a --- /dev/null +++ b/crates/ml/src/cuda_pipeline/sp20_stats_compute.rs @@ -0,0 +1,176 @@ +#![allow(unsafe_code)] // CUDA kernel launch via cudarc requires unsafe. +//! SP20 Phase 1.1 (2026-05-09) — `sp20_stats_compute_kernel` launcher. +//! +//! Component 5 / Kernel 3 of the SP20 design. Produces two scalar f32 +//! outputs from a `[B, 3]` `aux_logits` tensor: +//! +//! out[0] = aux_conf_p50 — median over the batch of +//! `max_c softmax(logits[i, *])[c] - 1/3` +//! out[1] = aux_conf_std — population std over the batch of the same +//! `aux_conf` row signal +//! +//! `aux_conf` measures peak class confidence above the K=3 uniform +//! baseline (so 0 ⇔ uniform, 2/3 ⇔ fully concentrated). The downstream +//! [`sp20_emas_compute_kernel`] (Phase 1.2) Wiener-blends these stats +//! into ISV slots [510..520); this kernel itself does NOT write to ISV. +//! +//! ── Hard rules covered ───────────────────────────────────────────── +//! - `feedback_no_atomicadd` — block tree-reduce + per-warp tile +//! histogram (the `sp4_histogram_p99` pattern adapted for p50). +//! - `feedback_no_cpu_compute_strict` — every operation is GPU-side. +//! - `feedback_no_htod_htoh_only_mapped_pinned` — output is a +//! `MappedF32Buffer<2>`; kernel emits `__threadfence_system()` for +//! PCIe-visible coherence. +//! - `pearl_no_host_branches_in_captured_graph` — single-block kernel, +//! captureable. +//! - `pearl_fused_per_group_statistics_oracle` — one fused kernel +//! streaming `aux_logits` once for both stats. +//! +//! Phase 1.1 wiring scope: **kernel + launcher + tests + build entry** +//! only. The training-loop call site lands in Phase 1.4 atomically with +//! the rest of the SP20 reward chain, per +//! `feedback_no_partial_refactor`. + +/// Row width of `aux_logits` — fixed by the SP14-C aux head's 3-class +/// (down/flat/up) direction output. Kernel reads `[B, AUX_K_CLASSES]`. +pub const AUX_K_CLASSES: usize = 3; + +/// Block size used at every internal kernel pass (max-reduce, sum, +/// sum-of-squares, histogram-of-warp-tiles). MUST stay 256: matches the +/// `sp4_histogram_p99<256>` template instantiation we re-use, and keeps +/// the per-warp tile shmem budget at `8 warps × 256 ints × 4 B = 8 KB` +/// (well under L40S's 48 KB / block budget). +pub const SP20_STATS_BLOCK: u32 = 256; + +/// 256-bin histogram (mirrors `SP4_HIST_BINS`). Re-exported for +/// shared-memory size accounting in callers and tests. +pub const SP20_STATS_HIST_BINS: usize = 256; + +/// Compute the dynamic shared-memory bytes the kernel needs for a given +/// batch size `b`. Layout (kernel-side): +/// +/// ```text +/// [warps × HIST_BINS × sizeof(int)] per-warp histogram tiles +/// [B × sizeof(float)] per-row aux_conf scratch +/// ``` +/// +/// With `BLOCK=256` (8 warps) the per-warp tiles consume +/// `8 × 256 × 4 = 8192 bytes`. The per-row scratch grows linearly with +/// `B` (e.g., `B=128` → 512 B, total 8704 B). +#[must_use] +pub fn dynamic_shmem_bytes(b: usize) -> u32 { + let warps = (SP20_STATS_BLOCK as usize) >> 5; // 8 + let tiles = warps + .checked_mul(SP20_STATS_HIST_BINS) + .and_then(|x| x.checked_mul(std::mem::size_of::())) + .expect("SP20 stats: per-warp tile size overflow"); + let aux_conf = b + .checked_mul(std::mem::size_of::()) + .expect("SP20 stats: aux_conf scratch size overflow"); + let total = tiles + .checked_add(aux_conf) + .expect("SP20 stats: total shmem size overflow"); + u32::try_from(total).expect("SP20 stats: shmem bytes exceed u32::MAX") +} + +/// Launch `sp20_stats_compute_kernel` on the producer's stream. +/// +/// Single-block, [`SP20_STATS_BLOCK`]-thread launch. `kernel` MUST be +/// the loaded `sp20_stats_compute_kernel` `CudaFunction` (from the +/// cubin emitted by `build.rs` — see manifest entry +/// `sp20_stats_compute_kernel.cu`). +/// +/// `aux_logits_dev` is a device pointer to a `[B, AUX_K_CLASSES]` +/// row-major f32 tensor; `out_dev` is the device pointer of a +/// `MappedF32Buffer<2>` ([p50, std]). +/// +/// # Safety +/// +/// All three inputs MUST be valid: +/// - `aux_logits_dev` points to at least `B × AUX_K_CLASSES` f32s. +/// - `out_dev` points to at least 2 f32s reachable from device space +/// (mapped-pinned per `feedback_no_htod_htoh_only_mapped_pinned`). +/// - `b >= 0`. Caller passes `0` only on empty-batch edge cases (the +/// kernel's degenerate guard writes `[0, 0]` and returns). +/// +/// Stream ordering with the aux-head forward (which writes +/// `aux_logits`) is the caller's responsibility — same-stream is +/// sufficient (CUDA stream-ordering invariant). Phase 1.4 will land +/// the production launch site enforcing that ordering. +pub unsafe fn launch_sp20_stats_compute( + stream: &cudarc::driver::CudaStream, + kernel: &cudarc::driver::CudaFunction, + aux_logits_dev: u64, + out_dev: u64, + b: i32, +) -> Result<(), crate::MLError> { + use cudarc::driver::{LaunchConfig, PushKernelArg}; + + debug_assert!(b >= 0, "launch_sp20_stats_compute: b must be non-negative, got {b}"); + debug_assert!(aux_logits_dev != 0, + "launch_sp20_stats_compute: aux_logits_dev must be a valid device pointer"); + debug_assert!(out_dev != 0, + "launch_sp20_stats_compute: out_dev must be a valid device pointer"); + + let shmem_bytes = dynamic_shmem_bytes(b.max(0) as usize); + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (SP20_STATS_BLOCK, 1, 1), + shared_mem_bytes: shmem_bytes, + }; + stream + .launch_builder(kernel) + .arg(&aux_logits_dev) + .arg(&out_dev) + .arg(&b) + .launch(cfg) + .map_err(|e| { + crate::MLError::ModelError(format!("sp20_stats_compute_kernel launch: {e}")) + })?; + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn aux_k_classes_is_three() { + // Locks the contract against silent changes — the SP14-C aux + // head emits 3 logits per row, and the kernel hardcodes that + // width via `AUX_K_CLASSES`. If the head ever moves to a + // different K, this constant + the kernel's `SP20_K_CLASSES` + // macro both must move together. + assert_eq!(AUX_K_CLASSES, 3); + } + + #[test] + fn block_size_is_warp_multiple() { + assert!(SP20_STATS_BLOCK >= 32); + assert_eq!(SP20_STATS_BLOCK % 32, 0); + } + + #[test] + fn shmem_budget_fits_l40s_block_limit() { + // L40S provides 48 KiB shmem/block by default; the kernel's + // shmem footprint MUST stay below that for the typical batch + // sizes the trainer uses (B ∈ [16, 1024]). + let max_b = 1024; + let bytes = dynamic_shmem_bytes(max_b); + assert!( + bytes < 48 * 1024, + "shmem footprint {bytes} bytes for B={max_b} exceeds L40S 48 KiB block budget", + ); + } + + #[test] + fn shmem_budget_handles_zero_batch() { + // Degenerate-batch path: kernel still needs the per-warp tile + // shmem so the histogram-zero pass does not OOB; aux_conf + // scratch can be 0-sized. + let bytes = dynamic_shmem_bytes(0); + let warps = (SP20_STATS_BLOCK as usize) >> 5; + let expected = warps * SP20_STATS_HIST_BINS * std::mem::size_of::(); + assert_eq!(bytes as usize, expected); + } +} diff --git a/crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu b/crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu new file mode 100644 index 000000000..1d5e09f07 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu @@ -0,0 +1,304 @@ +/* ══════════════════════════════════════════════════════════════════════════ + * SP20 Phase 1.1 (2026-05-09) — sp20_stats_compute kernel. + * + * Component 5 / Kernel 3 of the SP20 design (the *stats* producer of the + * three-kernel SP20 fused-producer chain; the other two — `sp20_emas_compute` + * and `sp20_controllers_compute` — land in Phase 1.2 and 1.3 and consume + * this kernel's outputs as their inputs). + * + * Reads a `[B, 3]` row-major `aux_logits` tensor (the SP14-C aux head's + * direction logits — `n_classes = 3` for {down, flat, up}) and produces + * two scalar f32 outputs: + * + * out[0] = aux_conf_p50 — median over the batch of + * aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/3 + * out[1] = aux_conf_std — population std over the batch of aux_conf[i] + * + * `aux_conf` is the per-row peak-confidence above the uniform baseline + * (uniform softmax at K=3 gives 1/3 per class, so subtracting 1/3 makes + * the signal zero when the head is maximally uncertain and approaches + * 2/3 when the head is fully concentrated on one class). + * + * Producer cadence: hot-path (per-training-step). Subsequent kernels + * (sp20_emas_compute, sp20_controllers_compute) Wiener-blend these into + * ISV slots [510..520) — this kernel itself does NOT write to ISV. + * + * ── Pearls + invariants ──────────────────────────────────────────────── + * - `feedback_no_atomicadd` — block tree-reduce only; no atomicAdd. + * p50 uses `sp4_histogram_p99<256>` with target_quantile = 0.50 + * (per-warp tile binning + cumulative-from-top, no atomicAdd). + * std uses two block tree-reductions (sum, sum-of-squares) sharing + * one shmem tile sequentially. + * - `pearl_fused_per_group_statistics_oracle` — one fused kernel + * reading `aux_logits` once for both stats; no separate p50 / std + * producers redundantly re-streaming the per-row aux_conf. + * - `pearl_no_host_branches_in_captured_graph` — single-block kernel, + * no host branches; safe to capture in the per-step CUDA Graph. + * - `feedback_no_htod_htoh_only_mapped_pinned` — outputs land in a + * `MappedF32Buffer<2>` written via `__threadfence_system()` for + * PCIe-visible coherence; host reads after stream sync. + * - `pearl_symmetric_clamp_audit` — the `aux_conf` formula is itself + * bounded by softmax composition (max ∈ [1/3, 1] ⇒ aux_conf ∈ + * [0, 2/3]); no runtime clamp needed on the per-row signal. + * - `feedback_no_cpu_compute_strict` — every operation lives in the + * kernel; no CPU fallback for the softmax / max / median / std math. + * + * Algorithm (single block, 256 threads): + * + * Pass A: compute `aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/3` + * for i in [0, B), strided across threads, into a flat + * per-row scratch `aux_conf_buf [B]` in dynamic shmem. + * Numerically-stable softmax: `m = max(l[i,0], l[i,1], l[i,2]); + * e_c = exp(l[i,c] - m); s = e_0 + e_1 + e_2; + * sm_c = e_c / s; max_sm = max_c sm_c`. Then + * aux_conf[i] = max_sm - 1.0/3.0. + * __syncthreads() after the pass. + * + * Pass B: block tree-reduce sum(aux_conf) and sum(aux_conf^2) for + * the std computation. Two reductions share the second half + * of the shmem tile sequentially. Thread 0 captures both. + * + * Pass C: single-block call to `sp4_histogram_p99<256>(aux_conf_buf, B)` + * with the percentile target overridden to 0.50 via the + * `cumul ≥ ceil(B/2)` cumulative-from-bottom variant + * (the `sp4_histogram_p99` cumulative-from-top variant + * returns the bin upper edge for the top 1% — for p50 we + * re-run pass 3 of the same histogram with a different + * cumulative target). To keep the shared `sp4_histogram_p99` + * contract intact, we inline pass 1+2 (which fill `s_bins`) + * and substitute pass 3 with a from-bottom cumulative scan. + * See per-pass code below for the detailed adaptation. + * + * Pass D: thread 0 writes: + * out[0] = p50 (median of aux_conf) + * out[1] = sqrt(var) (population std of aux_conf) + * via `__threadfence_system()` for mapped-pinned coherence. + * + * Launch contract: + * grid_dim = (1, 1, 1) + * block_dim = (SP20_STATS_BLOCK, 1, 1) — 256 threads + * shared mem (dynamic): + * [SP4_HIST_BINS warp tiles] 8 warps × 256 ints × 4B = 8192 bytes + * + B × sizeof(float) per-row aux_conf_buf + * + block tile reuses s_bins[0..256] (already 1024 bytes static) + * With B=128, shmem ≈ 8192 + 512 = 8704 bytes (well under L40S 48 KB). + * + * Args: + * aux_logits — `[B, 3]` row-major aux head logits. f32 device ptr. + * Row stride is 3 floats; contiguous per row. + * out — `MappedF32Buffer<2>` device ptr for [p50, std]. + * B — batch size (rows of aux_logits). + * ══════════════════════════════════════════════════════════════════════════ */ + +#include +#include "sp4_histogram_p99.cuh" + +#define SP20_STATS_BLOCK 256 +#define SP20_K_CLASSES 3 +#define SP20_UNIFORM_K3 0.333333343f /* 1.0f / 3.0f rounded to f32 */ + +extern "C" __global__ void sp20_stats_compute_kernel( + const float* __restrict__ aux_logits, /* [B, 3] row-major */ + float* __restrict__ out, /* MappedF32Buffer<2> */ + int B) +{ + /* Single-block launch contract. */ + if (blockIdx.x != 0) return; + + const int tid = (int)threadIdx.x; + const int bdim = (int)blockDim.x; + + /* Degenerate guard: empty batch — write zeros and return. The + * caller's downstream EMA producer (Phase 1.2) treats zero stats + * as a no-op via Pearl-A bootstrap (sentinel = 0). */ + if (B <= 0) { + if (tid == 0) { + out[0] = 0.0f; + out[1] = 0.0f; + __threadfence_system(); + } + return; + } + + /* Static shared: 256 ints for s_bins (1024 B) + 1 float s_step_max + * (4 B). Used by sp4_histogram_p99-style passes. */ + __shared__ int s_bins[SP4_HIST_BINS]; + __shared__ float s_step_max; + __shared__ float s_sum; + __shared__ float s_sumsq; + + /* Dynamic shared layout (caller MUST size accordingly): + * [0 .. warps × SP4_HIST_BINS × sizeof(int)) per-warp tiles + * [warps × SP4_HIST_BINS × sizeof(int) ..) aux_conf_buf [B] + * + * Both regions are accessed as `int*` then `float*` views; the + * caller passes one contiguous dynamic-shmem allocation sized + * (warps × SP4_HIST_BINS × sizeof(int)) + (B × sizeof(float)). + */ + extern __shared__ int s_dyn_int[]; + int* s_warp_tiles = s_dyn_int; + const int warps = bdim >> 5; + float* s_aux_conf = reinterpret_cast( + s_warp_tiles + warps * SP4_HIST_BINS); + + /* ── Pass A: per-row softmax-max → aux_conf[i] ──────────────────── */ + for (int i = tid; i < B; i += bdim) { + const int row_off = i * SP20_K_CLASSES; + const float l0 = aux_logits[row_off + 0]; + const float l1 = aux_logits[row_off + 1]; + const float l2 = aux_logits[row_off + 2]; + + /* Numerically-stable softmax: subtract row-max before exp. */ + const float m = fmaxf(l0, fmaxf(l1, l2)); + const float e0 = __expf(l0 - m); + const float e1 = __expf(l1 - m); + const float e2 = __expf(l2 - m); + const float denom = e0 + e1 + e2; + const float sm0 = e0 / denom; + const float sm1 = e1 / denom; + const float sm2 = e2 / denom; + const float max_sm = fmaxf(sm0, fmaxf(sm1, sm2)); + + s_aux_conf[i] = max_sm - SP20_UNIFORM_K3; + } + __syncthreads(); + + /* ── Pass B1: block tree-reduce sum(aux_conf) ───────────────────── */ + float local_sum = 0.0f; + for (int i = tid; i < B; i += bdim) { + local_sum += s_aux_conf[i]; + } + /* Reuse s_bins[0..bdim) as a float-via-int reinterpret reduce tile. + * BLOCK_SIZE=256 ≤ SP4_HIST_BINS=256 so the alias fits exactly. */ + s_bins[tid] = __float_as_int(local_sum); + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) { + const float a = __int_as_float(s_bins[tid]); + const float b = __int_as_float(s_bins[tid + s]); + s_bins[tid] = __float_as_int(a + b); + } + __syncthreads(); + } + if (tid == 0) s_sum = __int_as_float(s_bins[0]); + __syncthreads(); + + /* ── Pass B2: block tree-reduce sum(aux_conf^2) ─────────────────── */ + float local_sumsq = 0.0f; + for (int i = tid; i < B; i += bdim) { + const float x = s_aux_conf[i]; + local_sumsq += x * x; + } + s_bins[tid] = __float_as_int(local_sumsq); + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) { + const float a = __int_as_float(s_bins[tid]); + const float b = __int_as_float(s_bins[tid + s]); + s_bins[tid] = __float_as_int(a + b); + } + __syncthreads(); + } + if (tid == 0) s_sumsq = __int_as_float(s_bins[0]); + __syncthreads(); + + /* ── Pass C: histogram over aux_conf for p50 (cumulative-from-bottom) ── + * Mirrors sp4_histogram_p99's pass 1 (max-reduce of |aux_conf|) and + * pass 2 (per-warp tile binning) verbatim, then substitutes a + * cumulative-from-bottom scan over s_bins[] to recover the median's + * bin upper-edge. We can't call sp4_histogram_p99<256>() directly + * because its pass 3 hardcodes the top-1% (p99) cumulative target; + * the rest of the algorithm is identical, so the inlining cost is + * one extra `for (b)` block (we save a kernel boundary). + */ + + /* Pass C.1: max-reduce of |aux_conf| via the same tile alias. + * aux_conf is non-negative by construction (max_sm ≥ 1/3 ⇒ + * aux_conf ≥ 0), so |x| = x; we still use fabsf to guarantee the + * sign discipline that sp4_histogram_p99 demands. */ + float local_max = 0.0f; + for (int i = tid; i < B; i += bdim) { + local_max = fmaxf(local_max, fabsf(s_aux_conf[i])); + } + s_bins[tid] = __float_as_int(local_max); + __syncthreads(); + for (int s = bdim / 2; s > 0; s >>= 1) { + if (tid < s) { + const float a = __int_as_float(s_bins[tid]); + const float b = __int_as_float(s_bins[tid + s]); + s_bins[tid] = __float_as_int(fmaxf(a, b)); + } + __syncthreads(); + } + if (tid == 0) s_step_max = __int_as_float(s_bins[0]); + __syncthreads(); + + /* Degenerate-distribution branch: every aux_conf is 0 (e.g., + * uniform-logit input). p50 = 0; std also = 0 by construction. */ + if (s_step_max == 0.0f) { + if (tid == 0) { + out[0] = 0.0f; + out[1] = 0.0f; + __threadfence_system(); + } + return; + } + + const float step_max = s_step_max; + const float bin_width = step_max / (float)SP4_HIST_BINS; + + /* Pass C.2: per-warp tile binning (no atomicAdd). Verbatim from + * sp4_histogram_p99 pass 2 — the caller-provided dynamic shmem + * already reserves `warps × SP4_HIST_BINS` ints at the front. */ + const int warp_id = tid >> 5; + const int lane = tid & 31; + + /* Zero this warp's tile (32 lanes × 8 ints = 256). */ + for (int b = lane; b < SP4_HIST_BINS; b += 32) { + s_warp_tiles[warp_id * SP4_HIST_BINS + b] = 0; + } + __syncwarp(); + + for (int i = tid; i < B; i += bdim) { + int bin_idx = (int)floorf(fabsf(s_aux_conf[i]) / bin_width); + if (bin_idx >= SP4_HIST_BINS) bin_idx = SP4_HIST_BINS - 1; + s_warp_tiles[warp_id * SP4_HIST_BINS + bin_idx] += 1; + } + __syncthreads(); + + /* Tree-reduce warp tiles into s_bins[]. */ + for (int b = tid; b < SP4_HIST_BINS; b += bdim) { + int sum = 0; + for (int w = 0; w < warps; ++w) { + sum += s_warp_tiles[w * SP4_HIST_BINS + b]; + } + s_bins[b] = sum; + } + __syncthreads(); + + /* Pass C.3: cumulative-from-BOTTOM scan → p50 = first bin where + * cumul ≥ ⌈B/2⌉. Mirrors sp4_histogram_p99 pass 3 mechanically + * but inverts the scan direction (bottom-up for the median, top- + * down for the 99th percentile). */ + if (tid == 0) { + const int target = (B + 1) / 2; /* ceil(B / 2) */ + int cumul = 0; + int p50_bin = 0; + for (int b = 0; b < SP4_HIST_BINS; ++b) { + cumul += s_bins[b]; + if (cumul >= target) { p50_bin = b; break; } + } + /* p50 = upper edge of p50_bin = (p50_bin + 1) × bin_width. */ + const float p50 = (float)(p50_bin + 1) * bin_width; + + /* Population std from the captured sum / sumsq. */ + const float n_f = (float)B; + const float mean = s_sum / n_f; + const float var = fmaxf(0.0f, s_sumsq / n_f - mean * mean); + const float std = sqrtf(var); + + out[0] = p50; + out[1] = std; + __threadfence_system(); + } +} diff --git a/crates/ml/tests/sp20_stats_compute_test.rs b/crates/ml/tests/sp20_stats_compute_test.rs new file mode 100644 index 000000000..42f9f06a8 --- /dev/null +++ b/crates/ml/tests/sp20_stats_compute_test.rs @@ -0,0 +1,305 @@ +//! SP20 Phase 1.1 (2026-05-09) — `sp20_stats_compute_kernel` GPU oracle tests. +//! +//! Verifies the fused stats producer emits +//! `[aux_conf_p50, aux_conf_std]` consistent with the closed-form CPU +//! oracle for two characteristic distributions of `aux_logits [B, 3]`: +//! +//! 1. **Uniform**: `aux_logits = 0` everywhere → softmax = [1/3,1/3,1/3] +//! → `aux_conf = max - 1/3 = 0` → p50 ≈ 0, std ≈ 0. +//! 2. **Concentrated**: `aux_logits[*, 0] = 10`, rest = 0 → softmax +//! ≈ [1, 0, 0] (one-hot) → `aux_conf ≈ 2/3` → p50 ≈ 0.667, std ≈ 0. +//! +//! Tests also cover a heterogeneous distribution with non-trivial p50 +//! and std so the kernel is exercised on both the histogram path and +//! the running-sum/sum-of-squares paths together. +//! +//! Per `feedback_no_atomicadd` the kernel uses block tree-reduce + per- +//! warp tile binning. Per `feedback_no_htod_htoh_only_mapped_pinned` +//! all CPU↔GPU buffers are `MappedF32Buffer`. +//! +//! Run on a GPU host: +//! +//! SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 \ +//! cargo test -p ml --test sp20_stats_compute_test --features cuda \ +//! -- --ignored --nocapture + +#![allow(clippy::tests_outside_test_module)] + +#[cfg(feature = "cuda")] +#[allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory. +mod gpu { + use std::sync::Arc; + + use cudarc::driver::{CudaContext, CudaFunction, CudaStream}; + use ml::cuda_pipeline::mapped_pinned::MappedF32Buffer; + use ml::cuda_pipeline::sp20_stats_compute::{ + launch_sp20_stats_compute, AUX_K_CLASSES, + }; + + const SP20_STATS_COMPUTE_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/sp20_stats_compute_kernel.cubin")); + + fn make_test_stream() -> Arc { + let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?"); + ctx.default_stream() + } + + fn load_sp20_stats_compute(stream: &Arc) -> CudaFunction { + let module = stream + .context() + .load_cubin(SP20_STATS_COMPUTE_CUBIN.to_vec()) + .expect("load sp20_stats_compute_kernel cubin"); + module + .load_function("sp20_stats_compute_kernel") + .expect("load sp20_stats_compute_kernel function") + } + + /// CPU oracle: per-row softmax → max → minus 1/3, then median + std. + /// Mirrors the kernel's math exactly so the tolerances below pin + /// fp32 rounding only, not algorithmic deltas. + fn cpu_oracle(aux_logits: &[f32], b: usize) -> (f32, f32) { + assert_eq!(aux_logits.len(), b * AUX_K_CLASSES); + let one_third = 1.0_f32 / 3.0_f32; + + let mut conf: Vec = (0..b) + .map(|i| { + let off = i * AUX_K_CLASSES; + let l0 = aux_logits[off]; + let l1 = aux_logits[off + 1]; + let l2 = aux_logits[off + 2]; + let m = l0.max(l1).max(l2); + let e0 = (l0 - m).exp(); + let e1 = (l1 - m).exp(); + let e2 = (l2 - m).exp(); + let s = e0 + e1 + e2; + let max_sm = (e0 / s).max(e1 / s).max(e2 / s); + max_sm - one_third + }) + .collect(); + + // Sort for the median: matches the histogram-bin upper-edge + // semantics of the kernel (the sorted-array median is what the + // histogram converges to as bin_width → 0; for our two + // characteristic tests below the values are concentrated at + // 0 or 2/3 which the histogram resolves exactly). + conf.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal)); + let median_idx = (b + 1) / 2 - 1; // 0-based ⌈B/2⌉ - 1 + let median = conf[median_idx]; + + let mean: f32 = conf.iter().sum::() / (b as f32); + let var: f32 = conf + .iter() + .map(|x| (x - mean).powi(2)) + .sum::() / (b as f32); + let std = var.max(0.0).sqrt(); + + (median, std) + } + + fn run_kernel(b: usize, aux_logits: &[f32]) -> (f32, f32) { + let stream = make_test_stream(); + let kernel = load_sp20_stats_compute(&stream); + + let aux_buf = unsafe { MappedF32Buffer::new(aux_logits.len()) } + .expect("alloc aux_logits buf"); + aux_buf.write_from_slice(aux_logits); + let out_buf = unsafe { MappedF32Buffer::new(2) } + .expect("alloc out buf"); + out_buf.write_from_slice(&[0.0, 0.0]); + + unsafe { + launch_sp20_stats_compute( + &stream, + &kernel, + aux_buf.dev_ptr, + out_buf.dev_ptr, + b as i32, + ) + .expect("launch sp20_stats_compute_kernel"); + } + stream.synchronize().expect("sync after sp20_stats_compute_kernel"); + + let out = out_buf.read_all(); + (out[0], out[1]) + } + + /// Test 1 — uniform-logit distribution. + /// + /// `aux_logits = 0` everywhere ⇒ softmax = [1/3, 1/3, 1/3] for every + /// row ⇒ `aux_conf = 1/3 - 1/3 = 0` for every row. + /// Expected: p50 ≈ 0, std ≈ 0. The kernel's degenerate-distribution + /// branch (step_max == 0) writes both outputs to 0 directly. + #[test] + #[ignore = "requires GPU"] + fn uniform_logits_emit_zero_stats() { + const B: usize = 128; + let aux_logits = vec![0.0_f32; B * AUX_K_CLASSES]; + + let (p50, std) = run_kernel(B, &aux_logits); + + // Both outputs MUST be exactly zero — the kernel's + // step_max == 0 branch writes 0 directly. + assert!( + p50.abs() < 1e-6, + "uniform: p50 expected 0, got {p50}", + ); + assert!( + std.abs() < 1e-6, + "uniform: std expected 0, got {std}", + ); + } + + /// Test 2 — concentrated-logit distribution (varied confidence). + /// + /// Per-row class-0 logit varies linearly from 0 to 3 across the + /// batch, with classes 1 and 2 held at 0. softmax(class 0) varies + /// from 1/3 (logit 0, uniform) to ≈ 0.85 (logit 3), so `aux_conf` + /// varies from 0 to ≈ 0.52 — a 0.5 range spread across well over + /// half of the 256 histogram bins (per + /// `pearl_sp4_histogram_warp_tile_undercount`: lockstep-uniform + /// inputs cause warp-tile collisions in the non-atomic per-warp + /// histogram tiles, so test data MUST vary per row across many + /// bins, not just within one bin's saturating-softmax tail). + /// + /// Logits are kept under the saturating-softmax regime (max < 4) + /// because the softmax tail compresses 0.85 → 1.0 confidences + /// into a range smaller than one bin_width (≈ 0.4% of step_max), + /// re-creating the warp-tile collision pattern this test + /// explicitly avoids. + /// + /// Expected: p50 is the median of a monotonically-increasing + /// aux_conf vector — i.e., the value at the lower midpoint + /// (`(B+1)/2 - 1` 0-based). The CPU oracle computes it exactly; + /// the GPU kernel matches to within a couple of bin_widths. + #[test] + #[ignore = "requires GPU"] + fn varied_confidence_logits_match_cpu_oracle() { + const B: usize = 128; + let mut aux_logits = vec![0.0_f32; B * AUX_K_CLASSES]; + for i in 0..B { + // Linear ramp 0 → 3 on class 0 — confidence ramps from + // uniform (0) to a non-saturating peak (≈ 0.85). + let frac = (i as f32) / ((B - 1) as f32); + aux_logits[i * AUX_K_CLASSES] = 3.0 * frac; + } + + let (p50_gpu, std_gpu) = run_kernel(B, &aux_logits); + let (p50_cpu, std_cpu) = cpu_oracle(&aux_logits, B); + + // The GPU histogram bin_width ≈ step_max / 256 ≈ 0.52 / 256 ≈ + // 0.002; agreement to within a few bin_widths is the design + // tolerance for the histogram path. + let p50_tol = 0.02_f32; + assert!( + (p50_gpu - p50_cpu).abs() < p50_tol, + "concentrated: p50_gpu={p50_gpu} vs cpu_oracle={p50_cpu} (diff={:.3e}, tol={p50_tol})", + (p50_gpu - p50_cpu).abs(), + ); + // CPU oracle sanity: with logits ramping 0 → 3, the median + // aux_conf is the row at logit ≈ 1.5 → softmax ≈ 0.578 → + // aux_conf ≈ 0.245. We pin it to the (0.15, 0.40) range to + // catch contract drift while tolerating fp32 / softmax math. + assert!( + p50_cpu > 0.15_f32 && p50_cpu < 0.40_f32, + "CPU oracle: concentrated p50 should land in (0.15, 0.40), got {p50_cpu}", + ); + // std non-trivial; CPU and GPU within fp32 noise. + let std_tol = 5e-3_f32; + assert!( + (std_gpu - std_cpu).abs() < std_tol, + "concentrated: std_gpu={std_gpu} vs cpu_oracle={std_cpu} (diff={:.3e}, tol={std_tol})", + (std_gpu - std_cpu).abs(), + ); + // CPU oracle std: the aux_conf vector spans ≈ 0.52 with a + // monotonic ramp, so std ≈ 0.52 / sqrt(12) ≈ 0.15 for a + // uniform-distribution-style spread. + assert!( + std_cpu > 0.05_f32 && std_cpu < 0.25_f32, + "CPU oracle: concentrated std should land in (0.05, 0.25), got {std_cpu}", + ); + } + + /// Test 3 — heterogeneous distribution exercising both p50 and std. + /// + /// Half the rows get jittered hot logits (aux_conf ≈ 0.667) and + /// half get jittered uniform logits (aux_conf ≈ 0). The per-row + /// `aux_conf` is therefore a 50/50 mix of values around 0 and + /// values around 0.667. The median lands at the boundary (since + /// `(B+1)/2 - 1 = 64` indexes the upper end of the lower cluster), + /// and the std is `≈ 1/3` (half-distance between the two + /// concentration clusters). + /// + /// Per-row jitter is required to avoid the + /// `pearl_sp4_histogram_warp_tile_undercount` lockstep-uniform + /// trap (32 lanes in a warp racing on the same bin without + /// atomicAdd undercounts the histogram). Production aux logits + /// are naturally jittered by the per-sample net forward, so + /// jittered test data IS the realistic-input model. + #[test] + #[ignore = "requires GPU"] + fn heterogeneous_logits_match_cpu_oracle() { + const B: usize = 128; + let mut aux_logits = vec![0.0_f32; B * AUX_K_CLASSES]; + for i in 0..(B / 2) { + // Hot half: jitter ±0.5 around 10.0 on class 0. + let jitter = 0.5 * ((i as f32) / ((B / 2) as f32) - 0.5); + aux_logits[i * AUX_K_CLASSES] = 10.0 + jitter; + } + // Uniform half: jitter all three logits ±0.05 around 0 so + // each row's max varies slightly across bins. + for i in (B / 2)..B { + let k = (i - B / 2) as f32; + let scale = 0.05; + aux_logits[i * AUX_K_CLASSES + 0] = scale * (k * 0.10).sin(); + aux_logits[i * AUX_K_CLASSES + 1] = scale * (k * 0.13).cos(); + aux_logits[i * AUX_K_CLASSES + 2] = scale * (k * 0.07).sin(); + } + + let (p50_gpu, std_gpu) = run_kernel(B, &aux_logits); + let (p50_cpu, std_cpu) = cpu_oracle(&aux_logits, B); + + // p50 sanity: the sorted aux_conf vector has 64 small-near-zero + // values then 64 values ≈ 0.667; ⌈B/2⌉-th element (1-based) + // = 64th = the last "near-zero", so the median is small in + // both the CPU oracle and the GPU histogram. With uniform-half + // jitter ≈ 0.05, the per-row aux_conf there is at most + // (max softmax 0.34) - 1/3 ≈ 0.01, so p50 < 0.02. Tolerate + // up to a couple of bin_widths above that. + let p50_tol = 0.02_f32; + assert!( + (p50_gpu - p50_cpu).abs() < p50_tol, + "het: p50_gpu={p50_gpu} vs cpu_oracle={p50_cpu} (diff={:.3e}, tol={p50_tol})", + (p50_gpu - p50_cpu).abs(), + ); + // std sanity: with the hot cluster at ≈ 0.667 and the uniform + // cluster near 0, var = 0.5 × (0.333)² + 0.5 × (0.333)² ≈ 1/9 + // ⇒ std ≈ 1/3. Jitter perturbs this slightly; allow ±5% of + // 1/3 between the GPU result and the CPU oracle. + let std_tol = 5e-3_f32; + assert!( + (std_gpu - std_cpu).abs() < std_tol, + "het: std_gpu={std_gpu} vs cpu_oracle={std_cpu} (diff={:.3e}, tol={std_tol})", + (std_gpu - std_cpu).abs(), + ); + assert!( + (std_cpu - (1.0_f32 / 3.0_f32)).abs() < 0.02, + "CPU oracle: het std should be ≈ 1/3, got {std_cpu}", + ); + } + + /// Test 4 — empty-batch degenerate guard. + /// + /// `B = 0` ⇒ the kernel's degenerate branch writes `[0, 0]` and + /// returns without launching the full pipeline. This is the + /// trainer's pre-warmup path (no aux logits yet). + #[test] + #[ignore = "requires GPU"] + fn empty_batch_writes_zero_stats() { + // We still need a non-zero allocation for the device pointer to + // be valid; the kernel will not read it because B == 0. + let aux_logits = vec![0.0_f32; AUX_K_CLASSES]; + let (p50, std) = run_kernel(0, &aux_logits); + assert!(p50.abs() < 1e-6, "empty: p50 expected 0, got {p50}"); + assert!(std.abs() < 1e-6, "empty: std expected 0, got {std}"); + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index d9521044f..cd4cf15f3 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -10751,3 +10751,159 @@ bash scripts/audit_sp18_consumers.sh --check # exit round-trip equality. Direction sometimes flips at low Q values, unrelated to ISV slot layout. Verified flaky pre-Commit B by stashing the SP19 changes and re-running 3× — fails 2/3 runs. + +## 2026-05-09 — SP20 Phase 1.1: sp20_stats_compute kernel (additive) + +### Component + +- New CUDA kernel `crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu` +- New Rust launcher `crates/ml/src/cuda_pipeline/sp20_stats_compute.rs` + (re-exported from `cuda_pipeline::sp20_stats_compute`) +- New GPU oracle test `crates/ml/tests/sp20_stats_compute_test.rs` +- Cubin manifest entry in `crates/ml/build.rs` + +### Purpose + +Component 5 / Kernel 3 of the SP20 fused-producer chain (Phase 1.4 +will land Kernels 1+2 — `sp20_emas_compute` + `sp20_controllers_ +compute` — and the production launch site atomically with the +training-loop wire-up per `feedback_no_partial_refactor`). + +Reads `aux_logits [B, 3]` (the SP14-C aux head's 3-class direction +logits) and emits `[aux_conf_p50, aux_conf_std]` into a +`MappedF32Buffer<2>`, where the per-row signal is + +``` +aux_conf[i] = max_c softmax(logits[i, *])[c] - 1/3 +``` + +`aux_conf` measures peak class confidence above the K=3 uniform +baseline (0 ⇔ uniform, 2/3 ⇔ fully concentrated). The downstream +EMA producer (Phase 1.2) Wiener-blends p50 + std into ISV slots +within [510..520) reserved by SP20 (`f5eed1fa7`). + +### Algorithm + +Single-block, 256-thread kernel; one fused stream of `aux_logits` +for both stats per `pearl_fused_per_group_statistics_oracle`. + + - **Pass A**: per-row numerically-stable softmax → max → minus 1/3, + write `aux_conf[i]` into per-row scratch in dynamic shmem. + - **Pass B1**: block tree-reduce `sum(aux_conf)`. Reuses + `s_bins[0..256]` (the histogram tile) as a float-via-int reduce + tile, BLOCK=256 ≤ 256 bins so the alias fits exactly. + - **Pass B2**: block tree-reduce `sum(aux_conf²)` using the same + shmem tile sequentially. + - **Pass C.1**: max-reduce `|aux_conf|` to seed `step_max` for the + histogram (mirrors `sp4_histogram_p99` pass 1 verbatim). + - **Pass C.2**: per-warp tile binning into 256 linear bins (no + atomicAdd per `feedback_no_atomicadd`; mirrors + `sp4_histogram_p99` pass 2). + - **Pass C.3**: cumulative-from-BOTTOM scan (`cumul ≥ ⌈B/2⌉`) → + p50 = bin upper-edge. Inverts the direction of + `sp4_histogram_p99` pass 3 (which scans top-down for the 99th + percentile); the scan logic itself is otherwise identical. + +Thread 0 writes `out[0] = p50`, `out[1] = sqrt(var)` with +`__threadfence_system()` for mapped-pinned coherence. + +### Wiring contract + +Phase 1.1 wires kernel + launcher + tests + build entry **only**. +There is NO production caller in this commit — the kernel is dead +code awaiting Phase 1.4's atomic wire-up of: + +1. The training-loop call site (post-aux-head-forward, pre-EMA + producer) supplying the `aux_logits` device pointer. +2. A trainer-owned `MappedF32Buffer<2>` for the `[p50, std]` output. +3. Stream ordering with the aux-head forward (same-stream is + sufficient; CUDA stream-ordering invariant guarantees the + `aux_logits` write is visible before this kernel reads it). + +The `cuda_pipeline::sp20_stats_compute` module exports +`launch_sp20_stats_compute` + the contract constants (`AUX_K_CLASSES`, +`SP20_STATS_BLOCK`, `SP20_STATS_HIST_BINS`, `dynamic_shmem_bytes`) +the wire-up will consume. + +### Pearls + invariants honoured + +- `feedback_no_atomicadd` — block tree-reduce + per-warp tile + binning; no atomicAdd anywhere. +- `feedback_no_cpu_compute_strict` — every operation GPU-side. +- `feedback_no_htod_htoh_only_mapped_pinned` — output is a + `MappedF32Buffer<2>`; kernel emits `__threadfence_system()`. +- `feedback_no_partial_refactor` — kernel + launcher + tests + + build entry land in one commit; production wire-up lands + atomically in Phase 1.4 with the EMA + controller producers. +- `pearl_fused_per_group_statistics_oracle` — one kernel + computing both p50 and std from a single stream of `aux_logits`. +- `pearl_no_host_branches_in_captured_graph` — single-block + kernel, captureable. +- `pearl_first_observation_bootstrap` — degenerate-distribution + branch (step_max == 0) writes `[0, 0]` (sentinel) so the + downstream EMA producer can apply Pearl-A bootstrap on the + next non-degenerate observation. +- `pearl_sp4_histogram_warp_tile_undercount` (test data) — test + vectors use varied per-row confidences (logit ramps, jittered + half-hot/half-uniform) to avoid the lockstep-uniform trap that + causes warp-tile race undercounts in the non-atomic per-warp + histogram. + +### Tests + +`crates/ml/tests/sp20_stats_compute_test.rs` — GPU oracle tests +(all `#[ignore = "requires GPU"]`): + + 1. `uniform_logits_emit_zero_stats` — `aux_logits = 0` ⇒ + softmax = [1/3, 1/3, 1/3] ⇒ aux_conf = 0 ⇒ p50 = std = 0. + Exercises the kernel's degenerate-distribution branch. + 2. `varied_confidence_logits_match_cpu_oracle` — class-0 logit + ramps 0 → 3, others 0; CPU oracle matches GPU within + bin_width tolerance for p50 (≈ 0.002) and 5e-3 for std. + 3. `heterogeneous_logits_match_cpu_oracle` — half jittered + hot, half jittered uniform; CPU oracle matches GPU within + 0.02 for p50 and 5e-3 for std. + 4. `empty_batch_writes_zero_stats` — `B = 0` exercises the + pre-warmup degenerate guard; kernel writes `[0, 0]` and + returns without entering the histogram path. + +Plus 4 unit tests in the launcher module (`AUX_K_CLASSES = 3`, +block-size discipline, shmem budget under L40S 48 KiB, zero-batch +shmem accounting). + +### Files modified + +| File | Status | Purpose | +|------|--------|---------| +| `crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu` | NEW | Single-block stats producer kernel | +| `crates/ml/src/cuda_pipeline/sp20_stats_compute.rs` | NEW | Rust launcher + contract constants | +| `crates/ml/tests/sp20_stats_compute_test.rs` | NEW | 4 GPU oracle tests | +| `crates/ml/src/cuda_pipeline/mod.rs` | +pub mod | Module declaration | +| `crates/ml/build.rs` | +cubin entry | Compile entry for nvcc | +| `docs/dqn-wire-up-audit.md` | This entry | Audit log | + +### Verification + +``` +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo check -p ml --features cuda +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml \ + --test sp20_stats_compute_test --features cuda -- --ignored --nocapture +SQLX_OFFLINE=true CUDA_COMPUTE_CAP=86 cargo test -p ml \ + --features cuda --lib sp20_stats_compute +``` + +All 4 GPU oracle tests + 4 unit tests pass on RTX 3050 Ti (sm_86). +L40S verification deferred until Phase 1.4 lands the production +caller. + +### Phase 1.1 → Phase 1.4 forward references + +- Phase 1.2: `sp20_emas_compute_kernel.cu` consumes + `[aux_conf_p50, aux_conf_std]` (this kernel's output) + + per-trade `R_event` / WR-window observations to advance the + Wiener EMAs in ISV slots [510..520). +- Phase 1.3: `sp20_controllers_compute_kernel.cu` derives + `LOSS_CAP / TARGET_HOLD_PCT / N_STEP / AUX_CONF_THRESHOLD / + AUX_GATE_TEMP / HOLD_COST_SCALE` from the EMAs. +- Phase 1.4: training-loop wire-up calls all 3 producers in + sequence on the same stream as the aux-head forward.