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

feat(sp4): Task A13.0 — retrofit h_s2_rms_ema with Pearls A+D + buffer growth

Browse Source 
Replaces the kernel's hardcoded `ema_alpha` with the shared `pearls_ad_update`
host-side helper (Task A3). Buffer growth in same commit so subsequent A13.x
retrofits land on a stable scratch/Wiener layout.

  - Kernel `h_s2_rms_ema_update` signature: `(h_s2, B, SH2, scratch_buf,
    scratch_idx)`. Reduces RMS = sqrt(sum_sq/(B*SH2)) via the existing
    256-thread shmem tree (no atomicAdd) and writes step_obs to
    `producer_step_scratch_buf[40]` with `__threadfence_system()`.
  - Launcher `launch_h_s2_rms_ema(_ema_alpha_unused: f32)` syncs the stream,
    then applies Pearls A+D host-side via zero-copy mapped-pinned reads/
    writes of `isv_signals_pinned[H_S2_RMS_EMA_INDEX=96]` and
    `wiener_state_buf[120..123)` (= scratch slot 40 × 3). Degenerate-zero
    short-circuit before mutating ISV/Wiener state.
  - Buffer growth: `SP4_PRODUCER_COUNT 47 → 69` (40 SP4 + 29 Task A13
    retrofit producers); `wiener_state_buf 141 → 207` floats;
    `producer_step_scratch_buf` grows to 69 entries.
  - Stable-layout doc-comments updated: `producer_step_scratch_buf` field
    comment, `launch_sp4_target_q_p99` slot-table comment, 5 launcher
    `wiener_offset + 2 < 141` safety comments (now `< 207`),
    `reset_sp4_wiener_state` doc + body comment, and
    `state_reset_registry.rs::sp4_wiener_state` description.
  - Test cubin reference `SP4_PRODUCER_COUNT: usize = 47` in
    `tests/sp4_producer_unit_tests.rs` updated to 69 across all 6
    occurrences.

