feat(sp20): Phase 1.1 sp20_stats_compute kernel
Component 5 / Kernel 3 of the SP20 fused-producer chain. Single-block BLOCK=256 kernel 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`. p50 uses the inlined `sp4_histogram_p99` pattern (per-warp tile binning + cumulative-from-bottom, no atomicAdd per `feedback_no_atomicadd`); std uses two block tree-reductions sharing one shmem tile sequentially. One fused kernel streams `aux_logits` once for both stats per `pearl_fused_per_group_statistics_oracle`. Phase 1.4 wires the production launch site atomically with the rest of the SP20 reward chain per `feedback_no_partial_refactor`. This commit lands kernel + Rust launcher + GPU oracle tests + build entry + audit-doc entry together so the kernel is independently verifiable on RTX 3050 Ti (sm_86) and L40S (sm_89) before the EMA + controller producers (Phase 1.2 + 1.3) reference its outputs. Tests verify: - uniform logits → aux_conf = 0 → [p50, std] = [0, 0] - varied confidence (logit ramp 0 → 3) → matches CPU oracle - heterogeneous half-hot half-uniform → matches CPU oracle - empty batch → degenerate-guard writes [0, 0] All 4 GPU oracle tests + 4 launcher unit tests pass on RTX 3050 Ti. Test data uses per-row variance to avoid the `pearl_sp4_histogram_warp_tile_undercount` lockstep-uniform trap (concentrated values within one bin_width race the per-warp non-atomic increments). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
176
crates/ml/src/cuda_pipeline/sp20_stats_compute.rs
Normal file
176
crates/ml/src/cuda_pipeline/sp20_stats_compute.rs
Normal file
@@ -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::<i32>()))
|
||||
.expect("SP20 stats: per-warp tile size overflow");
|
||||
let aux_conf = b
|
||||
.checked_mul(std::mem::size_of::<f32>())
|
||||
.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::<i32>();
|
||||
assert_eq!(bytes as usize, expected);
|
||||
}
|
||||
}
|
||||
304
crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu
Normal file
304
crates/ml/src/cuda_pipeline/sp20_stats_compute_kernel.cu
Normal file
@@ -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 <cuda_runtime.h>
|
||||
#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<float*>(
|
||||
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();
|
||||
}
|
||||
}
|
||||
305
crates/ml/tests/sp20_stats_compute_test.rs
Normal file
305
crates/ml/tests/sp20_stats_compute_test.rs
Normal file
@@ -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<CudaStream> {
|
||||
let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
|
||||
ctx.default_stream()
|
||||
}
|
||||
|
||||
fn load_sp20_stats_compute(stream: &Arc<CudaStream>) -> 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<f32> = (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::<f32>() / (b as f32);
|
||||
let var: f32 = conf
|
||||
.iter()
|
||||
.map(|x| (x - mean).powi(2))
|
||||
.sum::<f32>() / (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}");
|
||||
}
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
Reference in New Issue
Block a user