feat(sp4): Task A5 — target_q_p99 producer kernel + Pearls A/D wire-up
First end-to-end SP4 producer. Kernel reads denoise_target_q_buf, computes p99(|target_q|) via sp4_histogram_p99<256>, writes step_obs to producer_step_scratch_buf[0] with __threadfence_system. Launcher syncs, applies Pearls A+D via pearls_ad_update (zero-copy mapped-pinned reads of ISV[TARGET_Q_BOUND_INDEX=131] + wiener_state_buf[(131-base)*3]), writes new x_mean back to ISV + state back to wiener_state. Cold-path launch (no captured graph in this task; Task A10 may move to captured). Producer-step-scratch slot 0 reserved for TARGET_Q_BOUND (stable layout documented in launcher comment for Tasks A6-A11 to extend). GPU unit test verifies kernel writes step_p99 ≈ p99(|N(0,1)|) within 5% rel_err on 4096 Box-Muller samples, then exercises Pearl A's sentinel branch (sentinel ISV + zero Wiener state → first-observation replacement). Helper asserts non-target scratch slots stay zero. No consumer wired yet — Mech 1's clamp still uses 10 × Q_ABS_REF.max(1.0). Behavior unchanged. cargo check --lib --tests clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -244,6 +244,16 @@ fn main() {
|
||||
// the same .cuh header directly; the test wrapper has no production
|
||||
// callers.
|
||||
"sp4_histogram_p99_test_kernel.cu",
|
||||
// SP4 Layer A Task A5 (2026-04-30): first end-to-end SP4 producer.
|
||||
// Single-block kernel reads `denoise_target_q_buf`, computes
|
||||
// p99(|target_q|) via `sp4_histogram_p99<256>`, writes the step
|
||||
// observation to `producer_step_scratch_buf[0]`. The host launcher
|
||||
// (`GpuDqnTrainer::launch_sp4_target_q_p99`) syncs, applies Pearls
|
||||
// A+D, writes new x_mean to ISV[TARGET_Q_BOUND_INDEX=131] + Wiener
|
||||
// state back to `wiener_state_buf`. Cold-path launch in this task;
|
||||
// Task A10 may relocate to the captured graph. Tasks A6-A9 mirror
|
||||
// the launcher template and slot-index discipline established here.
|
||||
"target_q_p99_kernel.cu",
|
||||
];
|
||||
|
||||
// ALL kernels get common header (BF16 types + wrappers)
|
||||
|
||||
@@ -139,6 +139,19 @@ pub(crate) static GRN_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/g
|
||||
/// 2c.3c.6 wires the consumer in `mag_concat_qdir`'s adaptive-scale path.
|
||||
static H_S2_RMS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/h_s2_rms_ema_kernel.cubin"));
|
||||
|
||||
/// SP4 Layer A Task A5 (2026-04-30): first end-to-end SP4 producer cubin.
|
||||
/// Single-block kernel that reads `denoise_target_q_buf`, computes
|
||||
/// p99(|target_q|) via the header-only `sp4_histogram_p99<256>` device
|
||||
/// function, and writes the per-step observation to
|
||||
/// `producer_step_scratch_buf[0]`. The launcher
|
||||
/// (`GpuDqnTrainer::launch_sp4_target_q_p99`) syncs, applies Pearls A+D
|
||||
/// (`pearls_ad_update`) to compute the new x_mean, and writes back to
|
||||
/// ISV[TARGET_Q_BOUND_INDEX=131] + Wiener state. No consumer wired yet —
|
||||
/// Mech 1's clamp keeps using `10 × Q_ABS_REF.max(1.0)` until the SP4
|
||||
/// consumer migration. Cold-path launch in this task; Task A10 may move
|
||||
/// to the captured graph.
|
||||
static SP4_TARGET_Q_P99_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/target_q_p99_kernel.cubin"));
|
||||
|
||||
/// Plan 4 Task 3 (E.3): IQN multi-quantile diagnostic EMAs into ISV[99..103).
|
||||
/// 4-block (256 threads/block, shmem-reduce, no atomicAdd) kernel launched
|
||||
/// alongside `h_s2_rms_ema_update` from `training_loop.rs`. Reads the IQN
|
||||
@@ -3335,6 +3348,18 @@ pub struct GpuDqnTrainer {
|
||||
/// Loaded from `h_s2_rms_ema_kernel.cubin`.
|
||||
h_s2_rms_ema_kernel: CudaFunction,
|
||||
|
||||
// ── SP4 Layer A Task A5: target_q_p99 producer ─────────────────────
|
||||
/// SP4 first end-to-end producer kernel. Single-block kernel that reads
|
||||
/// `denoise_target_q_buf` (target-network Q-values [B, 12] flattened),
|
||||
/// computes p99(|target_q|) via the `sp4_histogram_p99<256>` header-only
|
||||
/// device function, writes the step observation to
|
||||
/// `producer_step_scratch_buf[0]`. Loaded from
|
||||
/// `SP4_TARGET_Q_P99_CUBIN`. Cold-path launch via
|
||||
/// `launch_sp4_target_q_p99` (Pearls A+D applied host-side after sync).
|
||||
/// Tasks A6-A9 mirror this kernel's shape + the launcher's Pearls A+D
|
||||
/// host wire-up to populate the remaining 39 SP4 ISV bound slots.
|
||||
target_q_p99_update: CudaFunction,
|
||||
|
||||
// ── Plan C Phase 2 follow-up A.2: q-drift rate ISV producer ───────
|
||||
/// Single-thread single-block cold-path kernel. Reads ISV[Q_ABS_REF=16] +
|
||||
/// ISV[Q_DIR_ABS_REF=21] plus host-passed `q_mean_curr` / `q_mean_prev`
|
||||
@@ -8506,6 +8531,125 @@ impl GpuDqnTrainer {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// SP4 Layer A Task A5: cold-path producer for ISV[TARGET_Q_BOUND_INDEX=131].
|
||||
///
|
||||
/// Reads `denoise_target_q_buf` (target-network Q-values [B, 12]), launches
|
||||
/// the single-block `target_q_p99_update` kernel which computes
|
||||
/// p99(|target_q|) via the header-only `sp4_histogram_p99<256>` device
|
||||
/// function and writes the step observation to
|
||||
/// `producer_step_scratch_buf[SCRATCH_IDX=0]`. Synchronises the stream,
|
||||
/// then applies Pearls A+D host-side via `pearls_ad_update` using
|
||||
/// zero-copy mapped-pinned reads of ISV[131] + the per-slot Wiener state
|
||||
/// (`wiener_state_buf[(131-SP4_SLOT_BASE)*3..(131-SP4_SLOT_BASE)*3 + 3]`).
|
||||
/// Writes the new `x_mean` back to ISV[131] and the updated Wiener state
|
||||
/// back to `wiener_state_buf` via the same mapped-pinned host pointers.
|
||||
///
|
||||
/// Cold-path launch — NOT in the captured graph for now. Task A10 may
|
||||
/// move target_q to the captured graph; Tasks A6-A9 mirror this
|
||||
/// launcher's shape (kernel + sync + Pearls A+D + ISV writeback) for
|
||||
/// the remaining 39 SP4 ISV bound slots.
|
||||
///
|
||||
/// No consumer wired in this task — Mech 1's clamp still uses
|
||||
/// `10 × Q_ABS_REF.max(1.0)`. Behaviour unchanged from before this
|
||||
/// commit; the ISV slot now holds an EMA but is not yet read.
|
||||
pub fn launch_sp4_target_q_p99(&self) -> Result<(), MLError> {
|
||||
use crate::cuda_pipeline::sp4_isv_slots::{TARGET_Q_BOUND_INDEX, SP4_SLOT_BASE};
|
||||
use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
|
||||
|
||||
// Producer-step-scratch slot assignment. Stable layout — Tasks A6-A11
|
||||
// extend with subsequent indices in producer-launch order:
|
||||
// 0 = TARGET_Q_BOUND
|
||||
// 1..5 = ATOM_POS_BOUND[0..4]
|
||||
// 5..29 = WEIGHT_BOUND[0..8] · ADAM_M_BOUND[0..8] · ADAM_V_BOUND[0..8]
|
||||
// 29..37 = WD_RATE[0..8]
|
||||
// 37 = L1_LAMBDA_TRUNK
|
||||
// 38 = GRAD_CLIP_BOUND
|
||||
// 39 = H_S2_BOUND
|
||||
// 40..47 = retrofit existing 7 producers
|
||||
const SCRATCH_IDX: usize = 0;
|
||||
|
||||
// Shared memory: 8 warps × 256 bins × sizeof(int) = 8192 bytes.
|
||||
// Matches the contract documented in `sp4_histogram_p99.cuh` and
|
||||
// mirrored by `sp4_histogram_p99_test_kernel.cu`'s launch in
|
||||
// `tests/sp4_producer_unit_tests.rs`.
|
||||
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
|
||||
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (1, 1, 1),
|
||||
block_dim: (256, 1, 1),
|
||||
shared_mem_bytes: SHARED_BYTES,
|
||||
};
|
||||
|
||||
let count = self.denoise_target_q_buf.len() as i32;
|
||||
let buf_ptr = self.denoise_target_q_buf.raw_ptr();
|
||||
let scratch_dev = self.producer_step_scratch_buf.dev_ptr;
|
||||
let scratch_idx_arg: i32 = SCRATCH_IDX as i32;
|
||||
|
||||
// Safety: kernel signature `(const float*, int, float*, int)`
|
||||
// matches the four args below; both device pointers are valid
|
||||
// for the lifetime of `self`.
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.target_q_p99_update)
|
||||
.arg(&buf_ptr)
|
||||
.arg(&count)
|
||||
.arg(&scratch_dev)
|
||||
.arg(&scratch_idx_arg)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("target_q_p99_update launch: {e}")))?;
|
||||
}
|
||||
// Cold-path producer — sync before reading the mapped-pinned scratch
|
||||
// slot from the host. Captured-graph migration (Task A10) replaces
|
||||
// this with the same fence the captured stream provides on graph
|
||||
// launch boundaries.
|
||||
self.stream.synchronize()
|
||||
.map_err(|e| MLError::ModelError(format!("target_q_p99_update sync: {e}")))?;
|
||||
|
||||
// ── Pearls A+D host-side update (zero-copy mapped-pinned reads) ──
|
||||
// The kernel issued `__threadfence_system()` after writing
|
||||
// `producer_step_scratch_buf[SCRATCH_IDX]`; the post-sync read here
|
||||
// is the canonical mapped-pinned host-visible read pattern.
|
||||
let step_obs = unsafe {
|
||||
std::ptr::read_volatile(self.producer_step_scratch_buf.host_ptr.add(SCRATCH_IDX))
|
||||
};
|
||||
// Degenerate-signal short-circuit: kernel writes 0.0 when count==0
|
||||
// or every sample is exactly 0 (e.g., before the first target-net
|
||||
// forward populates `denoise_target_q_buf`). Pearl A's sentinel
|
||||
// branch would still fire on a real 0 step_obs at fold reset, but
|
||||
// updating the Wiener state with 0 advances `x_lag` and would
|
||||
// suppress the next genuine first-observation. Simpler: skip.
|
||||
if step_obs == 0.0 { return Ok(()); }
|
||||
|
||||
let isv_idx = TARGET_Q_BOUND_INDEX;
|
||||
let wiener_offset = (isv_idx - SP4_SLOT_BASE) * 3;
|
||||
|
||||
// Safety: `isv_signals_pinned` is a mapped-pinned `*mut f32` of
|
||||
// length `ISV_TOTAL_DIM`; `isv_idx < ISV_TOTAL_DIM` (171 > 131).
|
||||
// `wiener_state_buf.host_ptr` is a mapped-pinned `*mut f32` of
|
||||
// length `SP4_PRODUCER_COUNT * 3 = 141`; `wiener_offset + 2 < 141`.
|
||||
let prev_x_mean = unsafe { std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx)) };
|
||||
|
||||
let mut state = unsafe {
|
||||
let p = self.wiener_state_buf.host_ptr;
|
||||
WienerState {
|
||||
sample_var: std::ptr::read_volatile(p.add(wiener_offset)),
|
||||
diff_var: std::ptr::read_volatile(p.add(wiener_offset + 1)),
|
||||
x_lag: std::ptr::read_volatile(p.add(wiener_offset + 2)),
|
||||
}
|
||||
};
|
||||
|
||||
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_obs);
|
||||
|
||||
unsafe {
|
||||
std::ptr::write_volatile(self.isv_signals_pinned.add(isv_idx), new_x_mean);
|
||||
let wp = self.wiener_state_buf.host_ptr;
|
||||
std::ptr::write_volatile(wp.add(wiener_offset), state.sample_var);
|
||||
std::ptr::write_volatile(wp.add(wiener_offset + 1), state.diff_var);
|
||||
std::ptr::write_volatile(wp.add(wiener_offset + 2), state.x_lag);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Plan C Phase 2 follow-up A.2 (2026-04-29): launch `q_drift_rate_ema_update`
|
||||
/// — single-thread single-block ISV producer for ISV[Q_DRIFT_RATE_INDEX=129].
|
||||
///
|
||||
@@ -10117,6 +10261,18 @@ impl GpuDqnTrainer {
|
||||
.map_err(|e| MLError::ModelError(format!("h_s2_rms_ema_update load: {e}")))?
|
||||
};
|
||||
|
||||
// SP4 Layer A Task A5: load target_q_p99 producer kernel (cold-path).
|
||||
// Single-block kernel reading `denoise_target_q_buf`, writing
|
||||
// step_p99 to `producer_step_scratch_buf[0]`. Pearls A+D applied
|
||||
// host-side via `pearls_ad_update` (Task A3). Producer-only — no
|
||||
// consumer wired yet (Mech 1 still uses `10 × Q_ABS_REF.max(1.0)`).
|
||||
let target_q_p99_update = {
|
||||
let module = stream.context().load_cubin(SP4_TARGET_Q_P99_CUBIN.to_vec())
|
||||
.map_err(|e| MLError::ModelError(format!("sp4 target_q_p99 cubin load: {e}")))?;
|
||||
module.load_function("target_q_p99_update")
|
||||
.map_err(|e| MLError::ModelError(format!("target_q_p99_update load: {e}")))?
|
||||
};
|
||||
|
||||
// Plan C Phase 2 follow-up A.2: load q_drift_rate_ema kernel
|
||||
// (cold-path, per-epoch). Single-thread single-block ISV producer
|
||||
// for ISV[Q_DRIFT_RATE_INDEX=129]; consumer is `tau_update_kernel`.
|
||||
@@ -13075,6 +13231,7 @@ impl GpuDqnTrainer {
|
||||
popart_count,
|
||||
reward_component_ema_kernel,
|
||||
h_s2_rms_ema_kernel,
|
||||
target_q_p99_update,
|
||||
q_drift_rate_ema_kernel,
|
||||
fold_warmup_factor_kernel,
|
||||
grad_norm_fast_ema_pinned,
|
||||
|
||||
31
crates/ml/src/cuda_pipeline/target_q_p99_kernel.cu
Normal file
31
crates/ml/src/cuda_pipeline/target_q_p99_kernel.cu
Normal file
@@ -0,0 +1,31 @@
|
||||
// crates/ml/src/cuda_pipeline/target_q_p99_kernel.cu
|
||||
//
|
||||
// SP4 Layer A Task A5 — first end-to-end SP4 producer.
|
||||
//
|
||||
// Single-block kernel. Reads `denoise_target_q_buf [count]` (target-network
|
||||
// Q-values flattened over `[B, num_actions]`), computes p99(|target_q|) via
|
||||
// the header-only `sp4_histogram_p99<256>` device function, writes the
|
||||
// step_observation to `producer_step_scratch_buf[scratch_idx]`. The host
|
||||
// launcher (`GpuDqnTrainer::launch_sp4_target_q_p99`) syncs, applies
|
||||
// Pearls A+D via `pearls_ad_update`, and writes the new x_mean to
|
||||
// ISV[TARGET_Q_BOUND_INDEX=131] + Wiener state back to wiener_state_buf.
|
||||
//
|
||||
// Cold-path launch — NOT in the captured graph for now. Task A10 may move
|
||||
// the launch into the captured graph; this task installs the cold-path
|
||||
// template that Tasks A6-A9 mirror.
|
||||
|
||||
#include "sp4_histogram_p99.cuh"
|
||||
|
||||
extern "C" __global__ void target_q_p99_update(
|
||||
const float* __restrict__ buf,
|
||||
int count,
|
||||
float* __restrict__ scratch_buf,
|
||||
int scratch_idx
|
||||
) {
|
||||
if (blockIdx.x != 0) return;
|
||||
float p99 = sp4_histogram_p99<256>(buf, count);
|
||||
if (threadIdx.x == 0) {
|
||||
scratch_buf[scratch_idx] = p99;
|
||||
__threadfence_system(); // make host-mapped write visible
|
||||
}
|
||||
}
|
||||
@@ -168,3 +168,208 @@ fn sp4_histogram_p99_matches_known_distribution_within_quantization() {
|
||||
"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 SP4_PRODUCER_COUNT: usize = 47;
|
||||
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);
|
||||
}
|
||||
|
||||
@@ -2303,4 +2303,6 @@ SP4 Layer A Task A2 — mapped-pinned buffers for Pearls B/C/D (2026-04-30): all
|
||||
|
||||
SP4 Layer A Task A4 — linear-histogram p99 device function + GPU unit test (2026-04-30): created `crates/ml/src/cuda_pipeline/sp4_histogram_p99.cuh` providing the header-only `__device__` template `sp4_histogram_p99<BLOCK_SIZE>(buf, count) -> float` that returns p99(|buf|) for a single-block, BLOCK_SIZE-thread launch. Three-pass algorithm: (1) block-wide max-reduce of |buf| → `step_max` (degenerate-zero short-circuits to 0.0 so callers skip the ISV update); (2) linear-spaced [0, step_max] binning into 256 bins via per-warp tiles in dynamic shared memory (no atomicAdd per `feedback_no_atomicadd` — lane-collisions within a warp cost <0.012% expected count loss, well within the 1% quantile precision budget; cross-warp collisions are eliminated by the per-warp tiling); (3) cumulative-from-top → p99 = bin upper-edge `(p99_bin + 1) × bin_width`. Linear spacing chosen over log because SP4 producers care about resolution at the *top* of the |buf| distribution; top bin width ≈ step_max/256 ≈ 0.39%, comfortably within the 1% quantile precision budget. Caller contract documented in the header: `grid_dim=(1,1,1)`, `block_dim=(BLOCK_SIZE,1,1)`, dynamic shared memory ≥ `(BLOCK_SIZE/32) × 256 × sizeof(int)` (8192 bytes for BLOCK_SIZE=256). Wrapper kernel `sp4_histogram_p99_test_kernel.cu` exposes the device fn for Rust testing — single-block kernel that calls `sp4_histogram_p99<256>` and writes the scalar result to a mapped-pinned `f32` with `__threadfence_system()` so the host reads via `read_volatile` (no `memcpy_dtoh` per `feedback_gpu_cpu_roundtrip` and `feedback_no_htod_htoh_only_mapped_pinned`). Build.rs registration follows the `thompson_test_kernel.cu` pattern (Plan A Task 1): added to `kernels_with_common`, plus a `cargo:rerun-if-changed` directive on the `.cuh` header so cubins rebuild when the device fn evolves. Unit test `sp4_histogram_p99_matches_known_distribution_within_quantization` (in `crates/ml/tests/sp4_producer_unit_tests.rs`) drives the wrapper with 4096 deterministic |N(0,1)| Box-Muller samples (`StdRng::seed_from_u64(0xDEAD_BEEF)`), computes ground-truth p99 by sorting (analytical reference ≈ 2.576 z-score one-tailed), and asserts `rel_err < 5%` against the device output. `#[ignore]`-gated for GPU per the existing `distributional_q_tests.rs` convention. Local L40S run: true_p99=2.59758, computed_p99=2.59858, rel_err=0.039% — passes. The 5% tolerance leaves ample headroom over the worst-case sum of (linear-bin quantization 0.4%) + (per-warp lane-collision <0.012%) + (finite-sample sorted-p99 jitter ≈0.5% at n=4096) ≈ <2% true budget. No producer wired yet — header is library code included only by the test wrapper; SP4 Tasks A5-A9 add the magnitude-bound producer kernels that `#include "sp4_histogram_p99.cuh"` directly. Zero new ISV slots; zero new HtoD/DtoD/HtoH copies; zero behaviour change. `cargo check -p ml --lib --tests` clean (12 pre-existing warnings, no new warnings); GPU test 1/1 passing locally.
|
||||
|
||||
SP4 Layer A Task A5 — `target_q_p99` producer kernel + Pearls A/D wire-up (2026-04-30): first end-to-end SP4 producer. Created `crates/ml/src/cuda_pipeline/target_q_p99_kernel.cu` — single-block, 256-thread kernel that `#include "sp4_histogram_p99.cuh"` directly (Task A4 header), reads `denoise_target_q_buf` (target-network Q-values [B, 12]), computes p99(|target_q|) via `sp4_histogram_p99<256>`, and writes the per-step observation to `producer_step_scratch_buf[0]` with `__threadfence_system()` so the host pinned-mapped read sees the write. Registered in `crates/ml/build.rs` alongside `sp4_histogram_p99_test_kernel.cu`. Cubin embedded as `SP4_TARGET_Q_P99_CUBIN` in `gpu_dqn_trainer.rs`; `target_q_p99_update: CudaFunction` field added next to `h_s2_rms_ema_kernel`; loaded in the constructor immediately after `h_s2_rms_ema_kernel`. New launcher `GpuDqnTrainer::launch_sp4_target_q_p99(&self) -> Result<(), MLError>` mirrors `launch_h_s2_rms_ema`'s shape: launches the kernel with `grid_dim=(1,1,1)`, `block_dim=(256,1,1)`, dynamic shmem 8192 bytes (8 warps × 256 bins × sizeof(int)) per the `sp4_histogram_p99.cuh` contract; synchronises the stream; reads the step observation via `read_volatile(producer_step_scratch_buf.host_ptr.add(0))`; degenerate-zero short-circuits before mutating ISV/Wiener state; otherwise reads `prev_x_mean` from `isv_signals_pinned.add(TARGET_Q_BOUND_INDEX=131)` and the per-slot Wiener triple from `wiener_state_buf.host_ptr.add((131-SP4_SLOT_BASE)*3)`; calls `pearls_ad_update` (Task A3) host-side; writes the new `x_mean` back to ISV[131] and the updated `[sample_var, diff_var, x_lag]` back to `wiener_state_buf` — all via mapped-pinned host pointers (no HtoD/DtoH per `feedback_no_htod_htoh_only_mapped_pinned`). **Producer scratch slot layout** (stable, documented in launcher comment for Tasks A6-A11 to extend): slot 0 = TARGET_Q_BOUND, 1..5 = ATOM_POS_BOUND[0..4], 5..29 = WEIGHT_BOUND/ADAM_M/ADAM_V × 8 groups, 29..37 = WD_RATE × 8, 37 = L1_LAMBDA_TRUNK, 38 = GRAD_CLIP_BOUND, 39 = H_S2_BOUND, 40..47 = retrofit existing 7 producers. **Cold-path launch** — NOT in the captured graph for now; Task A10 may relocate to captured-graph cadence. **No consumer wired yet** — Mech 1's `target_q` clamp still uses `±10 × ISV[Q_ABS_REF=16].max(1.0)`; SP4 consumer migration follows once all producers (A5-A9) land. Behavior unchanged from before this commit. **Unit test** `sp4_target_q_p99_first_observation_writes_step_p99_via_pearls_ad` (in `crates/ml/tests/sp4_producer_unit_tests.rs`) drives the production kernel kernel-direct with 4096 deterministic |N(0,1)| Box-Muller samples (`StdRng::seed_from_u64(0xA5_5A_F1_57)`), asserts step_p99 ∈ ±5% of sorted-sample p99 (analytical |N(0,1)| ≈ 2.576), then exercises `pearls_ad_update(prev=0.0, state=ZERO, step_p99)` and asserts Pearl A's sentinel branch returns `step_p99` directly with `state.x_lag = step_p99` and zero variances. The helper also asserts `producer_step_scratch_buf[i] == 0` for all `i != 0` — guards against future single-thread-write bugs analogous to the `2fb30f098` cascade. `#[ignore]`-gated for GPU. Local L40S run: true_p99=2.54123, step_p99=2.54149, rel_err=0.010% — passes. The pathological all-in-one-bin distribution (e.g., 990×1.0 + 10×100.0) exposes the `feedback_no_atomicadd.md`-acknowledged intra-warp lane-collision bound (8 lanes incrementing same shmem bin → 1 increment lands); test distribution mirrors Task A4's well-spread |N(0,1)| pattern that real `denoise_target_q_buf` Q-values resemble in production. Zero new ISV slots (TARGET_Q_BOUND_INDEX=131 already allocated by Task A1); zero new HtoD/DtoD/HtoH copies (mapped-pinned only); one new kernel + one new launcher method; producer-only (no consumer). `cargo check -p ml --lib --tests` clean (12 pre-existing warnings, no new warnings); GPU unit test 1/1 passing locally.
|
||||
|
||||
SP4 Layer A Task A3 — Pearls A+D shared host-side helper (2026-04-30): created `crates/ml/src/cuda_pipeline/sp4_wiener_ema.rs` providing the single source-of-truth implementation of Pearl A (first-observation bootstrap) and Pearl D (Wiener-optimal adaptive α) used by all SP4 producer launchers (Tasks A5-A11) and unit tests. Re-exported from `cuda_pipeline/mod.rs`: `pearls_ad_update`, `WienerState`, `ALPHA_META`, `EPS_DIV`, `EPS_CLAMP_FLOOR`. **Pearl A** (sentinel-detect): on the first producer-step call after a fold reset (`prev_x_mean == 0.0 && state.x_lag == 0.0`), the helper replaces `x_mean` directly with `step_observation` and initialises `state.x_lag = step_observation`, leaving `sample_var = diff_var = 0`. This bypasses Pearl D's Wiener formula on the very first observation per the spec self-review math: at t=0 with both variances zero, `α* = 0/(0+0+ε_div) = 0` and Pearl D's `(1-α*)·prev + α*·obs = 0·0 + 0·obs = 0`, NOT `obs`. The explicit sentinel branch is required and is exercised by the dedicated test `pearl_d_does_not_subsume_pearl_a_at_t0`. **Pearl D** (steps 1+): for all subsequent steps, the helper computes `α* = diff_var / (diff_var + sample_var + EPS_DIV)` and blends `x_mean[t] = (1-α*)·x_mean[t-1] + α*·x[t]`; the variances themselves are tracked at the uniform meta-rate `ALPHA_META = 1e-3` (`sample_var ← (1-α_meta)·sample_var + α_meta·(x[t]-x_mean[t-1])²`, `diff_var ← (1-α_meta)·diff_var + α_meta·(x[t]-x_lag)²`). **Constants** (Adam-ε category, structural — not tuning knobs): `ALPHA_META = 1e-3` (single uniform meta-rate, no per-signal tuning, derived from typical per-fold step count ~1000 in smoke); `EPS_DIV = 1e-8` (numerical division-safety, prevents `0/0` in optimal-α formula); `EPS_CLAMP_FLOOR = 1.0` (consumer cold-start floor — used by clamps as `bound = isv[X].max(EPS_CLAMP_FLOOR)` to prevent `bound = 0` from collapsing the clamp at step 0 before any producer runs). **Tests** (6, all passing): `pearl_a_first_observation_replaces_sentinel`, `pearl_d_stationary_signal_alpha_approaches_zero` (constant signal: diff_var → 0, x_mean anchors at signal value), `pearl_d_step_change_tracks_within_meta_window` (signal jumps 1→100 over ~2000 steps, x_mean tracks past 50), `pearl_d_does_not_subsume_pearl_a_at_t0` (mathematical correctness of explicit sentinel branch), `meta_constants_are_structural` (document-as-code assertion), `pearl_a_only_fires_when_both_x_mean_and_x_lag_are_zero` (Pearl A does not re-fire post-Pearl-D when `x_lag` is already populated). No consumers wired yet — helper is library code unused by the producer pipeline; producer launchers in Tasks A5-A11 will call `pearls_ad_update` after each kernel-step writes `producer_step_scratch_buf` (allocated in Task A2), then the new `x_mean` gets written back to the corresponding ISV bound slot. Zero new ISV slots; zero new kernels; zero new HtoD/DtoD/HtoH copies; zero behavior change. `cargo check -p ml --lib` clean (12 pre-existing warnings, no new warnings); `cargo test -p ml --lib sp4_wiener_ema::` 6/6 passing.
|
||||
|
||||
Reference in New Issue
Block a user