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

feat(sp4): Task A8 — grad_norm_p99 producer (degenerate single-element)

Browse Source 
grad_norm_buf is single-element f32; p99 == the element. Kernel copies
to scratch slot 38 with __threadfence_system; host applies Pearls A+D
to update ISV[GRAD_CLIP_BOUND_INDEX=168]. Trivial pattern — establishes
the launcher template for Mech 6's eventual ISV read in Layer B.

GPU test verifies first-observation Pearl A replacement (input=42.0 →
scratch[38]=42.0 bit-identical → new_x_mean=42.0, x_lag=42.0, both
variances stay zero) and that all non-target scratch slots remain 0.

No consumer wired yet. Behaviour unchanged. cargo check --lib --tests clean.

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

11
crates/ml/build.rs
View File

@@ -285,6 +285,17 @@ fn main() {
// tuples, then applies Pearls A+D host-side per output via
// `pearls_ad_update`.
"param_group_oracle_kernel.cu",
// SP4 Layer A Task A8 (2026-04-30): grad_norm_p99 producer
// (degenerate single-element). `grad_norm_buf` is a single-element
// f32 device buffer; p99 over 1 element IS the element itself, so
// the kernel reduces to a single-thread copy of `grad_norm_buf[0]`
// into `producer_step_scratch_buf[scratch_idx=38]` with
// `__threadfence_system()`. The host launcher
// (`GpuDqnTrainer::launch_sp4_grad_norm_p99`) syncs, applies Pearls
// A+D via `pearls_ad_update`, writes the new x_mean to
// ISV[GRAD_CLIP_BOUND_INDEX=168] + Wiener state. Establishes the
// launcher template that Mech 6's eventual Layer-B ISV read consumes.
"grad_norm_p99_kernel.cu",
];
// ALL kernels get common header (BF16 types + wrappers)

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

