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

feat(sp4): Task A13.5 — retrofit reward_component_ema with Pearls A+D + cross-boundary wiring + orphan deletion

Browse Source 
Replaces the kernel's hardcoded `ema_alpha` with the shared `pearls_ad_update`
host-side helper, wires the host-side update across the trainer/collector
boundary (mirrors the A14/A15 Pearl C wiring path), and deletes the
trainer's orphan `launch_reward_component_ema` per
`feedback_wire_everything_up.md`.

Kernel + collector launcher:
  - Kernel `reward_component_ema` signature: drops `(isv_out, ema_alpha,
    isv_reward_base_slot)` for `(scratch_buf, scratch_first_index=63)`.
    Single-block 6-thread (one per component); each writes mean|r_c| to
    `scratch_buf[scratch_first_index + c]` with `__threadfence_system()`.
  - 6 ISV slots wired with Pearls A+D: ISV[63..69) (Wiener offsets
    189..207 — last slots in the post-A13 207-float wiener_state_buf).
  - `GpuExperienceCollector::launch_reward_component_ema_inplace` now:
    launches kernel → syncs stream → applies Pearls A+D in 6-iteration
    loop → memsets reward_components_per_sample to zero (preserves
    original behaviour). Degenerate-zero short-circuit per slot covers
    the always-zero placeholder components (c=2 trail, c=5 bonus) plus
    cold-start.

Cross-boundary wiring (mirrors A14/A15 precedent):
  - 4 new fields on `GpuExperienceCollector`:
    `reward_component_pearls_{wiener_host_ptr, scratch_dev_ptr,
    scratch_host_ptr, isv_pinned_ptr}` — all NULL/0 until wired.
  - New setter `set_reward_component_pearls_buffers(...)` on collector.
  - New accessors on `GpuDqnTrainer`: `wiener_state_buf_host_ptr()`,
    `producer_step_scratch_buf_dev_ptr()`,
    `producer_step_scratch_buf_host_ptr()`, `isv_signals_pinned_ptr()`.
  - New wire helper `FusedTrainingCtx::wire_reward_component_pearls_buffers`
    pulls all 4 pointers from trainer, pushes into collector.
  - Wired once in `training_loop.rs::init_gpu_experience_collector`
    immediately after `set_curiosity_pearl_c_buffers`.

Orphan deletion (per `feedback_wire_everything_up.md`):
  - Removed `GpuDqnTrainer::launch_reward_component_ema` (zero call sites
    pre-deletion).
  - Removed trainer-side `reward_component_ema_kernel: CudaFunction`
    field (zero consumers post-launcher-deletion).
  - Removed cubin loader + struct-init line.
  - `REWARD_COMPONENT_EMA_CUBIN` static remains because the collector
    still loads from it.

Tests:
  - `sp4_reward_component_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d`:
    drives kernel with N=128 reward_components where r[i*6+c] = sign(i) ×
    (c+1), asserts each slot ∈ ±1e-5 of (c+1), non-target slots remain 0,
    Pearl A bootstrap + Pearl D convergence verified.
  - The cross-boundary wiring path exercised in production by integration
    smoke harness; this kernel-direct test isolates kernel signature +
    Pearls A+D semantics.
  - `cargo test -p ml --lib sp4_wiener_ema --offline` 6/6 passing.