Behavior: stationary signals converge to the same RMS at adaptive rate
(Pearl D's α* derived from per-slot signal-vs-noise variance);
non-stationary signals respond Wiener-optimally faster. Cold-path producer
with no consumer-facing change beyond the EMA mechanism — slot 96 is read
by `mag_concat_qdir`'s adaptive-scale path and stays semantically identical
(RMS of `save_h_s2`).

Tests: new `sp4_h_s2_rms_ema_writes_step_rms_via_pearl_a_then_converges_pearl_d`
unit test (`#[ignore]`-gated for GPU) drives the production kernel kernel-
direct on a constant-5.0 stationary signal of B=4 SH2=64, asserts
step_rms ≈ analytical RMS=5.0 ± 1e-4, asserts non-target scratch slots
remain 0, then exercises Pearl A bootstrap (returns step_rms directly +
seeds x_lag) and Pearl D convergence (1000 stationary observations →
x_mean within 1% of 5.0).

Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`,
`feedback_no_partial_refactor.md`. Build: `cargo check -p ml --lib --tests
--offline` clean (11 pre-existing warnings, no new warnings); `cargo test
-p ml --lib sp4_wiener_ema --offline` 6/6 passing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-01 02:03:50 +02:00
parent 138d41b761
commit aada419de3
5 changed files with 266 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2936,12 +2936,23 @@ pub struct GpuDqnTrainer {
/// at fold boundary via the state-reset registry (Task A12).
pub(crate) pearl_c_rate_deficit_ema: [f32; SP4_PARAM_GROUP_COUNT],
/// SP4 producer step-observation scratch — each of the 47 producers
/// (40 SP4 + 7 retrofit) writes its per-step `step_observation` to
/// SP4 producer step-observation scratch — each of the 69 producers
/// (40 SP4 + 29 retrofit) writes its per-step `step_observation` to
/// its dedicated slot here. Host then applies Pearls A+D
/// (`pearls_ad_update`, Task A3) to compute the new `x_mean` and writes
/// to the corresponding ISV bound slot. Mapped-pinned f32, zero-copy
/// host reads via mapped-pinned device-ptr.
///
/// Layout (post Task A13):
/// [0..40) — SP4 producers (Tasks A5-A11)
/// [40] — h_s2_rms_ema → ISV[96] (A13.0)
/// [41..43) — aux_heads_loss_ema → ISV[113],ISV[114] (A13.1)
/// [43] — aux_label_scale_ema → ISV[117] (A13.1)
/// [44..52) — moe_expert_util_ema → ISV[118..126) (A13.2)
/// [52] — moe_gate_entropy_ema → ISV[126] (A13.2)
/// [53..59) — vsn_mask_ema → ISV[105..111) (A13.3)
/// [59..63) — iqn_quantile_ema → ISV[99..103]\med (A13.4)
/// [63..69) — reward_component_ema → ISV[63..69) (A13.5)
pub(crate) producer_step_scratch_buf: MappedF32Buffer,
/// SP4 Task A7 fix-up #2: oracle sub-buffer device-pointer table.
@@ -8784,39 +8795,53 @@ impl GpuDqnTrainer {
Ok(())
}
/// Plan 4 Task 2c.3c.5: launch `h_s2_rms_ema_update` (single-block,
/// 256-thread shmem-reduction kernel). Reduces RMS = sqrt(sum_sq / (B*SH2))
/// over the trainer's `save_h_s2 [B, SH2]` buffer (the online trunk's
/// post-GRN activation, valid after the most recent forward) and EMAs the
/// result into ISV[H_S2_RMS_EMA_INDEX=96].
/// Plan 4 Task 2c.3c.5; SP4 Layer A Task A13.0 retrofit (2026-05-01):
/// launch `h_s2_rms_ema_update` (single-block, 256-thread shmem-reduction
/// kernel). Reduces RMS = sqrt(sum_sq / (B*SH2)) over the trainer's
/// `save_h_s2 [B, SH2]` buffer and writes the step observation to
/// `producer_step_scratch_buf[SCRATCH_IDX=40]`. Synchronises the stream,
/// then applies Pearls A+D host-side via `pearls_ad_update` using
/// zero-copy mapped-pinned reads of ISV[H_S2_RMS_EMA_INDEX=96] +
/// `wiener_state_buf[(40)*3..(40)*3+3]` — slot 40 in the SP4 producer
/// scratch / Wiener-state layout (first retrofit slot after the 40 SP4
/// producers).
///
/// **α dropped per SP4 — Pearls A+D adapt α from per-slot signal-vs-noise
/// variance.** The `_ema_alpha_unused` argument is preserved so callers
/// (training_loop.rs) compile unchanged; α is now derived adaptively per
/// step inside `pearls_ad_update`.
///
/// Producer-only — 2c.3c.6 wires the consumer in `mag_concat_qdir`'s
/// adaptive-scale path. Cold-path-cadence: launched once per training step
/// alongside `launch_reward_component_ema`. No atomicAdd, no DtoH.
pub fn launch_h_s2_rms_ema(&self, ema_alpha: f32) -> Result<(), MLError> {
// `isv_signals_dev_ptr` is set by the constructor via
// `cuMemHostGetDevicePointer_v2` on a `cuMemAllocHost_v2`-backed pinned
// buffer (with its own `assert_eq!` on success) and is never reassigned,
// so non-zero is an invariant. Silent-skip on `== 0` would mask a
// misconfiguration; `debug_assert!` flags it in dev builds and release
// inherits the loud kernel-launch error rather than a no-op.
pub fn launch_h_s2_rms_ema(&self, _ema_alpha_unused: f32) -> Result<(), MLError> {
use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
// `isv_signals_dev_ptr` / `isv_signals_pinned` are set by the constructor
// via `cuMemHostGetDevicePointer_v2` on a `cuMemAllocHost_v2`-backed
// pinned buffer and are never reassigned. Silent-skip on `== 0` would
// mask a misconfiguration; `debug_assert!` flags it in dev builds.
debug_assert!(self.isv_signals_dev_ptr != 0,
"launch_h_s2_rms_ema: isv_signals_dev_ptr must be allocated by constructor");
// Task A13.0: scratch slot 40 — first retrofit slot.
const SCRATCH_IDX: usize = 40;
const BLOCK_DIM: u32 = 256;
let smem_bytes = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
let h_s2_ptr = self.save_h_s2.raw_ptr();
let b_i = self.config.batch_size as i32;
let sh2_i = self.config.shared_h2 as i32;
let isv_idx = H_S2_RMS_EMA_INDEX as i32;
let isv_dev_ptr = self.isv_signals_dev_ptr;
const BLOCK_DIM: u32 = 256;
let smem_bytes = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
let scratch_dev = self.producer_step_scratch_buf.dev_ptr;
let scratch_idx_arg: i32 = SCRATCH_IDX as i32;
unsafe {
self.stream.launch_builder(&self.h_s2_rms_ema_kernel)
.arg(&h_s2_ptr)
.arg(&b_i)
.arg(&sh2_i)
.arg(&isv_dev_ptr)
.arg(&isv_idx)
.arg(&ema_alpha)
.arg(&scratch_dev)
.arg(&scratch_idx_arg)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (BLOCK_DIM, 1, 1),
@@ -8824,6 +8849,51 @@ impl GpuDqnTrainer {
})
.map_err(|e| MLError::ModelError(format!("h_s2_rms_ema_update: {e}")))?;
}
// Cold-path producer — sync before reading the mapped-pinned scratch
// slot from the host. Captured-graph migration would replace this
// with the captured-stream's launch-boundary fence.
self.stream.synchronize()
.map_err(|e| MLError::ModelError(format!("h_s2_rms_ema 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 every
// sample is exactly 0 (e.g., before the first online forward
// populates `save_h_s2`). Updating the Wiener state with 0 would
// advance `x_lag` and suppress the next genuine first-observation.
if step_obs == 0.0 { return Ok(()); }
let isv_idx = H_S2_RMS_EMA_INDEX;
// Wiener offset: SCRATCH_IDX * 3 (one Wiener triple per scratch slot).
// SCRATCH_IDX=40 → wiener_offset=120, well within 207 floats.
let wiener_offset = SCRATCH_IDX * 3;
// Safety: `isv_signals_pinned` is mapped-pinned `*mut f32` of length
// ISV_TOTAL_DIM (>= 171); H_S2_RMS_EMA_INDEX=96 < ISV_TOTAL_DIM.
// `wiener_state_buf.host_ptr` is mapped-pinned `*mut f32` of length
// `SP4_PRODUCER_COUNT * 3 = 207`; `wiener_offset + 2 < 207`.
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(())
}
@@ -8861,7 +8931,15 @@ impl GpuDqnTrainer {
// 37 = L1_LAMBDA_TRUNK
// 38 = GRAD_CLIP_BOUND
// 39 = H_S2_BOUND
// 40..47 = retrofit existing 7 producers
// ── Task A13 retrofit producers (Pearls A+D for existing EMA kernels) ──
// 40 = H_S2_RMS_EMA → ISV[96] (A13.0)
// 41..43 = AUX_HEADS_LOSS_EMA → ISV[113],ISV[114] (A13.1)
// 43 = AUX_LABEL_SCALE_EMA → ISV[117] (A13.1)
// 44..52 = MOE_EXPERT_UTIL_EMA → ISV[118..126) (A13.2)
// 52 = MOE_GATE_ENTROPY_EMA → ISV[126] (A13.2)
// 53..59 = VSN_MASK_EMA → ISV[105..111) (A13.3)
// 59..63 = IQN_QUANTILE_EMA → ISV[99..103]\med (A13.4)
// 63..69 = REWARD_COMPONENT_EMA → ISV[63..69) (A13.5)
const SCRATCH_IDX: usize = 0;
// Shared memory: 8 warps × 256 bins × sizeof(int) = 8192 bytes.
@@ -8922,7 +9000,7 @@ impl GpuDqnTrainer {
// Safety: `isv_signals_pinned` is a mapped-pinned `*mut f32` of
// length `ISV_TOTAL_DIM`; `isv_idx < ISV_TOTAL_DIM` (171 > 131).
// `wiener_state_buf.host_ptr` is a mapped-pinned `*mut f32` of
// length `SP4_PRODUCER_COUNT * 3 = 141`; `wiener_offset + 2 < 141`.
// length `SP4_PRODUCER_COUNT * 3 = 207`; `wiener_offset + 2 < 207`.
let prev_x_mean = unsafe { std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx)) };
let mut state = unsafe {
@@ -9064,7 +9142,7 @@ impl GpuDqnTrainer {
// 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`.
// length `SP4_PRODUCER_COUNT * 3 = 207`; `wiener_offset + 2 < 207`.
let prev_x_mean = unsafe {
std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx))
};
@@ -9327,7 +9405,7 @@ impl GpuDqnTrainer {
// 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.
// length `SP4_PRODUCER_COUNT * 3 = 207`; wiener_offset + 2 < 207.
let prev_x_mean = unsafe {
std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx))
};
@@ -9550,7 +9628,7 @@ impl GpuDqnTrainer {
// 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`.
// length `SP4_PRODUCER_COUNT * 3 = 207`; `wiener_offset + 2 < 207`.
let prev_x_mean = unsafe { std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx)) };
let mut state = unsafe {
@@ -9660,7 +9738,7 @@ impl GpuDqnTrainer {
// Safety: `isv_signals_pinned` is a mapped-pinned `*mut f32` of
// length `ISV_TOTAL_DIM`; `isv_idx < ISV_TOTAL_DIM` (171 > 169).
// `wiener_state_buf.host_ptr` is a mapped-pinned `*mut f32` of
// length `SP4_PRODUCER_COUNT * 3 = 141`; `wiener_offset + 2 < 141`.
// length `SP4_PRODUCER_COUNT * 3 = 207`; `wiener_offset + 2 < 207`.
let prev_x_mean = unsafe { std::ptr::read_volatile(self.isv_signals_pinned.add(isv_idx)) };
let mut state = unsafe {
@@ -12808,9 +12886,9 @@ impl GpuDqnTrainer {
// Per `feedback_no_htod_htoh_only_mapped_pinned`: mapped pinned
// (`cuMemHostAlloc(DEVICEMAP|PORTABLE)`) is the only allowed CPU↔GPU
// path for these state buffers.
const SP4_PRODUCER_COUNT: usize = 47; // 40 SP4 + 7 retrofit existing producers
const SP4_PRODUCER_COUNT: usize = 69; // 40 SP4 + 29 Task-A13 retrofit producers
const SP4_WIENER_FLOATS_PER_SLOT: usize = 3; // sample_var, diff_var, x_lag
const SP4_WIENER_TOTAL_FLOATS: usize = SP4_PRODUCER_COUNT * SP4_WIENER_FLOATS_PER_SLOT; // 141
const SP4_WIENER_TOTAL_FLOATS: usize = SP4_PRODUCER_COUNT * SP4_WIENER_FLOATS_PER_SLOT; // 207
// SP4 Task A14: SP4_PARAM_GROUP_COUNT, MAX_BLOCKS_PER_ADAM, and
// SP4_ENGAGE_BUF_LEN are now public constants in `sp4_isv_slots`
// so launch sites + helpers can reference them without re-declaring.
@@ -22142,8 +22220,8 @@ impl GpuDqnTrainer {
/// SP4 Layer A Task A12 (2026-04-30): bulk-zero the Pearl D Wiener-EMA
/// state at fold boundary.
///
/// `wiener_state_buf` is a 141-float mapped-pinned buffer holding all
/// 47 producers' Wiener triples `[sample_var, diff_var, x_lag]`. Pearl
/// `wiener_state_buf` is a 207-float mapped-pinned buffer holding all
/// 69 producers' Wiener triples `[sample_var, diff_var, x_lag]`. Pearl
/// A's first-observation sentinel requires both `prev_x_mean` (the ISV
/// bound slot) AND `state.x_lag` (Wiener triple offset 2) to be exactly
/// 0 when a producer first fires in a new fold. Without this reset the
@@ -22161,7 +22239,7 @@ impl GpuDqnTrainer {
// Safety: `host_ptr` is a mapped-pinned `*mut f32` of length
// `self.wiener_state_buf.len` returned by `cuMemHostAlloc`.
// f32 zero == bytewise-zero per IEEE-754 +0.0 representation, so
// `write_bytes(.., 0, ..)` produces 141 valid f32 zeros.
// `write_bytes(.., 0, ..)` produces 207 valid f32 zeros.
unsafe {
std::ptr::write_bytes(
self.wiener_state_buf.host_ptr as *mut u8,

51
crates/ml/src/cuda_pipeline/h_s2_rms_ema_kernel.cu
View File

@@ -1,21 +1,29 @@
/* h_s2_rms_ema — GPU-driven RMS(save_h_s2) EMA into ISV[96].
/* h_s2_rms_ema — GPU-driven RMS(save_h_s2) step observation into
* `producer_step_scratch_buf[scratch_idx]` for Pearls A+D consumption.
*
* Plan 4 Task 2c.3c.5. GPU-drives-CPU-reads (spec §4.C.6) and
* Plan 4 Task 2c.3c.5; SP4 Layer A Task A13.0 retrofit (2026-05-01).
* GPU-drives-CPU-reads (spec §4.C.6) and
* pearl_cold_path_no_exception_to_gpu_drives.md compliant.
*
* Reduces RMS = sqrt(sum_sq / N) over the flattened `save_h_s2 [B, SH2]`
* buffer (the online trunk's post-GRN activation) and EMA-updates
* ISV[H_S2_RMS_EMA_INDEX] with rate `ema_alpha`.
* buffer (the online trunk's post-GRN activation). The reduction result
* is written to `producer_step_scratch_buf[scratch_idx]` with
* `__threadfence_system()`; the host launcher (`launch_h_s2_rms_ema`)
* synchronises the stream and then applies Pearls A+D via
* `pearls_ad_update`, mapping `step_observation` to ISV[H_S2_RMS_EMA_INDEX]
* + `wiener_state_buf` triple. **The hardcoded EMA α dropped per Task A13
* is replaced by per-slot signal-vs-noise variance ratio (Pearl D's
* Wiener-optimal blend); first-observation bootstrap (Pearl A) covers
* fold cold-start so prev=0 is no longer a structural problem.**
*
* Single-block kernel (256 threads, shmem-reduction). No atomicAdd
* (per `feedback_no_atomicadd.md`); the in-block tree reduction is
* standard `s >>= 1` with `__syncthreads()`. Only thread 0 writes the
* EMA result.
* step observation.
*
* Producer-only — no consumer reads ISV[96] in this commit. 2c.3c.6
* wires the consumer in `mag_concat_qdir`'s adaptive-scale path so the
* residual stack is normalised by the trunk-output RMS regardless of
* GRN drift across training.
* Producer-only — the consumer in `mag_concat_qdir`'s adaptive-scale path
* (wired in 2c.3c.6) continues to read ISV[H_S2_RMS_EMA_INDEX]; only the
* EMA mechanism behind the slot changes here.
*
* Cost: one launch per training step (cold-path cadence — same site as
* `reward_component_ema`). 256 threads × ceil(B*SH2/256) strided loads
@@ -26,9 +34,8 @@ extern "C" __global__ void h_s2_rms_ema_update(
const float* __restrict__ h_s2, /* [B, SH2] row-major */
int B,
int SH2,
float* __restrict__ isv, /* ISV bus [ISV_TOTAL_DIM] */
int isv_h_s2_rms_ema_index, /* H_S2_RMS_EMA_INDEX = 96 */
float ema_alpha /* per-batch EMA rate (0,1) */
float* __restrict__ scratch_buf, /* producer_step_scratch_buf [SP4_PRODUCER_COUNT] */
int scratch_idx /* per-producer scratch slot index */
) {
/* Single-block contract — match reward_component_ema_kernel's grid guard. */
if (blockIdx.x != 0) return;
@@ -53,17 +60,15 @@ extern "C" __global__ void h_s2_rms_ema_update(
__syncthreads();
}
/* Thread 0 writes the EMA. The host caller's trainer config requires
* batch_size > 0 and shared_h2 > 0, so N = B*SH2 > 0 is invariant.
* No defensive `(N > 0) ? N : 1` ternary here — masking the divide
* would silently EMA toward 0 if the invariant ever broke; instead the
* NaN that division-by-zero produces propagates to ISV[96] and
* downstream consumers (mag_concat_qdir in 2c.3c.6) where it is
* loudly visible. */
/* Thread 0 writes the step observation. The host caller's trainer
* config requires batch_size > 0 and shared_h2 > 0, so N = B*SH2 > 0
* is invariant. No defensive `(N > 0) ? N : 1` ternary here — masking
* the divide would silently produce 0 if the invariant ever broke;
* the NaN that division-by-zero produces propagates to scratch_buf
* and surfaces loudly in the host-side Pearls A+D update. */
if (tid == 0) {
const float rms = sqrtf(smem[0] / (float)N);
const int slot = isv_h_s2_rms_ema_index;
const float prev = isv[slot];
isv[slot] = (1.0f - ema_alpha) * prev + ema_alpha * rms;
const float rms = sqrtf(smem[0] / (float)N);
scratch_buf[scratch_idx] = rms;
__threadfence_system(); /* PCIe-visible write for mapped pinned host_ptr */
}
}

2
crates/ml/src/trainers/dqn/state_reset_registry.rs
View File

@@ -522,7 +522,7 @@ impl StateResetRegistry {
RegistryEntry {
name: "sp4_wiener_state",
category: ResetCategory::FoldReset,
description: "GpuDqnTrainer.wiener_state_buf [141 floats = 47 producers × {sample_var, diff_var, x_lag}] — SP4 Pearl D Wiener-EMA state. Mapped-pinned f32; bulk write_bytes(0) at fold boundary so Pearl A's sentinel branch fires on the new fold's first producer launch (`prev_x_mean == 0 AND state.x_lag == 0` → first-observation replacement). Companion to the 9 SP4 ISV bound entries above — both halves of the sentinel contract reset together per `feedback_no_partial_refactor.md` (Task A12)",
description: "GpuDqnTrainer.wiener_state_buf [207 floats = 69 producers × {sample_var, diff_var, x_lag}] — SP4 Pearl D Wiener-EMA state (40 SP4 + 29 Task-A13 retrofit producers). Mapped-pinned f32; bulk write_bytes(0) at fold boundary so Pearl A's sentinel branch fires on the new fold's first producer launch (`prev_x_mean == 0 AND state.x_lag == 0` → first-observation replacement). Companion to the 9 SP4 ISV bound entries above — both halves of the sentinel contract reset together per `feedback_no_partial_refactor.md` (Task A12; buffer grew 141→207 in Task A13.0 to cover the 29 retrofit producers — h_s2_rms, aux_heads_loss×2, aux_label_scale, moe_expert_util×8, moe_gate_entropy, vsn_mask×6, iqn_quantile×4, reward_component×6)",
},
RegistryEntry {
name: "sp4_clamp_engage_counters",

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

@@ -324,7 +324,7 @@ fn sp4_target_q_p99_first_observation_writes_step_p99_via_pearls_ad() {
// production-allocated TARGET_Q_BOUND slot; we use the same index so
// the kernel-direct test keeps the layout invariant the production
// launcher depends on.
const SP4_PRODUCER_COUNT: usize = 47;
const SP4_PRODUCER_COUNT: usize = 69;
const SCRATCH_IDX_TARGET_Q: usize = 0;
// ── Kernel ──
@@ -438,7 +438,7 @@ fn sp4_atom_pos_p99_per_branch_writes_distinct_isv_slots() {
// 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 SP4_PRODUCER_COUNT: usize = 69;
const SCRATCH_BASE: usize = 1;
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
const N_PER_BRANCH: usize = 4096;
@@ -652,7 +652,7 @@ fn launch_sp4_param_group_oracle_for_group(
k_in: i32,
h_dim: i32,
) -> Vec<f32> {
const SP4_PRODUCER_COUNT: usize = 47;
const SP4_PRODUCER_COUNT: usize = 69;
const SCRATCH_BASE_W: usize = 5;
const SCRATCH_BASE_M: usize = 13;
const SCRATCH_BASE_V: usize = 21;
@@ -802,7 +802,7 @@ fn sp4_param_group_oracle_per_group_writes_distinct_isv_slots() {
SP4_PARAM_GROUP_COUNT, SP4_SLOT_BASE, WD_RATE_BASE, WEIGHT_BOUND_BASE,
};
const SP4_PRODUCER_COUNT: usize = 47;
const SP4_PRODUCER_COUNT: usize = 69;
const SCRATCH_BASE_W: usize = 5;
const SCRATCH_BASE_M: usize = 13;
const SCRATCH_BASE_V: usize = 21;
@@ -1145,7 +1145,7 @@ fn sp4_grad_norm_p99_pearl_a_first_observation_replaces_scalar() {
// 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 SP4_PRODUCER_COUNT: usize = 69;
const SCRATCH_IDX_GRAD_CLIP: usize = 38;
// Single-element grad_norm scalar — known synthetic value chosen to be
@@ -1312,7 +1312,7 @@ fn sp4_h_s2_p99_writes_step_p99_to_scratch_via_pearl_a() {
// Production scratch shape (47 = SP4 producers count). Slot 39 is the
// production-allocated H_S2_BOUND slot per the documented stable layout
// in `launch_sp4_target_q_p99` / `launch_sp4_h_s2_p99`.
const SP4_PRODUCER_COUNT: usize = 47;
const SP4_PRODUCER_COUNT: usize = 69;
const SCRATCH_IDX_H_S2: usize = 39;
const SHARED_BYTES: u32 = (256 / 32) * 256 * 4;
@@ -1397,3 +1397,122 @@ fn sp4_h_s2_p99_writes_step_p99_to_scratch_via_pearl_a() {
assert_eq!(SP4_SLOT_BASE, 131);
assert_eq!((H_S2_BOUND_INDEX - SP4_SLOT_BASE) * 3, 114);
}
// ── SP4 Task A13.0: h_s2_rms_ema retrofit (Pearls A+D) ───────────────────────
/// Test-only cubin for the SP4 Task A13.0 retrofit `h_s2_rms_ema_update`
/// kernel. Same cubin as production (`H_S2_RMS_EMA_CUBIN`); this kernel-direct
/// test bypasses the trainer harness and exercises the kernel's scratch-slot
/// write contract on a controlled stationary signal.
const SP4_H_S2_RMS_EMA_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/h_s2_rms_ema_kernel.cubin"));
fn load_h_s2_rms_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
let module = stream
.context()
.load_cubin(SP4_H_S2_RMS_EMA_CUBIN.to_vec())
.expect("load h_s2_rms_ema_kernel cubin");
module
.load_function("h_s2_rms_ema_update")
.expect("load h_s2_rms_ema_update function")
}
/// SP4 Task A13.0: retrofit `h_s2_rms_ema_update` to write the per-step
/// `step_observation = sqrt(mean(h_s2²))` into `producer_step_scratch[40]`,
/// then verify Pearls A+D bootstrap (Pearl A) + stationary-convergence
/// (Pearl D) host-side using the production `pearls_ad_update` helper.
///
/// Stationary signal → constant `h_s2 = 5.0` over `B*SH2 = 4*64 = 256`
/// floats → analytical RMS = 5.0. After 1000 stationary observations the
/// converged ISV value must be within 1% of 5.0; first observation must
/// equal RMS (Pearl A bypasses Pearl D's variance-ratio at t=0).
#[test]
#[ignore = "requires GPU"]
fn sp4_h_s2_rms_ema_writes_step_rms_via_pearl_a_then_converges_pearl_d() {
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
const B: usize = 4;
const SH2: usize = 64;
const N: usize = B * SH2;
const SCRATCH_IDX: usize = 40;
const SP4_PRODUCER_COUNT: usize = 69;
const BLOCK_DIM: u32 = 256;
const SHARED_BYTES: u32 = BLOCK_DIM * std::mem::size_of::<f32>() as u32;
let stationary_signal: f32 = 5.0;
let samples: Vec<f32> = vec![stationary_signal; N];
let stream = make_test_stream();
let kernel = load_h_s2_rms_ema_kernel(&stream);
// Safety: CUDA context active on this thread.
let in_buf = unsafe { MappedF32Buffer::new(N) }
.expect("alloc h_s2 input buffer");
in_buf.write_from_slice(&samples);
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
.expect("alloc producer scratch buffer");
let b_i32 = B as i32;
let sh2_i32 = SH2 as i32;
let scratch_idx_arg: i32 = SCRATCH_IDX as i32;
let in_dev_ptr = in_buf.dev_ptr;
let scratch_dev_ptr = scratch_buf.dev_ptr;
unsafe {
stream
.launch_builder(&kernel)
.arg(&in_dev_ptr)
.arg(&b_i32)
.arg(&sh2_i32)
.arg(&scratch_dev_ptr)
.arg(&scratch_idx_arg)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (BLOCK_DIM, 1, 1),
shared_mem_bytes: SHARED_BYTES,
})
.expect("launch h_s2_rms_ema_update");
}
stream
.synchronize()
.expect("sync after h_s2_rms_ema_update launch");
// Verify ALL non-target slots remained zero — guards against write-cascade.
let host = scratch_buf.read_all();
for (i, &v) in host.iter().enumerate() {
if i != SCRATCH_IDX {
assert_eq!(
v, 0.0,
"h_s2_rms_ema wrote to scratch[{i}] outside target slot {SCRATCH_IDX}: {v}",
);
}
}
let step_rms = host[SCRATCH_IDX];
let expected_rms = stationary_signal; // sqrt(mean(5²)) = 5
assert!(
(step_rms - expected_rms).abs() < 1e-4,
"kernel step_rms {step_rms} vs analytical {expected_rms}",
);
// ── Pearl A bootstrap: first observation replaces sentinel directly ──
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_rms);
assert_eq!(
new_x_mean, step_rms,
"Pearl A: first observation must replace sentinel, got {new_x_mean}",
);
assert_eq!(state.x_lag, step_rms, "Pearl A seeds x_lag");
// ── Pearl D convergence: 1000 stationary observations of step_rms ──
let mut x_mean = new_x_mean;
for _ in 0..1000 {
x_mean = pearls_ad_update(x_mean, &mut state, step_rms);
}
let rel_err = ((x_mean - expected_rms) / expected_rms).abs();
assert!(
rel_err < 0.01,
"Pearl D: stationary signal converged to {x_mean} (expected {expected_rms}, rel_err {rel_err})",
);
}

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

File diff suppressed because one or more lines are too long