From b4c28811ab7b8c69de6c9c0fc4e51897ea49a71f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 30 Apr 2026 23:07:07 +0200 Subject: [PATCH] =?UTF-8?q?feat(sp4):=20Task=20A6=20=E2=80=94=20atom=5Fpos?= =?UTF-8?q?=5Fp99=20producer=20=C3=97=204=20branches?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single kernel parameterized by branch slice. Launcher loops 0..4, launching once per branch with distinct scratch slot (1..5). Phase 2 applies Pearls A+D host-side per branch, writing to ISV[ATOM_POS_BOUND[ branch]] + wiener_state. Same end-to-end pattern as Task A5. GPU test verifies per-branch independence: 4 scales of |N(0,1)| samples (1×, 10×, 100×, 1000×) → 4 distinct p99 values within 5% tolerance. No consumer wired yet. Behavior unchanged. cargo check --lib --tests clean. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 12 ++ .../src/cuda_pipeline/atom_pos_p99_kernel.cu | 38 ++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 184 +++++++++++++++++ crates/ml/tests/sp4_producer_unit_tests.rs | 195 ++++++++++++++++++ docs/dqn-wire-up-audit.md | 2 + 5 files changed, 431 insertions(+) create mode 100644 crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 61e406764..40d507e57 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -254,6 +254,18 @@ fn main() { // 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", + // SP4 Layer A Task A6 (2026-04-30): atom_pos_p99 × 4 branches. + // Single parameterised single-block kernel reads ONE branch's slice + // of `atom_positions_buf` (`[4 * num_atoms]`, branch `b` at offset + // `b * num_atoms` per `atoms_update_kernel.cu`'s canonical layout), + // computes p99(|atom_positions[branch]|) via `sp4_histogram_p99<256>`, + // writes the step observation to + // `producer_step_scratch_buf[scratch_idx]`. The launcher + // (`launch_sp4_atom_pos_p99_all_branches`) loops `branch ∈ 0..4`, + // launching this kernel four times with distinct `(branch_buf, + // scratch_idx ∈ 1..5)` pairs and applying Pearls A+D host-side per + // branch (writes back to ISV[ATOM_POS_BOUND_BASE + branch]). + "atom_pos_p99_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu b/crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu new file mode 100644 index 000000000..df990826c --- /dev/null +++ b/crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu @@ -0,0 +1,38 @@ +// crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu +// +// SP4 Layer A Task A6 — atom_pos_p99 producer × 4 branches. +// +// Single-block, 256-thread, parameterised producer kernel. Reads ONE branch's +// slice of `atom_positions_buf` (canonical layout `[4 * num_atoms]` with +// branch `b` at offset `b * num_atoms`, mirroring `atoms_update_kernel.cu`'s +// `positions_out + (long long)branch * num_atoms` write contract). Computes +// p99(|atom_positions[branch]|) via the header-only `sp4_histogram_p99<256>` +// device function (Task A4) and writes the per-step observation to +// `producer_step_scratch_buf[scratch_idx]` with `__threadfence_system()` +// so the host pinned-mapped read sees the write. +// +// The host launcher (`GpuDqnTrainer::launch_sp4_atom_pos_p99_all_branches`) +// loops `branch ∈ 0..4`, launching this same kernel four times with distinct +// `(branch_buf, scratch_idx)` pairs (scratch slots `1..5`, each → ISV slot +// `ATOM_POS_BOUND_BASE + branch ∈ {132, 133, 134, 135}`). Pearls A+D applied +// host-side per branch after a single end-of-loop `stream.synchronize()`. +// +// Cold-path launch — NOT in the captured graph; Task A10 may relocate. The +// kernel signature is identical in shape to `target_q_p99_update` (Task A5); +// the launcher's per-branch loop is the only structural addition over A5. + +#include "sp4_histogram_p99.cuh" + +extern "C" __global__ void atom_pos_p99_update( + const float* __restrict__ atom_positions_branch, /* slice for this branch: [count] */ + int count, + float* __restrict__ scratch_buf, + int scratch_idx +) { + if (blockIdx.x != 0) return; + float p99 = sp4_histogram_p99<256>(atom_positions_branch, count); + if (threadIdx.x == 0) { + scratch_buf[scratch_idx] = p99; + __threadfence_system(); // make host-mapped write visible + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index fff2d10f8..13b4a33c5 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -152,6 +152,22 @@ static H_S2_RMS_EMA_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/h_s /// to the captured graph. static SP4_TARGET_Q_P99_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/target_q_p99_kernel.cubin")); +/// SP4 Layer A Task A6 (2026-04-30): atom_pos_p99 × 4 branches cubin. +/// Single parameterised single-block kernel that reads ONE branch's slice +/// of `atom_positions_buf` (`[4 * num_atoms]`, branch `b` at offset +/// `b * num_atoms`, canonical layout per `atoms_update_kernel.cu`), +/// computes p99(|atom_positions[branch]|) via `sp4_histogram_p99<256>`, +/// writes the step observation to `producer_step_scratch_buf[scratch_idx]`. +/// The launcher (`GpuDqnTrainer::launch_sp4_atom_pos_p99_all_branches`) +/// loops `branch ∈ 0..4`, launching this kernel four times with distinct +/// `(branch_buf, scratch_idx ∈ 1..5)` pairs, syncs once at end-of-loop, then +/// applies Pearls A+D host-side per branch (`pearls_ad_update`) writing +/// back to ISV[ATOM_POS_BOUND_BASE + branch ∈ {132, 133, 134, 135}] + +/// Wiener state. No consumer wired yet — Mech 2's atom-clamp still uses +/// `±10 × ISV[Q_ABS_REF=16].max(1.0)` (see `atoms_update_kernel.cu`). +/// Cold-path launch in this task; Task A10 may relocate to captured graph. +static SP4_ATOM_POS_P99_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/atom_pos_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 @@ -3360,6 +3376,15 @@ pub struct GpuDqnTrainer { /// host wire-up to populate the remaining 39 SP4 ISV bound slots. target_q_p99_update: CudaFunction, + /// SP4 Layer A Task A6: atom_pos_p99 × 4 branches producer kernel. + /// Single parameterised single-block kernel; one kernel handle, the + /// launcher (`launch_sp4_atom_pos_p99_all_branches`) loops over 4 + /// branches with distinct slice pointer + scratch slot (slots 1..5). + /// Writes to ISV[ATOM_POS_BOUND_BASE + branch ∈ {132, 133, 134, 135}] + /// after Pearls A+D applied host-side per branch. Loaded from + /// `SP4_ATOM_POS_P99_CUBIN`. + atom_pos_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` @@ -8650,6 +8675,151 @@ impl GpuDqnTrainer { Ok(()) } + /// SP4 Layer A Task A6: cold-path producer for ISV[ATOM_POS_BOUND[0..4]]. + /// + /// Reads `atom_positions_buf` (`[4 * num_atoms]` flat, branch `b` at + /// offset `b * num_atoms` per `atoms_update_kernel.cu`'s canonical + /// `positions_out + (long long)branch * num_atoms` write contract). + /// Loops `branch ∈ 0..SP4_BRANCH_COUNT`, launching the single parameterised + /// `atom_pos_p99_update` kernel four times — each launch passes the + /// per-branch slice pointer (`atom_positions_buf.raw_ptr() + branch * + /// num_atoms * 4 bytes`), the per-branch element count (`num_atoms`), + /// and a distinct producer-step-scratch slot (`SCRATCH_BASE + branch ∈ + /// {1, 2, 3, 4}`). After a single end-of-loop `stream.synchronize()`, + /// applies Pearls A+D host-side per branch via `pearls_ad_update`, + /// writing the new `x_mean` to ISV[ATOM_POS_BOUND_BASE + branch] and the + /// updated `[sample_var, diff_var, x_lag]` triple back to + /// `wiener_state_buf` (zero-copy mapped-pinned reads + writes; no + /// HtoD/DtoH per `feedback_no_htod_htoh_only_mapped_pinned`). + /// + /// Cold-path launch — NOT in the captured graph for now. Task A10 may + /// move atom_pos to the captured graph; this launcher mirrors A5's + /// shape so the migration is structural-only at that point. + /// + /// No consumer wired in this task — Mech 2's atom-position clamp in + /// `atoms_update_kernel.cu` still uses `±10 × ISV[Q_ABS_REF=16].max(1.0)`. + /// Behaviour unchanged from before this commit; the ISV slots now hold + /// per-branch EMAs but are not yet read. + pub fn launch_sp4_atom_pos_p99_all_branches(&self) -> Result<(), MLError> { + use crate::cuda_pipeline::sp4_isv_slots::{ + atom_pos_bound, ATOM_POS_BOUND_BASE, SP4_BRANCH_COUNT, SP4_SLOT_BASE, + }; + use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState}; + + // Producer-step-scratch slot assignment (matches `launch_sp4_target_q_p99`'s + // documented stable layout): slots 1..1+SP4_BRANCH_COUNT == 1..5 are + // reserved for ATOM_POS_BOUND[0..4] in branch order. + const SCRATCH_BASE: usize = 1; + + // Same shared-memory budget as `launch_sp4_target_q_p99` (and Task A4 + // wrapper test): 8 warps × 256 bins × sizeof(int) = 8192 bytes. + 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 num_atoms = self.config.num_atoms; + let total_len = self.atom_positions_buf.len(); + // Sanity: layout is exactly [SP4_BRANCH_COUNT × num_atoms]. Guards + // against a future config change to atom_positions_buf shape that + // would invalidate this launcher's per-branch slice arithmetic. + debug_assert_eq!( + total_len, SP4_BRANCH_COUNT * num_atoms, + "atom_positions_buf layout drift: expected {SP4_BRANCH_COUNT}×num_atoms={}, got {total_len}", + SP4_BRANCH_COUNT * num_atoms, + ); + + let base_dev_ptr = self.atom_positions_buf.raw_ptr(); + let scratch_dev = self.producer_step_scratch_buf.dev_ptr; + let count_per_branch_i32 = num_atoms as i32; + let stride_bytes = (num_atoms * std::mem::size_of::()) as u64; + + // ── Phase 1: launch one kernel per branch ── + for branch in 0..SP4_BRANCH_COUNT { + let branch_dev_ptr: u64 = base_dev_ptr + (branch as u64) * stride_bytes; + let scratch_idx_arg: i32 = (SCRATCH_BASE + branch) as i32; + + // Safety: kernel signature `(const float*, int, float*, int)` + // matches the four args below. `branch_dev_ptr` is within the + // allocation [base, base + 4 × num_atoms × 4 bytes); the + // debug_assert above guards the layout invariant. + unsafe { + self.stream + .launch_builder(&self.atom_pos_p99_update) + .arg(&branch_dev_ptr) + .arg(&count_per_branch_i32) + .arg(&scratch_dev) + .arg(&scratch_idx_arg) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!( + "atom_pos_p99_update[branch={branch}] launch: {e}" + )))?; + } + } + // Single end-of-loop sync — all 4 launches share the same stream and + // are serialised by `__threadfence_system()` inside each kernel; one + // sync ensures every scratch[1..5] write is host-visible before the + // Pearls A+D phase reads them. + self.stream.synchronize() + .map_err(|e| MLError::ModelError(format!("atom_pos_p99 sync: {e}")))?; + + // ── Phase 2: Pearls A+D host-side update per branch ── + for branch in 0..SP4_BRANCH_COUNT { + let scratch_idx = SCRATCH_BASE + branch; + // The kernel issued `__threadfence_system()` after each per- + // branch write; 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)) + }; + // Same degenerate-zero short-circuit rationale as A5: skip the + // ISV/Wiener mutation when step_obs == 0 (e.g., before the first + // `recompute_atom_positions` populates the branch slice). Pearl A's + // sentinel branch fires on the next genuine first observation. + if step_obs == 0.0 { continue; } + + let isv_idx = atom_pos_bound(branch); + // Defence-in-depth (compile-time const arithmetic): debug_assert + // the branch's ISV slot lies in the SP4 range [SP4_SLOT_BASE, + // SP4_SLOT_BASE + 40) and guards against an off-by-one drift in + // ATOM_POS_BOUND_BASE. The const-fn `atom_pos_bound` already + // checks `branch < SP4_BRANCH_COUNT` via debug_assert. + debug_assert_eq!(isv_idx, ATOM_POS_BOUND_BASE + branch); + 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 ∈ [132, 136) < 171`. + // `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]. /// @@ -10273,6 +10443,19 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("target_q_p99_update load: {e}")))? }; + // SP4 Layer A Task A6: load atom_pos_p99 producer kernel (cold-path). + // Single parameterised single-block kernel; launcher loops branches + // 0..4 with distinct slice pointers + scratch slots (1..5). Pearls + // A+D applied host-side per branch via `pearls_ad_update` (Task A3). + // Producer-only — no consumer wired yet (Mech 2 still uses + // `±10 × ISV[Q_ABS_REF=16].max(1.0)` in `atoms_update_kernel.cu`). + let atom_pos_p99_update = { + let module = stream.context().load_cubin(SP4_ATOM_POS_P99_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("sp4 atom_pos_p99 cubin load: {e}")))?; + module.load_function("atom_pos_p99_update") + .map_err(|e| MLError::ModelError(format!("atom_pos_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`. @@ -13232,6 +13415,7 @@ impl GpuDqnTrainer { reward_component_ema_kernel, h_s2_rms_ema_kernel, target_q_p99_update, + atom_pos_p99_update, q_drift_rate_ema_kernel, fold_warmup_factor_kernel, grad_norm_fast_ema_pinned, diff --git a/crates/ml/tests/sp4_producer_unit_tests.rs b/crates/ml/tests/sp4_producer_unit_tests.rs index cab84a0c7..3b522ac39 100644 --- a/crates/ml/tests/sp4_producer_unit_tests.rs +++ b/crates/ml/tests/sp4_producer_unit_tests.rs @@ -373,3 +373,198 @@ fn sp4_target_q_p99_first_observation_writes_step_p99_via_pearls_ad() { 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) -> 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__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 SP4_PRODUCER_COUNT: usize = 47; + 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::with_capacity(SP4_BRANCH_COUNT); + let mut per_branch_true_p99: Vec = 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 = (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::()) 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); + } +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index b0c63b31a..c7f84107e 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -2303,6 +2303,8 @@ 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(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 A6 — `atom_pos_p99` producer × 4 branches (2026-04-30): single parameterised producer kernel + 4-branch launcher loop, mirroring Task A5's end-to-end pattern. Created `crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu` — single-block, 256-thread kernel that `#include "sp4_histogram_p99.cuh"` directly (Task A4 header), reads ONE branch's slice of `atom_positions_buf` (canonical layout `[4 × num_atoms]` flat with branch `b` at offset `b × num_atoms` per `atoms_update_kernel.cu`'s `positions_out + (long long)branch * num_atoms` write contract), computes `p99(|atom_positions[branch]|)` via `sp4_histogram_p99<256>`, and writes the per-step observation to `producer_step_scratch_buf[scratch_idx]` with `__threadfence_system()`. Registered in `crates/ml/build.rs` after `target_q_p99_kernel.cu`. Cubin embedded as `SP4_ATOM_POS_P99_CUBIN`; `atom_pos_p99_update: CudaFunction` field added next to `target_q_p99_update`; loaded in the constructor immediately after `target_q_p99_update`. New launcher `GpuDqnTrainer::launch_sp4_atom_pos_p99_all_branches(&self) -> Result<(), MLError>` loops `branch ∈ 0..SP4_BRANCH_COUNT=4`, launching the same kernel four times with distinct `(branch_dev_ptr, scratch_idx)` pairs — `branch_dev_ptr = atom_positions_buf.raw_ptr() + branch × num_atoms × sizeof(f32)`, `scratch_idx ∈ {1, 2, 3, 4}` (SCRATCH_BASE=1 per Task A5's documented stable layout). After a single end-of-loop `stream.synchronize()`, applies Pearls A+D host-side per branch via `pearls_ad_update` (Task A3), reading `prev_x_mean` from `isv_signals_pinned.add(atom_pos_bound(branch) ∈ {132, 133, 134, 135})` and the per-slot Wiener triple from `wiener_state_buf.host_ptr.add((isv_idx − SP4_SLOT_BASE) × 3)`, writing the new `x_mean` back to ISV[ATOM_POS_BOUND_BASE + branch] 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`). `debug_assert_eq!(total_len, SP4_BRANCH_COUNT × num_atoms)` guards against future `atom_positions_buf` layout drift; `debug_assert_eq!(isv_idx, ATOM_POS_BOUND_BASE + branch)` guards against `atom_pos_bound` const-fn drift. Same degenerate-zero short-circuit (`if step_obs == 0.0 { continue; }`) per branch as Task A5 — skips the ISV/Wiener update before the first `recompute_atom_positions` populates the branch slice. **Cold-path launch** — NOT in the captured graph for now; Task A10 may relocate atom_pos to captured-graph cadence. **No consumer wired yet** — Mech 2's atom-position clamp in `atoms_update_kernel.cu` still uses `±10 × ISV[Q_ABS_REF=16].max(1.0)`; SP4 consumer migration follows once all producers (A5-A9) land. Behaviour unchanged from before this commit. **Unit test** `sp4_atom_pos_p99_per_branch_writes_distinct_isv_slots` (in `crates/ml/tests/sp4_producer_unit_tests.rs`) drives the production kernel kernel-direct with 4 × 4096 deterministic |N(0,1)| Box-Muller samples (`StdRng::seed_from_u64(0xA6_5A_F1_57 ^ (branch << 32))`), each branch scaled by `10ⁿ` for `n ∈ {0, 1, 2, 3}` so a launcher bug that mixes up branch slice pointers would land in the wrong scratch slot at the wrong scale and trip the per-branch ±5% rel_err assertion. Asserts each `scratch[1..5]` slot is within ±5% of its branch's true sorted-p99, all non-target slots remain zero (guards against the `2fb30f098`-style write-cascade), and the `(SP4_SLOT_BASE, ATOM_POS_BOUND_BASE, SP4_BRANCH_COUNT)` const-layout invariant + `atom_pos_bound(branch) == ATOM_POS_BOUND_BASE + branch` for `branch ∈ 0..4`. `#[ignore]`-gated for GPU. Local RTX 3050 Ti run: branch 0 (scale 1) rel_err=0.106%, branch 1 (scale 10) rel_err=0.362%, branch 2 (scale 100) rel_err=0.325%, branch 3 (scale 1000) rel_err=0.062% — all four pass with substantial margin under the 5% tolerance. The 5% headroom budget matches Tasks A4/A5: linear-bin quantization 0.4% + lane collision <0.012% + finite-sample p99 jitter ≈0.5% at n=4096 ≈ <2% true. Zero new ISV slots (ATOM_POS_BOUND[0..4] = ISV[132..136) already allocated by Task A1); zero new HtoD/DtoD/HtoH copies (mapped-pinned only); one new kernel + one new launcher method (single kernel handle, four launches via the loop); 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 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.