test(alpha): kill-criteria GPU smoke matching hand-computed observations
Adds the second of the two Phase E.1 kernel smoke tests in
alpha_kernels.rs (companion to the munchausen_target smoke from
91d1a52b9). End-to-end exercises the alpha_kill_criteria_compute_kernel
launcher with synthetic inputs covering all 4 outputs:
q_values = [[1, 2, 3], [5, 5, 5]] → q_spread ≈ 0.2041
action_counts = [10, 30, 60] → action_entropy ≈ 0.8980
rollout_R=100, isv[547]=-5185,
isv[548]=4953 → return_vs_random ≈ 1.0670
q_init=50, q_early=55 → early_movement = 0.1000
ISV buffer is sized to 552 floats with slots 547/548 populated using the
committed Task 7c baseline values — this exercises the production
slot-indexing path through the kernel's
`isv[random_baseline_mean_slot]` / `isv[random_baseline_std_slot]` reads,
not just isolated kernel arithmetic.
Does NOT chain apply_pearls_ad_kernel afterward — the smoothing path is
canonical SP4 applicator territory already covered elsewhere. This test
isolates the kill-criteria producer arithmetic.
Tolerance 0.01 on all four observations; passes on RTX 3050 Ti in <2s
including kernel JIT.
`cargo test -p ml --lib alpha_kernels`: 3 pass (compile witness + both
GPU smokes). Audit doc docs/isv-slots.md updated per Invariant 7.
This commit is contained in:
@@ -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<f32> = vec![1.0, 2.0, 3.0, 5.0, 5.0, 5.0];
|
||||
// Non-uniform action counts → non-maximal entropy.
|
||||
let action_counts: Vec<i32> = vec![10, 30, 60];
|
||||
// scalar_inputs: [rollout_reward_mean, q_init_norm, q_early_step_norm]
|
||||
let scalar_inputs: Vec<f32> = 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<f32> = 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::<f32>(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:?})"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user