feat(sp4): Task A6 — atom_pos_p99 producer × 4 branches
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -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)
|
||||
|
||||
38
crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu
Normal file
38
crates/ml/src/cuda_pipeline/atom_pos_p99_kernel.cu
Normal file
@@ -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
|
||||
}
|
||||
}
|
||||
@@ -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::<f32>()) 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,
|
||||
|
||||
@@ -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<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 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<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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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<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 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.
|
||||
|
||||
Reference in New Issue
Block a user