From 0f664cf361b0831ae86010eb9d1c3cf683629fdf Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 21 May 2026 00:34:06 +0200 Subject: [PATCH] test(crt-train): invariant tests for smoothness_lambda_controller --- ...smoothness_lambda_controller_invariants.rs | 254 ++++++++++++++++++ 1 file changed, 254 insertions(+) create mode 100644 crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs diff --git a/crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs b/crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs new file mode 100644 index 000000000..d73072c05 --- /dev/null +++ b/crates/ml-alpha/tests/smoothness_lambda_controller_invariants.rs @@ -0,0 +1,254 @@ +//! Invariant tests for the ISV-driven λ controller. +//! +//! Per `feedback_no_cpu_test_fallbacks`: GPU oracle only. Per +//! `pearl_tests_must_prove_not_lock_observations`: assertions check +//! invariants (sentinel bootstrap, steady-state at target, excess +//! proportionality, floor activation under zero amplitude) — not +//! arbitrary numeric snapshots. +//! +//! Per `feedback_no_htod_htoh_only_mapped_pinned`: all CPU↔GPU +//! transfers go via `MappedF32Buffer` + `memcpy_dtod_async`. The i32 +//! sentinel is shuttled by reinterpreting the same mapped-pinned +//! buffer as i32 (canonical pattern, see `loss.rs::download_i32`). + +use std::sync::Arc; + +use anyhow::{Context, Result}; +use approx::assert_relative_eq; +use cudarc::driver::{CudaSlice, CudaStream, DevicePtr, DevicePtrMut, LaunchConfig, PushKernelArg}; +use ml_alpha::heads::N_HORIZONS; +use ml_alpha::pinned_mem::MappedF32Buffer; +use ml_core::device::MlDevice; + +const CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/smoothness_lambda_controller.cubin")); +const LAMBDA_FLOOR: f32 = 1.0e-4; + +struct ControllerOut { + jitter_ema: Vec, + first_obs: i32, + lambda: Vec, +} + +fn run_controller( + dev: &MlDevice, + raw_per_h: &[f32; N_HORIZONS], + jitter_ema_init: &[f32; N_HORIZONS], + first_obs_init: i32, + base_lambda: f32, +) -> Result { + let stream: &Arc = dev.cuda_stream().context("controller stream")?; + let ctx = dev.cuda_context().context("controller ctx")?; + let module = ctx.load_cubin(CUBIN.to_vec()).context("load controller cubin")?; + let func = module.load_function("smoothness_lambda_controller").context("load fn")?; + + let raw_d = upload_f32(stream, raw_per_h.as_slice())?; + let mut jitter_d = upload_f32(stream, jitter_ema_init.as_slice())?; + let mut first_obs_d = upload_i32(stream, &[first_obs_init])?; + let mut lambda_d = stream.alloc_zeros::(N_HORIZONS).context("lambda alloc")?; + + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (N_HORIZONS as u32, 1, 1), + shared_mem_bytes: 0, + }; + let mut launch = stream.launch_builder(&func); + launch + .arg(&raw_d) + .arg(&mut jitter_d) + .arg(&mut first_obs_d) + .arg(&base_lambda) + .arg(&mut lambda_d); + unsafe { + launch.launch(cfg).context("controller launch")?; + } + stream.synchronize().context("controller sync")?; + + Ok(ControllerOut { + jitter_ema: download_f32(stream, &jitter_d)?, + first_obs: download_i32(stream, &first_obs_d)?[0], + lambda: download_f32(stream, &lambda_d)?, + }) +} + +fn upload_f32(stream: &Arc, host: &[f32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("upload staging: {e}"))?; + staging.write_from_slice(host); + let mut dst = stream.alloc_zeros::(n).context("upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + staging.dev_ptr, + nbytes, + stream.cu_stream(), + ) + .context("upload DtoD")?; + } + } + Ok(dst) +} + +/// Upload an `i32` slice via the canonical mapped-pinned path by +/// reinterpreting `MappedF32Buffer`'s host pointer as `*mut i32` +/// (same pattern as `trainer/loss.rs::download_i32`). One stream sync +/// barrier inside the kernel launch makes the host write visible +/// to the device through the same mapped page. +fn upload_i32(stream: &Arc, host: &[i32]) -> Result> { + let n = host.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("i32 upload staging: {e}"))?; + // Write through the host pointer reinterpreted as i32. + unsafe { + let host_ptr = staging.host_ptr.cast::(); + for (i, &v) in host.iter().enumerate() { + std::ptr::write_volatile(host_ptr.add(i), v); + } + } + let mut dst = stream.alloc_zeros::(n).context("i32 upload alloc")?; + if n > 0 { + let nbytes = n * std::mem::size_of::(); + unsafe { + let (dst_ptr, _g) = dst.device_ptr_mut(stream); + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + staging.dev_ptr, + nbytes, + stream.cu_stream(), + ) + .context("i32 upload DtoD")?; + } + } + // Keep staging alive until the DtoD has been issued onto the stream. + // `dst` carries an internal record; the subsequent kernel launch and + // sync barrier in `run_controller` flushes the copy before staging + // drops here. + stream.synchronize().context("i32 upload sync")?; + drop(staging); + Ok(dst) +} + +fn download_f32(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, + src_ptr, + nbytes, + stream.cu_stream(), + ) + .context("download DtoD")?; + } + stream.synchronize().context("download sync")?; + Ok(staging.read_all()) +} + +/// Download an `i32` slice via the canonical mapped-pinned path +/// (matches `trainer/loss.rs::download_i32`). +fn download_i32(stream: &Arc, src: &CudaSlice) -> Result> { + let n = src.len(); + let staging = unsafe { MappedF32Buffer::new(n) } + .map_err(|e| anyhow::anyhow!("i32 download staging: {e}"))?; + let nbytes = n * std::mem::size_of::(); + unsafe { + let (src_ptr, _g) = src.device_ptr(stream); + cudarc::driver::result::memcpy_dtod_async( + staging.dev_ptr, + src_ptr, + nbytes, + stream.cu_stream(), + ) + .context("i32 download DtoD")?; + } + stream.synchronize().context("i32 download sync")?; + let host_ptr = staging.host_ptr.cast::(); + let mut out = Vec::with_capacity(n); + unsafe { + for i in 0..n { + out.push(std::ptr::read_volatile(host_ptr.add(i))); + } + } + Ok(out) +} + +fn test_device() -> MlDevice { + MlDevice::cuda(0).expect("CUDA 0 required for ml-alpha tests") +} + +// ───── Tests ───── + +#[test] +fn first_observation_bootstraps_ema_and_sets_sentinel() { + let dev = test_device(); + // raw_per_h satisfies the target ratio exactly (raw[h]/raw[0] = K_h30/K_h), + // so on the bootstrap step jitter_ema == raw → target == jitter → excess_ratio = 0. + let raw = [0.5_f32, 0.15, 0.05, 0.015, 0.0025]; + let zero_ema = [0.0_f32; N_HORIZONS]; + let base = 0.01_f32; + let out = run_controller(&dev, &raw, &zero_ema, 0, base).unwrap(); + assert_eq!( + out.first_obs, 1, + "sentinel should advance to 1 on first observation" + ); + for h in 0..N_HORIZONS { + assert_relative_eq!(out.jitter_ema[h], raw[h], epsilon = 1e-6); + // excess_ratio = 0 → lambda = base; floor doesn't engage (base > floor). + assert_relative_eq!(out.lambda[h], base, epsilon = 1e-6); + } +} + +#[test] +fn steady_state_at_target_yields_base_lambda() { + let dev = test_device(); + let ema = [0.5_f32, 0.15, 0.05, 0.015, 0.0025]; + let raw = ema; // identical → EMA update is a no-op (0.5*old + 0.5*old = old) + let base = 0.01_f32; + let out = run_controller(&dev, &raw, &ema, 1, base).unwrap(); + for h in 0..N_HORIZONS { + assert_relative_eq!(out.jitter_ema[h], ema[h], epsilon = 1e-6); + assert_relative_eq!(out.lambda[h], base, epsilon = 1e-5); + } +} + +#[test] +fn excess_at_h6000_lifts_lambda_proportionally() { + let dev = test_device(); + // h6000 EMA at 0.025 = 10× target (= jitter_ema[h30] * 30/6000 = 0.0025) + let ema = [0.5_f32, 0.15, 0.05, 0.015, 0.025]; + let raw = ema; // EMA update is no-op + let base = 0.01_f32; + let out = run_controller(&dev, &raw, &ema, 1, base).unwrap(); + // h30: at-target → lambda = base + assert_relative_eq!(out.lambda[0], base, epsilon = 1e-5); + // h100, h300, h1000: at-target → lambda = base + assert_relative_eq!(out.lambda[1], base, epsilon = 1e-5); + assert_relative_eq!(out.lambda[2], base, epsilon = 1e-5); + assert_relative_eq!(out.lambda[3], base, epsilon = 1e-5); + // h6000: excess_ratio = 0.025/0.0025 - 1 = 9 → lambda = base * (1 + 9) = 0.10 + let expected_lambda_h6000 = base * 10.0; + assert_relative_eq!(out.lambda[4], expected_lambda_h6000, max_relative = 1e-3); +} + +#[test] +fn lambda_floor_when_base_lambda_zero() { + let dev = test_device(); + let ema = [0.5_f32; N_HORIZONS]; + let raw = ema; + let base = 0.0_f32; + let out = run_controller(&dev, &raw, &ema, 1, base).unwrap(); + // base = 0 → lambda_new = 0 (regardless of excess_ratio) → floor activates + for h in 0..N_HORIZONS { + assert!( + (out.lambda[h] - LAMBDA_FLOOR).abs() < 1e-9, + "lambda[{h}] should equal LAMBDA_FLOOR (={LAMBDA_FLOOR:.1e}); got {}", + out.lambda[h] + ); + } +}