diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 01a06e9fb..7f27f7d8d 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -3576,12 +3576,6 @@ pub struct GpuDqnTrainer { /// v8: PopArt sample count [1]. popart_count: CudaSlice, - // ── 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. diff --git a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs index c92e68aea..f0862a6e3 100644 --- a/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_experience_collector.rs @@ -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; diff --git a/crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu b/crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu index 8c9c63597..d6a02b29f 100644 --- a/crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu +++ b/crates/ml/src/cuda_pipeline/reward_component_ema_kernel.cu @@ -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 */ } diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 957470742..dfd9f8953 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -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 diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index c2a643486..671acb58a 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -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); } diff --git a/crates/ml/tests/sp4_producer_unit_tests.rs b/crates/ml/tests/sp4_producer_unit_tests.rs index 651fe542c..ea0b6c61a 100644 --- a/crates/ml/tests/sp4_producer_unit_tests.rs +++ b/crates/ml/tests/sp4_producer_unit_tests.rs @@ -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) -> 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 = (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)", + ); +} diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index ef5fbbf7a..c65a811e3 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -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.