1389d1c810
Layer A L40S smoke smoke-test-v9kjv revealed CUDA Graph capture failure
at aux_label_scale_ema mid-step host sync inside aux_heads_forward Step 2b
(CUDA_ERROR_STREAM_CAPTURE_UNSUPPORTED). Per feedback_no_cpu_compute_strict
and pearl_cold_path_no_exception_to_gpu_drives: CPU compute path is
strictly forbidden — any formula (EMA / reduction / Pearls A+D / adaptive α)
belongs on GPU regardless of frequency.
This commit migrates ALL 11 Pearls A+D applications + 1 inline application
to a single GPU apply_pearls_ad_kernel:
- 5 SP4 producers (target_q, atom_pos×4, param_group_oracle, grad_norm, h_s2)
- 6 A13 retrofits (h_s2_rms, aux_heads_loss, moe_expert_util,
vsn_mask, iqn_quantile, reward_component)
- 1 inline aux_label_scale block in aux_heads_forward Step 2b
Eliminates: stream.synchronize() + host_ptr read_volatile + host arithmetic
+ host_ptr write_volatile pattern from all 11 launchers + 1 inline.
apply_pearls_to_slot host helper deleted. pearls_ad_update kept as
test-only reference implementation (also retained for the post-Adam
pearl_c_post_adam_engagement_check host-side diagnostic which reduces
mapped-pinned i32 counters outside any captured graph). New GPU unit
test asserts kernel output matches reference within fp32 ULP for 5
representative cases plus a multi-slot batch sanity check.
The refactor is graph-capture-compatible by construction: the kernel
runs single-thread in the same stream as the producer kernel; no host
synchronisation needed between producer and applicator.
Build clean (cargo check --workspace), 6 host sp4_wiener_ema tests pass,
13 SP4 GPU tests + 1 new oracle test pass on RTX 3050 Ti. L40S smoke
re-validation deferred to A17 redo.
Refs: smoke-test-v9kjv graph-capture failure (terminated), commit 4c231fa81.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2401 lines
99 KiB
Rust
2401 lines
99 KiB
Rust
#![allow(unsafe_code)] // CUDA kernel launch + mapped-pinned memory.
|
||
|
||
//! SP4 producer kernel unit tests.
|
||
//!
|
||
//! Each producer added by Tasks A5-A9 has at least one test in this file
|
||
//! validating its output against an analytical or synthetic-distribution
|
||
//! ground truth. Task A4 lands the first test: the linear-histogram p99
|
||
//! `__device__` function (`sp4_histogram_p99` in `sp4_histogram_p99.cuh`),
|
||
//! exercised through the standalone wrapper kernel
|
||
//! `sp4_histogram_p99_test_kernel.cu`.
|
||
//!
|
||
//! All tests in this file are `#[ignore]`-gated for GPU. Run with:
|
||
//!
|
||
//! cargo test -p ml --tests sp4_producer_unit_tests -- --ignored --nocapture
|
||
//!
|
||
//! No `memcpy_dtoh` anywhere — the wrapper kernel writes its scalar `f32`
|
||
//! result through a mapped-pinned allocation (`cuMemHostAlloc` with
|
||
//! `DEVICEMAP|PORTABLE`), and the host reads via `read_volatile` after a
|
||
//! stream sync, mirroring `gpu_training_guard.rs::MappedBuffer` and
|
||
//! `distributional_q_tests.rs`.
|
||
|
||
use std::sync::Arc;
|
||
|
||
use cudarc::driver::{CudaContext, CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
|
||
use ml::cuda_pipeline::mapped_pinned::{MappedF32Buffer, MappedI32Buffer, MappedU64Buffer};
|
||
use ml::cuda_pipeline::SP4_PRODUCER_COUNT;
|
||
|
||
/// Test-only cubin for the SP4 histogram-p99 wrapper kernel. Built by
|
||
/// `crates/ml/build.rs`; the cubin path is the standard `OUT_DIR` slot
|
||
/// used by every other CUDA kernel in this crate.
|
||
const SP4_HISTOGRAM_P99_TEST_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/sp4_histogram_p99_test_kernel.cubin"));
|
||
|
||
/// Build a CUDA stream against device 0. Mirrors `make_test_stream` in
|
||
/// `distributional_q_tests.rs`. If no GPU is available the call panics —
|
||
/// the test is `#[ignore]`-gated so CPU-only CI never reaches this path.
|
||
fn make_test_stream() -> Arc<CudaStream> {
|
||
let ctx = CudaContext::new(0).expect("CUDA context — is a GPU available?");
|
||
ctx.default_stream()
|
||
}
|
||
|
||
/// Resolve the `sp4_histogram_p99_test_kernel` function handle from the
|
||
/// embedded cubin. Per-test resolution is fine here because the file
|
||
/// holds a single producer test today; the Plan A pattern's `OnceLock`
|
||
/// caching is overkill until Tasks A5-A9 add more wrappers.
|
||
fn load_sp4_histogram_p99_test_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_HISTOGRAM_P99_TEST_CUBIN.to_vec())
|
||
.expect("load sp4_histogram_p99_test_kernel cubin");
|
||
module
|
||
.load_function("sp4_histogram_p99_test_kernel")
|
||
.expect("load sp4_histogram_p99_test_kernel function")
|
||
}
|
||
|
||
/// Run the histogram-p99 wrapper kernel on `samples` and return the
|
||
/// device-computed p99. Single-block launch, BLOCK_SIZE=256 (matches the
|
||
/// `sp4_histogram_p99<256>` instantiation inside the wrapper kernel),
|
||
/// 8 warp tiles × 256 ints × 4 bytes = 8192 bytes dynamic shared memory.
|
||
fn launch_sp4_histogram_p99(stream: &Arc<CudaStream>, samples: &[f32]) -> f32 {
|
||
let kernel = load_sp4_histogram_p99_test_kernel(stream);
|
||
|
||
// Safety: CUDA context is active on this thread (resolved through
|
||
// the stream's context just above).
|
||
let in_buf = unsafe { MappedF32Buffer::new(samples.len()) }
|
||
.expect("alloc MappedF32Buffer for SP4 histogram input");
|
||
in_buf.write_from_slice(samples);
|
||
|
||
let out_buf = unsafe { MappedF32Buffer::new(1) }
|
||
.expect("alloc MappedF32Buffer for SP4 histogram output");
|
||
|
||
let count_i32 = i32::try_from(samples.len()).expect("sample count fits in i32");
|
||
let in_dev_ptr = in_buf.dev_ptr;
|
||
let out_dev_ptr = out_buf.dev_ptr;
|
||
|
||
// Per-warp tiles: 8 warps (BLOCK_SIZE=256, warpSize=32) × 256 bins ×
|
||
// sizeof(int). Matches the contract documented in
|
||
// `sp4_histogram_p99.cuh`.
|
||
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
|
||
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&in_dev_ptr)
|
||
.arg(&count_i32)
|
||
.arg(&out_dev_ptr)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch sp4_histogram_p99_test_kernel");
|
||
}
|
||
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after sp4_histogram_p99_test_kernel launch");
|
||
|
||
let host = out_buf.read_all();
|
||
*host.first().expect("SP4 histogram p99 output empty")
|
||
}
|
||
|
||
/// Touch `MappedI32Buffer` so the import (carried for symmetry with
|
||
/// `distributional_q_tests.rs` and ready for the count-bucket producer in
|
||
/// Task A8) does not provoke an `unused_imports` warning before that
|
||
/// task lands.
|
||
#[allow(dead_code)]
|
||
fn _mapped_i32_import_anchor() {
|
||
let _: Option<Result<MappedI32Buffer, String>> = None;
|
||
}
|
||
|
||
/// SP4 Task A4 unit test: the GPU-side linear-histogram p99 estimator
|
||
/// agrees with a sorted-sample p99 within the 1%-quantile precision
|
||
/// budget of 256 linear bins.
|
||
///
|
||
/// Sample distribution: 4096 deterministic |N(0,1)| draws (Box-Muller
|
||
/// on a fixed-seed `StdRng`). Folded standard-normal magnitudes, where
|
||
/// the analytical p99 ≈ 2.576 (the one-tailed 99th-percentile z-score)
|
||
/// — this synthetic-input ground truth keeps the test independent of any
|
||
/// CPU fallback per `feedback_no_cpu_test_fallbacks.md`.
|
||
///
|
||
/// Tolerance: 5% relative error. Worst-case headroom budget:
|
||
/// - linear-bin quantization at top of distribution: 1/256 ≈ 0.4%
|
||
/// - per-warp lane-collision count loss: <0.012% expected
|
||
/// - finite-sample sorted-p99 vs population p99 jitter: ≈0.5% at n=4096
|
||
/// Total <2%; the 5% threshold is comfortably above that and well above
|
||
/// the analytical ≈2.576 reference.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_histogram_p99_matches_known_distribution_within_quantization() {
|
||
use rand::rngs::StdRng;
|
||
use rand::{Rng, SeedableRng};
|
||
|
||
let mut rng = StdRng::seed_from_u64(0xDEAD_BEEF);
|
||
let n: usize = 4096;
|
||
let samples: Vec<f32> = (0..n)
|
||
.map(|_| {
|
||
// Box-Muller for one f32 N(0,1) sample, then fold to the
|
||
// half-line via `abs()` so we get |N(0,1)|.
|
||
let u1: f32 = rng.gen_range(1e-8_f32..1.0_f32);
|
||
let u2: f32 = rng.gen_range(0.0_f32..1.0_f32);
|
||
let z = (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos();
|
||
z.abs()
|
||
})
|
||
.collect();
|
||
|
||
// Sample-sorted ground truth.
|
||
let mut sorted = samples.clone();
|
||
sorted.sort_by(|a, b| {
|
||
a.partial_cmp(b)
|
||
.expect("|N(0,1)| samples are finite (no NaN)")
|
||
});
|
||
let true_p99 = sorted[(n * 99) / 100];
|
||
|
||
let stream = make_test_stream();
|
||
let computed_p99 = launch_sp4_histogram_p99(&stream, &samples);
|
||
|
||
assert!(
|
||
computed_p99 > 0.0,
|
||
"computed_p99 should be positive (step_max > 0); got {computed_p99}"
|
||
);
|
||
|
||
let rel_err = ((computed_p99 - true_p99) / true_p99).abs();
|
||
println!(
|
||
"SP4 histogram p99 — true_p99={true_p99:.5}, computed_p99={computed_p99:.5}, rel_err={rel_err:.5}",
|
||
);
|
||
assert!(
|
||
rel_err < 0.05,
|
||
"computed_p99={computed_p99}, true_p99={true_p99}, rel_err={rel_err} (> 5%)"
|
||
);
|
||
}
|
||
|
||
// ── SP4 Task A5: target_q_p99 producer kernel ─────────────────────────────────
|
||
|
||
/// Test-only cubin for the SP4 target_q_p99 producer kernel. The same cubin
|
||
/// is loaded by `GpuDqnTrainer::new` in production via `SP4_TARGET_Q_P99_CUBIN`;
|
||
/// this kernel-direct test bypasses the full trainer harness and exercises
|
||
/// the kernel's scratch-slot write contract on a controlled distribution.
|
||
const SP4_TARGET_Q_P99_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/target_q_p99_kernel.cubin"));
|
||
|
||
/// Resolve the `target_q_p99_update` function handle from its cubin.
|
||
fn load_sp4_target_q_p99_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_TARGET_Q_P99_CUBIN.to_vec())
|
||
.expect("load target_q_p99_kernel cubin");
|
||
module
|
||
.load_function("target_q_p99_update")
|
||
.expect("load target_q_p99_update function")
|
||
}
|
||
|
||
/// Run the production producer kernel on `samples` and return the device-
|
||
/// computed step_p99 written to `scratch_buf[scratch_idx]`. Single-block
|
||
/// launch, BLOCK_SIZE=256 (matches the `sp4_histogram_p99<256>` instantiation
|
||
/// inside the kernel), 8 warp tiles × 256 ints × 4 bytes = 8192 bytes
|
||
/// dynamic shared memory. `scratch_idx` mirrors the production layout
|
||
/// (slot 0 = TARGET_Q_BOUND in `launch_sp4_target_q_p99`).
|
||
fn launch_sp4_target_q_p99_into_scratch(
|
||
stream: &Arc<CudaStream>,
|
||
samples: &[f32],
|
||
scratch_idx: usize,
|
||
scratch_len: usize,
|
||
) -> f32 {
|
||
assert!(scratch_idx < scratch_len);
|
||
let kernel = load_sp4_target_q_p99_kernel(stream);
|
||
|
||
// Safety: CUDA context active on this thread (resolved via stream).
|
||
let in_buf = unsafe { MappedF32Buffer::new(samples.len()) }
|
||
.expect("alloc MappedF32Buffer for SP4 target_q_p99 input");
|
||
in_buf.write_from_slice(samples);
|
||
|
||
// Production-shape scratch buffer (47 slots in `GpuDqnTrainer`); the test
|
||
// sizes it to whatever the caller wants but writes via the same indexed
|
||
// contract as the production launcher.
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(scratch_len) }
|
||
.expect("alloc MappedF32Buffer for SP4 producer scratch");
|
||
|
||
let count_i32 = i32::try_from(samples.len()).expect("sample count fits in i32");
|
||
let scratch_idx_i32 = i32::try_from(scratch_idx).expect("scratch idx fits in i32");
|
||
let in_dev_ptr = in_buf.dev_ptr;
|
||
let scratch_dev_ptr = scratch_buf.dev_ptr;
|
||
|
||
// Same shared-memory budget as the production launcher (and the
|
||
// Task A4 wrapper test): 8 warps × 256 bins × 4 bytes.
|
||
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
|
||
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&in_dev_ptr)
|
||
.arg(&count_i32)
|
||
.arg(&scratch_dev_ptr)
|
||
.arg(&scratch_idx_i32)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch target_q_p99_update");
|
||
}
|
||
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after target_q_p99_update launch");
|
||
|
||
// Verify ALL non-target slots remained zero — the kernel must only
|
||
// touch the slot at `scratch_idx`. This guards against a future
|
||
// single-thread-write bug analogous to the `2fb30f098` cascade.
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
if i != scratch_idx {
|
||
assert_eq!(
|
||
v, 0.0,
|
||
"SP4 producer kernel wrote to scratch[{i}] outside target slot {scratch_idx}: {v}"
|
||
);
|
||
}
|
||
}
|
||
host[scratch_idx]
|
||
}
|
||
|
||
/// SP4 Task A5 unit test: the production `target_q_p99_update` kernel
|
||
/// writes step_p99 to `producer_step_scratch_buf[scratch_idx]`, and the
|
||
/// host-side `pearls_ad_update` helper (Task A3) replaces the ISV
|
||
/// sentinel directly with that step observation on the first call after a
|
||
/// fold-boundary reset (Pearl A bootstrap).
|
||
///
|
||
/// Test distribution: 4096 deterministic |N(0,1)| samples (Box-Muller on
|
||
/// fixed-seed `StdRng`) — the same well-spread synthetic-input pattern
|
||
/// used by the Task A4 reference test. Real `denoise_target_q_buf` values
|
||
/// (target-network Q across [B, 12] actions) are similarly well-spread by
|
||
/// the time the kernel runs in production; degenerate single-bin masses
|
||
/// (e.g., the all-1.0/all-100.0 split) hit the `feedback_no_atomicadd.md`
|
||
/// lane-collision bound where 8 lanes in one warp increment the same
|
||
/// shared-memory bin in the same instruction → only 1 increment lands.
|
||
/// The Box-Muller distribution avoids that pathology naturally.
|
||
///
|
||
/// Analytical p99 of |N(0,1)| ≈ 2.576 (the one-tailed 99th-percentile
|
||
/// z-score). Tolerance: 5% relative — same headroom budget as Task A4
|
||
/// (linear-bin quantization 0.4% + lane collision <0.012% + finite-sample
|
||
/// p99 jitter ≈0.5% at n=4096; total <2%).
|
||
///
|
||
/// Pearl A semantics tested explicitly:
|
||
/// - `prev_x_mean = 0.0` (fold-reset sentinel) and `WienerState::ZERO`
|
||
/// (untouched) → `pearls_ad_update` returns `step_observation`
|
||
/// directly and seeds `state.x_lag = step_observation`.
|
||
/// - Subsequent Pearl D calls are exercised in
|
||
/// `crates/ml/src/cuda_pipeline/sp4_wiener_ema.rs::tests`; this test
|
||
/// verifies the kernel-side step_p99 reaches the host and threads
|
||
/// through Pearl A's sentinel branch end-to-end.
|
||
///
|
||
/// Production launcher contract validated:
|
||
/// - Kernel writes step_p99 to `scratch[SCRATCH_IDX_TARGET_Q=0]`.
|
||
/// - All other scratch slots remain zero (helper checks this in
|
||
/// `launch_sp4_target_q_p99_into_scratch`).
|
||
/// - `TARGET_Q_BOUND_INDEX == SP4_SLOT_BASE == 131` (slot-layout guard).
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_target_q_p99_first_observation_writes_step_p99_via_pearls_ad() {
|
||
use ml::cuda_pipeline::sp4_isv_slots::{SP4_SLOT_BASE, TARGET_Q_BOUND_INDEX};
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
use rand::rngs::StdRng;
|
||
use rand::{Rng, SeedableRng};
|
||
|
||
// 4096 |N(0,1)| samples via Box-Muller, fixed seed for determinism.
|
||
// Mirrors the Task A4 test's distribution pattern; analytical p99 of
|
||
// |N(0,1)| ≈ 2.576. Well-spread across bins → no lane-collision
|
||
// pathology.
|
||
let mut rng = StdRng::seed_from_u64(0xA5_5A_F1_57);
|
||
let n: usize = 4096;
|
||
let samples: Vec<f32> = (0..n)
|
||
.map(|_| {
|
||
let u1: f32 = rng.gen_range(1e-8_f32..1.0_f32);
|
||
let u2: f32 = rng.gen_range(0.0_f32..1.0_f32);
|
||
let z = (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos();
|
||
z.abs()
|
||
})
|
||
.collect();
|
||
|
||
let mut sorted = samples.clone();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).expect("|N(0,1)| samples are finite"));
|
||
let true_p99 = sorted[(n * 99) / 100];
|
||
|
||
// Production scratch shape (47 = SP4 producers count). Slot 0 is the
|
||
// production-allocated TARGET_Q_BOUND slot; we use the same index so
|
||
// the kernel-direct test keeps the layout invariant the production
|
||
// launcher depends on.
|
||
const SCRATCH_IDX_TARGET_Q: usize = 0;
|
||
|
||
// ── Kernel ──
|
||
let stream = make_test_stream();
|
||
let step_p99 = launch_sp4_target_q_p99_into_scratch(
|
||
&stream,
|
||
&samples,
|
||
SCRATCH_IDX_TARGET_Q,
|
||
SP4_PRODUCER_COUNT,
|
||
);
|
||
|
||
let rel_err = ((step_p99 - true_p99) / true_p99).abs();
|
||
println!(
|
||
"SP4 target_q_p99 — true_p99={true_p99:.5}, step_p99 (kernel)={step_p99:.5}, rel_err={rel_err:.5}",
|
||
);
|
||
assert!(
|
||
step_p99 > 0.0,
|
||
"kernel step_p99 (slot {SCRATCH_IDX_TARGET_Q}) should be positive; got {step_p99}",
|
||
);
|
||
assert!(
|
||
rel_err < 0.05,
|
||
"kernel step_p99 ({step_p99}) vs true p99 ({true_p99}) rel_err {rel_err} > 5%",
|
||
);
|
||
|
||
// ── Pearl A bootstrap (host-side, mirrors `launch_sp4_target_q_p99`) ──
|
||
// Sentinel ISV slot + zero Wiener state — replicates the cold-start
|
||
// condition immediately after a fold-boundary reset.
|
||
let prev_x_mean: f32 = 0.0;
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_p99);
|
||
|
||
assert_eq!(
|
||
new_x_mean, step_p99,
|
||
"Pearl A: first observation must replace sentinel directly; got {new_x_mean} vs step_p99={step_p99}",
|
||
);
|
||
assert_eq!(
|
||
state.x_lag, step_p99,
|
||
"Pearl A: x_lag must seed to first observation",
|
||
);
|
||
assert_eq!(state.sample_var, 0.0, "Pearl A: sample_var stays zero");
|
||
assert_eq!(state.diff_var, 0.0, "Pearl A: diff_var stays zero");
|
||
|
||
// Sanity: the slot index this producer writes is the documented
|
||
// TARGET_Q_BOUND slot at the SP4 base. Guards against an off-by-one
|
||
// SP4_SLOT_BASE drift in the ISV layout.
|
||
assert_eq!(TARGET_Q_BOUND_INDEX, SP4_SLOT_BASE);
|
||
assert_eq!(TARGET_Q_BOUND_INDEX, 131);
|
||
}
|
||
|
||
// ── SP4 Task A6: atom_pos_p99 producer × 4 branches ───────────────────────────
|
||
|
||
/// Test-only cubin for the SP4 atom_pos_p99 producer kernel. The same cubin
|
||
/// is loaded by `GpuDqnTrainer::new` in production via `SP4_ATOM_POS_P99_CUBIN`;
|
||
/// this kernel-direct test bypasses the full trainer harness and exercises
|
||
/// the kernel's per-branch scratch-slot write contract by launching it 4
|
||
/// times with distinct slice pointers + scratch slots, mirroring the
|
||
/// production `launch_sp4_atom_pos_p99_all_branches` launcher.
|
||
const SP4_ATOM_POS_P99_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/atom_pos_p99_kernel.cubin"));
|
||
|
||
/// Resolve the `atom_pos_p99_update` function handle from its cubin.
|
||
fn load_sp4_atom_pos_p99_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_ATOM_POS_P99_CUBIN.to_vec())
|
||
.expect("load atom_pos_p99_kernel cubin");
|
||
module
|
||
.load_function("atom_pos_p99_update")
|
||
.expect("load atom_pos_p99_update function")
|
||
}
|
||
|
||
/// SP4 Task A6 unit test: the production `atom_pos_p99_update` kernel
|
||
/// computes p99(|atom_positions[branch]|) per branch and writes each branch's
|
||
/// step_p99 to its distinct producer-step-scratch slot. Verifies per-branch
|
||
/// independence by giving each branch a different scale of |N(0,1)| samples
|
||
/// (`1×, 10×, 100×, 1000×`) — each branch's slot must be within ±5% of its
|
||
/// own scale × analytical p99 ≈ 2.576, while non-target slots must remain
|
||
/// zero.
|
||
///
|
||
/// Distribution: 4096 deterministic |N(0,1)| Box-Muller samples per branch
|
||
/// (`StdRng::seed_from_u64(0xA6_<branch>_F1_57)`), scaled by the per-branch
|
||
/// factor. The well-spread base distribution avoids the
|
||
/// `feedback_no_atomicadd.md`-acknowledged intra-warp lane-collision bound;
|
||
/// the per-branch scaling factors (10ⁿ for n=0..3) span four orders of
|
||
/// magnitude so a launcher bug that re-uses the wrong slice pointer would
|
||
/// land in the wrong scratch slot at the wrong scale, causing a >5% rel_err
|
||
/// in at least one of the four assertions.
|
||
///
|
||
/// Tolerance: 5% relative — same headroom budget as Tasks A4 and A5
|
||
/// (linear-bin quantization 0.4% + lane collision <0.012% + finite-sample
|
||
/// p99 jitter ≈0.5% at n=4096; total <2%).
|
||
///
|
||
/// Production launcher contract validated:
|
||
/// - Each kernel launch writes step_p99 to `scratch[SCRATCH_BASE + branch]`
|
||
/// (SCRATCH_BASE=1, slots 1..5 per `launch_sp4_atom_pos_p99_all_branches`).
|
||
/// - All other scratch slots remain zero (slot 0 is reserved for
|
||
/// TARGET_Q_BOUND in the production layout; slots 5..47 reserved for
|
||
/// Tasks A7-A9).
|
||
/// - `ATOM_POS_BOUND_BASE == 132 == SP4_SLOT_BASE + 1` (slot-layout guard).
|
||
/// - `atom_pos_bound(branch) == ATOM_POS_BOUND_BASE + branch` for branch
|
||
/// ∈ 0..SP4_BRANCH_COUNT (convenience-fn invariant).
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_atom_pos_p99_per_branch_writes_distinct_isv_slots() {
|
||
use ml::cuda_pipeline::sp4_isv_slots::{
|
||
atom_pos_bound, ATOM_POS_BOUND_BASE, SP4_BRANCH_COUNT, SP4_SLOT_BASE,
|
||
};
|
||
use rand::rngs::StdRng;
|
||
use rand::{Rng, SeedableRng};
|
||
|
||
// Same scratch shape as `launch_sp4_atom_pos_p99_all_branches` writes
|
||
// into in production. SCRATCH_BASE=1 mirrors the launcher's documented
|
||
// slot assignment (slot 0 = TARGET_Q_BOUND, slots 1..5 = ATOM_POS_BOUND[0..4]).
|
||
const SCRATCH_BASE: usize = 1;
|
||
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
|
||
const N_PER_BRANCH: usize = 4096;
|
||
|
||
// Per-branch scaling factors span 10⁰..10³ — a bug that mixes up branch
|
||
// slices would hit wrong-scale results well outside the 5% tolerance on
|
||
// at least one branch's assertion.
|
||
let scales: [f32; 4] = [1.0, 10.0, 100.0, 1000.0];
|
||
assert_eq!(scales.len(), SP4_BRANCH_COUNT);
|
||
|
||
// Generate per-branch |N(0,1)| × scale samples, deterministic seeds.
|
||
let mut per_branch_samples: Vec<Vec<f32>> = Vec::with_capacity(SP4_BRANCH_COUNT);
|
||
let mut per_branch_true_p99: Vec<f32> = Vec::with_capacity(SP4_BRANCH_COUNT);
|
||
for branch in 0..SP4_BRANCH_COUNT {
|
||
let seed: u64 = 0xA6_5A_F1_57 ^ ((branch as u64) << 32);
|
||
let mut rng = StdRng::seed_from_u64(seed);
|
||
let scale = scales[branch];
|
||
let samples: Vec<f32> = (0..N_PER_BRANCH)
|
||
.map(|_| {
|
||
let u1: f32 = rng.gen_range(1e-8_f32..1.0_f32);
|
||
let u2: f32 = rng.gen_range(0.0_f32..1.0_f32);
|
||
let z = (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos();
|
||
z.abs() * scale
|
||
})
|
||
.collect();
|
||
let mut sorted = samples.clone();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).expect("|N(0,1)|×scale samples are finite"));
|
||
per_branch_true_p99.push(sorted[(N_PER_BRANCH * 99) / 100]);
|
||
per_branch_samples.push(samples);
|
||
}
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_sp4_atom_pos_p99_kernel(&stream);
|
||
|
||
// ── Allocate ONE flat input buffer matching production layout
|
||
// [SP4_BRANCH_COUNT × N_PER_BRANCH], branch b at offset b*N_PER_BRANCH.
|
||
// The kernel takes a per-branch slice pointer; pointer arithmetic in
|
||
// the loop below mirrors the launcher's `base + branch * stride`.
|
||
let total_len = SP4_BRANCH_COUNT * N_PER_BRANCH;
|
||
let in_buf = unsafe { MappedF32Buffer::new(total_len) }
|
||
.expect("alloc MappedF32Buffer for SP4 atom_pos_p99 input");
|
||
{
|
||
// Pack all 4 branches into the flat buffer (host-side write via
|
||
// mapped-pinned slice — same `write_from_slice` API as A4/A5).
|
||
let mut flat = Vec::with_capacity(total_len);
|
||
for branch in 0..SP4_BRANCH_COUNT {
|
||
flat.extend_from_slice(&per_branch_samples[branch]);
|
||
}
|
||
in_buf.write_from_slice(&flat);
|
||
}
|
||
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc MappedF32Buffer for SP4 producer scratch");
|
||
|
||
let count_per_branch = i32::try_from(N_PER_BRANCH).expect("N_PER_BRANCH fits in i32");
|
||
let scratch_dev_ptr = scratch_buf.dev_ptr;
|
||
let stride_bytes = (N_PER_BRANCH * std::mem::size_of::<f32>()) as u64;
|
||
|
||
// ── Phase 1: launch the same kernel 4 times with distinct slice + slot ──
|
||
for branch in 0..SP4_BRANCH_COUNT {
|
||
let branch_dev_ptr: u64 = in_buf.dev_ptr + (branch as u64) * stride_bytes;
|
||
let scratch_idx_arg: i32 = (SCRATCH_BASE + branch) as i32;
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&branch_dev_ptr)
|
||
.arg(&count_per_branch)
|
||
.arg(&scratch_dev_ptr)
|
||
.arg(&scratch_idx_arg)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch atom_pos_p99_update");
|
||
}
|
||
}
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after 4 atom_pos_p99_update launches");
|
||
|
||
// ── Phase 2: assertions ──
|
||
let host = scratch_buf.read_all();
|
||
|
||
// Non-target slots (anything outside 1..5) must remain zero — guards
|
||
// against any writes leaking outside the documented per-branch slot
|
||
// assignment.
|
||
for (i, &v) in host.iter().enumerate() {
|
||
let in_range = (SCRATCH_BASE..SCRATCH_BASE + SP4_BRANCH_COUNT).contains(&i);
|
||
if !in_range {
|
||
assert_eq!(
|
||
v, 0.0,
|
||
"atom_pos_p99 producer wrote to scratch[{i}] outside ATOM_POS_BOUND slots {SCRATCH_BASE}..{}: {v}",
|
||
SCRATCH_BASE + SP4_BRANCH_COUNT,
|
||
);
|
||
}
|
||
}
|
||
|
||
// Per-branch slot must hold step_p99 within ±5% of the scaled p99.
|
||
for branch in 0..SP4_BRANCH_COUNT {
|
||
let scratch_idx = SCRATCH_BASE + branch;
|
||
let step_p99 = host[scratch_idx];
|
||
let true_p99 = per_branch_true_p99[branch];
|
||
let rel_err = ((step_p99 - true_p99) / true_p99).abs();
|
||
println!(
|
||
"SP4 atom_pos_p99[branch={branch}, scale={}] — true_p99={true_p99:.5}, step_p99={step_p99:.5}, rel_err={rel_err:.5}",
|
||
scales[branch],
|
||
);
|
||
assert!(
|
||
step_p99 > 0.0,
|
||
"branch {branch} step_p99 (slot {scratch_idx}) should be positive; got {step_p99}",
|
||
);
|
||
assert!(
|
||
rel_err < 0.05,
|
||
"branch {branch} step_p99 ({step_p99}) vs true p99 ({true_p99}) rel_err {rel_err} > 5%",
|
||
);
|
||
}
|
||
|
||
// Slot-layout invariants. Guards against an off-by-one drift in
|
||
// ATOM_POS_BOUND_BASE or SP4_BRANCH_COUNT that would silently route
|
||
// writes to the wrong ISV indices.
|
||
assert_eq!(SP4_SLOT_BASE, 131);
|
||
assert_eq!(ATOM_POS_BOUND_BASE, SP4_SLOT_BASE + 1);
|
||
assert_eq!(ATOM_POS_BOUND_BASE, 132);
|
||
assert_eq!(SP4_BRANCH_COUNT, 4);
|
||
for branch in 0..SP4_BRANCH_COUNT {
|
||
assert_eq!(atom_pos_bound(branch), ATOM_POS_BOUND_BASE + branch);
|
||
}
|
||
}
|
||
|
||
// ── SP4 Task A7: Pearl B fused per-param-group oracle kernel ──────────────────
|
||
|
||
/// Test-only cubin for the SP4 Pearl B fused per-param-group oracle. The same
|
||
/// cubin is loaded by `GpuDqnTrainer::new` in production via
|
||
/// `SP4_PARAM_GROUP_ORACLE_CUBIN`. This kernel-direct test bypasses the full
|
||
/// trainer harness and exercises the kernel's 4-5 fused passes (max-reduce
|
||
/// + p99 histogram per (params, adam_m, adam_v); `Σ w·g / Σ w²` weight-decay
|
||
/// rate; trunk-only L1-lambda from gradient-direction entropy deficit) on
|
||
/// controlled distributions with analytically-known ground truth.
|
||
const SP4_PARAM_GROUP_ORACLE_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/param_group_oracle_kernel.cubin"));
|
||
|
||
/// Resolve the `param_group_oracle_update` function handle from its cubin.
|
||
fn load_sp4_param_group_oracle_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_PARAM_GROUP_ORACLE_CUBIN.to_vec())
|
||
.expect("load param_group_oracle_kernel cubin");
|
||
module
|
||
.load_function("param_group_oracle_update")
|
||
.expect("load param_group_oracle_update function")
|
||
}
|
||
|
||
/// Box-Muller deterministic |N(0, σ²)| sample generator. Produces folded
|
||
/// half-normal magnitudes so the per-pass p99 lookups are well-spread across
|
||
/// the 256 linear-histogram bins (avoids the lane-collision pathology
|
||
/// flagged in `feedback_no_atomicadd.md`).
|
||
fn boxmuller_abs_normal(seed: u64, n: usize, sigma: f32) -> Vec<f32> {
|
||
use rand::rngs::StdRng;
|
||
use rand::{Rng, SeedableRng};
|
||
let mut rng = StdRng::seed_from_u64(seed);
|
||
(0..n)
|
||
.map(|_| {
|
||
let u1: f32 = rng.gen_range(1e-8_f32..1.0_f32);
|
||
let u2: f32 = rng.gen_range(0.0_f32..1.0_f32);
|
||
let z = (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos();
|
||
(z * sigma).abs()
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
/// Sample-sorted p99 of `xs` (positive samples). Mirrors the reference
|
||
/// computation used by the Task A4/A5/A6 tests — `sorted[(n * 99) / 100]`.
|
||
fn sample_p99(xs: &[f32]) -> f32 {
|
||
let mut sorted: Vec<f32> = xs.to_vec();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).expect("samples are finite"));
|
||
sorted[(sorted.len() * 99) / 100]
|
||
}
|
||
|
||
/// One sub-buffer's host-side data for a kernel-direct test launch.
|
||
/// Mirrors the production `Sp4SubBuffer` Rust type but holds host
|
||
/// vectors instead of device pointers. The test launcher allocates
|
||
/// matching mapped-pinned device-visible buffers and populates them.
|
||
struct TestSubBuffer {
|
||
params_host: Vec<f32>,
|
||
grads_host: Vec<f32>,
|
||
adam_m_host: Vec<f32>,
|
||
adam_v_host: Vec<f32>,
|
||
}
|
||
|
||
impl TestSubBuffer {
|
||
fn count(&self) -> usize {
|
||
let n = self.params_host.len();
|
||
assert_eq!(self.grads_host.len(), n);
|
||
assert_eq!(self.adam_m_host.len(), n);
|
||
assert_eq!(self.adam_v_host.len(), n);
|
||
n
|
||
}
|
||
}
|
||
|
||
/// Launch the production `param_group_oracle_update` kernel kernel-direct
|
||
/// across one or more sub-buffers (groups 0-6: 1 sub-buffer; group 7: 4
|
||
/// sub-buffers). Returns the host-readable scratch buffer for assertion.
|
||
///
|
||
/// `k_in`/`h_dim` non-zero only for `g_idx == 0` — Pass E (L1 lambda)
|
||
/// gating on `l1_lambda_scratch_idx >= 0`. For other groups, pass 0/0/-1.
|
||
fn launch_sp4_param_group_oracle_for_group(
|
||
stream: &Arc<CudaStream>,
|
||
g_idx: usize,
|
||
sub_buffers: &[TestSubBuffer],
|
||
k_in: i32,
|
||
h_dim: i32,
|
||
) -> Vec<f32> {
|
||
const SCRATCH_BASE_W: usize = 5;
|
||
const SCRATCH_BASE_M: usize = 13;
|
||
const SCRATCH_BASE_V: usize = 21;
|
||
const SCRATCH_BASE_WD: usize = 29;
|
||
const SCRATCH_L1_TRUNK: usize = 37;
|
||
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
|
||
const K_MAX: usize = 4; // matches production `SP4_ORACLE_TABLE_MAX_SUB`
|
||
|
||
let n_sub = sub_buffers.len();
|
||
assert!(n_sub >= 1 && n_sub <= K_MAX,
|
||
"n_sub={n_sub} must be in [1, {K_MAX}]");
|
||
let total_count: usize = sub_buffers.iter().map(|sb| sb.count()).sum();
|
||
assert!(total_count > 0, "total_count must be positive");
|
||
|
||
let kernel = load_sp4_param_group_oracle_kernel(stream);
|
||
|
||
// Allocate per-sub-buffer mapped-pinned device-visible storage and
|
||
// write the host data through. We hold the buffers alive in a Vec
|
||
// for the duration of the kernel launch.
|
||
let mut buf_owners: Vec<(MappedF32Buffer, MappedF32Buffer, MappedF32Buffer, MappedF32Buffer)>
|
||
= Vec::with_capacity(n_sub);
|
||
for sb in sub_buffers {
|
||
let n = sb.count();
|
||
let p_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc params");
|
||
let g_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc grads");
|
||
let m_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc adam_m");
|
||
let v_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc adam_v");
|
||
p_buf.write_from_slice(&sb.params_host);
|
||
g_buf.write_from_slice(&sb.grads_host);
|
||
m_buf.write_from_slice(&sb.adam_m_host);
|
||
v_buf.write_from_slice(&sb.adam_v_host);
|
||
buf_owners.push((p_buf, g_buf, m_buf, v_buf));
|
||
}
|
||
|
||
// Build the device-pointer table: 4 ptr-arrays of length K_MAX,
|
||
// packed contiguously in a single mapped-pinned u64 buffer. Counts
|
||
// sit in a parallel mapped-pinned i32 buffer of length K_MAX.
|
||
let table_buf = unsafe { MappedU64Buffer::new(4 * K_MAX) }.expect("alloc table");
|
||
let counts_buf = unsafe { MappedI32Buffer::new(K_MAX) }.expect("alloc counts");
|
||
{
|
||
// Build host-side images then write them through.
|
||
let mut table_img: Vec<u64> = vec![0; 4 * K_MAX];
|
||
let mut counts_img: Vec<i32> = vec![0; K_MAX];
|
||
for (s, owners) in buf_owners.iter().enumerate() {
|
||
table_img[0 * K_MAX + s] = owners.0.dev_ptr;
|
||
table_img[1 * K_MAX + s] = owners.1.dev_ptr;
|
||
table_img[2 * K_MAX + s] = owners.2.dev_ptr;
|
||
table_img[3 * K_MAX + s] = owners.3.dev_ptr;
|
||
counts_img[s] = sub_buffers[s].count() as i32;
|
||
}
|
||
table_buf.write_from_slice(&table_img);
|
||
counts_buf.write_from_slice(&counts_img);
|
||
}
|
||
|
||
let scratch_buf =
|
||
unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }.expect("alloc scratch");
|
||
|
||
let n_sub_i32 = n_sub as i32;
|
||
let total_count_i32 = total_count as i32;
|
||
let weight_idx = (SCRATCH_BASE_W + g_idx) as i32;
|
||
let adam_m_idx = (SCRATCH_BASE_M + g_idx) as i32;
|
||
let adam_v_idx = (SCRATCH_BASE_V + g_idx) as i32;
|
||
let wd_rate_idx = (SCRATCH_BASE_WD + g_idx) as i32;
|
||
let l1_idx: i32 = if g_idx == 0 { SCRATCH_L1_TRUNK as i32 } else { -1 };
|
||
|
||
let table_dev = table_buf.dev_ptr;
|
||
let stride_bytes = (K_MAX * std::mem::size_of::<u64>()) as u64;
|
||
let params_ptrs_dev = table_dev;
|
||
let grads_ptrs_dev = table_dev + stride_bytes;
|
||
let adam_m_ptrs_dev = table_dev + 2 * stride_bytes;
|
||
let adam_v_ptrs_dev = table_dev + 3 * stride_bytes;
|
||
let counts_dev = counts_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(¶ms_ptrs_dev)
|
||
.arg(&grads_ptrs_dev)
|
||
.arg(&adam_m_ptrs_dev)
|
||
.arg(&adam_v_ptrs_dev)
|
||
.arg(&counts_dev)
|
||
.arg(&n_sub_i32)
|
||
.arg(&total_count_i32)
|
||
.arg(&k_in)
|
||
.arg(&h_dim)
|
||
.arg(&scratch_dev)
|
||
.arg(&weight_idx)
|
||
.arg(&adam_m_idx)
|
||
.arg(&adam_v_idx)
|
||
.arg(&wd_rate_idx)
|
||
.arg(&l1_idx)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch param_group_oracle_update");
|
||
}
|
||
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after param_group_oracle_update launch");
|
||
|
||
scratch_buf.read_all()
|
||
}
|
||
|
||
/// SP4 Task A7 unit test: parameterised across all 8 param-groups. Each
|
||
/// group launches the production kernel kernel-direct with a distinct
|
||
/// `(seed, sigma_w, sigma_g, sigma_m, sigma_v, count)` tuple — the per-group
|
||
/// scaling factors span 1×..1000× across the 4 buffers so a kernel bug
|
||
/// that mixes up buffer arguments would land in the wrong scratch slot at
|
||
/// the wrong scale and trip a per-output rel_err assertion.
|
||
///
|
||
/// **Distributions:** 4096 deterministic |N(0, σ²)| Box-Muller samples per
|
||
/// buffer per group (well-spread → no lane-collision pathology). Per-buffer
|
||
/// seeds: `0xA7_<group>_<buf>_F1_57`. Per-group `count` chosen to keep the
|
||
/// test runtime small while exceeding 1024 (the kernel's per-thread strided
|
||
/// reduction needs ≥ 1 stride for thread-256's accumulator to be exercised).
|
||
///
|
||
/// **Tolerances:**
|
||
/// - WEIGHT/M/V p99: 5% relative error (linear-bin quantization 0.4% +
|
||
/// lane collision <0.012% + finite-sample p99 jitter ≈0.5% at n=4096
|
||
/// ≈ <2% true; 5% is generous).
|
||
/// - WD_RATE: 2% relative error (analytical formula `|Σ w·g| / Σ w²`
|
||
/// reduces in `f32` with no histogram quantization; 2% headroom for
|
||
/// summation order non-determinism).
|
||
/// - L1 (group 0 only): 5% relative error — composes (mean|g|/mean|w|)
|
||
/// × deficit, both `f32` reductions plus a per-feature L2 norm chain
|
||
/// with two `logf` calls. Tolerance matches the p99 family.
|
||
///
|
||
/// **Production launcher contract validated:**
|
||
/// - The kernel writes step observations to scratch slots
|
||
/// `5 + g`, `13 + g`, `21 + g`, `29 + g` per group (and slot 37 for
|
||
/// group 0's L1 trunk lambda).
|
||
/// - All non-target slots remain zero (guards against the
|
||
/// `2fb30f098`-class write-cascade where a kernel silently writes
|
||
/// outside its assigned slot).
|
||
/// - SP4 ISV slot conventions hold: `weight_bound(g)`, `adam_m_bound(g)`,
|
||
/// `adam_v_bound(g)`, `wd_rate(g)`, `L1_LAMBDA_TRUNK_INDEX = 170`.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_param_group_oracle_per_group_writes_distinct_isv_slots() {
|
||
use ml::cuda_pipeline::sp4_isv_slots::{
|
||
adam_m_bound, adam_v_bound, weight_bound, wd_rate, ParamGroup,
|
||
ADAM_M_BOUND_BASE, ADAM_V_BOUND_BASE, L1_LAMBDA_TRUNK_INDEX,
|
||
SP4_PARAM_GROUP_COUNT, SP4_SLOT_BASE, WD_RATE_BASE, WEIGHT_BOUND_BASE,
|
||
};
|
||
|
||
const SCRATCH_BASE_W: usize = 5;
|
||
const SCRATCH_BASE_M: usize = 13;
|
||
const SCRATCH_BASE_V: usize = 21;
|
||
const SCRATCH_BASE_WD: usize = 29;
|
||
const SCRATCH_L1_TRUNK: usize = 37;
|
||
const N: usize = 4096;
|
||
|
||
// Per-group input shape:
|
||
// group 0 (DqnTrunk): k_in × h_dim = 64 × 64 = 4096 (Pass E exercised)
|
||
// groups 1..7: flat 4096 element slice (Pass E gated off)
|
||
// sigmas chosen so p99(|N(0,σ²)|) lands at ~2.576 × σ:
|
||
// sigma_w = 0.1 × (g+1) → WEIGHT_BOUND[g] ≈ 0.258 × (g+1)
|
||
// sigma_g = 0.01 × (g+1) → grads at 1/10 weight scale
|
||
// sigma_m = 0.05 × (g+1) → adam_m at 1/2 weight scale
|
||
// sigma_v = 0.001 × (g+1) → adam_v second-moment small (mostly squared g)
|
||
let stream = make_test_stream();
|
||
|
||
// Layout invariants — guard against off-by-one drift in the SP4 slot map
|
||
// that would silently route writes to wrong ISV indices.
|
||
assert_eq!(SP4_PARAM_GROUP_COUNT, 8);
|
||
assert_eq!(SP4_SLOT_BASE, 131);
|
||
assert_eq!(WEIGHT_BOUND_BASE, 136);
|
||
assert_eq!(ADAM_M_BOUND_BASE, 144);
|
||
assert_eq!(ADAM_V_BOUND_BASE, 152);
|
||
assert_eq!(WD_RATE_BASE, 160);
|
||
assert_eq!(L1_LAMBDA_TRUNK_INDEX, 170);
|
||
for g in 0..SP4_PARAM_GROUP_COUNT {
|
||
assert_eq!(weight_bound(g), 136 + g);
|
||
assert_eq!(adam_m_bound(g), 144 + g);
|
||
assert_eq!(adam_v_bound(g), 152 + g);
|
||
assert_eq!(wd_rate(g), 160 + g);
|
||
}
|
||
// ParamGroup::ALL ordering matches our seed-keying.
|
||
let groups = ParamGroup::ALL;
|
||
assert_eq!(groups.len(), SP4_PARAM_GROUP_COUNT);
|
||
|
||
for g_idx in 0..SP4_PARAM_GROUP_COUNT {
|
||
let scale = (g_idx as f32) + 1.0;
|
||
let sigma_w = 0.1 * scale;
|
||
let sigma_g = 0.01 * scale;
|
||
let sigma_m = 0.05 * scale;
|
||
let sigma_v = 0.001 * scale;
|
||
|
||
// Group 7 (Curiosity) exercises the multi-sub-buffer path with 4
|
||
// sub-buffers of distinct shapes (w1/b1/w2/b2-like sizes). All
|
||
// other groups use a single contiguous N-element sub-buffer.
|
||
let sub_shapes: Vec<usize> = if g_idx == 7 {
|
||
// Mirror Curiosity production layout: w1=[128×45]=5760, b1=128,
|
||
// w2=[42×128]=5376, b2=42 (total 11306) — but scale to fit the
|
||
// test budget. Use 1024/32/1024/32 to total 2112 elements (kept
|
||
// small for fast tests; still exercises the multi-sub-buffer
|
||
// iteration in Pass A/B/C/D).
|
||
vec![1024, 32, 1024, 32]
|
||
} else {
|
||
vec![N]
|
||
};
|
||
|
||
let mut sub_buffers: Vec<TestSubBuffer> = Vec::with_capacity(sub_shapes.len());
|
||
for (s, &shape_n) in sub_shapes.iter().enumerate() {
|
||
// Per-sub-buffer seed offsets so distinct sub-buffers have
|
||
// distinct distributions (catches buffer-mixup bugs in the
|
||
// kernel's sub-buffer iteration).
|
||
let seed_w = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ ((s as u64) << 16) ^ 0x01;
|
||
let seed_g = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ ((s as u64) << 16) ^ 0x02;
|
||
let seed_m = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ ((s as u64) << 16) ^ 0x03;
|
||
let seed_v = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ ((s as u64) << 16) ^ 0x04;
|
||
|
||
let params_host = boxmuller_abs_normal(seed_w, shape_n, sigma_w);
|
||
// grads: signed (mix of positive/negative) so Σ w·g doesn't
|
||
// trivially collapse to Σ w·|g|. Alternate signs by index parity.
|
||
let grads_mag = boxmuller_abs_normal(seed_g, shape_n, sigma_g);
|
||
let grads_host: Vec<f32> = grads_mag
|
||
.iter()
|
||
.enumerate()
|
||
.map(|(i, &g)| if i % 2 == 0 { g } else { -g })
|
||
.collect();
|
||
let adam_m_host = boxmuller_abs_normal(seed_m, shape_n, sigma_m);
|
||
let adam_v_host = boxmuller_abs_normal(seed_v, shape_n, sigma_v);
|
||
|
||
sub_buffers.push(TestSubBuffer {
|
||
params_host,
|
||
grads_host,
|
||
adam_m_host,
|
||
adam_v_host,
|
||
});
|
||
}
|
||
|
||
// For group 0 (DqnTrunk) the kernel uses k_in × h_dim total
|
||
// elements (= sub_buffers[0].count() since group 0 is single-sub-
|
||
// buffer) to drive Pass E. We pick k_in = 64, h_dim = 64 so the
|
||
// matrix has N=4096 elements.
|
||
let (k_in, h_dim) = if g_idx == 0 { (64_i32, 64_i32) } else { (0_i32, 0_i32) };
|
||
|
||
let scratch = launch_sp4_param_group_oracle_for_group(
|
||
&stream,
|
||
g_idx,
|
||
&sub_buffers,
|
||
k_in,
|
||
h_dim,
|
||
);
|
||
|
||
// ── Slot-write isolation: every slot outside this group's targets
|
||
// must remain zero. Targets are [5+g, 13+g, 21+g, 29+g] and (for
|
||
// g==0) slot 37.
|
||
let mut is_target = [false; SP4_PRODUCER_COUNT];
|
||
is_target[SCRATCH_BASE_W + g_idx] = true;
|
||
is_target[SCRATCH_BASE_M + g_idx] = true;
|
||
is_target[SCRATCH_BASE_V + g_idx] = true;
|
||
is_target[SCRATCH_BASE_WD + g_idx] = true;
|
||
if g_idx == 0 {
|
||
is_target[SCRATCH_L1_TRUNK] = true;
|
||
}
|
||
for (i, &v) in scratch.iter().enumerate() {
|
||
if !is_target[i] {
|
||
assert_eq!(
|
||
v, 0.0,
|
||
"param_group_oracle[group={g_idx}] wrote to non-target scratch[{i}] = {v}",
|
||
);
|
||
}
|
||
}
|
||
|
||
// Build the union distributions for reference computation.
|
||
// p99 is computed over the union of all sub-buffer entries — this
|
||
// matches the kernel's `sp4_histogram_p99_multi` semantics.
|
||
let union_params: Vec<f32> = sub_buffers.iter()
|
||
.flat_map(|sb| sb.params_host.iter().copied()).collect();
|
||
let union_grads: Vec<f32> = sub_buffers.iter()
|
||
.flat_map(|sb| sb.grads_host.iter().copied()).collect();
|
||
let union_adam_m: Vec<f32> = sub_buffers.iter()
|
||
.flat_map(|sb| sb.adam_m_host.iter().copied()).collect();
|
||
let union_adam_v: Vec<f32> = sub_buffers.iter()
|
||
.flat_map(|sb| sb.adam_v_host.iter().copied()).collect();
|
||
let union_total_count: usize = union_params.len();
|
||
|
||
// ── Pass A: WEIGHT_BOUND[g] p99 ──
|
||
let true_p99_w = sample_p99(&union_params);
|
||
let kernel_p99_w = scratch[SCRATCH_BASE_W + g_idx];
|
||
let rel_w = ((kernel_p99_w - true_p99_w) / true_p99_w).abs();
|
||
println!(
|
||
"SP4 param_group[{g_idx}] WEIGHT — true_p99={true_p99_w:.5}, kernel_p99={kernel_p99_w:.5}, rel_err={rel_w:.5} (n_sub={} of len {})",
|
||
sub_buffers.len(),
|
||
sub_buffers.iter().map(|sb| sb.count().to_string()).collect::<Vec<_>>().join("+"),
|
||
);
|
||
assert!(kernel_p99_w > 0.0, "WEIGHT p99 must be positive");
|
||
assert!(
|
||
rel_w < 0.05,
|
||
"group {g_idx} WEIGHT rel_err {rel_w} > 5% (kernel={kernel_p99_w}, true={true_p99_w})",
|
||
);
|
||
|
||
// ── Pass B: ADAM_M_BOUND[g] p99 ──
|
||
let true_p99_m = sample_p99(&union_adam_m);
|
||
let kernel_p99_m = scratch[SCRATCH_BASE_M + g_idx];
|
||
let rel_m = ((kernel_p99_m - true_p99_m) / true_p99_m).abs();
|
||
println!(
|
||
"SP4 param_group[{g_idx}] ADAM_M — true_p99={true_p99_m:.5}, kernel_p99={kernel_p99_m:.5}, rel_err={rel_m:.5}",
|
||
);
|
||
assert!(kernel_p99_m > 0.0, "ADAM_M p99 must be positive");
|
||
assert!(
|
||
rel_m < 0.05,
|
||
"group {g_idx} ADAM_M rel_err {rel_m} > 5% (kernel={kernel_p99_m}, true={true_p99_m})",
|
||
);
|
||
|
||
// ── Pass C: ADAM_V_BOUND[g] p99 ──
|
||
let true_p99_v = sample_p99(&union_adam_v);
|
||
let kernel_p99_v = scratch[SCRATCH_BASE_V + g_idx];
|
||
let rel_v = ((kernel_p99_v - true_p99_v) / true_p99_v).abs();
|
||
println!(
|
||
"SP4 param_group[{g_idx}] ADAM_V — true_p99={true_p99_v:.5}, kernel_p99={kernel_p99_v:.5}, rel_err={rel_v:.5}",
|
||
);
|
||
assert!(kernel_p99_v > 0.0, "ADAM_V p99 must be positive");
|
||
assert!(
|
||
rel_v < 0.05,
|
||
"group {g_idx} ADAM_V rel_err {rel_v} > 5% (kernel={kernel_p99_v}, true={true_p99_v})",
|
||
);
|
||
|
||
// ── Pass D: WD_RATE[g] = |Σ w·g| / max(Σ w², EPS_DIV) ──
|
||
// Analytically computed in f64 host-side then cast to f32 to bound
|
||
// accumulation order error. The kernel does the same reduction in
|
||
// f32 only over the union of sub-buffers.
|
||
let mut sum_wg: f64 = 0.0;
|
||
let mut sum_w2: f64 = 0.0;
|
||
for i in 0..union_total_count {
|
||
let w = union_params[i] as f64;
|
||
let g = union_grads[i] as f64;
|
||
sum_wg += w * g;
|
||
sum_w2 += w * w;
|
||
}
|
||
const EPS_DIV: f64 = 1.0e-8;
|
||
let true_wd_rate = (sum_wg.abs() / sum_w2.max(EPS_DIV)) as f32;
|
||
let kernel_wd_rate = scratch[SCRATCH_BASE_WD + g_idx];
|
||
let rel_wd = ((kernel_wd_rate - true_wd_rate) / true_wd_rate).abs();
|
||
println!(
|
||
"SP4 param_group[{g_idx}] WD_RATE — true={true_wd_rate:.6e}, kernel={kernel_wd_rate:.6e}, rel_err={rel_wd:.5}",
|
||
);
|
||
assert!(
|
||
kernel_wd_rate > 0.0,
|
||
"WD_RATE must be positive (signed grads × signed-aligned weights yields non-zero |Σ w·g|)",
|
||
);
|
||
assert!(
|
||
rel_wd < 0.02,
|
||
"group {g_idx} WD_RATE rel_err {rel_wd} > 2% (kernel={kernel_wd_rate}, true={true_wd_rate})",
|
||
);
|
||
|
||
// ── Pass E (group 0 only): L1_LAMBDA_TRUNK ──
|
||
// Analytical reference computes per-feature L2 norms of the
|
||
// [h_dim × k_in] gradient matrix, normalises them to a discrete
|
||
// distribution, then computes (mean|g| / max(mean|w|, EPS_DIV)) ×
|
||
// (log K − H) / log K. Mirrors the kernel's Pass E exactly.
|
||
if g_idx == 0 {
|
||
let k_in_us = k_in as usize;
|
||
let h_dim_us = h_dim as usize;
|
||
// Group 0 (DqnTrunk) is single-sub-buffer by construction; the
|
||
// trunk grad-w matrix sits in `sub_buffers[0]` and `union_*`
|
||
// are identical.
|
||
assert_eq!(sub_buffers.len(), 1);
|
||
assert_eq!(k_in_us * h_dim_us, union_total_count);
|
||
let grads_host = &sub_buffers[0].grads_host;
|
||
let params_host = &sub_buffers[0].params_host;
|
||
|
||
// Per-feature L2 norms (column-wise across the [h_dim × k_in]
|
||
// row-major grad matrix → feature `i` indexed at
|
||
// [h*k_in + i] for h ∈ 0..h_dim).
|
||
let mut g_feat: Vec<f64> = vec![0.0; k_in_us];
|
||
for h in 0..h_dim_us {
|
||
for i in 0..k_in_us {
|
||
let gv = grads_host[h * k_in_us + i] as f64;
|
||
g_feat[i] += gv * gv;
|
||
}
|
||
}
|
||
for v in g_feat.iter_mut() {
|
||
*v = v.sqrt();
|
||
}
|
||
let total: f64 = g_feat.iter().sum();
|
||
let mut deficit: f64 = 0.0;
|
||
if total > EPS_DIV && k_in_us > 1 {
|
||
let log_k = (k_in_us as f64).ln();
|
||
let mut entropy = 0.0_f64;
|
||
for &v in g_feat.iter() {
|
||
let p = v / total;
|
||
if p > EPS_DIV {
|
||
entropy -= p * p.ln();
|
||
}
|
||
}
|
||
deficit = ((log_k - entropy) / log_k).clamp(0.0, 1.0);
|
||
}
|
||
let total_count_f64 = union_total_count as f64;
|
||
let mean_abs_g: f64 =
|
||
grads_host.iter().map(|&g| (g as f64).abs()).sum::<f64>() / total_count_f64;
|
||
let mean_abs_w: f64 =
|
||
params_host.iter().map(|&w| (w as f64).abs()).sum::<f64>() / total_count_f64;
|
||
let scale = mean_abs_g / mean_abs_w.max(EPS_DIV);
|
||
let true_l1 = (scale * deficit) as f32;
|
||
|
||
let kernel_l1 = scratch[SCRATCH_L1_TRUNK];
|
||
// Both numerator and deficit can be small (near-uniform feature
|
||
// gradient distribution → deficit ≈ 0, l1 ≈ 0). Use the same
|
||
// 5% rel_err budget as the p99 family for non-degenerate cases;
|
||
// when the analytical reference hits the natural-zero floor of
|
||
// the formula, the kernel's `__threadfence_system()` still
|
||
// wrote 0 to the slot but the rel_err formula is undefined —
|
||
// assert kernel_l1 ≈ 0 directly in that branch.
|
||
println!(
|
||
"SP4 param_group[0] L1_LAMBDA_TRUNK — true={true_l1:.5e}, kernel={kernel_l1:.5e}, deficit={deficit:.5}",
|
||
);
|
||
if true_l1.abs() < 1e-6 {
|
||
assert!(
|
||
kernel_l1.abs() < 1e-3,
|
||
"group 0 L1 ≈ 0 expected, kernel returned {kernel_l1}",
|
||
);
|
||
} else {
|
||
let rel_l1 = ((kernel_l1 - true_l1) / true_l1).abs();
|
||
assert!(
|
||
rel_l1 < 0.05,
|
||
"group 0 L1 rel_err {rel_l1} > 5% (kernel={kernel_l1}, true={true_l1})",
|
||
);
|
||
}
|
||
}
|
||
}
|
||
|
||
// Final layout sanity — `weight_bound(g)` etc. correctly resolve to the
|
||
// ISV slots the host launcher reads/writes. The kernel-direct test does
|
||
// NOT touch ISV (no GpuDqnTrainer), so this is purely a const-fn drift
|
||
// guard mirroring Tasks A5/A6.
|
||
for g in 0..SP4_PARAM_GROUP_COUNT {
|
||
assert!(
|
||
weight_bound(g) >= 136 && weight_bound(g) < 144,
|
||
"weight_bound({g}) = {} out of [136, 144)", weight_bound(g),
|
||
);
|
||
assert!(adam_m_bound(g) >= 144 && adam_m_bound(g) < 152);
|
||
assert!(adam_v_bound(g) >= 152 && adam_v_bound(g) < 160);
|
||
assert!(wd_rate(g) >= 160 && wd_rate(g) < 168);
|
||
}
|
||
}
|
||
|
||
// ── SP4 Task A8: grad_norm_p99 producer (degenerate single-element) ───────────
|
||
|
||
/// Test-only cubin for the SP4 grad_norm_p99 producer kernel. The same cubin
|
||
/// is loaded by `GpuDqnTrainer::new` in production via `SP4_GRAD_NORM_P99_CUBIN`;
|
||
/// this kernel-direct test bypasses the full trainer harness and exercises
|
||
/// the kernel's degenerate single-thread copy contract.
|
||
const SP4_GRAD_NORM_P99_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/grad_norm_p99_kernel.cubin"));
|
||
|
||
/// Resolve the `grad_norm_p99_update` function handle from its cubin.
|
||
fn load_sp4_grad_norm_p99_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_GRAD_NORM_P99_CUBIN.to_vec())
|
||
.expect("load grad_norm_p99_kernel cubin");
|
||
module
|
||
.load_function("grad_norm_p99_update")
|
||
.expect("load grad_norm_p99_update function")
|
||
}
|
||
|
||
/// SP4 Task A8 unit test: the production `grad_norm_p99_update` kernel
|
||
/// performs a single-thread copy of `grad_norm_buf[0]` to
|
||
/// `producer_step_scratch_buf[scratch_idx=38]` and the host-side
|
||
/// `pearls_ad_update` helper (Task A3) replaces the ISV sentinel directly
|
||
/// with that step observation on the first call after a fold-boundary
|
||
/// reset (Pearl A bootstrap).
|
||
///
|
||
/// `grad_norm_buf` is a single-element f32 device buffer in production;
|
||
/// p99 over 1 element IS the element itself, so the kernel is degenerate
|
||
/// (no histogram). This test sets that single element to a known scalar
|
||
/// (42.0) and verifies the kernel writes exactly that value to slot 38
|
||
/// while leaving every other slot zero.
|
||
///
|
||
/// Production launcher contract validated:
|
||
/// - Kernel writes the input scalar to `scratch[SCRATCH_IDX_GRAD_CLIP=38]`.
|
||
/// - All other scratch slots remain zero (helper assertion guards the
|
||
/// `2fb30f098`-style write-cascade pattern).
|
||
/// - `GRAD_CLIP_BOUND_INDEX == 168` and `(168 - SP4_SLOT_BASE) * 3 == 111`
|
||
/// (Wiener state offset for this slot).
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_grad_norm_p99_pearl_a_first_observation_replaces_scalar() {
|
||
use ml::cuda_pipeline::sp4_isv_slots::{GRAD_CLIP_BOUND_INDEX, SP4_SLOT_BASE};
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
// Production scratch shape (47 = SP4 producers count). Slot 38 is the
|
||
// production-allocated GRAD_CLIP_BOUND slot (per the documented stable
|
||
// layout in `launch_sp4_target_q_p99` / `launch_sp4_grad_norm_p99`).
|
||
const SCRATCH_IDX_GRAD_CLIP: usize = 38;
|
||
|
||
// Single-element grad_norm scalar — known synthetic value chosen to be
|
||
// distinguishable from the Pearl-A sentinel (0.0) and from a stale
|
||
// initialisation (1.0) so a "did the kernel actually copy?" question
|
||
// is answered unambiguously.
|
||
let grad_norm_scalar: f32 = 42.0;
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_sp4_grad_norm_p99_kernel(&stream);
|
||
|
||
// Safety: CUDA context active on this thread (resolved via stream).
|
||
let in_buf = unsafe { MappedF32Buffer::new(1) }
|
||
.expect("alloc MappedF32Buffer for SP4 grad_norm_p99 input");
|
||
in_buf.write_from_slice(&[grad_norm_scalar]);
|
||
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc MappedF32Buffer for SP4 producer scratch");
|
||
|
||
let scratch_idx_arg: i32 = SCRATCH_IDX_GRAD_CLIP as i32;
|
||
let in_dev_ptr = in_buf.dev_ptr;
|
||
let scratch_dev_ptr = scratch_buf.dev_ptr;
|
||
|
||
// Single-thread degenerate kernel — minimal launch config, no shared
|
||
// memory. Mirrors the production `launch_sp4_grad_norm_p99` config.
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&in_dev_ptr)
|
||
.arg(&scratch_dev_ptr)
|
||
.arg(&scratch_idx_arg)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.expect("launch grad_norm_p99_update");
|
||
}
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after grad_norm_p99_update launch");
|
||
|
||
// Verify ALL non-target slots remained zero — the kernel must only
|
||
// touch the slot at `scratch_idx=38`.
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
if i != SCRATCH_IDX_GRAD_CLIP {
|
||
assert_eq!(
|
||
v, 0.0,
|
||
"grad_norm_p99 producer wrote to scratch[{i}] outside target slot {SCRATCH_IDX_GRAD_CLIP}: {v}",
|
||
);
|
||
}
|
||
}
|
||
let step_obs = host[SCRATCH_IDX_GRAD_CLIP];
|
||
println!(
|
||
"SP4 grad_norm_p99 — input={grad_norm_scalar:.5}, step_obs={step_obs:.5}",
|
||
);
|
||
// Degenerate copy — bit-identical equality (no histogram quantization).
|
||
assert_eq!(
|
||
step_obs, grad_norm_scalar,
|
||
"kernel must copy grad_norm_buf[0] verbatim to scratch[{SCRATCH_IDX_GRAD_CLIP}]; got {step_obs} vs {grad_norm_scalar}",
|
||
);
|
||
|
||
// ── Pearl A bootstrap (host-side, mirrors `launch_sp4_grad_norm_p99`) ──
|
||
// Sentinel ISV slot + zero Wiener state — replicates the cold-start
|
||
// condition immediately after a fold-boundary reset.
|
||
let prev_x_mean: f32 = 0.0;
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_obs);
|
||
|
||
assert_eq!(
|
||
new_x_mean, step_obs,
|
||
"Pearl A: first observation must replace sentinel directly; got {new_x_mean} vs step_obs={step_obs}",
|
||
);
|
||
assert_eq!(
|
||
state.x_lag, step_obs,
|
||
"Pearl A: x_lag must seed to first observation",
|
||
);
|
||
assert_eq!(state.sample_var, 0.0, "Pearl A: sample_var stays zero");
|
||
assert_eq!(state.diff_var, 0.0, "Pearl A: diff_var stays zero");
|
||
|
||
// Slot-layout invariants — guard against off-by-one drift in
|
||
// GRAD_CLIP_BOUND_INDEX or SP4_SLOT_BASE.
|
||
assert_eq!(GRAD_CLIP_BOUND_INDEX, 168);
|
||
assert_eq!(SP4_SLOT_BASE, 131);
|
||
assert_eq!((GRAD_CLIP_BOUND_INDEX - SP4_SLOT_BASE) * 3, 111);
|
||
}
|
||
|
||
// ── SP4 Task A9: h_s2_p99 producer for ISV[H_S2_BOUND=169] ───────────────────
|
||
|
||
/// Test-only cubin for the SP4 h_s2_p99 producer kernel. The same cubin is
|
||
/// loaded by `GpuDqnTrainer::new` in production via `SP4_H_S2_P99_CUBIN`;
|
||
/// this kernel-direct test bypasses the full trainer harness and exercises
|
||
/// the kernel's scratch-slot write contract on a controlled distribution.
|
||
const SP4_H_S2_P99_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/h_s2_p99_kernel.cubin"));
|
||
|
||
/// Resolve the `h_s2_p99_update` function handle from its cubin.
|
||
fn load_sp4_h_s2_p99_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_H_S2_P99_CUBIN.to_vec())
|
||
.expect("load h_s2_p99_kernel cubin");
|
||
module
|
||
.load_function("h_s2_p99_update")
|
||
.expect("load h_s2_p99_update function")
|
||
}
|
||
|
||
/// SP4 Task A9 unit test: the production `h_s2_p99_update` kernel writes
|
||
/// step_p99 to `producer_step_scratch_buf[scratch_idx=39]`, and the host-side
|
||
/// `pearls_ad_update` helper (Task A3) replaces the ISV sentinel directly
|
||
/// with that step observation on the first call after a fold-boundary
|
||
/// reset (Pearl A bootstrap).
|
||
///
|
||
/// Test distribution: 4096 deterministic |N(0,1)| samples (Box-Muller on
|
||
/// fixed-seed `StdRng`) — the same well-spread synthetic-input pattern
|
||
/// used by Tasks A4/A5/A6/A7. Real `save_h_s2` activations (post-GRN trunk
|
||
/// output of shape `[B, SH2]`) are similarly well-spread by the time the
|
||
/// kernel runs in production.
|
||
///
|
||
/// Analytical p99 of |N(0,1)| ≈ 2.576 (the one-tailed 99th-percentile
|
||
/// z-score). Tolerance: 5% relative — same headroom budget as Tasks
|
||
/// A4/A5/A6 (linear-bin quantization 0.4% + lane collision <0.012% +
|
||
/// finite-sample p99 jitter ≈0.5% at n=4096; total <2%).
|
||
///
|
||
/// Pearl A semantics tested explicitly:
|
||
/// - `prev_x_mean = 0.0` (fold-reset sentinel) and `WienerState::ZERO`
|
||
/// (untouched) → `pearls_ad_update` returns `step_observation`
|
||
/// directly and seeds `state.x_lag = step_observation`.
|
||
///
|
||
/// Production launcher contract validated:
|
||
/// - Kernel writes step_p99 to `scratch[SCRATCH_IDX_H_S2=39]`.
|
||
/// - All other scratch slots remain zero (helper checks this against the
|
||
/// `2fb30f098`-style write-cascade).
|
||
/// - `H_S2_BOUND_INDEX == 169`, `SP4_SLOT_BASE == 131`,
|
||
/// `(169 - SP4_SLOT_BASE) * 3 == 114` (Wiener state offset).
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_h_s2_p99_writes_step_p99_to_scratch_via_pearl_a() {
|
||
use ml::cuda_pipeline::sp4_isv_slots::{H_S2_BOUND_INDEX, SP4_SLOT_BASE};
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
use rand::rngs::StdRng;
|
||
use rand::{Rng, SeedableRng};
|
||
|
||
// 4096 |N(0,1)| samples via Box-Muller, fixed seed for determinism.
|
||
// Mirrors Tasks A4/A5/A7's distribution pattern; analytical p99 of
|
||
// |N(0,1)| ≈ 2.576. Well-spread across bins → no lane-collision
|
||
// pathology.
|
||
let mut rng = StdRng::seed_from_u64(0xA9_5A_F1_57);
|
||
let n: usize = 4096;
|
||
let samples: Vec<f32> = (0..n)
|
||
.map(|_| {
|
||
let u1: f32 = rng.gen_range(1e-8_f32..1.0_f32);
|
||
let u2: f32 = rng.gen_range(0.0_f32..1.0_f32);
|
||
let z = (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos();
|
||
z.abs()
|
||
})
|
||
.collect();
|
||
|
||
let mut sorted = samples.clone();
|
||
sorted.sort_by(|a, b| a.partial_cmp(b).expect("|N(0,1)| samples are finite"));
|
||
let true_p99 = sorted[(n * 99) / 100];
|
||
|
||
// Production scratch shape (47 = SP4 producers count). Slot 39 is the
|
||
// production-allocated H_S2_BOUND slot per the documented stable layout
|
||
// in `launch_sp4_target_q_p99` / `launch_sp4_h_s2_p99`.
|
||
const SCRATCH_IDX_H_S2: usize = 39;
|
||
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_sp4_h_s2_p99_kernel(&stream);
|
||
|
||
// Safety: CUDA context active on this thread (resolved via stream).
|
||
let in_buf = unsafe { MappedF32Buffer::new(samples.len()) }
|
||
.expect("alloc MappedF32Buffer for SP4 h_s2_p99 input");
|
||
in_buf.write_from_slice(&samples);
|
||
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc MappedF32Buffer for SP4 producer scratch");
|
||
|
||
let count_i32 = i32::try_from(samples.len()).expect("sample count fits in i32");
|
||
let scratch_idx_arg: i32 = SCRATCH_IDX_H_S2 as i32;
|
||
let in_dev_ptr = in_buf.dev_ptr;
|
||
let scratch_dev_ptr = scratch_buf.dev_ptr;
|
||
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&in_dev_ptr)
|
||
.arg(&count_i32)
|
||
.arg(&scratch_dev_ptr)
|
||
.arg(&scratch_idx_arg)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (256, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch h_s2_p99_update");
|
||
}
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after h_s2_p99_update launch");
|
||
|
||
// Verify ALL non-target slots remained zero — guards the
|
||
// `2fb30f098`-style write-cascade pattern.
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
if i != SCRATCH_IDX_H_S2 {
|
||
assert_eq!(
|
||
v, 0.0,
|
||
"h_s2_p99 producer wrote to scratch[{i}] outside target slot {SCRATCH_IDX_H_S2}: {v}",
|
||
);
|
||
}
|
||
}
|
||
let step_p99 = host[SCRATCH_IDX_H_S2];
|
||
let rel_err = ((step_p99 - true_p99) / true_p99).abs();
|
||
println!(
|
||
"SP4 h_s2_p99 — true_p99={true_p99:.5}, step_p99 (kernel)={step_p99:.5}, rel_err={rel_err:.5}",
|
||
);
|
||
assert!(
|
||
step_p99 > 0.0,
|
||
"kernel step_p99 (slot {SCRATCH_IDX_H_S2}) should be positive; got {step_p99}",
|
||
);
|
||
assert!(
|
||
rel_err < 0.05,
|
||
"kernel step_p99 ({step_p99}) vs true p99 ({true_p99}) rel_err {rel_err} > 5%",
|
||
);
|
||
|
||
// ── Pearl A bootstrap (host-side, mirrors `launch_sp4_h_s2_p99`) ──
|
||
let prev_x_mean: f32 = 0.0;
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_p99);
|
||
|
||
assert_eq!(
|
||
new_x_mean, step_p99,
|
||
"Pearl A: first observation must replace sentinel directly; got {new_x_mean} vs step_p99={step_p99}",
|
||
);
|
||
assert_eq!(
|
||
state.x_lag, step_p99,
|
||
"Pearl A: x_lag must seed to first observation",
|
||
);
|
||
assert_eq!(state.sample_var, 0.0, "Pearl A: sample_var stays zero");
|
||
assert_eq!(state.diff_var, 0.0, "Pearl A: diff_var stays zero");
|
||
|
||
// Slot-layout invariants — guard against off-by-one drift in
|
||
// H_S2_BOUND_INDEX or SP4_SLOT_BASE.
|
||
assert_eq!(H_S2_BOUND_INDEX, 169);
|
||
assert_eq!(SP4_SLOT_BASE, 131);
|
||
assert_eq!((H_S2_BOUND_INDEX - SP4_SLOT_BASE) * 3, 114);
|
||
}
|
||
|
||
// ── SP4 Task A13.0: h_s2_rms_ema retrofit (Pearls A+D) ───────────────────────
|
||
|
||
/// Test-only cubin for the SP4 Task A13.0 retrofit `h_s2_rms_ema_update`
|
||
/// kernel. Same cubin as production (`H_S2_RMS_EMA_CUBIN`); this kernel-direct
|
||
/// test bypasses the trainer harness and exercises the kernel's scratch-slot
|
||
/// write contract on a controlled stationary signal.
|
||
const SP4_H_S2_RMS_EMA_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/h_s2_rms_ema_kernel.cubin"));
|
||
|
||
fn load_h_s2_rms_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_H_S2_RMS_EMA_CUBIN.to_vec())
|
||
.expect("load h_s2_rms_ema_kernel cubin");
|
||
module
|
||
.load_function("h_s2_rms_ema_update")
|
||
.expect("load h_s2_rms_ema_update function")
|
||
}
|
||
|
||
/// SP4 Task A13.0: retrofit `h_s2_rms_ema_update` to write the per-step
|
||
/// `step_observation = sqrt(mean(h_s2²))` into `producer_step_scratch[40]`,
|
||
/// then verify Pearls A+D bootstrap (Pearl A) + stationary-convergence
|
||
/// (Pearl D) host-side using the production `pearls_ad_update` helper.
|
||
///
|
||
/// Stationary signal → constant `h_s2 = 5.0` over `B*SH2 = 4*64 = 256`
|
||
/// floats → analytical RMS = 5.0. After 1000 stationary observations the
|
||
/// converged ISV value must be within 1% of 5.0; first observation must
|
||
/// equal RMS (Pearl A bypasses Pearl D's variance-ratio at t=0).
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_h_s2_rms_ema_writes_step_rms_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const B: usize = 4;
|
||
const SH2: usize = 64;
|
||
const N: usize = B * SH2;
|
||
const SCRATCH_IDX: usize = 40;
|
||
const BLOCK_DIM: u32 = 256;
|
||
const SHARED_BYTES: u32 = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
|
||
|
||
let stationary_signal: f32 = 5.0;
|
||
let samples: Vec<f32> = vec![stationary_signal; N];
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_h_s2_rms_ema_kernel(&stream);
|
||
|
||
// Safety: CUDA context active on this thread.
|
||
let in_buf = unsafe { MappedF32Buffer::new(N) }
|
||
.expect("alloc h_s2 input buffer");
|
||
in_buf.write_from_slice(&samples);
|
||
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc producer scratch buffer");
|
||
|
||
let b_i32 = B as i32;
|
||
let sh2_i32 = SH2 as i32;
|
||
let scratch_idx_arg: i32 = SCRATCH_IDX as i32;
|
||
let in_dev_ptr = in_buf.dev_ptr;
|
||
let scratch_dev_ptr = scratch_buf.dev_ptr;
|
||
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&in_dev_ptr)
|
||
.arg(&b_i32)
|
||
.arg(&sh2_i32)
|
||
.arg(&scratch_dev_ptr)
|
||
.arg(&scratch_idx_arg)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (BLOCK_DIM, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch h_s2_rms_ema_update");
|
||
}
|
||
stream
|
||
.synchronize()
|
||
.expect("sync after h_s2_rms_ema_update launch");
|
||
|
||
// Verify ALL non-target slots remained zero — guards against write-cascade.
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
if i != SCRATCH_IDX {
|
||
assert_eq!(
|
||
v, 0.0,
|
||
"h_s2_rms_ema wrote to scratch[{i}] outside target slot {SCRATCH_IDX}: {v}",
|
||
);
|
||
}
|
||
}
|
||
let step_rms = host[SCRATCH_IDX];
|
||
let expected_rms = stationary_signal; // sqrt(mean(5²)) = 5
|
||
assert!(
|
||
(step_rms - expected_rms).abs() < 1e-4,
|
||
"kernel step_rms {step_rms} vs analytical {expected_rms}",
|
||
);
|
||
|
||
// ── Pearl A bootstrap: first observation replaces sentinel directly ──
|
||
let prev_x_mean: f32 = 0.0;
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_rms);
|
||
assert_eq!(
|
||
new_x_mean, step_rms,
|
||
"Pearl A: first observation must replace sentinel, got {new_x_mean}",
|
||
);
|
||
assert_eq!(state.x_lag, step_rms, "Pearl A seeds x_lag");
|
||
|
||
// ── Pearl D convergence: 1000 stationary observations of step_rms ──
|
||
let mut x_mean = new_x_mean;
|
||
for _ in 0..1000 {
|
||
x_mean = pearls_ad_update(x_mean, &mut state, step_rms);
|
||
}
|
||
let rel_err = ((x_mean - expected_rms) / expected_rms).abs();
|
||
assert!(
|
||
rel_err < 0.01,
|
||
"Pearl D: stationary signal converged to {x_mean} (expected {expected_rms}, rel_err {rel_err})",
|
||
);
|
||
}
|
||
|
||
// ── SP4 Task A13.1: aux_heads_loss + aux_label_scale retrofit ────────────────
|
||
|
||
const SP4_AUX_HEADS_LOSS_EMA_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/aux_heads_loss_ema_kernel.cubin"));
|
||
|
||
fn load_aux_heads_loss_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_AUX_HEADS_LOSS_EMA_CUBIN.to_vec())
|
||
.expect("load aux_heads_loss_ema cubin");
|
||
module
|
||
.load_function("aux_heads_loss_ema_update")
|
||
.expect("load aux_heads_loss_ema_update function")
|
||
}
|
||
|
||
fn load_aux_label_scale_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_AUX_HEADS_LOSS_EMA_CUBIN.to_vec())
|
||
.expect("load aux_heads_loss_ema cubin (label scale)");
|
||
module
|
||
.load_function("aux_label_scale_ema_update")
|
||
.expect("load aux_label_scale_ema_update function")
|
||
}
|
||
|
||
/// SP4 Task A13.1: `aux_heads_loss_ema_update` writes both scalar loss
|
||
/// values to scratch[41]/[42], and Pearls A+D bootstrap (Pearl A) replaces
|
||
/// the ISV sentinels with those step observations on first call.
|
||
/// Stationary signals (`nb_loss=0.5`, `rg_loss=1.2`) converge to the same
|
||
/// values within 1% after 1000 observations.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_aux_heads_loss_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const SCRATCH_IDX_NEXT_BAR: usize = 41;
|
||
const SCRATCH_IDX_REGIME: usize = 42;
|
||
|
||
let nb_loss: f32 = 0.5;
|
||
let rg_loss: f32 = 1.2;
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_aux_heads_loss_ema_kernel(&stream);
|
||
|
||
let nb_buf = unsafe { MappedF32Buffer::new(1) }.expect("alloc nb buf");
|
||
nb_buf.write_from_slice(&[nb_loss]);
|
||
let rg_buf = unsafe { MappedF32Buffer::new(1) }.expect("alloc rg buf");
|
||
rg_buf.write_from_slice(&[rg_loss]);
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc scratch");
|
||
|
||
let nb_dev = nb_buf.dev_ptr;
|
||
let rg_dev = rg_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let nb_idx: i32 = SCRATCH_IDX_NEXT_BAR as i32;
|
||
let rg_idx: i32 = SCRATCH_IDX_REGIME as i32;
|
||
|
||
unsafe {
|
||
stream.launch_builder(&kernel)
|
||
.arg(&nb_dev)
|
||
.arg(&rg_dev)
|
||
.arg(&scratch_dev)
|
||
.arg(&nb_idx)
|
||
.arg(&rg_idx)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.expect("launch aux_heads_loss_ema_update");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
if i != SCRATCH_IDX_NEXT_BAR && i != SCRATCH_IDX_REGIME {
|
||
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
|
||
}
|
||
}
|
||
assert!((host[SCRATCH_IDX_NEXT_BAR] - nb_loss).abs() < 1e-6);
|
||
assert!((host[SCRATCH_IDX_REGIME] - rg_loss).abs() < 1e-6);
|
||
|
||
// Pearl A bootstrap for both slots.
|
||
let mut nb_state = WienerState::ZERO;
|
||
let mut rg_state = WienerState::ZERO;
|
||
let nb_x = pearls_ad_update(0.0, &mut nb_state, host[SCRATCH_IDX_NEXT_BAR]);
|
||
let rg_x = pearls_ad_update(0.0, &mut rg_state, host[SCRATCH_IDX_REGIME]);
|
||
assert_eq!(nb_x, nb_loss);
|
||
assert_eq!(rg_x, rg_loss);
|
||
assert_eq!(nb_state.x_lag, nb_loss);
|
||
assert_eq!(rg_state.x_lag, rg_loss);
|
||
|
||
// Pearl D convergence: 1000 stationary observations.
|
||
let mut nb_mean = nb_x;
|
||
let mut rg_mean = rg_x;
|
||
for _ in 0..1000 {
|
||
nb_mean = pearls_ad_update(nb_mean, &mut nb_state, nb_loss);
|
||
rg_mean = pearls_ad_update(rg_mean, &mut rg_state, rg_loss);
|
||
}
|
||
assert!(((nb_mean - nb_loss) / nb_loss).abs() < 0.01,
|
||
"nb stationary should converge: got {nb_mean}, expected {nb_loss}");
|
||
assert!(((rg_mean - rg_loss) / rg_loss).abs() < 0.01,
|
||
"rg stationary should converge: got {rg_mean}, expected {rg_loss}");
|
||
}
|
||
|
||
/// SP4 Task A13.1: `aux_label_scale_ema_update` writes mean(|label|) to
|
||
/// scratch[43], and Pearls A+D bootstrap+convergence behave correctly on
|
||
/// stationary input.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_aux_label_scale_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const SCRATCH_IDX: usize = 43;
|
||
const B: usize = 256;
|
||
const BLOCK_DIM: u32 = 256;
|
||
const SHARED_BYTES: u32 = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
|
||
|
||
// Stationary mixed-sign labels with absolute mean = 3.0.
|
||
let labels: Vec<f32> = (0..B)
|
||
.map(|i| if i % 2 == 0 { 3.0 } else { -3.0 })
|
||
.collect();
|
||
let expected_mean_abs: f32 = 3.0;
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_aux_label_scale_ema_kernel(&stream);
|
||
|
||
let label_buf = unsafe { MappedF32Buffer::new(B) }.expect("alloc label buf");
|
||
label_buf.write_from_slice(&labels);
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc scratch");
|
||
|
||
let label_dev = label_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let b_i32 = B as i32;
|
||
let scratch_idx_arg: i32 = SCRATCH_IDX as i32;
|
||
|
||
unsafe {
|
||
stream.launch_builder(&kernel)
|
||
.arg(&label_dev)
|
||
.arg(&b_i32)
|
||
.arg(&scratch_dev)
|
||
.arg(&scratch_idx_arg)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (BLOCK_DIM, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch aux_label_scale_ema_update");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
if i != SCRATCH_IDX {
|
||
assert_eq!(v, 0.0, "wrote outside target slot: scratch[{i}]={v}");
|
||
}
|
||
}
|
||
let step_obs = host[SCRATCH_IDX];
|
||
assert!(
|
||
(step_obs - expected_mean_abs).abs() < 1e-4,
|
||
"step_obs {step_obs} vs expected {expected_mean_abs}",
|
||
);
|
||
|
||
// Pearl A bootstrap.
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(0.0, &mut state, step_obs);
|
||
assert_eq!(new_x_mean, step_obs);
|
||
assert_eq!(state.x_lag, step_obs);
|
||
|
||
// Pearl D convergence.
|
||
let mut x_mean = new_x_mean;
|
||
for _ in 0..1000 {
|
||
x_mean = pearls_ad_update(x_mean, &mut state, step_obs);
|
||
}
|
||
assert!(((x_mean - expected_mean_abs) / expected_mean_abs).abs() < 0.01);
|
||
}
|
||
|
||
// ── SP4 Task A13.2: moe_expert_util_ema retrofit ─────────────────────────────
|
||
|
||
const SP4_MOE_KERNELS_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/moe_kernels.cubin"));
|
||
|
||
fn load_moe_expert_util_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_MOE_KERNELS_CUBIN.to_vec())
|
||
.expect("load moe_kernels cubin");
|
||
module
|
||
.load_function("moe_expert_util_ema_update")
|
||
.expect("load moe_expert_util_ema_update function")
|
||
}
|
||
|
||
/// SP4 Task A13.2: `moe_expert_util_ema_update` writes K col_means
|
||
/// to scratch[44..52) plus gate Shannon entropy to scratch[52]. Pearls A+D
|
||
/// bootstrap + convergence verified host-side via `pearls_ad_update`.
|
||
///
|
||
/// Test gate: B=128 samples, K=8 experts, perfect uniform softmax →
|
||
/// each col_mean = 1/8, entropy = ln(8) ≈ 2.0794.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_moe_expert_util_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const B: usize = 128;
|
||
const K: usize = 8;
|
||
const SCRATCH_IDX_UTIL_BASE: usize = 44;
|
||
const SCRATCH_IDX_ENTROPY: usize = 52;
|
||
|
||
// Uniform gate: each row sums to 1, all entries equal 1/K.
|
||
let uniform: f32 = 1.0 / K as f32;
|
||
let gate: Vec<f32> = vec![uniform; B * K];
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_moe_expert_util_ema_kernel(&stream);
|
||
|
||
let gate_buf = unsafe { MappedF32Buffer::new(B * K) }.expect("alloc gate");
|
||
gate_buf.write_from_slice(&gate);
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc scratch");
|
||
|
||
let gate_dev = gate_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let b_i32 = B as i32;
|
||
let k_i32 = K as i32;
|
||
let util_base_i32 = SCRATCH_IDX_UTIL_BASE as i32;
|
||
let entropy_idx_i32 = SCRATCH_IDX_ENTROPY as i32;
|
||
|
||
unsafe {
|
||
stream.launch_builder(&kernel)
|
||
.arg(&gate_dev)
|
||
.arg(&scratch_dev)
|
||
.arg(&b_i32)
|
||
.arg(&k_i32)
|
||
.arg(&util_base_i32)
|
||
.arg(&entropy_idx_i32)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.expect("launch moe_expert_util_ema_update");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
|
||
let host = scratch_buf.read_all();
|
||
|
||
// Verify only target slots written.
|
||
for (i, &v) in host.iter().enumerate() {
|
||
let in_util = i >= SCRATCH_IDX_UTIL_BASE && i < SCRATCH_IDX_UTIL_BASE + K;
|
||
let is_entropy = i == SCRATCH_IDX_ENTROPY;
|
||
if !in_util && !is_entropy {
|
||
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
|
||
}
|
||
}
|
||
|
||
// Each col_mean should equal 1/8.
|
||
for k in 0..K {
|
||
let col_mean = host[SCRATCH_IDX_UTIL_BASE + k];
|
||
assert!(
|
||
(col_mean - uniform).abs() < 1e-6,
|
||
"expert {k} col_mean {col_mean} vs expected {uniform}",
|
||
);
|
||
}
|
||
|
||
// Entropy = ln(K).
|
||
let entropy = host[SCRATCH_IDX_ENTROPY];
|
||
let expected_entropy: f32 = (K as f32).ln();
|
||
assert!(
|
||
(entropy - expected_entropy).abs() < 1e-5,
|
||
"entropy {entropy} vs expected {expected_entropy}",
|
||
);
|
||
|
||
// Pearl A bootstrap on the entropy slot (representative — same logic
|
||
// for util slots).
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(0.0, &mut state, entropy);
|
||
assert_eq!(new_x_mean, entropy, "Pearl A returns step_obs directly");
|
||
assert_eq!(state.x_lag, entropy);
|
||
|
||
// Pearl D convergence — 1000 stationary observations.
|
||
let mut x_mean = new_x_mean;
|
||
for _ in 0..1000 {
|
||
x_mean = pearls_ad_update(x_mean, &mut state, entropy);
|
||
}
|
||
assert!(
|
||
((x_mean - expected_entropy) / expected_entropy).abs() < 0.01,
|
||
"Pearl D: stationary entropy converged to {x_mean} (expected {expected_entropy})",
|
||
);
|
||
}
|
||
|
||
// ── SP4 Task A13.3: vsn_mask_ema retrofit ────────────────────────────────────
|
||
|
||
const SP4_VSN_MASK_EMA_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/vsn_mask_ema_kernel.cubin"));
|
||
|
||
fn load_vsn_mask_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_VSN_MASK_EMA_CUBIN.to_vec())
|
||
.expect("load vsn_mask_ema cubin");
|
||
module
|
||
.load_function("vsn_mask_ema_update")
|
||
.expect("load vsn_mask_ema_update function")
|
||
}
|
||
|
||
/// SP4 Task A13.3: `vsn_mask_ema_update` writes per-group mean to
|
||
/// scratch[53..59) (6 groups). Pearls A+D bootstrap+convergence verified.
|
||
///
|
||
/// Test mask: B=128 samples × num_groups=6, where mask[b,g] = (g+1)/21
|
||
/// so each row sums to 1 and each group's mean = (g+1)/21.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_vsn_mask_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const B: usize = 128;
|
||
const NUM_GROUPS: usize = 6;
|
||
const SCRATCH_FIRST: usize = 53;
|
||
const BLOCK_DIM: u32 = 256;
|
||
const SHARED_BYTES: u32 = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
|
||
|
||
// mask[b, g] = (g+1)/21 → rows sum = (1+2+3+4+5+6)/21 = 21/21 = 1.
|
||
// Group g's mean = (g+1)/21.
|
||
let mask: Vec<f32> = (0..B)
|
||
.flat_map(|_| (0..NUM_GROUPS).map(|g| (g as f32 + 1.0) / 21.0))
|
||
.collect();
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_vsn_mask_ema_kernel(&stream);
|
||
|
||
let mask_buf = unsafe { MappedF32Buffer::new(B * NUM_GROUPS) }
|
||
.expect("alloc mask buf");
|
||
mask_buf.write_from_slice(&mask);
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc scratch");
|
||
|
||
let mask_dev = mask_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let b_i32 = B as i32;
|
||
let num_groups_i32 = NUM_GROUPS as i32;
|
||
let scratch_first_i32 = SCRATCH_FIRST as i32;
|
||
|
||
unsafe {
|
||
stream.launch_builder(&kernel)
|
||
.arg(&mask_dev)
|
||
.arg(&b_i32)
|
||
.arg(&num_groups_i32)
|
||
.arg(&scratch_dev)
|
||
.arg(&scratch_first_i32)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (BLOCK_DIM, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch vsn_mask_ema_update");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
let in_target = i >= SCRATCH_FIRST && i < SCRATCH_FIRST + NUM_GROUPS;
|
||
if !in_target {
|
||
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
|
||
}
|
||
}
|
||
for g in 0..NUM_GROUPS {
|
||
let actual = host[SCRATCH_FIRST + g];
|
||
let expected = (g as f32 + 1.0) / 21.0;
|
||
assert!(
|
||
(actual - expected).abs() < 1e-5,
|
||
"group {g} mean {actual} vs expected {expected}",
|
||
);
|
||
}
|
||
|
||
// Pearl A bootstrap on group 0 (representative — same logic per group).
|
||
let group_0_obs = host[SCRATCH_FIRST];
|
||
let expected_group_0: f32 = 1.0 / 21.0;
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(0.0, &mut state, group_0_obs);
|
||
assert_eq!(new_x_mean, group_0_obs);
|
||
assert_eq!(state.x_lag, group_0_obs);
|
||
|
||
// Pearl D convergence.
|
||
let mut x_mean = new_x_mean;
|
||
for _ in 0..1000 {
|
||
x_mean = pearls_ad_update(x_mean, &mut state, group_0_obs);
|
||
}
|
||
assert!(
|
||
((x_mean - expected_group_0) / expected_group_0).abs() < 0.01,
|
||
"Pearl D convergence: got {x_mean}, expected {expected_group_0}",
|
||
);
|
||
}
|
||
|
||
// ── SP4 Task A13.4: iqn_quantile_ema retrofit ────────────────────────────────
|
||
|
||
const SP4_IQN_QUANTILE_EMA_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/iqn_quantile_ema_kernel.cubin"));
|
||
|
||
fn load_iqn_quantile_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_IQN_QUANTILE_EMA_CUBIN.to_vec())
|
||
.expect("load iqn_quantile_ema cubin");
|
||
module
|
||
.load_function("iqn_quantile_ema_update")
|
||
.expect("load iqn_quantile_ema_update function")
|
||
}
|
||
|
||
/// SP4 Task A13.4: `iqn_quantile_ema_update` writes 4 mean-|Q| step
|
||
/// observations to scratch[59..63) (one per non-median quantile). Pearls
|
||
/// A+D bootstrap+convergence verified host-side.
|
||
///
|
||
/// Test surface: `save_q_online [TBA, B*Q]` col-major where
|
||
/// Q[a, (b*Q + tau_idx)] = (tau_idx + 1) × 0.7. Each tau's mean|Q|
|
||
/// matches its (tau_idx + 1) × 0.7 setpoint. Skipped median (tau_idx=2)
|
||
/// is consequently ignored by the kernel.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_iqn_quantile_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const B: usize = 32;
|
||
const Q: usize = 5;
|
||
const TBA: usize = 12;
|
||
const SCRATCH_FIRST: usize = 59;
|
||
const BLOCK_DIM: u32 = 256;
|
||
const SHARED_BYTES: u32 = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
|
||
|
||
// save_q_online[a + (b*Q + tau_idx)*TBA] = (tau_idx+1) * 0.7
|
||
let mut q_data = vec![0.0_f32; TBA * B * Q];
|
||
for b in 0..B {
|
||
for tau in 0..Q {
|
||
let val = (tau as f32 + 1.0) * 0.7;
|
||
for a in 0..TBA {
|
||
let col = b * Q + tau;
|
||
let flat = a + col * TBA;
|
||
q_data[flat] = val;
|
||
}
|
||
}
|
||
}
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_iqn_quantile_ema_kernel(&stream);
|
||
|
||
let q_buf = unsafe { MappedF32Buffer::new(q_data.len()) }.expect("alloc q buf");
|
||
q_buf.write_from_slice(&q_data);
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc scratch");
|
||
|
||
let q_dev = q_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let b_i32 = B as i32;
|
||
let q_i32 = Q as i32;
|
||
let tba_i32 = TBA as i32;
|
||
let scratch_first_i32 = SCRATCH_FIRST as i32;
|
||
|
||
unsafe {
|
||
stream.launch_builder(&kernel)
|
||
.arg(&q_dev)
|
||
.arg(&b_i32)
|
||
.arg(&q_i32)
|
||
.arg(&tba_i32)
|
||
.arg(&scratch_dev)
|
||
.arg(&scratch_first_i32)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (4, 1, 1),
|
||
block_dim: (BLOCK_DIM, 1, 1),
|
||
shared_mem_bytes: SHARED_BYTES,
|
||
})
|
||
.expect("launch iqn_quantile_ema_update");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
let in_target = i >= SCRATCH_FIRST && i < SCRATCH_FIRST + 4;
|
||
if !in_target {
|
||
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
|
||
}
|
||
}
|
||
|
||
// Slot order: +0 → p05 (tau_idx=0, expected 0.7),
|
||
// +1 → p25 (tau_idx=1, expected 1.4),
|
||
// +2 → p75 (tau_idx=3, expected 2.8),
|
||
// +3 → p95 (tau_idx=4, expected 3.5).
|
||
let expected: [f32; 4] = [0.7, 1.4, 2.8, 3.5];
|
||
for (off, &exp) in expected.iter().enumerate() {
|
||
let actual = host[SCRATCH_FIRST + off];
|
||
assert!(
|
||
(actual - exp).abs() < 1e-5,
|
||
"slot offset {off}: got {actual}, expected {exp}",
|
||
);
|
||
}
|
||
|
||
// Pearl A bootstrap on slot +0 (p05).
|
||
let p05_obs = host[SCRATCH_FIRST];
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(0.0, &mut state, p05_obs);
|
||
assert_eq!(new_x_mean, p05_obs);
|
||
assert_eq!(state.x_lag, p05_obs);
|
||
|
||
// Pearl D convergence on the same slot.
|
||
let mut x_mean = new_x_mean;
|
||
for _ in 0..1000 {
|
||
x_mean = pearls_ad_update(x_mean, &mut state, p05_obs);
|
||
}
|
||
assert!(
|
||
((x_mean - 0.7) / 0.7).abs() < 0.01,
|
||
"Pearl D: stationary p05 converged to {x_mean} (expected 0.7)",
|
||
);
|
||
}
|
||
|
||
// ── SP4 Task A13.5: reward_component_ema retrofit + cross-boundary wiring ────
|
||
|
||
const SP4_REWARD_COMPONENT_EMA_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/reward_component_ema_kernel.cubin"));
|
||
|
||
fn load_reward_component_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(SP4_REWARD_COMPONENT_EMA_CUBIN.to_vec())
|
||
.expect("load reward_component_ema cubin");
|
||
module
|
||
.load_function("reward_component_ema")
|
||
.expect("load reward_component_ema function")
|
||
}
|
||
|
||
/// SP4 Task A13.5: `reward_component_ema` writes 6 mean-|r_c| step
|
||
/// observations to scratch[63..69) (one per reward component). Pearls A+D
|
||
/// bootstrap+convergence verified host-side via `pearls_ad_update`. The
|
||
/// production cross-boundary wiring (FusedTrainingCtx →
|
||
/// GpuExperienceCollector::set_reward_component_pearls_buffers) is
|
||
/// exercised by the integration smoke harness; this kernel-direct test
|
||
/// validates the kernel signature retrofit + Pearls A+D semantics in
|
||
/// isolation.
|
||
///
|
||
/// Test surface: n_samples=128 with reward_components[i*6 + c] = c+1
|
||
/// (constant per component, mixed sign on alternating samples). Each
|
||
/// component's mean|r_c| = c+1.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_reward_component_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
const N: usize = 128;
|
||
const NUM_COMPONENTS: usize = 6;
|
||
const SCRATCH_FIRST: usize = 63;
|
||
|
||
// reward_components[i*6 + c] = (c+1) with sign flip on odd i, so
|
||
// mean(|r_c|) = c+1 ∈ {1, 2, 3, 4, 5, 6}.
|
||
let rewards: Vec<f32> = (0..N)
|
||
.flat_map(|i| {
|
||
let sign: f32 = if i % 2 == 0 { 1.0 } else { -1.0 };
|
||
(0..NUM_COMPONENTS).map(move |c| sign * (c as f32 + 1.0))
|
||
})
|
||
.collect();
|
||
|
||
let stream = make_test_stream();
|
||
let kernel = load_reward_component_ema_kernel(&stream);
|
||
|
||
let rewards_buf = unsafe { MappedF32Buffer::new(rewards.len()) }
|
||
.expect("alloc rewards buf");
|
||
rewards_buf.write_from_slice(&rewards);
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
|
||
.expect("alloc scratch");
|
||
|
||
let rewards_dev = rewards_buf.dev_ptr;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let n_i32 = N as i32;
|
||
let scratch_first_i32 = SCRATCH_FIRST as i32;
|
||
|
||
unsafe {
|
||
stream.launch_builder(&kernel)
|
||
.arg(&rewards_dev)
|
||
.arg(&n_i32)
|
||
.arg(&scratch_dev)
|
||
.arg(&scratch_first_i32)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (NUM_COMPONENTS as u32, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.expect("launch reward_component_ema");
|
||
}
|
||
stream.synchronize().expect("sync");
|
||
|
||
let host = scratch_buf.read_all();
|
||
for (i, &v) in host.iter().enumerate() {
|
||
let in_target = i >= SCRATCH_FIRST && i < SCRATCH_FIRST + NUM_COMPONENTS;
|
||
if !in_target {
|
||
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
|
||
}
|
||
}
|
||
for c in 0..NUM_COMPONENTS {
|
||
let actual = host[SCRATCH_FIRST + c];
|
||
let expected = c as f32 + 1.0;
|
||
assert!(
|
||
(actual - expected).abs() < 1e-5,
|
||
"component {c} mean|r| {actual} vs expected {expected}",
|
||
);
|
||
}
|
||
|
||
// Pearl A bootstrap on component 0 (representative — same logic per
|
||
// component).
|
||
let c0_obs = host[SCRATCH_FIRST];
|
||
let mut state = WienerState::ZERO;
|
||
let new_x_mean = pearls_ad_update(0.0, &mut state, c0_obs);
|
||
assert_eq!(new_x_mean, c0_obs);
|
||
assert_eq!(state.x_lag, c0_obs);
|
||
|
||
// Pearl D convergence on component 0.
|
||
let mut x_mean = new_x_mean;
|
||
for _ in 0..1000 {
|
||
x_mean = pearls_ad_update(x_mean, &mut state, c0_obs);
|
||
}
|
||
assert!(
|
||
((x_mean - 1.0) / 1.0).abs() < 0.01,
|
||
"Pearl D: stationary component 0 converged to {x_mean} (expected 1.0)",
|
||
);
|
||
}
|
||
|
||
// ── GPU-Pearls A+D applicator kernel oracle test (2026-05-01) ────────────────
|
||
|
||
/// Test-only cubin for the GPU-Pearls A+D applicator. Same cubin loaded by
|
||
/// `GpuDqnTrainer::new` in production via `APPLY_PEARLS_AD_CUBIN`; this
|
||
/// kernel-direct test bypasses the full trainer harness and exercises the
|
||
/// kernel's per-slot Pearl-A-or-Pearl-D writeback contract on representative
|
||
/// trajectories, asserting bit-equivalent agreement (within fp32 ULP) with
|
||
/// the host-side `pearls_ad_update` reference implementation in
|
||
/// `crates/ml/src/cuda_pipeline/sp4_wiener_ema.rs`.
|
||
const APPLY_PEARLS_AD_TEST_CUBIN: &[u8] =
|
||
include_bytes!(concat!(env!("OUT_DIR"), "/apply_pearls_kernel.cubin"));
|
||
|
||
/// Resolve the `apply_pearls_ad_kernel` function handle from its cubin.
|
||
fn load_apply_pearls_ad_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
|
||
let module = stream
|
||
.context()
|
||
.load_cubin(APPLY_PEARLS_AD_TEST_CUBIN.to_vec())
|
||
.expect("load apply_pearls_kernel cubin");
|
||
module
|
||
.load_function("apply_pearls_ad_kernel")
|
||
.expect("load apply_pearls_ad_kernel function")
|
||
}
|
||
|
||
/// Run the GPU `apply_pearls_ad_kernel` on `(scratch, isv_init, wiener_init)`
|
||
/// for `n_slots` consecutive slots and return `(isv_after, wiener_after)`.
|
||
/// All inputs/outputs ride on a single mapped-pinned allocation per buffer
|
||
/// (zero-copy, no DtoH/HtoD per `feedback_no_htod_htoh_only_mapped_pinned`).
|
||
/// Single-slot per ISV/scratch index (`*_idx_base = 0`); wiener offset 0.
|
||
fn launch_apply_pearls_ad_oracle(
|
||
stream: &Arc<CudaStream>,
|
||
scratch: &[f32],
|
||
isv_init: &[f32],
|
||
wiener_init: &[f32],
|
||
) -> (Vec<f32>, Vec<f32>) {
|
||
let n = scratch.len();
|
||
assert_eq!(isv_init.len(), n);
|
||
assert_eq!(wiener_init.len(), n * 3);
|
||
let kernel = load_apply_pearls_ad_kernel(stream);
|
||
|
||
// Allocate mapped-pinned buffers; copy host-side seeds in via
|
||
// write_from_slice (HtoH on mapped-pinned is allowed per
|
||
// feedback_no_htod_htoh_only_mapped_pinned.md).
|
||
let scratch_buf = unsafe { MappedF32Buffer::new(n) }
|
||
.expect("alloc scratch_buf");
|
||
scratch_buf.write_from_slice(scratch);
|
||
let isv_buf = unsafe { MappedF32Buffer::new(n) }
|
||
.expect("alloc isv_buf");
|
||
isv_buf.write_from_slice(isv_init);
|
||
let wiener_buf = unsafe { MappedF32Buffer::new(n * 3) }
|
||
.expect("alloc wiener_buf");
|
||
wiener_buf.write_from_slice(wiener_init);
|
||
|
||
let scratch_idx_base: i32 = 0;
|
||
let isv_idx_base: i32 = 0;
|
||
let wiener_offset_base: i32 = 0;
|
||
let n_slots: i32 = n as i32;
|
||
let scratch_dev = scratch_buf.dev_ptr;
|
||
let isv_dev = isv_buf.dev_ptr;
|
||
let wiener_dev = wiener_buf.dev_ptr;
|
||
|
||
unsafe {
|
||
stream
|
||
.launch_builder(&kernel)
|
||
.arg(&scratch_dev)
|
||
.arg(&scratch_idx_base)
|
||
.arg(&isv_dev)
|
||
.arg(&isv_idx_base)
|
||
.arg(&wiener_dev)
|
||
.arg(&wiener_offset_base)
|
||
.arg(&n_slots)
|
||
.launch(LaunchConfig {
|
||
grid_dim: (1, 1, 1),
|
||
block_dim: (1, 1, 1),
|
||
shared_mem_bytes: 0,
|
||
})
|
||
.expect("launch apply_pearls_ad_kernel");
|
||
}
|
||
stream
|
||
.synchronize()
|
||
.expect("synchronize after apply_pearls_ad_kernel launch");
|
||
|
||
(isv_buf.read_all(), wiener_buf.read_all())
|
||
}
|
||
|
||
/// SP4 GPU-Pearls A+D applicator oracle test (2026-05-01).
|
||
///
|
||
/// Asserts the GPU `apply_pearls_ad_kernel` produces bit-equivalent
|
||
/// (fp32 ULP) outputs to the host-side `pearls_ad_update` reference for a
|
||
/// representative set of step observations covering all five trajectory
|
||
/// classes the production callers exercise:
|
||
///
|
||
/// 1. **Pearl A bootstrap** — first observation, both `prev_x_mean` and
|
||
/// `x_lag` start at zero. Kernel must replace x_mean directly with
|
||
/// step_obs (Pearl A sentinel branch), seeding `x_lag = step_obs`
|
||
/// and keeping `sample_var = diff_var = 0`.
|
||
/// 2. **Stationary convergence** — after a Pearl A bootstrap, repeated
|
||
/// identical step observations should drive `diff_var → 0` and the
|
||
/// EMA toward the stationary value. Kernel and host should reach
|
||
/// the same final state after many steps.
|
||
/// 3. **Step-function response** — large abrupt change after a long
|
||
/// stationary period. Kernel must follow Pearl D's adaptive α toward
|
||
/// the new value at the same rate as the host.
|
||
/// 4. **Degenerate-zero short-circuit** — step_obs == 0 should leave
|
||
/// ISV + Wiener state untouched (kernel skip path matching the
|
||
/// pre-refactor host-side `if step_obs == 0.0 { continue; }`).
|
||
/// 5. **Large-magnitude observation** — verify the Pearl D variance
|
||
/// formula doesn't suffer fp32 precision loss vs the host reference
|
||
/// on observations with large absolute value.
|
||
///
|
||
/// Per `feedback_no_cpu_test_fallbacks.md`: the host reference is
|
||
/// analytically correct (Pearl D's Wiener-optimal blend has a closed
|
||
/// form), not a CPU fallback — the GPU kernel mirrors the same formula
|
||
/// bit-for-bit, so identity (within fp32 ULP) is the strict assertion.
|
||
#[test]
|
||
#[ignore = "requires GPU"]
|
||
fn sp4_apply_pearls_ad_kernel_matches_host_reference_within_fp32_ulp() {
|
||
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||
|
||
let stream = make_test_stream();
|
||
|
||
// Build a 5-slot trajectory covering the five representative cases.
|
||
// Each slot starts at the Pearl A sentinel (prev=0, x_lag=0, vars=0)
|
||
// and receives ONE step observation in this kernel launch — so for
|
||
// tests that need multiple steps (cases 2/3) we run iterated launches
|
||
// until convergence, comparing slot-by-slot against the host. For
|
||
// single-step cases (1/4/5), one launch suffices.
|
||
|
||
// ── Case 1: Pearl A bootstrap (first observation = 42.0) ──
|
||
{
|
||
let scratch = [42.0_f32];
|
||
let isv_init = [0.0_f32];
|
||
let wiener_init = [0.0_f32; 3];
|
||
let (isv_gpu, wiener_gpu) =
|
||
launch_apply_pearls_ad_oracle(&stream, &scratch, &isv_init, &wiener_init);
|
||
|
||
let mut state = WienerState::ZERO;
|
||
let isv_host = pearls_ad_update(0.0, &mut state, 42.0);
|
||
|
||
assert_eq!(isv_gpu[0], isv_host,
|
||
"Pearl A bootstrap: GPU={} host={}", isv_gpu[0], isv_host);
|
||
assert_eq!(wiener_gpu[0], state.sample_var,
|
||
"Pearl A: sample_var GPU={} host={}", wiener_gpu[0], state.sample_var);
|
||
assert_eq!(wiener_gpu[1], state.diff_var,
|
||
"Pearl A: diff_var GPU={} host={}", wiener_gpu[1], state.diff_var);
|
||
assert_eq!(wiener_gpu[2], state.x_lag,
|
||
"Pearl A: x_lag GPU={} host={}", wiener_gpu[2], state.x_lag);
|
||
}
|
||
|
||
// ── Case 2: Stationary convergence after bootstrap (5.0 × 200 steps) ──
|
||
{
|
||
// Drive both the GPU kernel and the host reference identically over
|
||
// 200 iterations. Each iteration is a fresh launch with the prior
|
||
// (isv, wiener) as input. Verify they agree at the end.
|
||
let mut isv = vec![0.0_f32];
|
||
let mut wiener = vec![0.0_f32; 3];
|
||
let mut host_state = WienerState::ZERO;
|
||
let mut host_x = 0.0_f32;
|
||
|
||
for step in 0..200 {
|
||
let scratch = [5.0_f32];
|
||
let (isv_next, wiener_next) =
|
||
launch_apply_pearls_ad_oracle(&stream, &scratch, &isv, &wiener);
|
||
isv = isv_next;
|
||
wiener = wiener_next;
|
||
host_x = pearls_ad_update(host_x, &mut host_state, 5.0);
|
||
|
||
// Per-iteration ULP check — catches drift early if it appears.
|
||
assert!((isv[0] - host_x).abs() <= 4.0 * f32::EPSILON * host_x.abs().max(1.0),
|
||
"step {step}: stationary x_mean drift GPU={} host={}", isv[0], host_x);
|
||
}
|
||
assert!((isv[0] - 5.0).abs() < 1e-2,
|
||
"stationary final x_mean should anchor at 5.0, got {}", isv[0]);
|
||
}
|
||
|
||
// ── Case 3: Step-function response (1.0 → 100.0 transition) ──
|
||
{
|
||
let mut isv = vec![0.0_f32];
|
||
let mut wiener = vec![0.0_f32; 3];
|
||
let mut host_state = WienerState::ZERO;
|
||
let mut host_x = 0.0_f32;
|
||
|
||
// Stationary at 1.0 for 100 steps.
|
||
for _ in 0..100 {
|
||
let (isv_next, wiener_next) =
|
||
launch_apply_pearls_ad_oracle(&stream, &[1.0_f32], &isv, &wiener);
|
||
isv = isv_next;
|
||
wiener = wiener_next;
|
||
host_x = pearls_ad_update(host_x, &mut host_state, 1.0);
|
||
}
|
||
// Step change to 100.0 for 200 steps.
|
||
for _ in 0..200 {
|
||
let (isv_next, wiener_next) =
|
||
launch_apply_pearls_ad_oracle(&stream, &[100.0_f32], &isv, &wiener);
|
||
isv = isv_next;
|
||
wiener = wiener_next;
|
||
host_x = pearls_ad_update(host_x, &mut host_state, 100.0);
|
||
}
|
||
// After 200 post-step iterations, Pearl D should have tracked
|
||
// toward 100; both implementations agree to fp32 ULP.
|
||
assert!((isv[0] - host_x).abs() <= 16.0 * f32::EPSILON * host_x.abs().max(1.0),
|
||
"step-change x_mean drift GPU={} host={}", isv[0], host_x);
|
||
assert!(isv[0] > 50.0,
|
||
"step-change should track toward 100, got {}", isv[0]);
|
||
}
|
||
|
||
// ── Case 4: Degenerate-zero short-circuit ──
|
||
{
|
||
// Sentinel ISV+wiener with non-zero scratch followed by a 0
|
||
// scratch — second iteration should leave ISV/wiener untouched.
|
||
let mut isv = vec![0.0_f32];
|
||
let mut wiener = vec![0.0_f32; 3];
|
||
let (isv_a, wiener_a) =
|
||
launch_apply_pearls_ad_oracle(&stream, &[7.5_f32], &isv, &wiener);
|
||
isv = isv_a.clone();
|
||
wiener = wiener_a.clone();
|
||
// step_obs = 0.0 → no-op (kernel `if step_obs == 0.0f { continue; }`).
|
||
let (isv_b, wiener_b) =
|
||
launch_apply_pearls_ad_oracle(&stream, &[0.0_f32], &isv, &wiener);
|
||
assert_eq!(isv_b, isv_a, "zero step_obs must not mutate ISV");
|
||
assert_eq!(wiener_b, wiener_a, "zero step_obs must not mutate Wiener state");
|
||
}
|
||
|
||
// ── Case 5: Large-magnitude observation (≈1e5) ──
|
||
{
|
||
let scratch = [1.0e5_f32];
|
||
let isv_init = [0.0_f32];
|
||
let wiener_init = [0.0_f32; 3];
|
||
let (isv_gpu, wiener_gpu) =
|
||
launch_apply_pearls_ad_oracle(&stream, &scratch, &isv_init, &wiener_init);
|
||
|
||
let mut state = WienerState::ZERO;
|
||
let isv_host = pearls_ad_update(0.0, &mut state, 1.0e5);
|
||
// Bootstrap: GPU should match host bit-for-bit (no arithmetic).
|
||
assert_eq!(isv_gpu[0], isv_host,
|
||
"large-magnitude bootstrap: GPU={} host={}", isv_gpu[0], isv_host);
|
||
assert_eq!(wiener_gpu[2], state.x_lag);
|
||
}
|
||
|
||
// ── Multi-slot batched launch sanity (n_slots=3, mixed cases) ──
|
||
// Validates that the device-side loop correctly increments the
|
||
// per-slot indices without aliasing across slots.
|
||
{
|
||
let scratch = [10.0_f32, 0.0, 50.0]; // bootstrap, zero, bootstrap
|
||
let isv_init = [0.0_f32, 5.0, 0.0]; // sentinel, populated, sentinel
|
||
let wiener_init = [
|
||
0.0, 0.0, 0.0, // slot 0: sentinel (Pearl A path)
|
||
1.0, 1.0, 5.0, // slot 1: established Wiener state, but step_obs=0 → skip
|
||
0.0, 0.0, 0.0, // slot 2: sentinel (Pearl A path)
|
||
];
|
||
let (isv_gpu, wiener_gpu) =
|
||
launch_apply_pearls_ad_oracle(&stream, &scratch, &isv_init, &wiener_init);
|
||
|
||
// Slot 0: Pearl A bootstrap → ISV[0] = 10, x_lag = 10.
|
||
assert_eq!(isv_gpu[0], 10.0);
|
||
assert_eq!(wiener_gpu[2], 10.0);
|
||
// Slot 1: zero step_obs → untouched.
|
||
assert_eq!(isv_gpu[1], 5.0);
|
||
assert_eq!(wiener_gpu[3], 1.0);
|
||
assert_eq!(wiener_gpu[4], 1.0);
|
||
assert_eq!(wiener_gpu[5], 5.0);
|
||
// Slot 2: Pearl A bootstrap → ISV[2] = 50, x_lag = 50.
|
||
assert_eq!(isv_gpu[2], 50.0);
|
||
assert_eq!(wiener_gpu[8], 50.0);
|
||
}
|
||
}
|