Per `feedback_no_atomicadd.md`,
`feedback_no_htod_htoh_only_mapped_pinned.md`,
`feedback_no_partial_refactor.md`, `feedback_wire_everything_up.md`.
Build: `cargo check -p ml --lib --tests --offline` clean (11 pre-existing
warnings, no new warnings).

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-01 02:27:50 +02:00
parent 5f800fe5c8
commit 4f82b74a51
7 changed files with 383 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -3576,12 +3576,6 @@ pub struct GpuDqnTrainer {
/// v8: PopArt sample count [1].
popart_count: CudaSlice<f32>,
// ── C.2 Reward-component attribution EMA (Plan 3 Task 1) ──────────
/// Single-block kernel (6 threads). Reads reward_components_per_sample [N*L, 6]
/// from experience_env_step, reduces mean |r_c| per component, EMA into ISV[63..69).
/// Cold-path: one launch per training step. Loaded from reward_component_ema_kernel.cubin.
reward_component_ema_kernel: CudaFunction,
// ── Plan 4 Task 2c.3c.5: H_S2 RMS EMA producer ────────────────────
/// Single-block kernel (256 threads, shmem-reduction, no atomicAdd). Reduces
/// RMS = sqrt(sum_sq / (B*SH2)) over `save_h_s2 [B, SH2]` and EMA-updates
@@ -8765,35 +8759,16 @@ impl GpuDqnTrainer {
Ok(())
}
/// C.2 Reward-component attribution (Plan 3 Task 1, spec §4.C.2).
///
/// Single-block kernel (6 threads). Reads `reward_components` [n_samples * 6],
/// reduces mean |r_c| per component, adaptive-rate EMA into ISV[63..69).
///
/// Called from `GpuExperienceCollector` after each `collect_experiences_gpu` epoch.
/// The caller passes the ISV device pointer and the number of live samples.
pub fn launch_reward_component_ema(
&self,
reward_components_ptr: u64, // [n_samples * 6] f32 device ptr
n_samples: usize,
isv_dev_ptr: u64,
ema_alpha: f32,
) -> Result<(), MLError> {
if n_samples == 0 { return Ok(()); }
let n_i32 = n_samples as i32;
let base_slot = REWARD_POPART_EMA_INDEX as i32;
unsafe {
self.stream.launch_builder(&self.reward_component_ema_kernel)
.arg(&reward_components_ptr)
.arg(&n_i32)
.arg(&isv_dev_ptr)
.arg(&ema_alpha)
.arg(&base_slot)
.launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (6, 1, 1), shared_mem_bytes: 0 })
.map_err(|e| MLError::ModelError(format!("reward_component_ema: {e}")))?;
}
Ok(())
}
// SP4 Layer A Task A13.5 (2026-05-01): the orphan
// `launch_reward_component_ema` previously here was deleted per
// `feedback_wire_everything_up.md`. The active launcher lives on
// `GpuExperienceCollector::launch_reward_component_ema_inplace`,
// which now takes the trainer's mapped-pinned Pearls A+D buffers
// via `set_reward_component_pearls_buffers` (wired from
// `FusedTrainingCtx::wire_reward_component_pearls_buffers`).
// The trainer-side `reward_component_ema_kernel` field is also
// deleted because no trainer launcher consumes it; the experience
// collector loads its own copy from the same cubin.
/// 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
@@ -11612,13 +11587,15 @@ impl GpuDqnTrainer {
.map_err(|e| MLError::ModelError(format!("q_quantile_reduce load: {e}")))?
};
// Plan 3 Task 1 C.2: load reward_component_ema kernel (cold-path, per-step).
let reward_component_ema_kernel = {
let module = stream.context().load_cubin(REWARD_COMPONENT_EMA_CUBIN.to_vec())
.map_err(|e| MLError::ModelError(format!("reward_component_ema cubin load: {e}")))?;
module.load_function("reward_component_ema")
.map_err(|e| MLError::ModelError(format!("reward_component_ema load: {e}")))?
};
// SP4 Task A13.5 (2026-05-01): the trainer-side
// `reward_component_ema_kernel` loader was deleted along with the
// orphan `launch_reward_component_ema`. The active kernel is loaded
// by `GpuExperienceCollector` from the same cubin
// (REWARD_COMPONENT_EMA_CUBIN) — the collector owns the launcher
// (`launch_reward_component_ema_inplace`) and threads the trainer's
// mapped-pinned Pearls A+D buffers via
// `set_reward_component_pearls_buffers` (called from
// `FusedTrainingCtx::wire_reward_component_pearls_buffers`).
// Plan 4 Task 2c.3c.5: load h_s2_rms_ema kernel (cold-path, per-step).
// Single-block 256-thread shmem-reduction kernel; no atomicAdd. Producer
@@ -14697,7 +14674,6 @@ impl GpuDqnTrainer {
popart_mean,
popart_var,
popart_count,
reward_component_ema_kernel,
h_s2_rms_ema_kernel,
target_q_p99_update,
atom_pos_p99_update,
@@ -22583,6 +22559,39 @@ impl GpuDqnTrainer {
self.clamp_engage_per_block_buf.dev_ptr
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
/// `wiener_state_buf`. Used by `GpuExperienceCollector` to share
/// the Pearl D Wiener state for `launch_reward_component_ema_inplace`
/// — the 6 reward-component slots occupy `wiener_state_buf[(63+c)*3..(63+c)*3 + 3)`
/// for c in 0..6.
pub fn wiener_state_buf_host_ptr(&self) -> *mut f32 {
self.wiener_state_buf.host_ptr
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned device pointer for
/// `producer_step_scratch_buf`. Used by `GpuExperienceCollector`
/// to launch the Pearls-A+D-retrofitted `reward_component_ema`
/// kernel writing slots [63..69).
pub fn producer_step_scratch_buf_dev_ptr(&self) -> u64 {
self.producer_step_scratch_buf.dev_ptr
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
/// `producer_step_scratch_buf`. Used by `GpuExperienceCollector`
/// to read step observations after the kernel + sync (slots [63..69)
/// for the 6 reward components).
pub fn producer_step_scratch_buf_host_ptr(&self) -> *mut f32 {
self.producer_step_scratch_buf.host_ptr
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
/// `isv_signals_pinned`. Used by `GpuExperienceCollector` to read
/// `prev_x_mean` and write `new_x_mean` for ISV[REWARD_POPART_EMA_INDEX..+6)
/// during the host-side Pearls A+D update.
pub fn isv_signals_pinned_ptr(&self) -> *mut f32 {
self.isv_signals_pinned
}
/// Read-only accessor for diagnostics / HEALTH_DIAG mirror.
pub fn grad_norm_fast_ema_value(&self) -> f32 { unsafe { *self.grad_norm_fast_ema_pinned } }
/// Read-only accessor for diagnostics / HEALTH_DIAG mirror.

146
crates/ml/src/cuda_pipeline/gpu_experience_collector.rs
View File

@@ -832,6 +832,22 @@ pub struct GpuExperienceCollector {
/// Set via set_isv_signals_ptr(). 0 = NULL (static hold).
isv_signals_dev_ptr: u64,
/// SP4 Layer A Task A13.5 (2026-05-01) — cross-boundary Pearls A+D
/// wiring for `launch_reward_component_ema_inplace`. The kernel writes
/// the 6 step observations to `producer_step_scratch_buf[63..69)` on
/// the trainer's mapped-pinned scratch, and the host-side Pearls A+D
/// update reads/writes the trainer's `wiener_state_buf` + `isv_signals_pinned`
/// directly via these pointers (mirrors the A14/A15 Pearl C precedent).
/// All four pointers are set together by
/// `set_reward_component_pearls_buffers`; `None` means the wiring has
/// not yet been threaded by `FusedTrainingCtx::wire_reward_component_pearls_buffers`,
/// in which case `launch_reward_component_ema_inplace` is a no-op
/// (loud `debug_assert!` flags the misconfiguration in dev builds).
reward_component_pearls_wiener_host_ptr: *mut f32,
reward_component_pearls_scratch_dev_ptr: u64,
reward_component_pearls_scratch_host_ptr: *mut f32,
reward_component_pearls_isv_pinned_ptr: *mut f32,
/// Trade plan params dev_ptr from fused training context.
/// Set via set_plan_params_ptr(). 0 = NULL (no plan).
plan_params_dev_ptr: u64,
@@ -1691,6 +1707,13 @@ impl GpuExperienceCollector {
cost_anneal_pinned,
cost_anneal_dev_ptr,
isv_signals_dev_ptr: 0, // NULL until trainer sets it
// SP4 Task A13.5: Pearls A+D buffers for reward_component_ema.
// All four NULL until FusedTrainingCtx wires them post-collector
// construction (mirrors A14/A15 Pearl C wiring path).
reward_component_pearls_wiener_host_ptr: std::ptr::null_mut(),
reward_component_pearls_scratch_dev_ptr: 0,
reward_component_pearls_scratch_host_ptr: std::ptr::null_mut(),
reward_component_pearls_isv_pinned_ptr: std::ptr::null_mut(),
plan_params_dev_ptr: 0, // NULL until trainer sets it
readiness_dev_ptr: 0, // NULL until trainer sets it
exploration_scale_pinned,
@@ -2244,42 +2267,112 @@ impl GpuExperienceCollector {
Ok([0.0, cf_rate, trail_r, micro_r])
}
/// C.2 Plan 3 Task 1 (spec §4.C.2): GPU-driven reward-component EMA.
/// C.2 Plan 3 Task 1 (spec §4.C.2); SP4 Layer A Task A13.5 retrofit
/// (2026-05-01): GPU-driven reward-component step-observation
/// producer + host-side Pearls A+D update.
///
/// Launches `reward_component_ema` (single-block, 6 threads) on this
/// collector's stream. Reads `reward_components_per_sample [n_samples * 6]`
/// written by the last `collect_experiences_gpu` epoch, reduces mean |r_c|
/// per component, and adaptive-rate EMA into ISV[63..69).
/// written by the last `collect_experiences_gpu` epoch, reduces mean
/// |r_c| per component, writes step observations to
/// `producer_step_scratch_buf[63..69)` (mapped-pinned scratch owned
/// by the trainer; host pointer + dev pointer threaded via
/// `set_reward_component_pearls_buffers`). Then syncs the stream and
/// applies Pearls A+D host-side for each of the 6 slots, reading/
/// writing `wiener_state_buf` + `isv_signals_pinned` directly via
/// the threaded mapped-pinned host pointers (no HtoD/DtoH per
/// `feedback_no_htod_htoh_only_mapped_pinned.md`).
///
/// After the kernel fires, resets `reward_components_per_sample` to zeros
/// via memset (same as other per-sample buffers).
/// **α 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.
///
/// After the Pearls A+D update completes, resets
/// `reward_components_per_sample` to zeros via memset (same as other
/// per-sample buffers).
///
/// `isv_dev_ptr` must be non-zero (the trainer's ISV device pointer).
/// `n_samples` should be `alloc_episodes * alloc_timesteps` (the full
/// per-epoch batch, not the CF-doubled buffer). `ema_alpha` ∈ (0, 1).
/// per-epoch batch, not the CF-doubled buffer).
pub fn launch_reward_component_ema_inplace(
&mut self,
n_samples: usize,
ema_alpha: f32,
_ema_alpha_unused: f32,
) -> Result<(), crate::MLError> {
use crate::cuda_pipeline::gpu_dqn_trainer::REWARD_POPART_EMA_INDEX;
use crate::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
if n_samples == 0 || self.isv_signals_dev_ptr == 0 {
return Ok(());
}
// SP4 Task A13.5: scratch slot 63 — first of 6 contiguous reward-
// component slots. Wiener offsets: (63+c)*3 = 189..207 for c in 0..6.
const SCRATCH_FIRST: usize = 63;
// The trainer's Pearls A+D buffers MUST be wired before the first
// launch. `FusedTrainingCtx::wire_reward_component_pearls_buffers`
// is the canonical wire site (called from training_loop.rs's
// `init_gpu_experience_collector`).
debug_assert!(
self.reward_component_pearls_scratch_dev_ptr != 0
&& !self.reward_component_pearls_scratch_host_ptr.is_null()
&& !self.reward_component_pearls_wiener_host_ptr.is_null()
&& !self.reward_component_pearls_isv_pinned_ptr.is_null(),
"launch_reward_component_ema_inplace: Pearls A+D buffers not wired \
(call FusedTrainingCtx::wire_reward_component_pearls_buffers)"
);
let n_i32 = n_samples as i32;
let base_slot = REWARD_POPART_EMA_INDEX as i32;
let scratch_first_i32 = SCRATCH_FIRST as i32;
let rc_ptr = self.reward_components_per_sample.raw_ptr();
let isv_ptr = self.isv_signals_dev_ptr;
let scratch_dev = self.reward_component_pearls_scratch_dev_ptr;
unsafe {
self.stream.launch_builder(&self.reward_component_ema_kernel)
.arg(&rc_ptr)
.arg(&n_i32)
.arg(&isv_ptr)
.arg(&ema_alpha)
.arg(&base_slot)
.arg(&scratch_dev)
.arg(&scratch_first_i32)
.launch(LaunchConfig { grid_dim: (1, 1, 1), block_dim: (6, 1, 1), shared_mem_bytes: 0 })
.map_err(|e| crate::MLError::ModelError(format!("reward_component_ema: {e}")))?;
}
// Cold-path producer — sync before reading the mapped-pinned scratch
// slots from the host.
self.stream.synchronize()
.map_err(|e| crate::MLError::ModelError(format!("reward_component_ema sync: {e}")))?;
// ── Pearls A+D host-side update for all 6 component slots ──
let scratch_host = self.reward_component_pearls_scratch_host_ptr;
let wiener_host = self.reward_component_pearls_wiener_host_ptr;
let isv_host = self.reward_component_pearls_isv_pinned_ptr;
for c in 0..6 {
let scratch_idx = SCRATCH_FIRST + c;
let isv_idx = REWARD_POPART_EMA_INDEX + c;
let step_obs = unsafe { std::ptr::read_volatile(scratch_host.add(scratch_idx)) };
// Reward components 2/5 (trail/bonus) are always 0.0 placeholders
// in this design — they should not advance Wiener state. The
// generic short-circuit also covers cold-start before the first
// collect_experiences_gpu populates `reward_components_per_sample`.
if step_obs == 0.0 { continue; }
let wiener_offset = scratch_idx * 3;
let prev_x_mean = unsafe { std::ptr::read_volatile(isv_host.add(isv_idx)) };
let mut state = unsafe {
WienerState {
sample_var: std::ptr::read_volatile(wiener_host.add(wiener_offset)),
diff_var: std::ptr::read_volatile(wiener_host.add(wiener_offset + 1)),
x_lag: std::ptr::read_volatile(wiener_host.add(wiener_offset + 2)),
}
};
let new_x_mean = pearls_ad_update(prev_x_mean, &mut state, step_obs);
unsafe {
std::ptr::write_volatile(isv_host.add(isv_idx), new_x_mean);
std::ptr::write_volatile(wiener_host.add(wiener_offset), state.sample_var);
std::ptr::write_volatile(wiener_host.add(wiener_offset + 1), state.diff_var);
std::ptr::write_volatile(wiener_host.add(wiener_offset + 2), state.x_lag);
}
}
self.stream.memset_zeros(&mut self.reward_components_per_sample)
.map_err(|e| crate::MLError::ModelError(format!("memset reward_components_per_sample: {e}")))?;
Ok(())
@@ -4291,6 +4384,33 @@ impl GpuExperienceCollector {
}
}
/// SP4 Layer A Task A13.5 (2026-05-01): thread the main DQN trainer's
/// mapped-pinned Pearls A+D buffers (`producer_step_scratch_buf`
/// host+dev pointers, `wiener_state_buf` host pointer, and
/// `isv_signals_pinned`) into this experience collector. Called once
/// post-construction by `FusedTrainingCtx::wire_reward_component_pearls_buffers`,
/// mirroring the A14/A15 Pearl C wiring path.
///
/// The collector's `launch_reward_component_ema_inplace` uses these
/// to (1) tell the kernel where to write step observations
/// (`scratch_dev_ptr`), (2) read the per-step scalars host-side via
/// `scratch_host_ptr`, (3) fetch the prior x_mean from
/// `isv_pinned_ptr`, and (4) update the per-slot Wiener triple in
/// `wiener_state_host_ptr` — all via mapped-pinned host pointers
/// (no HtoD/DtoH per `feedback_no_htod_htoh_only_mapped_pinned.md`).
pub fn set_reward_component_pearls_buffers(
&mut self,
wiener_state_host_ptr: *mut f32,
scratch_dev_ptr: u64,
scratch_host_ptr: *mut f32,
isv_pinned_ptr: *mut f32,
) {
self.reward_component_pearls_wiener_host_ptr = wiener_state_host_ptr;
self.reward_component_pearls_scratch_dev_ptr = scratch_dev_ptr;
self.reward_component_pearls_scratch_host_ptr = scratch_host_ptr;
self.reward_component_pearls_isv_pinned_ptr = isv_pinned_ptr;
}
/// Disable curiosity training (e.g. after an async kernel crash).
pub fn disable_curiosity(&mut self) {
self.curiosity_trainer = None;

34
crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu
View File

@@ -1,6 +1,8 @@
/* reward_component_ema — GPU-driven per-component |reward| EMA into ISV.
/* reward_component_ema — GPU-driven per-component |reward| step
* observation into producer_step_scratch_buf for Pearls A+D consumption.
*
* Plan 3 Task 1, spec §4.C.2. GPU-drives-CPU-reads (spec §4.C.6).
* Plan 3 Task 1, spec §4.C.2; SP4 Layer A Task A13.5 retrofit (2026-05-01).
* GPU-drives-CPU-reads (spec §4.C.6).
*
* Reads per-sample per-component rewards from `reward_components_per_sample`
* [n_samples * 6], where the 6 components are laid out as:
@@ -12,19 +14,29 @@
* [c=4] opp_cost — Flat opportunity cost (Plan 3 Task 2 B.1; ISV[21]-scaled)
* [c=5] bonus — 0.0 placeholder (future trade-attempt/timing bonus)
*
* Reduces mean |r_c| over the batch per component, then applies an
* adaptive-rate EMA into ISV slots [REWARD_POPART_EMA_INDEX..+6).
* Reduces mean |r_c| over the batch per component, then writes the step
* observation to `scratch_buf[scratch_first_index + c]` for c in 0..6.
* The host launcher (`launch_reward_component_ema_inplace` on
* `GpuExperienceCollector`) syncs the stream and applies Pearls A+D via
* `pearls_ad_update`, mapping each scratch slot to ISV[REWARD_POPART_EMA_INDEX + c].
*
* **α dropped per SP4 — Pearls A+D adapt α from per-slot signal-vs-noise
* variance.** Cross-boundary wiring: the trainer's mapped-pinned
* `producer_step_scratch_buf` + `wiener_state_buf` + `isv_signals_pinned`
* are threaded into the collector via
* `set_reward_component_pearls_buffers`, called once by
* `FusedTrainingCtx::wire_reward_component_pearls_buffers` after
* collector construction (mirrors A14/A15 Pearl C wiring path).
*
* Single-block kernel (one thread per component). Cold-path — once per
* training step. Cost: 6 threads × n_samples additions + 6 global writes.
* training step. Cost: 6 threads × n_samples additions + 6 scratch writes.
*/
extern "C" __global__ void reward_component_ema(
const float* __restrict__ reward_components, /* [n_samples * 6] row-major */
int n_samples,
float* __restrict__ isv_out, /* ISV bus [ISV_TOTAL_DIM] */
float ema_alpha, /* adaptive EMA rate (0,1) */
int isv_reward_base_slot /* REWARD_POPART_EMA_INDEX = 63 */
float* __restrict__ scratch_buf, /* producer_step_scratch_buf */
int scratch_first_index /* slot 63 — first of 6 contiguous slots */
) {
int c = (int)threadIdx.x;
if (c >= 6 || blockIdx.x != 0) return;
@@ -36,8 +48,6 @@ extern "C" __global__ void reward_component_ema(
}
float mean_abs = acc / fmaxf(1.0f, (float)n_samples);
/* Adaptive-rate EMA into ISV[isv_reward_base_slot + c]. */
int slot = isv_reward_base_slot + c;
float prev = isv_out[slot];
isv_out[slot] = (1.0f - ema_alpha) * prev + ema_alpha * mean_abs;
scratch_buf[scratch_first_index + c] = mean_abs;
__threadfence_system(); /* PCIe-visible to mapped-pinned host_ptr */
}

54
crates/ml/src/trainers/dqn/fused_training.rs
View File

@@ -3302,6 +3302,60 @@ impl FusedTrainingCtx {
self.trainer.clamp_engage_per_block_buf_dev_ptr()
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
/// `wiener_state_buf`. Used by `init_gpu_experience_collector`
/// to wire `set_reward_component_pearls_buffers` on the
/// experience collector (cross-boundary Pearls A+D wiring).
pub(crate) fn wiener_state_buf_host_ptr(&self) -> *mut f32 {
self.trainer.wiener_state_buf_host_ptr()
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned device pointer for
/// `producer_step_scratch_buf` (used by the Pearls-A+D-retrofitted
/// `reward_component_ema` kernel as its scratch destination).
pub(crate) fn producer_step_scratch_buf_dev_ptr(&self) -> u64 {
self.trainer.producer_step_scratch_buf_dev_ptr()
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
/// `producer_step_scratch_buf` (used by `GpuExperienceCollector`
/// to read the 6 reward-component step observations after kernel
/// + sync).
pub(crate) fn producer_step_scratch_buf_host_ptr(&self) -> *mut f32 {
self.trainer.producer_step_scratch_buf_host_ptr()
}
/// SP4 Task A13.5 (2026-05-01): mapped-pinned host pointer for
/// `isv_signals_pinned` (used by `GpuExperienceCollector` to
/// read/write ISV[REWARD_POPART_EMA_INDEX..+6) during Pearls A+D
/// host-side update).
pub(crate) fn isv_signals_pinned_ptr(&self) -> *mut f32 {
self.trainer.isv_signals_pinned_ptr()
}
/// SP4 Layer A Task A13.5 (2026-05-01): wire the main DQN trainer's
/// mapped-pinned Pearls A+D buffers into the experience collector
/// for `launch_reward_component_ema_inplace`. Mirrors the A14/A15
/// `wire_aux_trainer_pearl_c_buffers` precedent — idempotent, safe
/// to call multiple times. Per `feedback_wire_everything_up.md`,
/// the trainer-side orphan `launch_reward_component_ema` was
/// deleted; this is the canonical (and only) wiring path now.
pub(crate) fn wire_reward_component_pearls_buffers(
&mut self,
collector: &mut crate::cuda_pipeline::gpu_experience_collector::GpuExperienceCollector,
) {
let wiener_host = self.trainer.wiener_state_buf_host_ptr();
let scratch_dev = self.trainer.producer_step_scratch_buf_dev_ptr();
let scratch_host = self.trainer.producer_step_scratch_buf_host_ptr();
let isv_pinned = self.trainer.isv_signals_pinned_ptr();
collector.set_reward_component_pearls_buffers(
wiener_host,
scratch_dev,
scratch_host,
isv_pinned,
);
}
/// SP4 Task A14: wire the main DQN trainer's mapped-pinned Pearl C
/// buffers into all aux trainers owned by this `FusedTrainingCtx`.
/// Idempotent — safe to call multiple times. Curiosity is wired

10
crates/ml/src/trainers/dqn/trainer/training_loop.rs
View File

@@ -1438,6 +1438,16 @@ impl DQNTrainer {
let engage_buf_ptr = fused_ctx.clamp_engage_per_block_buf_dev_ptr();
collector.set_curiosity_pearl_c_buffers(nan_flags_ptr, engage_buf_ptr);
info!("SP4 Pearl C: aux trainer engagement-counter buffers wired");
// SP4 Task A13.5 (2026-05-01): wire the trainer's
// mapped-pinned Pearls A+D buffers (producer_step_scratch_buf,
// wiener_state_buf, isv_signals_pinned) into the
// experience collector so launch_reward_component_ema_inplace
// can run the host-side Pearls A+D update for slots 63..69.
// Per `feedback_wire_everything_up.md`, the trainer-side
// orphan `launch_reward_component_ema` is deleted —
// this is the canonical (only) launch path now.
fused_ctx.wire_reward_component_pearls_buffers(&mut collector);
info!("SP4 Task A13.5: reward_component_ema Pearls A+D buffers wired");
}
self.gpu_experience_collector = Some(collector);
}

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

@@ -2034,3 +2034,113 @@ fn sp4_iqn_quantile_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
"Pearl D: stationary p05 converged to {x_mean} (expected 0.7)",
);
}
// ── SP4 Task A13.5: reward_component_ema retrofit + cross-boundary wiring ────
const SP4_REWARD_COMPONENT_EMA_CUBIN: &[u8] =
include_bytes!(concat!(env!("OUT_DIR"), "/reward_component_ema_kernel.cubin"));
fn load_reward_component_ema_kernel(stream: &Arc<CudaStream>) -> CudaFunction {
let module = stream
.context()
.load_cubin(SP4_REWARD_COMPONENT_EMA_CUBIN.to_vec())
.expect("load reward_component_ema cubin");
module
.load_function("reward_component_ema")
.expect("load reward_component_ema function")
}
/// SP4 Task A13.5: `reward_component_ema` writes 6 mean-|r_c| step
/// observations to scratch[63..69) (one per reward component). Pearls A+D
/// bootstrap+convergence verified host-side via `pearls_ad_update`. The
/// production cross-boundary wiring (FusedTrainingCtx →
/// GpuExperienceCollector::set_reward_component_pearls_buffers) is
/// exercised by the integration smoke harness; this kernel-direct test
/// validates the kernel signature retrofit + Pearls A+D semantics in
/// isolation.
///
/// Test surface: n_samples=128 with reward_components[i*6 + c] = c+1
/// (constant per component, mixed sign on alternating samples). Each
/// component's mean|r_c| = c+1.
#[test]
#[ignore = "requires GPU"]
fn sp4_reward_component_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d() {
use ml::cuda_pipeline::sp4_wiener_ema::{pearls_ad_update, WienerState};
const N: usize = 128;
const NUM_COMPONENTS: usize = 6;
const SCRATCH_FIRST: usize = 63;
const SP4_PRODUCER_COUNT: usize = 69;
// reward_components[i*6 + c] = (c+1) with sign flip on odd i, so
// mean(|r_c|) = c+1 ∈ {1, 2, 3, 4, 5, 6}.
let rewards: Vec<f32> = (0..N)
.flat_map(|i| {
let sign: f32 = if i % 2 == 0 { 1.0 } else { -1.0 };
(0..NUM_COMPONENTS).map(move |c| sign * (c as f32 + 1.0))
})
.collect();
let stream = make_test_stream();
let kernel = load_reward_component_ema_kernel(&stream);
let rewards_buf = unsafe { MappedF32Buffer::new(rewards.len()) }
.expect("alloc rewards buf");
rewards_buf.write_from_slice(&rewards);
let scratch_buf = unsafe { MappedF32Buffer::new(SP4_PRODUCER_COUNT) }
.expect("alloc scratch");
let rewards_dev = rewards_buf.dev_ptr;
let scratch_dev = scratch_buf.dev_ptr;
let n_i32 = N as i32;
let scratch_first_i32 = SCRATCH_FIRST as i32;
unsafe {
stream.launch_builder(&kernel)
.arg(&rewards_dev)
.arg(&n_i32)
.arg(&scratch_dev)
.arg(&scratch_first_i32)
.launch(LaunchConfig {
grid_dim: (1, 1, 1),
block_dim: (NUM_COMPONENTS as u32, 1, 1),
shared_mem_bytes: 0,
})
.expect("launch reward_component_ema");
}
stream.synchronize().expect("sync");
let host = scratch_buf.read_all();
for (i, &v) in host.iter().enumerate() {
let in_target = i >= SCRATCH_FIRST && i < SCRATCH_FIRST + NUM_COMPONENTS;
if !in_target {
assert_eq!(v, 0.0, "wrote outside target slots: scratch[{i}]={v}");
}
}
for c in 0..NUM_COMPONENTS {
let actual = host[SCRATCH_FIRST + c];
let expected = c as f32 + 1.0;
assert!(
(actual - expected).abs() < 1e-5,
"component {c} mean|r| {actual} vs expected {expected}",
);
}
// Pearl A bootstrap on component 0 (representative — same logic per
// component).
let c0_obs = host[SCRATCH_FIRST];
let mut state = WienerState::ZERO;
let new_x_mean = pearls_ad_update(0.0, &mut state, c0_obs);
assert_eq!(new_x_mean, c0_obs);
assert_eq!(state.x_lag, c0_obs);
// Pearl D convergence on component 0.
let mut x_mean = new_x_mean;
for _ in 0..1000 {
x_mean = pearls_ad_update(x_mean, &mut state, c0_obs);
}
assert!(
((x_mean - 1.0) / 1.0).abs() < 0.01,
"Pearl D: stationary component 0 converged to {x_mean} (expected 1.0)",
);
}

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

@@ -2334,3 +2334,5 @@ SP4 Layer A Task A13.2 — moe_expert_util_ema retrofit (Pearls A+D) (2026-05-01
SP4 Layer A Task A13.3 — vsn_mask_ema retrofit (Pearls A+D) (2026-05-01): retrofit the existing per-step VSN feature-group mask producer in `crates/ml/src/cuda_pipeline/vsn_mask_ema_kernel.cu`. Kernel signature changed: drops `(isv, isv_first_index, ema_alpha)` for `(scratch_buf, scratch_first_index=53)`. Single-block, 256-thread shmem-tree reduce per group; `num_groups` separate reductions sharing the same 256-float scratchpad — total smem footprint identical to the prior implementation. Per-group mean is written to `scratch_buf[scratch_first_index + g]` for g in 0..num_groups; single `__threadfence_system()` after all 6 groups write (PCIe-visible to mapped-pinned host_ptr). **6 ISV slots wired with Pearls A+D**: `ISV[VSN_MASK_GROUP_0_EMA_INDEX..+6) = ISV[105..111)`, with Wiener offsets `(53+g)*3` = 159..177 for g in 0..6. Wrapper `GpuDqnTrainer::launch_vsn_mask_ema(_ema_alpha_unused: f32)` syncs the stream after launch and applies Pearls A+D in a `for` loop over the 6 groups. `debug_assert_eq!(num_groups, 6)` guards against `SL_NUM_FEATURE_GROUPS` changes silently breaking the contiguous scratch layout. Degenerate-zero short-circuit at every slot prevents Wiener-state advancement when the kernel hasn't run (cold start before VSN forward populates `vsn_mask_buf`). **Unit test** `sp4_vsn_mask_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` (`#[ignore]`-gated for GPU): drives kernel with B=128 num_groups=6 mask `mask[b,g] = (g+1)/21` (rows sum to 1, per-group mean = (g+1)/21). Asserts each scratch slot ∈ ±1e-5, non-target slots remain 0; verifies Pearl A bootstrap on group 0 + Pearl D convergence (1000 stationary observations within 1% of analytical). Behaviour: cold-path producer with no consumer-facing change beyond the EMA mechanism — slots [105..111) stay semantically identical (per-group mean of VSN softmax mask) and continue to be read by HEALTH_DIAG mirror + VSN focus monitor. Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`. `cargo check -p ml --lib --tests --offline` clean (11 pre-existing warnings, no new warnings).
SP4 Layer A Task A13.4 — iqn_quantile_ema retrofit (Pearls A+D) (2026-05-01): retrofit the existing per-step IQN off-median-quantile diagnostic producer in `crates/ml/src/cuda_pipeline/iqn_quantile_ema_kernel.cu`. Kernel signature changed: drops `(isv, isv_q_p05_idx, isv_q_p25_idx, isv_q_p75_idx, isv_q_p95_idx, ema_alpha)` for `(scratch_buf, scratch_first_index=59)`. Block-dispatch unchanged: 4 blocks × 256 threads, blockIdx.x ∈ {0,1,2,3} maps to (τ_idx ∈ {0,1,3,4}, slot_offset ∈ {0,1,2,3}); the median (τ_idx=2) is intentionally skipped. Each block reduces mean |Q(s,a;τ)| over (B × TBA) entries via shmem-tree (no atomicAdd) and writes the step observation to `scratch_buf[scratch_first_index + slot_offset]` with `__threadfence_system()`. **4 ISV slots wired with Pearls A+D**: ISV[IQN_Q_P05_EMA_INDEX=99] (Wiener offset 177), ISV[IQN_Q_P25_EMA_INDEX=100] (Wiener offset 180), ISV[IQN_Q_P75_EMA_INDEX=101] (Wiener offset 183), ISV[IQN_Q_P95_EMA_INDEX=102] (Wiener offset 186). Wrapper `GpuDqnTrainer::launch_iqn_quantile_ema(save_q_online_ptr, tba, num_quantiles, _ema_alpha_unused: f32)` retains its early-return `if save_q_online_ptr == 0` guard (no IQN active in this trainer config), and now syncs the stream after launch and applies Pearls A+D in a `for` loop over the 4 SLOT_PAIRS const array. Degenerate-zero short-circuit at every slot prevents Wiener-state advancement when the kernel hasn't run (cold start before IQN forward populates `save_q_online`). **Unit test** `sp4_iqn_quantile_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` (`#[ignore]`-gated for GPU): drives kernel with B=32 Q=5 TBA=12 controlled `save_q_online` where `Q[a, b*Q+tau_idx] = (tau_idx+1) × 0.7` (analytical mean per tau = (tau_idx+1) × 0.7). Asserts each slot ∈ ±1e-5 of expected (0.7/1.4/2.8/3.5 for p05/p25/p75/p95), median (tau_idx=2 expected 2.1) intentionally absent, non-target scratch slots remain 0; verifies Pearl A bootstrap on p05 + Pearl D convergence (1000 stationary observations within 1% of 0.7). Behaviour: cold-path producer with no consumer-facing change beyond the EMA mechanism — slots [99..103) remain diagnostic-only (HEALTH_DIAG / risk-monitoring surface). Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`. `cargo check -p ml --lib --tests --offline` clean (11 pre-existing warnings, no new warnings).
SP4 Layer A Task A13.5 — reward_component_ema retrofit (Pearls A+D) + cross-boundary wiring + orphan deletion (2026-05-01): retrofit the existing per-step reward-component EMA producer in `crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu` AND wire the host-side Pearls A+D update across the trainer/collector boundary (mirroring the A14/A15 Pearl C wiring path) AND delete the trainer's orphan `launch_reward_component_ema` per `feedback_wire_everything_up.md`. Kernel signature changed: drops `(isv_out, ema_alpha, isv_reward_base_slot)` for `(scratch_buf, scratch_first_index=63)`. Single-block 6-thread (one per component); each thread reduces mean |r_c| over `n_samples` samples and writes the step observation to `scratch_buf[scratch_first_index + c]` for c in 0..6 with `__threadfence_system()`. **6 ISV slots wired with Pearls A+D**: ISV[REWARD_POPART_EMA_INDEX..+6) = ISV[63..69), with Wiener offsets `(63+c)*3 = 189..207` for c in 0..6 — these are the LAST slots in the post-A13 207-float `wiener_state_buf`. **Cross-boundary wiring (mirrors A14/A15 Pearl C precedent)**: 4 new fields on `GpuExperienceCollector` — `reward_component_pearls_wiener_host_ptr: *mut f32`, `reward_component_pearls_scratch_dev_ptr: u64`, `reward_component_pearls_scratch_host_ptr: *mut f32`, `reward_component_pearls_isv_pinned_ptr: *mut f32` — all `NULL`/`0` until wired. New setter `set_reward_component_pearls_buffers(wiener_host, scratch_dev, scratch_host, isv_pinned)` on the collector. New accessors on `GpuDqnTrainer`: `wiener_state_buf_host_ptr()`, `producer_step_scratch_buf_dev_ptr()`, `producer_step_scratch_buf_host_ptr()`, `isv_signals_pinned_ptr()`. New wire helper `FusedTrainingCtx::wire_reward_component_pearls_buffers(&mut self, &mut GpuExperienceCollector)` pulls all four pointers from the trainer and pushes them into the collector. Wired once in `training_loop.rs::init_gpu_experience_collector` immediately after `set_curiosity_pearl_c_buffers` (mirrors the A14/A15 wire call site). The retrofitted `launch_reward_component_ema_inplace` on the collector now: launches the kernel with `(rewards_ptr, n_samples, scratch_dev, scratch_first_index=63)`, syncs the stream, applies Pearls A+D in a `for c in 0..6` loop (read prev_x_mean from `isv_pinned`, read Wiener triple from `wiener_host`, call `pearls_ad_update`, write back ISV + Wiener) — all via mapped-pinned host pointers (no HtoD/DtoH per `feedback_no_htod_htoh_only_mapped_pinned.md`), then memsets `reward_components_per_sample` to zero (preserving original behaviour). Degenerate-zero short-circuit per slot — covers the always-zero placeholder components (c=2 trail, c=5 bonus) plus cold-start before first `collect_experiences_gpu`. **Orphan deletion** per `feedback_wire_everything_up.md`: removed `GpuDqnTrainer::launch_reward_component_ema` (lines 8775-8796 — zero call sites pre-deletion), removed the trainer-side `reward_component_ema_kernel: CudaFunction` field (zero consumers post-launcher-deletion), removed the cubin loader at line 11591-11596, removed the field from the constructor's struct-init at line 14677. The `REWARD_COMPONENT_EMA_CUBIN` static remains because the collector still loads from it. **Unit test** `sp4_reward_component_ema_writes_step_obs_via_pearl_a_then_converges_pearl_d` (`#[ignore]`-gated for GPU): drives kernel with N=128 reward_components where `r[i*6+c] = sign(i) × (c+1)` (analytical mean|r_c| = c+1 for c in 0..6). Asserts each scratch slot ∈ ±1e-5 of (c+1), non-target slots remain 0; verifies Pearl A bootstrap on component 0 + Pearl D convergence (1000 stationary observations within 1% of 1.0). The cross-boundary wiring path is exercised in production by the integration smoke harness; this kernel-direct test isolates the kernel signature retrofit + Pearls A+D semantics. Per `feedback_no_atomicadd.md`, `feedback_no_htod_htoh_only_mapped_pinned.md`, `feedback_no_partial_refactor.md`, `feedback_wire_everything_up.md`. `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.