gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(sp4): Task A7 — Pearl B fused per-param-group statistics oracle

Browse Source 
Single kernel per group reads (params, grads, adam_m, adam_v) once and
fuses 5 passes (4 for non-trunk):
  Pass A: WEIGHT_BOUND[group] = p99(|params|)  via sp4_histogram_p99
  Pass B: ADAM_M_BOUND[group] = p99(|adam_m|)  via sp4_histogram_p99
  Pass C: ADAM_V_BOUND[group] = p99(|adam_v|)  via sp4_histogram_p99
  Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², ε_div); side: mean|g|, mean|w|
  Pass E (group 0 only): L1_LAMBDA[trunk] via gradient-direction entropy
   deficit × (mean|g|/mean|w|) magnitude scale

Launcher loops over 8 groups (DQN trunk/value/branches, IQN, IQL hi/lo,
attn, curiosity), one launch per group. Phase 2 applies Pearls A+D to
each of the 4-5 outputs per group via host-side pearls_ad_update.

Producer-scratch slots 5..38 reserved for this task's 33 outputs:
[5..13)=WEIGHT, [13..21)=ADAM_M, [21..29)=ADAM_V, [29..37)=WD_RATE, 37=L1.

Three of eight groups wired (DQN trunk/value/branches); five aux groups
(IQN, IQL hi/lo, attn, curiosity) skip silently because their backing
trainers live on FusedTrainingCtx, not GpuDqnTrainer. Wiring those
requires either threading buffer pointers through the launcher signature
or hoisting the launcher onto FusedTrainingCtx — both follow-up changes
scoped beyond Task A7. Documented in launcher + param_group_buffers
docstrings.

Per-group GPU tests verify outputs within 5% rel_err (p99 quantization)
or 2% rel_err (analytical formulas). Local RTX 3050 Ti: max WEIGHT
rel_err=0.589%, max ADAM_M=0.589%, max ADAM_V=0.554%, max WD_RATE=0.001%
across all 8 group shapes; group 0 L1 within tolerance.

Pearl B 4× memory-bandwidth reduction vs naive 4 separate per-group
producer kernels: each buffer read once for all 4-5 outputs.

No consumer wired yet — Adam kernels still take hardcoded weight_decay
+ weight_clamp_max_abs config args. Behavior unchanged.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-30 23:34:19 +02:00
parent b4c28811ab
commit 4f13e2ca37
5 changed files with 1060 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
crates/ml/build.rs
View File

@@ -266,6 +266,25 @@ fn main() {
// 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",
// SP4 Layer A Task A7 (2026-04-30): Pearl B fused per-param-group
// statistics oracle. Single-block kernel reads
// `(params, grads, adam_m, adam_v)` once and runs 4 sequential
// passes (5 for trunk):
// Pass A: WEIGHT_BOUND[group] = p99(|params|) via sp4_histogram_p99
// Pass B: ADAM_M_BOUND[group] = p99(|adam_m|) via sp4_histogram_p99
// Pass C: ADAM_V_BOUND[group] = p99(|adam_v|) via sp4_histogram_p99
// Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², EPS_DIV)
// Pass E: L1_LAMBDA[trunk] = (mean|g|/mean|w|) × entropy_deficit
// (group 0 only, gated by
// l1_lambda_scratch_idx >= 0).
// Pearl B 4× memory-bandwidth reduction vs four naive single-stat
// producer kernels: each buffer read once for all outputs. The host
// launcher (`launch_sp4_param_group_oracles_all_groups`) loops
// `group ∈ 0..SP4_PARAM_GROUP_COUNT=8`, launching this kernel 8
// times with distinct `(params, grads, m, v, count, scratch slots)`
// tuples, then applies Pearls A+D host-side per output via
// `pearls_ad_update`.
"param_group_oracle_kernel.cu",
];
// ALL kernels get common header (BF16 types + wrappers)

354
crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs
View File