@@ -189,6 +189,18 @@ static SP4_ATOM_POS_P99_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "
static SP4_PARAM_GROUP_ORACLE_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/param_group_oracle_kernel.cubin"));
/// SP4 Layer A Task A8 (2026-04-30): grad_norm_p99 producer cubin.
/// `grad_norm_buf` is a single-element f32 device buffer; p99 over 1 element
/// is the element itself, so the kernel reduces to a single-thread copy from
/// `grad_norm_buf[0]` into `producer_step_scratch_buf[scratch_idx]` with
/// `__threadfence_system()`. The launcher
/// (`GpuDqnTrainer::launch_sp4_grad_norm_p99`) syncs, applies Pearls A+D
/// (`pearls_ad_update`), and writes the new x_mean to
/// ISV[GRAD_CLIP_BOUND_INDEX=168] + Wiener state. No consumer wired yet —
/// Mech 6's adaptive_clip upper-bound consumer migrates in Layer B. Cold-path
/// launch in this task; Task A10 may move to the captured graph.
static SP4_GRAD_NORM_P99_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/grad_norm_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
@@ -3424,6 +3436,17 @@ pub struct GpuDqnTrainer {
/// `SP4_PARAM_GROUP_ORACLE_CUBIN`.
param_group_oracle_update: CudaFunction,
/// SP4 Layer A Task A8: grad_norm_p99 producer (degenerate single-element).
/// `grad_norm_buf` is a single-element f32 device buffer (populated by
/// `launch_grad_norm_finalize`). The kernel reduces to a single-thread
/// copy of `grad_norm_buf[0]` → `producer_step_scratch_buf[scratch_idx=38]`
/// with `__threadfence_system()`. The launcher
/// (`launch_sp4_grad_norm_p99`) syncs, applies Pearls A+D via
/// `pearls_ad_update` and writes back to ISV[GRAD_CLIP_BOUND_INDEX=168] +
/// Wiener state at `wiener_state_buf[(168-SP4_SLOT_BASE)*3 = 111..114)`.
/// Loaded from `SP4_GRAD_NORM_P99_CUBIN`.
grad_norm_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`
@@ -9158,6 +9181,122 @@ impl GpuDqnTrainer {
}
}
/// SP4 Layer A Task A8: cold-path producer for ISV[GRAD_CLIP_BOUND_INDEX=168].
///
/// `grad_norm_buf` is a single-element f32 device buffer populated by
/// `launch_grad_norm_finalize` with the current step's L2 grad norm.
/// p99 over a single element is the element itself — the kernel
/// (`grad_norm_p99_update`) reduces to a single-thread copy of
/// `grad_norm_buf[0]` → `producer_step_scratch_buf[SCRATCH_IDX=38]` with
/// `__threadfence_system()`. No histogram pass is needed; the launcher
/// keeps the same shape as `launch_sp4_target_q_p99` for consistency
/// across SP4 producers, and Mech 6's eventual Layer-B ISV read in the
/// adaptive_clip path consumes ISV[168] directly.
///
/// Synchronises the stream then applies Pearls A+D host-side via
/// `pearls_ad_update`, using zero-copy mapped-pinned reads of ISV[168]
/// + the per-slot Wiener state at `wiener_state_buf[(168 -
/// SP4_SLOT_BASE) * 3 = 111..114)`. Writes the new `x_mean` back to
/// ISV[168] and the updated Wiener state back to `wiener_state_buf` via
/// the same mapped-pinned host pointers — no HtoD/DtoH per
/// `feedback_no_htod_htoh_only_mapped_pinned`.
///
/// **Cold-path launch** — NOT in the captured graph for now. Task A10 may
/// relocate to captured-graph cadence; this launcher mirrors A5's shape
/// so the migration is structural-only at that point.
///
/// **No consumer wired in this task** — Mech 6's `adaptive_clip` upper
/// bound still uses the existing fast/slow grad-norm EMA path; SP4
/// consumer migration follows once all producers (A5-A9) land.
pub fn launch_sp4_grad_norm_p99(&self) -> Result<(), MLError> {
use crate::cuda_pipeline::sp4_isv_slots::{GRAD_CLIP_BOUND_INDEX, SP4_SLOT_BASE};
use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
// Producer-step-scratch slot 38 is reserved for GRAD_CLIP_BOUND per
// the stable layout documented on `launch_sp4_target_q_p99`:
// 0 = TARGET_Q_BOUND (Task A5)
// 1..5 = ATOM_POS_BOUND[0..4] (Task A6)
// 5..29 = WEIGHT/ADAM_M/ADAM_V × 8 (Task A7)
// 29..37 = WD_RATE × 8 (Task A7)
// 37 = L1_LAMBDA_TRUNK (Task A7)
// 38 = GRAD_CLIP_BOUND (this task)
// 39 = H_S2_BOUND (Task A9)
// 40..47 = retrofit existing 7 producers
const SCRATCH_IDX: usize = 38;
// Single-thread degenerate kernel — minimal launch config, no shared
// memory. Mirrors the `q_drift_rate_ema_kernel.cu` /
// `moe_lambda_eff_kernel.cu` cold-path single-thread footprint.
let cfg = LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
};
let grad_norm_ptr = self.ptrs.grad_norm_buf;
let scratch_dev = self.producer_step_scratch_buf.dev_ptr;
let scratch_idx_arg: i32 = SCRATCH_IDX as i32;
// Safety: kernel signature `(const float*, float*, int)` matches
// the three args below; both device pointers are valid for the
// lifetime of `self`.
unsafe {
self.stream
.launch_builder(&self.grad_norm_p99_update)
.arg(&grad_norm_ptr)
.arg(&scratch_dev)
.arg(&scratch_idx_arg)
.launch(cfg)
.map_err(|e| MLError::ModelError(format!("grad_norm_p99_update launch: {e}")))?;
}
// Cold-path producer — sync before reading the mapped-pinned scratch
// slot from the host. Captured-graph migration (Task A10) replaces
// this with the same fence the captured stream provides on graph
// launch boundaries.
self.stream.synchronize()
.map_err(|e| MLError::ModelError(format!("grad_norm_p99_update sync: {e}")))?;
// ── Pearls A+D host-side update (zero-copy mapped-pinned reads) ──
let step_obs = unsafe {
std::ptr::read_volatile(self.producer_step_scratch_buf.host_ptr.add(SCRATCH_IDX))
};
// Degenerate-signal short-circuit: kernel writes 0.0 if grad_norm
// hasn't been populated yet (e.g., before first backward). Avoid
// advancing Wiener state with a zero observation, which would
// suppress the next genuine first-observation. Same semantics as
// Tasks A5/A6/A7.
if step_obs == 0.0 { return Ok(()); }
let isv_idx = GRAD_CLIP_BOUND_INDEX;
let wiener_offset = (isv_idx - SP4_SLOT_BASE) * 3;
// Safety: `isv_signals_pinned` is a mapped-pinned `*mut f32` of
// length `ISV_TOTAL_DIM`; `isv_idx < ISV_TOTAL_DIM` (171 > 168).
// `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].
///
@@ -10809,6 +10948,21 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("param_group_oracle_update load: {e}")))?
};
// SP4 Layer A Task A8: load grad_norm_p99 producer kernel (cold-path).
// Single-thread single-block kernel that copies `grad_norm_buf[0]`
// (single-element scalar populated by `launch_grad_norm_finalize`) to
// `producer_step_scratch_buf[38]` with `__threadfence_system()`.
// p99 over a single element IS the element itself, hence the
// degenerate-copy form (no histogram). Pearls A+D applied host-side
// via `pearls_ad_update` (Task A3). Producer-only — Mech 6's
// adaptive_clip upper bound consumer migrates in Layer B.
let grad_norm_p99_update = {
let module = stream.context().load_cubin(SP4_GRAD_NORM_P99_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("sp4 grad_norm_p99 cubin load: {e}")))?;
module.load_function("grad_norm_p99_update")
.map_err(|e| MLError::ModelError(format!("grad_norm_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`.
@@ -13770,6 +13924,7 @@ impl GpuDqnTrainer {
target_q_p99_update,
atom_pos_p99_update,
param_group_oracle_update,
grad_norm_p99_update,
q_drift_rate_ema_kernel,
fold_warmup_factor_kernel,
grad_norm_fast_ema_pinned,

30
crates/ml/src/cuda_pipeline/grad_norm_p99_kernel.cu Normal file
View File

@@ -0,0 +1,30 @@
// crates/ml/src/cuda_pipeline/grad_norm_p99_kernel.cu
//
// SP4 Layer A Task A8 — grad_norm_p99 producer (degenerate single-element).
//
// `grad_norm_buf` is a single-element f32 device buffer containing the current
// step's L2 grad norm (populated by `launch_grad_norm_finalize`). p99 over a
// single element IS the element itself, so the "histogram" reduces to a copy:
// the single-thread kernel reads `grad_norm_buf[0]` and writes it 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_grad_norm_p99`) syncs, applies
// Pearls A+D via `pearls_ad_update`, and writes the new x_mean to
// ISV[GRAD_CLIP_BOUND_INDEX=168] + Wiener state at slot 38 of
// `wiener_state_buf` (offset `(168 - SP4_SLOT_BASE) * 3 = 111`).
//
// Cold-path launch — NOT in the captured graph for now. Task A10 may relocate.
// Task A8 establishes the launcher template that Mech 6's eventual ISV read
// in Layer B will consume. Mirrors Task A5's launcher shape end-to-end with
// the kernel itself reduced to a copy.
extern "C" __global__ void grad_norm_p99_update(
const float* __restrict__ grad_norm_buf, /* [1] */
float* __restrict__ scratch_buf,
int scratch_idx
) {
if (blockIdx.x != 0 || threadIdx.x != 0) return;
scratch_buf[scratch_idx] = grad_norm_buf[0];
__threadfence_system(); // make host-mapped write visible
}

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

@@ -998,3 +998,140 @@ fn sp4_param_group_oracle_per_group_writes_distinct_isv_slots() {
assert!(wd_rate(g) >= 160 && wd_rate(g) < 168);
}
}
// ── SP4 Task A8: grad_norm_p99 producer (degenerate single-element) ───────────
/// Test-only cubin for the SP4 grad_norm_p99 producer kernel. The same cubin
/// is loaded by `GpuDqnTrainer::new` in production via `SP4_GRAD_NORM_P99_CUBIN`;
/// this kernel-direct test bypasses the full trainer harness and exercises
/// the kernel's degenerate single-thread copy contract.
const SP4_GRAD_NORM_P99_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/grad_norm_p99_kernel.cubin"));
/// Resolve the `grad_norm_p99_update` function handle from its cubin.
fn load_sp4_grad_norm_p99_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
let module = stream
.context()
.load_cubin(SP4_GRAD_NORM_P99_CUBIN.to_vec())
.expect("load grad_norm_p99_kernel cubin");
module
.load_function("grad_norm_p99_update")
.expect("load grad_norm_p99_update function")
}
/// SP4 Task A8 unit test: the production `grad_norm_p99_update` kernel
/// performs a single-thread copy of `grad_norm_buf[0]` to
/// `producer_step_scratch_buf[scratch_idx=38]` and the host-side
/// `pearls_ad_update` helper (Task A3) replaces the ISV sentinel directly
/// with that step observation on the first call after a fold-boundary
/// reset (Pearl A bootstrap).
///
/// `grad_norm_buf` is a single-element f32 device buffer in production;
/// p99 over 1 element IS the element itself, so the kernel is degenerate
/// (no histogram). This test sets that single element to a known scalar
/// (42.0) and verifies the kernel writes exactly that value to slot 38
/// while leaving every other slot zero.
///
/// Production launcher contract validated:
/// - Kernel writes the input scalar to `scratch[SCRATCH_IDX_GRAD_CLIP=38]`.
/// - All other scratch slots remain zero (helper assertion guards the
/// `2fb30f098`-style write-cascade pattern).
/// - `GRAD_CLIP_BOUND_INDEX == 168` and `(168 - SP4_SLOT_BASE) * 3 == 111`
/// (Wiener state offset for this slot).
#[test]
#[ignore = "requires GPU"]
fn sp4_grad_norm_p99_pearl_a_first_observation_replaces_scalar() {
use ml::cuda_pipeline::sp4_isv_slots::{GRAD_CLIP_BOUND_INDEX, SP4_SLOT_BASE};
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
// Production scratch shape (47 = SP4 producers count). Slot 38 is the
// production-allocated GRAD_CLIP_BOUND slot (per the documented stable
// layout in `launch_sp4_target_q_p99` / `launch_sp4_grad_norm_p99`).
const SP4_PRODUCER_COUNT: usize = 47;
const SCRATCH_IDX_GRAD_CLIP: usize = 38;
// Single-element grad_norm scalar — known synthetic value chosen to be
// distinguishable from the Pearl-A sentinel (0.0) and from a stale
// initialisation (1.0) so a "did the kernel actually copy?" question
// is answered unambiguously.
let grad_norm_scalar: f32 = 42.0;
let stream = make_test_stream();
let kernel = load_sp4_grad_norm_p99_kernel(&stream);
// Safety: CUDA context active on this thread (resolved via stream).
let in_buf = unsafe { MappedF32Buffer::new(1) }
.expect("alloc MappedF32Buffer for SP4 grad_norm_p99 input");
in_buf.write_from_slice(&[grad_norm_scalar]);
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
.expect("alloc MappedF32Buffer for SP4 producer scratch");
let scratch_idx_arg: i32 = SCRATCH_IDX_GRAD_CLIP as i32;
let in_dev_ptr = in_buf.dev_ptr;
let scratch_dev_ptr = scratch_buf.dev_ptr;
// Single-thread degenerate kernel — minimal launch config, no shared
// memory. Mirrors the production `launch_sp4_grad_norm_p99` config.
unsafe {
stream
.launch_builder(&kernel)
.arg(&in_dev_ptr)
.arg(&scratch_dev_ptr)
.arg(&scratch_idx_arg)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (1, 1, 1),
shared_mem_bytes: 0,
})
.expect("launch grad_norm_p99_update");
}
stream
.synchronize()
.expect("synchronize after grad_norm_p99_update launch");
// Verify ALL non-target slots remained zero — the kernel must only
// touch the slot at `scratch_idx=38`.
let host = scratch_buf.read_all();
for (i, &v) in host.iter().enumerate() {
if i != SCRATCH_IDX_GRAD_CLIP {
assert_eq!(
v, 0.0,
"grad_norm_p99 producer wrote to scratch[{i}] outside target slot {SCRATCH_IDX_GRAD_CLIP}: {v}",
);
}
}
let step_obs = host[SCRATCH_IDX_GRAD_CLIP];
println!(
"SP4 grad_norm_p99 — input={grad_norm_scalar:.5}, step_obs={step_obs:.5}",
);
// Degenerate copy — bit-identical equality (no histogram quantization).
assert_eq!(
step_obs, grad_norm_scalar,
"kernel must copy grad_norm_buf[0] verbatim to scratch[{SCRATCH_IDX_GRAD_CLIP}]; got {step_obs} vs {grad_norm_scalar}",
);
// ── Pearl A bootstrap (host-side, mirrors `launch_sp4_grad_norm_p99`) ──
// Sentinel ISV slot + zero Wiener state — replicates the cold-start
// condition immediately after a fold-boundary reset.
let prev_x_mean: f32 = 0.0;
let mut state = WienerState::ZERO;
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_obs);
assert_eq!(
new_x_mean, step_obs,
"Pearl A: first observation must replace sentinel directly; got {new_x_mean} vs step_obs={step_obs}",
);
assert_eq!(
state.x_lag, step_obs,
"Pearl A: x_lag must seed to first observation",
);
assert_eq!(state.sample_var, 0.0, "Pearl A: sample_var stays zero");
assert_eq!(state.diff_var, 0.0, "Pearl A: diff_var stays zero");
// Slot-layout invariants — guard against off-by-one drift in
// GRAD_CLIP_BOUND_INDEX or SP4_SLOT_BASE.
assert_eq!(GRAD_CLIP_BOUND_INDEX, 168);
assert_eq!(SP4_SLOT_BASE, 131);
assert_eq!((GRAD_CLIP_BOUND_INDEX - SP4_SLOT_BASE) * 3, 111);
}

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

File diff suppressed because one or more lines are too long