diff --git a/crates/ml/src/cuda_pipeline/alpha_kernels.rs b/crates/ml/src/cuda_pipeline/alpha_kernels.rs index b4c99ffd4..56a7b839f 100644 --- a/crates/ml/src/cuda_pipeline/alpha_kernels.rs +++ b/crates/ml/src/cuda_pipeline/alpha_kernels.rs @@ -171,6 +171,10 @@ pub(crate) unsafe fn launch_alpha_munchausen_target( static ALPHA_MUNCHAUSEN_TARGET_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/alpha_munchausen_target.cubin")); +/// Precompiled kill-criteria producer cubin, embedded at compile time by build.rs. +static ALPHA_KILL_CRITERIA_CUBIN: &[u8] = + include_bytes!(concat!(env!("OUT_DIR"), "/alpha_kill_criteria.cubin")); + #[cfg(test)] mod tests { use super::*; @@ -282,4 +286,124 @@ mod tests { ); } } + + /// End-to-end GPU smoke for the kill-criteria producer launcher. + /// + /// Loads the cubin, allocates synthetic inputs covering ALL four output + /// observations (with hand-computed expected values), launches the + /// kernel via `launch_alpha_kill_criteria`, and validates the raw + /// scratch outputs against the hand math. Does NOT chain + /// `apply_pearls_ad_kernel` — the smoothing path is the canonical SP4 + /// applicator and is already covered elsewhere; this test isolates the + /// kill-criteria arithmetic. + /// + /// Hand-math (batch=2, n_actions=3): + /// - q_values = [[1, 2, 3], [5, 5, 5]] + /// row 0: mean=2, var=2/3, std≈0.8165, ratio=0.8165/2=0.4082 + /// row 1: mean=5, var=0, std=0, ratio=0/5=0 + /// batch_mean(ratio) = (0.4082 + 0)/2 = 0.2041 + /// - action_counts = [10, 30, 60], total=100 → p=[0.1, 0.3, 0.6] + /// entropy = -[0.1·ln(0.1) + 0.3·ln(0.3) + 0.6·ln(0.6)] + /// = -[(0.1)(-2.3026) + (0.3)(-1.2040) + (0.6)(-0.5108)] + /// = -[-0.2303 - 0.3612 - 0.3065] = 0.8980 + /// - rollout_R=100, isv[547]=-5185, isv[548]=4953 + /// return_vs_random = (100 - (-5185)) / 4953 = 5285/4953 = 1.0670 + /// - q_init=50, q_early=55 + /// early_movement = |55 - 50| / |50| = 0.1 + /// + /// Expected scratch_out = [0.2041, 0.8980, 1.0670, 0.1]. + #[test] + fn kill_criteria_smoke_matches_hand_computation() { + use crate::cuda_pipeline::alpha_isv_slots::{ + RANDOM_BASELINE_MEAN_INDEX, RANDOM_BASELINE_STD_INDEX, + }; + + let ctx = match CudaContext::new(0) { + Ok(c) => c, + Err(e) => { + eprintln!("Skipping GPU smoke — no CUDA device 0: {e}"); + return; + } + }; + let stream = ctx.default_stream(); + + let module = ctx + .load_cubin(ALPHA_KILL_CRITERIA_CUBIN.to_vec()) + .expect("kill_criteria cubin module load"); + let kernel = module + .load_function("alpha_kill_criteria_compute_kernel") + .expect("kill_criteria kernel load"); + + let batch: i32 = 2; + let n_actions: i32 = 3; + + // q_values: row 0 spreads 1..3; row 1 is constant 5. + let q_values: Vec = vec![1.0, 2.0, 3.0, 5.0, 5.0, 5.0]; + // Non-uniform action counts → non-maximal entropy. + let action_counts: Vec = vec![10, 30, 60]; + // scalar_inputs: [rollout_reward_mean, q_init_norm, q_early_step_norm] + let scalar_inputs: Vec = vec![100.0, 50.0, 55.0]; + + // Construct an ISV buffer with the random-baseline anchors at slots + // 547 and 548. Using the Task 7c committed values so the test + // exercises the production slot-indexing path with real numbers. + let mut isv_host: Vec = vec![0.0; 552]; + isv_host[RANDOM_BASELINE_MEAN_INDEX] = -5185.0; + isv_host[RANDOM_BASELINE_STD_INDEX] = 4953.0; + + let q_values_dev = stream.clone_htod(&q_values).expect("upload q_values"); + let action_counts_dev = stream + .clone_htod(&action_counts) + .expect("upload action_counts"); + let scalar_inputs_dev = stream + .clone_htod(&scalar_inputs) + .expect("upload scalar_inputs"); + let isv_dev = stream.clone_htod(&isv_host).expect("upload isv"); + let mut scratch_out_dev = stream.alloc_zeros::(4).expect("alloc scratch_out"); + + // Scoped block so device-ptr guards drop before the download. + { + let (q_values_ptr, _g0) = q_values_dev.device_ptr(&stream); + let (action_counts_ptr, _g1) = action_counts_dev.device_ptr(&stream); + let (scalar_inputs_ptr, _g2) = scalar_inputs_dev.device_ptr(&stream); + let (isv_ptr, _g3) = isv_dev.device_ptr(&stream); + let (scratch_out_ptr, _g4) = scratch_out_dev.device_ptr_mut(&stream); + + unsafe { + launch_alpha_kill_criteria( + &stream, + &kernel, + q_values_ptr, + action_counts_ptr, + scalar_inputs_ptr, + isv_ptr, + batch, + n_actions, + scratch_out_ptr, + ) + .expect("kill_criteria launch"); + } + } + stream.synchronize().expect("stream sync"); + + let scratch_out = stream + .clone_dtoh(&scratch_out_dev) + .expect("download scratch_out"); + assert_eq!(scratch_out.len(), 4); + + let expected = [0.2041_f32, 0.8980_f32, 1.0670_f32, 0.1_f32]; + let labels = ["q_spread", "action_entropy", "return_vs_random", "early_movement"]; + for (i, ((got, exp), label)) in scratch_out + .iter() + .zip(expected.iter()) + .zip(labels.iter()) + .enumerate() + { + assert!( + (got - exp).abs() < 0.01, + "scratch_out[{i}] = {label}: expected {exp:.4}, got {got:.4} \ + (full output: {scratch_out:?})" + ); + } + } } diff --git a/docs/isv-slots.md b/docs/isv-slots.md index e28838922..d1b697ab0 100644 --- a/docs/isv-slots.md +++ b/docs/isv-slots.md @@ -605,3 +605,16 @@ Both launchers follow `launch_apply_pearls` (in `sp4_wiener_ema.rs`) — pre-loa - `crates/ml/src/cuda_pipeline/mod.rs` — registers `pub(crate) mod alpha_kernels;` **Tests:** `cargo test -p ml --lib alpha_kernels` — 1 compile-witness passes. Real GPU integration test lands in Phase E.1 Task 12 (kernel smoke binary). + +## Phase E.1 — kill-criteria GPU smoke (2026-05-15) + +`crates/ml/src/cuda_pipeline/alpha_kernels.rs::tests::kill_criteria_smoke_matches_hand_computation` exercises the `alpha_kill_criteria_compute_kernel` end-to-end on the local GPU with synthetic inputs covering all 4 outputs. Hand-math: + +- `q_values = [[1, 2, 3], [5, 5, 5]]` → row 0 spread / |mean| = 0.4082; row 1 = 0 → batch mean = **0.2041** +- `action_counts = [10, 30, 60]` → H(p=[0.1, 0.3, 0.6]) = **0.8980** nats +- `rollout_R = 100`, `ISV[547] = -5185`, `ISV[548] = 4953` → `(100 + 5185) / 4953` = **1.0670** σ above random +- `q_init = 50`, `q_early = 55` → `5 / 50` = **0.1000** + +ISV buffer is sized to 552 floats with the production slot indices populated using the committed Task 7c baseline values; this exercises the slot-indexing path through `isv[random_baseline_mean_slot]` / `isv[random_baseline_std_slot]` reads. + +Companion to the `munchausen_target_smoke_matches_hand_computation` test (91d1a52b9). Both pass on RTX 3050 Ti.