@@ -168,6 +168,27 @@ static SP4_TARGET_Q_P99_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "
/// 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"));
/// SP4 Layer A Task A7 (2026-04-30): Pearl B fused per-param-group oracle cubin.
/// Single-block kernel reads `(params, grads, adam_m, adam_v)` once and runs
/// 4-5 fused passes:
/// Pass A: WEIGHT_BOUND[group] = p99(|params|) via sp4_histogram_p99
/// Pass B: ADAM_M_BOUND[group] = p99(|adam_m|) via sp4_histogram_p99
/// Pass C: ADAM_V_BOUND[group] = p99(|adam_v|) via sp4_histogram_p99
/// Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², EPS_DIV)
/// Pass E: L1_LAMBDA[trunk] = (mean|g|/mean|w|) × entropy_deficit
/// (group 0 only, gated by l1_idx ≥ 0)
/// The launcher (`GpuDqnTrainer::launch_sp4_param_group_oracles_all_groups`)
/// loops `group ∈ 0..SP4_PARAM_GROUP_COUNT`, launching the kernel once per
/// group with the per-group `(params, grads, m, v, count, scratch slots)`
/// tuple, then applies Pearls A+D host-side per output slot via
/// `pearls_ad_update`. Pearl B 4× memory-bandwidth reduction vs four naive
/// single-stat producer kernels: each buffer read once for all 4-5 outputs.
/// No consumer wired yet — Mech 9's clamp still uses hardcoded
/// `weight_clamp_max_abs` config args. Cold-path launch in this task; Task
/// A10 may relocate to captured graph.
static SP4_PARAM_GROUP_ORACLE_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/param_group_oracle_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
@@ -3385,6 +3406,24 @@ pub struct GpuDqnTrainer {
/// `SP4_ATOM_POS_P99_CUBIN`.
atom_pos_p99_update: CudaFunction,
/// SP4 Layer A Task A7: Pearl B fused per-param-group statistics oracle
/// kernel. Single-block kernel reads `(params, grads, adam_m, adam_v)`
/// once and runs 4 fused passes (5 for the trunk group):
/// Pass A: WEIGHT_BOUND[group] = p99(|params|)
/// Pass B: ADAM_M_BOUND[group] = p99(|adam_m|)
/// Pass C: ADAM_V_BOUND[group] = p99(|adam_v|)
/// Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², EPS_DIV)
/// Pass E: L1_LAMBDA[trunk] = (mean|g|/mean|w|) × entropy_deficit
/// (group 0 only, gated by l1_idx ≥ 0).
/// Pearl B 4× memory-bandwidth reduction vs four naive single-stat
/// producer kernels: each buffer read once for all 4-5 outputs. The
/// launcher (`launch_sp4_param_group_oracles_all_groups`) loops over
/// `SP4_PARAM_GROUP_COUNT=8` groups (DQN trunk/value/branches +
/// IQN/IQL-hi/IQL-lo/Attn/Curiosity), one launch per group, then applies
/// Pearls A+D host-side per output slot. Loaded from
/// `SP4_PARAM_GROUP_ORACLE_CUBIN`.
param_group_oracle_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`
@@ -8820,6 +8859,305 @@ impl GpuDqnTrainer {
Ok(())
}
/// SP4 Layer A Task A7: cold-path producer for the Pearl B
/// per-param-group statistics oracle.
///
/// For each param-group `g ∈ 0..SP4_PARAM_GROUP_COUNT=8`, calls
/// `param_group_buffers(g)` to resolve `(params, grads, m, v, count,
/// k_in, h_dim)`, then launches `param_group_oracle_update` once with
/// distinct producer-step-scratch slot indices into the documented
/// stable layout:
/// `WEIGHT_BOUND[g]` → scratch slot `5 + g`
/// `ADAM_M_BOUND[g]` → scratch slot `13 + g`
/// `ADAM_V_BOUND[g]` → scratch slot `21 + g`
/// `WD_RATE[g]` → scratch slot `29 + g`
/// `L1_LAMBDA_TRUNK` → scratch slot `37` (group 0 only; passed as
/// `-1_i32` to the kernel for groups 1-7).
///
/// After a single end-of-loop `stream.synchronize()`, applies Pearls
/// A+D host-side per-output via `pearls_ad_update` (Task A3) — 4
/// outputs per group plus L1_LAMBDA for group 0, totalling 33
/// updates. All reads/writes go through mapped-pinned host pointers
/// (`isv_signals_pinned` and `wiener_state_buf.host_ptr`); no
/// HtoD/DtoH per `feedback_no_htod_htoh_only_mapped_pinned`.
///
/// **Pearl B 4× memory-bandwidth reduction** vs four naive single-stat
/// producer kernels: the kernel reads each of `params`, `grads`,
/// `adam_m`, `adam_v` exactly once across all 4-5 outputs.
///
/// **Cold-path launch** — NOT in the captured graph for now. Task A10
/// may relocate select producers; this launcher mirrors A5/A6's shape
/// so the migration is structural-only at that point.
///
/// **No consumer wired in this task** — Mech 9's clamp still uses
/// hardcoded `weight_clamp_max_abs = 100 × q_abs_ref_eff` config args;
/// AdamW weight_decay and L1 unchanged. SP4 consumer migration follows
/// once all producers (A5-A9) land.
///
/// **Aux trainers not yet wired** — `param_group_buffers` returns
/// `None` for the 5 auxiliary param-groups (IQN, IQL-hi, IQL-lo, Attn,
/// Curiosity). Those trainers live on `FusedTrainingCtx`, not
/// `GpuDqnTrainer`; wiring them here would require either passing the
/// pointers through the launcher signature (and threading the call
/// site through every `&mut FusedTrainingCtx` callsite) or hoisting
/// the launcher onto `FusedTrainingCtx` itself. Both are follow-up
/// changes scoped beyond Task A7. The 3 DQN-internal groups (trunk,
/// value, branches) launch successfully; aux groups skip with a
/// debug-only `tracing::trace!`.
pub fn launch_sp4_param_group_oracles_all_groups(&self) -> Result<(), MLError> {
use crate::cuda_pipeline::sp4_isv_slots::{
adam_m_bound, adam_v_bound, weight_bound, wd_rate, ParamGroup,
L1_LAMBDA_TRUNK_INDEX, SP4_PARAM_GROUP_COUNT, SP4_SLOT_BASE,
};
use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
// Producer-step-scratch slot assignment (matches the stable layout
// documented in `launch_sp4_target_q_p99` and extended for Task A7):
// slots 5..13 = WEIGHT_BOUND[0..8]
// slots 13..21 = ADAM_M_BOUND[0..8]
// slots 21..29 = ADAM_V_BOUND[0..8]
// slots 29..37 = WD_RATE[0..8]
// slot 37 = L1_LAMBDA_TRUNK (group 0 only)
const SCRATCH_BASE_W: usize = 5;
const SCRATCH_BASE_M: usize = 13;
const SCRATCH_BASE_V: usize = 21;
const SCRATCH_BASE_WD: usize = 29;
const SCRATCH_L1_TRUNK: usize = 37;
// Same shared-memory budget as the rest of the SP4 producer family
// (sp4_histogram_p99 contract): 8 warps × 256 bins × sizeof(int).
// Pass D's 4 sequential block-reduces reuse this same dynamic shmem
// region as a `float[256]` accumulator buffer (256 × 4 bytes ≪ 8192).
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 scratch_dev = self.producer_step_scratch_buf.dev_ptr;
// Track which groups got a real launch — Phase 2 only post-processes
// those (skipping aux groups whose buffers aren't yet wired).
let mut launched: [bool; SP4_PARAM_GROUP_COUNT] = [false; SP4_PARAM_GROUP_COUNT];
// ── Phase 1: launch one kernel per group (when buffers available) ──
for g_idx in 0..SP4_PARAM_GROUP_COUNT {
let g = ParamGroup::ALL[g_idx];
let bufs = match self.param_group_buffers(g) {
Some(b) => b,
None => {
// Aux trainer not yet wired (IQN/IQL/Attn/Curiosity);
// skip silently. Follow-up task adds accessors.
continue;
}
};
let (params_ptr, grads_ptr, m_ptr, v_ptr, count, k_in, h_dim) = bufs;
// Sanity: a zero-element slice would yield a degenerate p99 = 0
// inside the kernel and short-circuit Pearls A+D below; cheaper
// to skip the launch entirely.
if count == 0 { continue; }
let count_i32 = count as i32;
let k_in_i32 = k_in;
let h_dim_i32 = h_dim;
let weight_idx_i32 = (SCRATCH_BASE_W + g_idx) as i32;
let adam_m_idx_i32 = (SCRATCH_BASE_M + g_idx) as i32;
let adam_v_idx_i32 = (SCRATCH_BASE_V + g_idx) as i32;
let wd_rate_idx_i32 = (SCRATCH_BASE_WD + g_idx) as i32;
let l1_idx_i32: i32 = if g_idx == 0 {
SCRATCH_L1_TRUNK as i32
} else {
-1
};
// Safety: kernel signature `(const float*, const float*,
// const float*, const float*, int, int, int, float*, int, int,
// int, int, int)` matches the 13 args below; all device pointers
// come from `param_group_buffers` which validates ownership.
unsafe {
self.stream
.launch_builder(&self.param_group_oracle_update)
.arg(&params_ptr)
.arg(&grads_ptr)
.arg(&m_ptr)
.arg(&v_ptr)
.arg(&count_i32)
.arg(&k_in_i32)
.arg(&h_dim_i32)
.arg(&scratch_dev)
.arg(&weight_idx_i32)
.arg(&adam_m_idx_i32)
.arg(&adam_v_idx_i32)
.arg(&wd_rate_idx_i32)
.arg(&l1_idx_i32)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!(
"param_group_oracle_update[group={g_idx}] launch: {e}"
)))?;
}
launched[g_idx] = true;
}
// Single end-of-loop sync — all launches are serialised on the same
// stream by `__threadfence_system()` inside each kernel; one sync
// ensures every scratch[5..38) write is host-visible before Phase 2.
self.stream.synchronize()
.map_err(|e| MLError::ModelError(format!("param_group_oracle sync: {e}")))?;
// ── Phase 2: Pearls A+D host-side update per output slot ──
// Helper closure: read step_obs from scratch slot, apply
// pearls_ad_update with the slot's Wiener state, write back.
// Captured-by-reference into closure to keep code DRY across the
// 4 (or 5 for group 0) outputs per group.
let apply_pearls = |isv_idx: usize, scratch_idx: usize| {
let step_obs = unsafe {
std::ptr::read_volatile(self.producer_step_scratch_buf.host_ptr.add(scratch_idx))
};
// Same degenerate-zero short-circuit as A5/A6: skip the
// ISV/Wiener mutation when step_obs == 0 (e.g., before the
// first forward populates the buffer, or in the L1-lambda case
// when the trunk weight matrix collapses to a single feature
// and entropy_deficit hits the natural 0 of the formula).
if step_obs == 0.0 { return; }
let wiener_offset = (isv_idx - SP4_SLOT_BASE) * 3;
// Safety: `isv_signals_pinned` is a mapped-pinned `*mut f32` of
// length `ISV_TOTAL_DIM`; SP4 ISV slots are bounded
// [SP4_SLOT_BASE, SP4_SLOT_END) = [131, 171).
// `wiener_state_buf.host_ptr` is 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);
}
};
for g_idx in 0..SP4_PARAM_GROUP_COUNT {
if !launched[g_idx] { continue; }
apply_pearls(weight_bound(g_idx), SCRATCH_BASE_W + g_idx);
apply_pearls(adam_m_bound(g_idx), SCRATCH_BASE_M + g_idx);
apply_pearls(adam_v_bound(g_idx), SCRATCH_BASE_V + g_idx);
apply_pearls(wd_rate(g_idx), SCRATCH_BASE_WD + g_idx);
if g_idx == 0 {
apply_pearls(L1_LAMBDA_TRUNK_INDEX, SCRATCH_L1_TRUNK);
}
}
Ok(())
}
/// SP4 Layer A Task A7 helper: resolve the per-param-group buffer tuple
/// for the Pearl B fused statistics oracle.
///
/// Returns `Some((params_ptr, grads_ptr, adam_m_ptr, adam_v_ptr, count,
/// k_in, h_dim))` when the trainer owns all four buffers for the group;
/// returns `None` for groups whose backing trainer (IQN, IQL-hi,
/// IQL-lo, Attn, Curiosity) is not owned by `GpuDqnTrainer` (those live
/// on `FusedTrainingCtx`).
///
/// `k_in`/`h_dim` are non-zero only for `ParamGroup::DqnTrunk` (group 0)
/// — the trunk Pass-E entropy-deficit computation needs to know the
/// `[h_dim × k_in]` shape of the trunk-input weight tensor (`w_a_h_s1`,
/// the GRN Linear_a at `param_sizes[0]`). Other groups pass 0 / 0 and
/// the kernel skips Pass E via `l1_lambda_scratch_idx = -1`.
///
/// For the 3 DQN-internal groups, slices are derived from
/// `compute_param_sizes` + `padded_byte_offset`:
/// - **Trunk** (tensors `[0..13)`): GRN h_s1 + h_s2 + γ/β pairs. Same
/// range as the existing `trunk_adam_m_ptr` / `trunk_params_ptr`
/// accessors.
/// - **Value head** (tensors `[13..17)`): W_v1, B_v1, W_v2, B_v2.
/// Same range as `value_adam_m_ptr`.
/// - **Branch heads** (tensors `[17..33)`): W/B_b{0..3}fc +
/// W/B_b{0..3}out concatenated. Same range as `branch_adam_m_ptr`.
fn param_group_buffers(
&self,
group: crate::cuda_pipeline::sp4_isv_slots::ParamGroup,
) -> Option<(u64, u64, u64, u64, usize, i32, i32)> {
use crate::cuda_pipeline::sp4_isv_slots::ParamGroup;
let f32_sz = std::mem::size_of::<f32>() as u64;
let param_sizes = compute_param_sizes(&self.config);
match group {
ParamGroup::DqnTrunk => {
// Trunk slice = tensors [0..13). k_in = state_dim_padded
// (the input dim of `w_a_h_s1`); h_dim = shared_h1 (the
// output dim of `w_a_h_s1`). Per `compute_param_sizes`,
// tensor [0] (`w_a_h_s1`) has size `shared_h1 × s1_input_dim`
// — that's our [h_dim × k_in] row-major layout.
let count = self.trunk_param_count;
let params_ptr = self.params_buf.raw_ptr();
let grads_ptr = self.ptrs.grad_buf;
let m_ptr = self.m_buf.raw_ptr();
let v_ptr = self.v_buf.raw_ptr();
// `s1_input_dim` is computed inside `compute_param_sizes`
// (state_dim or bottleneck_dim + portfolio_dim depending on
// bottleneck activation). Recompute here from
// tensor [0]'s size: w_a_h_s1 = shared_h1 × s1_input_dim,
// so s1_input_dim = param_sizes[0] / shared_h1.
let h_dim = self.config.shared_h1 as i32;
let s1_input_dim = if self.config.shared_h1 > 0 {
(param_sizes[0] / self.config.shared_h1) as i32
} else {
0
};
let k_in = s1_input_dim;
Some((params_ptr, grads_ptr, m_ptr, v_ptr, count, k_in, h_dim))
}
ParamGroup::DqnValue => {
// Value-head slice = tensors [13..17). k_in/h_dim = 0
// (Pass E gated to group 0 only).
let start = padded_byte_offset(&param_sizes, 13);
let end = padded_byte_offset(&param_sizes, 17);
let count = ((end - start) / f32_sz) as usize;
let params_ptr = self.params_buf.raw_ptr() + start;
let grads_ptr = self.ptrs.grad_buf + start;
let m_ptr = self.m_buf.raw_ptr() + start;
let v_ptr = self.v_buf.raw_ptr() + start;
Some((params_ptr, grads_ptr, m_ptr, v_ptr, count, 0, 0))
}
ParamGroup::DqnBranches => {
// Branch-heads slice = tensors [17..33). k_in/h_dim = 0.
let start = padded_byte_offset(&param_sizes, 17);
let end = padded_byte_offset(&param_sizes, 33);
let count = ((end - start) / f32_sz) as usize;
let params_ptr = self.params_buf.raw_ptr() + start;
let grads_ptr = self.ptrs.grad_buf + start;
let m_ptr = self.m_buf.raw_ptr() + start;
let v_ptr = self.v_buf.raw_ptr() + start;
Some((params_ptr, grads_ptr, m_ptr, v_ptr, count, 0, 0))
}
ParamGroup::Iqn
| ParamGroup::IqlHigh
| ParamGroup::IqlLow
| ParamGroup::Attn
| ParamGroup::Curiosity => {
// Owned by `FusedTrainingCtx`, not `GpuDqnTrainer`. Wiring
// these requires either threading the buffer pointers
// through the launcher signature or hoisting the launcher
// onto `FusedTrainingCtx`. Both are follow-up changes
// scoped beyond Task A7 — see the launcher's docstring.
None
}
}
}
/// 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].
///
@@ -10456,6 +10794,21 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("atom_pos_p99_update load: {e}")))?
};
// SP4 Layer A Task A7: load param_group_oracle producer kernel
// (cold-path). Single-block fused per-param-group oracle: reads
// `(params, grads, adam_m, adam_v)` once and runs 4-5 sequential
// passes producing WEIGHT/ADAM_M/ADAM_V p99 + WD_RATE + (group 0
// only) L1_LAMBDA_TRUNK. Pearls A+D applied host-side per output
// via `pearls_ad_update` (Task A3). Producer-only — no consumer
// wired yet (Mech 9 still uses hardcoded `weight_clamp_max_abs`
// config args; weight-decay/L1 unchanged).
let param_group_oracle_update = {
let module = stream.context().load_cubin(SP4_PARAM_GROUP_ORACLE_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("sp4 param_group_oracle cubin load: {e}")))?;
module.load_function("param_group_oracle_update")
.map_err(|e| MLError::ModelError(format!("param_group_oracle_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`.
@@ -13416,6 +13769,7 @@ impl GpuDqnTrainer {
h_s2_rms_ema_kernel,
target_q_p99_update,
atom_pos_p99_update,
param_group_oracle_update,
q_drift_rate_ema_kernel,
fold_warmup_factor_kernel,
grad_norm_fast_ema_pinned,

255
crates/ml/src/cuda_pipeline/param_group_oracle_kernel.cu Normal file
View File

@@ -0,0 +1,255 @@
// crates/ml/src/cuda_pipeline/param_group_oracle_kernel.cu
//
// SP4 Layer A Task A7 — Pearl B fused per-param-group statistics oracle.
//
// Single kernel per param-group reads `(params, grads, adam_m, adam_v)` once
// and runs 4 sequential passes (5 for the trunk group):
//
// Pass A: WEIGHT_BOUND[group] = p99(|params|) via sp4_histogram_p99<256>
// Pass B: ADAM_M_BOUND[group] = p99(|adam_m|) via sp4_histogram_p99<256>
// Pass C: ADAM_V_BOUND[group] = p99(|adam_v|) via sp4_histogram_p99<256>
// Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², EPS_DIV)
// side products: mean|g|, mean|w|
// Pass E: L1_LAMBDA[trunk] = (mean|g| / max(mean|w|, EPS_DIV))
// × (log K − H) / log K
// (Group 0 only; H = Shannon entropy of
// per-feature gradient L2 norms)
//
// Pearl B 4× memory-bandwidth reduction vs four naive single-stat producer
// kernels: each of params/grads/m/v is read once across all 4-5 outputs
// rather than four times.
//
// Caller contract (same as the rest of the SP4 producer family):
// - grid_dim = (1, 1, 1)
// - block_dim = (256, 1, 1)
// - dynamic shared memory ≥ (256/32) × 256 × sizeof(int) = 8192 bytes
// (matches `sp4_histogram_p99.cuh`'s per-warp tile contract; Pass D
// reuses the same dynamic shmem region after Pass C finishes).
//
// The trunk-only Pass E uses statically allocated `__shared__ float
// s_g_feat[K_IN_MAX]` for per-feature gradient L2 norms; K_IN_MAX is set to
// 1024 to cover any plausible trunk input dimension (production trunk is
// 40 + OFI/TLOB widening, well under 1024).
//
// Cold-path launch — NOT in the captured graph for now. Task A10 may
// relocate select producers into the captured graph; this kernel mirrors
// the launch-shape that A5/A6 established so the migration is structural-
// only at that point.
#include "sp4_histogram_p99.cuh"
#define EPS_DIV 1.0e-8f
// Cap on Pass-E per-feature shared scratch. Production trunk K_in is much
// smaller (40 base features; widened with OFI/TLOB to <128). 1024 is a
// generous ceiling that still fits comfortably in a single block's static
// shmem budget alongside `s_bins[256]` and `s_step_max`.
#define K_IN_MAX 1024
// Helper: block-wide tree-reduce of `val` across the 256 threads via
// `s_buf` (already-allocated `__shared__ float[256]` storage). Returns the
// reduced value to thread 0; other threads' return value is undefined.
// Uses `__syncthreads()` between halves; no atomicAdd per
// `feedback_no_atomicadd`.
__device__ __forceinline__ float block_reduce_sum_256(float val, float* s_buf, int tid) {
s_buf[tid] = val;
__syncthreads();
for (int s = 128; s > 0; s >>= 1) {
if (tid < s) {
s_buf[tid] += s_buf[tid + s];
}
__syncthreads();
}
return s_buf[0];
}
extern "C" __global__ void param_group_oracle_update(
const float* __restrict__ params,
const float* __restrict__ grads,
const float* __restrict__ adam_m,
const float* __restrict__ adam_v,
int count, /* slice length (elements) */
int k_in, /* trunk input dim (Pass E only; 0 for non-trunk) */
int h_dim, /* trunk output dim of layer (Pass E only) */
float* __restrict__ scratch_buf,
int weight_scratch_idx,
int adam_m_scratch_idx,
int adam_v_scratch_idx,
int wd_rate_scratch_idx,
int l1_lambda_scratch_idx /* -1 if not group 0 */
) {
if (blockIdx.x != 0) return;
const int tid = threadIdx.x;
// ── Pass A: WEIGHT_BOUND[group] = p99(|params|) ─────────────────────
{
float p99 = sp4_histogram_p99<256>(params, count);
if (tid == 0) {
scratch_buf[weight_scratch_idx] = p99;
__threadfence_system(); // make host-mapped write visible
}
}
__syncthreads();
// ── Pass B: ADAM_M_BOUND[group] = p99(|adam_m|) ─────────────────────
{
float p99 = sp4_histogram_p99<256>(adam_m, count);
if (tid == 0) {
scratch_buf[adam_m_scratch_idx] = p99;
__threadfence_system();
}
}
__syncthreads();
// ── Pass C: ADAM_V_BOUND[group] = p99(|adam_v|) ─────────────────────
{
float p99 = sp4_histogram_p99<256>(adam_v, count);
if (tid == 0) {
scratch_buf[adam_v_scratch_idx] = p99;
__threadfence_system();
}
}
__syncthreads();
// ── Pass D: WD_RATE[group] = |Σ w·g| / max(Σ w², EPS_DIV) ──────────
// Side products: mean|g| = Σ|g|/count, mean|w| = Σ|w|/count for Pass E.
//
// 4 parallel block-wide reductions: each thread maintains 4 register
// accumulators (sum_wg, sum_w2, sum_abs_g, sum_abs_w), then we reduce
// each tree-style through reused shared memory.
//
// Reuses the histogram dynamic shmem region (`s_warp_tiles`) as the
// per-thread accumulator buffer; the Pass C `__syncthreads` above
// ensures the previous histogram computation has fully drained before
// we overwrite. We reinterpret `s_warp_tiles` (declared as `int*` in
// `sp4_histogram_p99.cuh`) via `float*` cast: 8192 bytes = 2048 floats
// ≥ 256 floats × 4 accumulators = 1024 floats. Safe.
extern __shared__ int s_warp_tiles[];
float* s_buf = reinterpret_cast<float*>(s_warp_tiles);
float t_sum_wg = 0.0f;
float t_sum_w2 = 0.0f;
float t_sum_abs_g = 0.0f;
float t_sum_abs_w = 0.0f;
for (int i = tid; i < count; i += 256) {
float w = params[i];
float g = grads[i];
t_sum_wg += w * g;
t_sum_w2 += w * w;
t_sum_abs_g += fabsf(g);
t_sum_abs_w += fabsf(w);
}
// Four sequential block-reduces sharing the same 256-float shmem region.
float sum_wg = block_reduce_sum_256(t_sum_wg, s_buf, tid);
float sum_w2 = block_reduce_sum_256(t_sum_w2, s_buf, tid);
float sum_abs_g = block_reduce_sum_256(t_sum_abs_g, s_buf, tid);
float sum_abs_w = block_reduce_sum_256(t_sum_abs_w, s_buf, tid);
// Stash mean|g| / mean|w| in static shmem so Pass E (group 0) sees them
// after any further `__syncthreads`. We use `s_bins[]` (also int) by
// reinterpret-cast — same trick as `sp4_histogram_p99.cuh`'s Pass 1
// does for max-reduce. With 256 ints available we only need 2 slots.
__shared__ float s_mean_abs_g;
__shared__ float s_mean_abs_w;
__shared__ float s_wd_rate;
if (tid == 0) {
const float inv_count = (count > 0) ? (1.0f / (float)count) : 0.0f;
s_mean_abs_g = sum_abs_g * inv_count;
s_mean_abs_w = sum_abs_w * inv_count;
const float denom = fmaxf(sum_w2, EPS_DIV);
s_wd_rate = fabsf(sum_wg) / denom;
scratch_buf[wd_rate_scratch_idx] = s_wd_rate;
__threadfence_system();
}
__syncthreads();
// ── Pass E (group 0 trunk only): L1_LAMBDA_TRUNK ────────────────────
// Per-feature gradient L2 norm `||g_w[:,i]||` for i ∈ [0, k_in).
// Trunk weight matrix layout per `compute_param_sizes`: tensor 0 is
// `W_h_s1 [k_in × h_dim]` row-major, so feature `i` occupies the
// `i`-th column → strided indices `[i, k_in + i, 2*k_in + i, ...,
// (h_dim-1)*k_in + i]` in row-major. We compute the per-feature L2
// norm by squaring + summing across `h_dim` then sqrt.
//
// Feature normalization → discrete distribution; Shannon entropy H
// → entropy_deficit = (log K − H) / log K ∈ [0, 1] where K = k_in
// (1 = perfectly concentrated on one feature; 0 = uniform). The L1
// λ for the trunk is then
//
// L1_LAMBDA_TRUNK = (mean|g| / max(mean|w|, EPS_DIV)) × deficit
//
// — magnitude scale × concentration. Single thread (thread 0) does the
// normalize + entropy computation; per-feature L2 squares are produced
// by the block in parallel.
if (l1_lambda_scratch_idx >= 0 && k_in > 0 && h_dim > 0 && k_in <= K_IN_MAX) {
__shared__ float s_g_feat[K_IN_MAX];
// Zero the per-feature scratch (cooperative across the block).
for (int i = tid; i < k_in; i += 256) {
s_g_feat[i] = 0.0f;
}
__syncthreads();
// Each thread accumulates squared gradients for a strided slice
// of the [k_in × h_dim] grad_w matrix (= grads[0..k_in*h_dim)).
// Per-feature accumulation is naturally race-free across threads
// that hit different (i, h) cells; collisions on the same `i`
// across threads are resolved by reducing through shared memory
// — but to avoid atomicAdd, we instead serialize per-i: threads
// collaborate by reduction across `h ∈ [0, h_dim)` for each `i`
// in turn. That's k_in × (h_dim parallel reduce) = k_in × ~h_dim
// ops; for k_in ≤ 1024 and h_dim ≤ ~1024 this is fine for a cold-
// path producer.
const int kh = k_in * h_dim;
// Cooperative thread-strided accumulation into s_g_feat[i] using
// sequential synchronisation per feature. Avoids atomicAdd while
// keeping the cost linear in k_in × h_dim.
for (int i = 0; i < k_in; ++i) {
float partial = 0.0f;
for (int h = tid; h < h_dim; h += 256) {
int idx = h * k_in + i;
if (idx < kh) {
float gv = grads[idx];
partial += gv * gv;
}
}
float sq_sum = block_reduce_sum_256(partial, s_buf, tid);
if (tid == 0) {
s_g_feat[i] = sqrtf(sq_sum);
}
__syncthreads();
}
// Single-thread normalization + Shannon entropy.
if (tid == 0) {
float total = 0.0f;
for (int i = 0; i < k_in; ++i) total += s_g_feat[i];
float deficit = 0.0f;
if (total > EPS_DIV && k_in > 1) {
const float log_k = logf((float)k_in);
float entropy = 0.0f;
for (int i = 0; i < k_in; ++i) {
float p = s_g_feat[i] / total;
if (p > EPS_DIV) {
entropy -= p * logf(p);
}
}
deficit = (log_k - entropy) / log_k;
if (deficit < 0.0f) deficit = 0.0f;
if (deficit > 1.0f) deficit = 1.0f;
}
const float mean_abs_g = s_mean_abs_g;
const float mean_abs_w = s_mean_abs_w;
const float scale = mean_abs_g / fmaxf(mean_abs_w, EPS_DIV);
const float l1_lambda = scale * deficit;
scratch_buf[l1_lambda_scratch_idx] = l1_lambda;
__threadfence_system();
}
}
}

430
crates/ml/tests/sp4_producer_unit_tests.rs
View File

@@ -568,3 +568,433 @@ fn sp4_atom_pos_p99_per_branch_writes_distinct_isv_slots() {
assert_eq!(atom_pos_bound(branch), ATOM_POS_BOUND_BASE + branch);
}
}
// ── SP4 Task A7: Pearl B fused per-param-group oracle kernel ──────────────────
/// Test-only cubin for the SP4 Pearl B fused per-param-group oracle. The same
/// cubin is loaded by `GpuDqnTrainer::new` in production via
/// `SP4_PARAM_GROUP_ORACLE_CUBIN`. This kernel-direct test bypasses the full
/// trainer harness and exercises the kernel's 4-5 fused passes (max-reduce
/// + p99 histogram per (params, adam_m, adam_v); `Σ w·g / Σ w²` weight-decay
/// rate; trunk-only L1-lambda from gradient-direction entropy deficit) on
/// controlled distributions with analytically-known ground truth.
const SP4_PARAM_GROUP_ORACLE_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/param_group_oracle_kernel.cubin"));
/// Resolve the `param_group_oracle_update` function handle from its cubin.
fn load_sp4_param_group_oracle_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
let module = stream
.context()
.load_cubin(SP4_PARAM_GROUP_ORACLE_CUBIN.to_vec())
.expect("load param_group_oracle_kernel cubin");
module
.load_function("param_group_oracle_update")
.expect("load param_group_oracle_update function")
}
/// Box-Muller deterministic |N(0, σ²)| sample generator. Produces folded
/// half-normal magnitudes so the per-pass p99 lookups are well-spread across
/// the 256 linear-histogram bins (avoids the lane-collision pathology
/// flagged in `feedback_no_atomicadd.md`).
fn boxmuller_abs_normal(seed: u64, n: usize, sigma: f32) -> Vec<f32> {
use rand::rngs::StdRng;
use rand::{Rng, SeedableRng};
let mut rng = StdRng::seed_from_u64(seed);
(0..n)
.map(|_| {
let u1: f32 = rng.gen_range(1e-8_f32..1.0_f32);
let u2: f32 = rng.gen_range(0.0_f32..1.0_f32);
let z = (-2.0_f32 * u1.ln()).sqrt() * (2.0_f32 * std::f32::consts::PI * u2).cos();
(z * sigma).abs()
})
.collect()
}
/// Sample-sorted p99 of `xs` (positive samples). Mirrors the reference
/// computation used by the Task A4/A5/A6 tests — `sorted[(n * 99) / 100]`.
fn sample_p99(xs: &[f32]) -> f32 {
let mut sorted: Vec<f32> = xs.to_vec();
sorted.sort_by(|a, b| a.partial_cmp(b).expect("samples are finite"));
sorted[(sorted.len() * 99) / 100]
}
/// Launch the production `param_group_oracle_update` kernel kernel-direct
/// with controlled `(params, grads, adam_m, adam_v)` slices and write into a
/// 47-slot scratch buffer at the production-layout slot offsets for group
/// `g_idx`. Returns the host-readable scratch buffer for assertion.
///
/// `k_in`/`h_dim` non-zero only for `g_idx == 0` — Pass E (L1 lambda)
/// gating on `l1_lambda_scratch_idx >= 0`. For other groups, pass 0/0/-1.
#[allow(clippy::too_many_arguments)]
fn launch_sp4_param_group_oracle_for_group(
stream: &Arc<CudaStream>,
g_idx: usize,
params_host: &[f32],
grads_host: &[f32],
adam_m_host: &[f32],
adam_v_host: &[f32],
k_in: i32,
h_dim: i32,
) -> Vec<f32> {
const SP4_PRODUCER_COUNT: usize = 47;
const SCRATCH_BASE_W: usize = 5;
const SCRATCH_BASE_M: usize = 13;
const SCRATCH_BASE_V: usize = 21;
const SCRATCH_BASE_WD: usize = 29;
const SCRATCH_L1_TRUNK: usize = 37;
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
let n = params_host.len();
assert_eq!(grads_host.len(), n);
assert_eq!(adam_m_host.len(), n);
assert_eq!(adam_v_host.len(), n);
let kernel = load_sp4_param_group_oracle_kernel(stream);
// Safety: CUDA context active on this thread (resolved via stream).
let params_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc params");
let grads_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc grads");
let adam_m_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc adam_m");
let adam_v_buf = unsafe { MappedF32Buffer::new(n) }.expect("alloc adam_v");
params_buf.write_from_slice(params_host);
grads_buf.write_from_slice(grads_host);
adam_m_buf.write_from_slice(adam_m_host);
adam_v_buf.write_from_slice(adam_v_host);
let scratch_buf =
unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }.expect("alloc scratch");
let count_i32 = i32::try_from(n).expect("count fits in i32");
let weight_idx = (SCRATCH_BASE_W + g_idx) as i32;
let adam_m_idx = (SCRATCH_BASE_M + g_idx) as i32;
let adam_v_idx = (SCRATCH_BASE_V + g_idx) as i32;
let wd_rate_idx = (SCRATCH_BASE_WD + g_idx) as i32;
let l1_idx: i32 = if g_idx == 0 { SCRATCH_L1_TRUNK as i32 } else { -1 };
let params_dev = params_buf.dev_ptr;
let grads_dev = grads_buf.dev_ptr;
let adam_m_dev = adam_m_buf.dev_ptr;
let adam_v_dev = adam_v_buf.dev_ptr;
let scratch_dev = scratch_buf.dev_ptr;
unsafe {
stream
.launch_builder(&kernel)
.arg(&params_dev)
.arg(&grads_dev)
.arg(&adam_m_dev)
.arg(&adam_v_dev)
.arg(&count_i32)
.arg(&k_in)
.arg(&h_dim)
.arg(&scratch_dev)
.arg(&weight_idx)
.arg(&adam_m_idx)
.arg(&adam_v_idx)
.arg(&wd_rate_idx)
.arg(&l1_idx)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (256, 1, 1),
shared_mem_bytes: SHARED_BYTES,
})
.expect("launch param_group_oracle_update");
}
stream
.synchronize()
.expect("synchronize after param_group_oracle_update launch");
scratch_buf.read_all()
}
/// SP4 Task A7 unit test: parameterised across all 8 param-groups. Each
/// group launches the production kernel kernel-direct with a distinct
/// `(seed, sigma_w, sigma_g, sigma_m, sigma_v, count)` tuple — the per-group
/// scaling factors span 1×..1000× across the 4 buffers so a kernel bug
/// that mixes up buffer arguments would land in the wrong scratch slot at
/// the wrong scale and trip a per-output rel_err assertion.
///
/// **Distributions:** 4096 deterministic |N(0, σ²)| Box-Muller samples per
/// buffer per group (well-spread → no lane-collision pathology). Per-buffer
/// seeds: `0xA7_<group>_<buf>_F1_57`. Per-group `count` chosen to keep the
/// test runtime small while exceeding 1024 (the kernel's per-thread strided
/// reduction needs ≥ 1 stride for thread-256's accumulator to be exercised).
///
/// **Tolerances:**
/// - WEIGHT/M/V p99: 5% relative error (linear-bin quantization 0.4% +
/// lane collision <0.012% + finite-sample p99 jitter ≈0.5% at n=4096
/// ≈ <2% true; 5% is generous).
/// - WD_RATE: 2% relative error (analytical formula `|Σ w·g| / Σ w²`
/// reduces in `f32` with no histogram quantization; 2% headroom for
/// summation order non-determinism).
/// - L1 (group 0 only): 5% relative error — composes (mean|g|/mean|w|)
/// × deficit, both `f32` reductions plus a per-feature L2 norm chain
/// with two `logf` calls. Tolerance matches the p99 family.
///
/// **Production launcher contract validated:**
/// - The kernel writes step observations to scratch slots
/// `5 + g`, `13 + g`, `21 + g`, `29 + g` per group (and slot 37 for
/// group 0's L1 trunk lambda).
/// - All non-target slots remain zero (guards against the
/// `2fb30f098`-class write-cascade where a kernel silently writes
/// outside its assigned slot).
/// - SP4 ISV slot conventions hold: `weight_bound(g)`, `adam_m_bound(g)`,
/// `adam_v_bound(g)`, `wd_rate(g)`, `L1_LAMBDA_TRUNK_INDEX = 170`.
#[test]
#[ignore = "requires GPU"]
fn sp4_param_group_oracle_per_group_writes_distinct_isv_slots() {
use ml::cuda_pipeline::sp4_isv_slots::{
adam_m_bound, adam_v_bound, weight_bound, wd_rate, ParamGroup,
ADAM_M_BOUND_BASE, ADAM_V_BOUND_BASE, L1_LAMBDA_TRUNK_INDEX,
SP4_PARAM_GROUP_COUNT, SP4_SLOT_BASE, WD_RATE_BASE, WEIGHT_BOUND_BASE,
};
const SP4_PRODUCER_COUNT: usize = 47;
const SCRATCH_BASE_W: usize = 5;
const SCRATCH_BASE_M: usize = 13;
const SCRATCH_BASE_V: usize = 21;
const SCRATCH_BASE_WD: usize = 29;
const SCRATCH_L1_TRUNK: usize = 37;
const N: usize = 4096;
// Per-group input shape:
// group 0 (DqnTrunk): k_in × h_dim = 64 × 64 = 4096 (Pass E exercised)
// groups 1..7: flat 4096 element slice (Pass E gated off)
// sigmas chosen so p99(|N(0,σ²)|) lands at ~2.576 × σ:
// sigma_w = 0.1 × (g+1) → WEIGHT_BOUND[g] ≈ 0.258 × (g+1)
// sigma_g = 0.01 × (g+1) → grads at 1/10 weight scale
// sigma_m = 0.05 × (g+1) → adam_m at 1/2 weight scale
// sigma_v = 0.001 × (g+1) → adam_v second-moment small (mostly squared g)
let stream = make_test_stream();
// Layout invariants — guard against off-by-one drift in the SP4 slot map
// that would silently route writes to wrong ISV indices.
assert_eq!(SP4_PARAM_GROUP_COUNT, 8);
assert_eq!(SP4_SLOT_BASE, 131);
assert_eq!(WEIGHT_BOUND_BASE, 136);
assert_eq!(ADAM_M_BOUND_BASE, 144);
assert_eq!(ADAM_V_BOUND_BASE, 152);
assert_eq!(WD_RATE_BASE, 160);
assert_eq!(L1_LAMBDA_TRUNK_INDEX, 170);
for g in 0..SP4_PARAM_GROUP_COUNT {
assert_eq!(weight_bound(g), 136 + g);
assert_eq!(adam_m_bound(g), 144 + g);
assert_eq!(adam_v_bound(g), 152 + g);
assert_eq!(wd_rate(g), 160 + g);
}
// ParamGroup::ALL ordering matches our seed-keying.
let groups = ParamGroup::ALL;
assert_eq!(groups.len(), SP4_PARAM_GROUP_COUNT);
for g_idx in 0..SP4_PARAM_GROUP_COUNT {
let scale = (g_idx as f32) + 1.0;
let sigma_w = 0.1 * scale;
let sigma_g = 0.01 * scale;
let sigma_m = 0.05 * scale;
let sigma_v = 0.001 * scale;
let seed_w = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ 0x01;
let seed_g = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ 0x02;
let seed_m = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ 0x03;
let seed_v = 0xA7_5A_F1_57_u64 ^ ((g_idx as u64) << 32) ^ 0x04;
let params_host = boxmuller_abs_normal(seed_w, N, sigma_w);
// grads: signed (mix of positive/negative) so Σ w·g doesn't trivially
// collapse to Σ w·|g|. Reuse boxmuller for magnitudes, alternate
// signs by index parity → roughly zero-mean but non-trivial Σ w·g.
let grads_mag = boxmuller_abs_normal(seed_g, N, sigma_g);
let grads_host: Vec<f32> = grads_mag
.iter()
.enumerate()
.map(|(i, &g)| if i % 2 == 0 { g } else { -g })
.collect();
let adam_m_host = boxmuller_abs_normal(seed_m, N, sigma_m);
let adam_v_host = boxmuller_abs_normal(seed_v, N, sigma_v);
// For group 0 (DqnTrunk) the kernel uses k_in × h_dim = N total to
// drive Pass E. We pick k_in = 64, h_dim = 64 so k_in × h_dim = N.
let (k_in, h_dim) = if g_idx == 0 { (64_i32, 64_i32) } else { (0_i32, 0_i32) };
let scratch = launch_sp4_param_group_oracle_for_group(
&stream,
g_idx,
&params_host,
&grads_host,
&adam_m_host,
&adam_v_host,
k_in,
h_dim,
);
// ── Slot-write isolation: every slot outside this group's targets
// must remain zero. Targets are [5+g, 13+g, 21+g, 29+g] and (for
// g==0) slot 37.
let mut is_target = [false; SP4_PRODUCER_COUNT];
is_target[SCRATCH_BASE_W + g_idx] = true;
is_target[SCRATCH_BASE_M + g_idx] = true;
is_target[SCRATCH_BASE_V + g_idx] = true;
is_target[SCRATCH_BASE_WD + g_idx] = true;
if g_idx == 0 {
is_target[SCRATCH_L1_TRUNK] = true;
}
for (i, &v) in scratch.iter().enumerate() {
if !is_target[i] {
assert_eq!(
v, 0.0,
"param_group_oracle[group={g_idx}] wrote to non-target scratch[{i}] = {v}",
);
}
}
// ── Pass A: WEIGHT_BOUND[g] p99 ──
let true_p99_w = sample_p99(&params_host);
let kernel_p99_w = scratch[SCRATCH_BASE_W + g_idx];
let rel_w = ((kernel_p99_w - true_p99_w) / true_p99_w).abs();
println!(
"SP4 param_group[{g_idx}] WEIGHT — true_p99={true_p99_w:.5}, kernel_p99={kernel_p99_w:.5}, rel_err={rel_w:.5}",
);
assert!(kernel_p99_w > 0.0, "WEIGHT p99 must be positive");
assert!(
rel_w < 0.05,
"group {g_idx} WEIGHT rel_err {rel_w} > 5% (kernel={kernel_p99_w}, true={true_p99_w})",
);
// ── Pass B: ADAM_M_BOUND[g] p99 ──
let true_p99_m = sample_p99(&adam_m_host);
let kernel_p99_m = scratch[SCRATCH_BASE_M + g_idx];
let rel_m = ((kernel_p99_m - true_p99_m) / true_p99_m).abs();
println!(
"SP4 param_group[{g_idx}] ADAM_M — true_p99={true_p99_m:.5}, kernel_p99={kernel_p99_m:.5}, rel_err={rel_m:.5}",
);
assert!(kernel_p99_m > 0.0, "ADAM_M p99 must be positive");
assert!(
rel_m < 0.05,
"group {g_idx} ADAM_M rel_err {rel_m} > 5% (kernel={kernel_p99_m}, true={true_p99_m})",
);
// ── Pass C: ADAM_V_BOUND[g] p99 ──
let true_p99_v = sample_p99(&adam_v_host);
let kernel_p99_v = scratch[SCRATCH_BASE_V + g_idx];
let rel_v = ((kernel_p99_v - true_p99_v) / true_p99_v).abs();
println!(
"SP4 param_group[{g_idx}] ADAM_V — true_p99={true_p99_v:.5}, kernel_p99={kernel_p99_v:.5}, rel_err={rel_v:.5}",
);
assert!(kernel_p99_v > 0.0, "ADAM_V p99 must be positive");
assert!(
rel_v < 0.05,
"group {g_idx} ADAM_V rel_err {rel_v} > 5% (kernel={kernel_p99_v}, true={true_p99_v})",
);
// ── Pass D: WD_RATE[g] = |Σ w·g| / max(Σ w², EPS_DIV) ──
// Analytically computed in f64 host-side then cast to f32 to bound
// accumulation order error. The kernel does the same reduction in
// f32 only, so a 2% tolerance comfortably covers summation drift
// across 4096 elements.
let mut sum_wg: f64 = 0.0;
let mut sum_w2: f64 = 0.0;
for i in 0..N {
let w = params_host[i] as f64;
let g = grads_host[i] as f64;
sum_wg += w * g;
sum_w2 += w * w;
}
const EPS_DIV: f64 = 1.0e-8;
let true_wd_rate = (sum_wg.abs() / sum_w2.max(EPS_DIV)) as f32;
let kernel_wd_rate = scratch[SCRATCH_BASE_WD + g_idx];
let rel_wd = ((kernel_wd_rate - true_wd_rate) / true_wd_rate).abs();
println!(
"SP4 param_group[{g_idx}] WD_RATE — true={true_wd_rate:.6e}, kernel={kernel_wd_rate:.6e}, rel_err={rel_wd:.5}",
);
assert!(
kernel_wd_rate > 0.0,
"WD_RATE must be positive (signed grads × signed-aligned weights yields non-zero |Σ w·g|)",
);
assert!(
rel_wd < 0.02,
"group {g_idx} WD_RATE rel_err {rel_wd} > 2% (kernel={kernel_wd_rate}, true={true_wd_rate})",
);
// ── Pass E (group 0 only): L1_LAMBDA_TRUNK ──
// Analytical reference computes per-feature L2 norms of the
// [h_dim × k_in] gradient matrix, normalises them to a discrete
// distribution, then computes (mean|g| / max(mean|w|, EPS_DIV)) ×
// (log K − H) / log K. Mirrors the kernel's Pass E exactly.
if g_idx == 0 {
let k_in_us = k_in as usize;
let h_dim_us = h_dim as usize;
assert_eq!(k_in_us * h_dim_us, N);
// Per-feature L2 norms (column-wise across the [h_dim × k_in]
// row-major grad matrix → feature `i` indexed at
// [h*k_in + i] for h ∈ 0..h_dim).
let mut g_feat: Vec<f64> = vec![0.0; k_in_us];
for h in 0..h_dim_us {
for i in 0..k_in_us {
let gv = grads_host[h * k_in_us + i] as f64;
g_feat[i] += gv * gv;
}
}
for v in g_feat.iter_mut() {
*v = v.sqrt();
}
let total: f64 = g_feat.iter().sum();
let mut deficit: f64 = 0.0;
if total > EPS_DIV && k_in_us > 1 {
let log_k = (k_in_us as f64).ln();
let mut entropy = 0.0_f64;
for &v in g_feat.iter() {
let p = v / total;
if p > EPS_DIV {
entropy -= p * p.ln();
}
}
deficit = ((log_k - entropy) / log_k).clamp(0.0, 1.0);
}
let mean_abs_g: f64 =
grads_host.iter().map(|&g| (g as f64).abs()).sum::<f64>() / (N as f64);
let mean_abs_w: f64 =
params_host.iter().map(|&w| (w as f64).abs()).sum::<f64>() / (N as f64);
let scale = mean_abs_g / mean_abs_w.max(EPS_DIV);
let true_l1 = (scale * deficit) as f32;
let kernel_l1 = scratch[SCRATCH_L1_TRUNK];
// Both numerator and deficit can be small (near-uniform feature
// gradient distribution → deficit ≈ 0, l1 ≈ 0). Use the same
// 5% rel_err budget as the p99 family for non-degenerate cases;
// when the analytical reference hits the natural-zero floor of
// the formula, the kernel's `__threadfence_system()` still
// wrote 0 to the slot but the rel_err formula is undefined —
// assert kernel_l1 ≈ 0 directly in that branch.
println!(
"SP4 param_group[0] L1_LAMBDA_TRUNK — true={true_l1:.5e}, kernel={kernel_l1:.5e}, deficit={deficit:.5}",
);
if true_l1.abs() < 1e-6 {
assert!(
kernel_l1.abs() < 1e-3,
"group 0 L1 ≈ 0 expected, kernel returned {kernel_l1}",
);
} else {
let rel_l1 = ((kernel_l1 - true_l1) / true_l1).abs();
assert!(
rel_l1 < 0.05,
"group 0 L1 rel_err {rel_l1} > 5% (kernel={kernel_l1}, true={true_l1})",
);
}
}
}
// Final layout sanity — `weight_bound(g)` etc. correctly resolve to the
// ISV slots the host launcher reads/writes. The kernel-direct test does
// NOT touch ISV (no GpuDqnTrainer), so this is purely a const-fn drift
// guard mirroring Tasks A5/A6.
for g in 0..SP4_PARAM_GROUP_COUNT {
assert!(
weight_bound(g) >= 136 && weight_bound(g) < 144,
"weight_bound({g}) = {} out of [136, 144)", weight_bound(g),
);
assert!(adam_m_bound(g) >= 144 && adam_m_bound(g) < 152);
assert!(adam_v_bound(g) >= 152 && adam_v_bound(g) < 160);
assert!(wd_rate(g) >= 160 && wd_rate(g) < 168);
}
}

2
docs/dqn-wire-up-audit.md
View File

File diff suppressed because one or more lines are too long