From 048c78d5f74ff4654dfe358301106c33dc8e1f5a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Fri, 24 Apr 2026 18:28:19 +0200 Subject: [PATCH] =?UTF-8?q?feat(dqn):=20Plan=201=20Task=209=20=E2=80=94=20?= =?UTF-8?q?atoms=5Fupdate=20GPU=20kernel=20+=20AtomsMonitor=20+=20delete?= =?UTF-8?q?=20CPU=20paths?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit atoms_update_kernel.cu: 4-block kernel (grid=(4,1,1), block=(256,1,1)) replaces the CPU loop that launched adaptive_atom_positions 4× per branch. Each block reads ISV v-range slots [23..31) and spacing_raw params[52..56); computes softmax + cumsum; writes atom_positions_buf[branch * num_atoms]. Same formula as the existing adaptive_atom_positions kernel in experience_kernels.cu. Deleted recompute_atom_positions and warm_start_atom_positions from GpuDqnTrainer and their fused_training.rs delegates. step_atom_positions now calls launch_atoms_update. training_loop.rs: recompute_atom_positions → launch_atoms_update; warm_start_atom_positions + compute_reward_quantiles block removed. AtomsMonitor reads ISV[V_CENTER_DIR_INDEX=23] as representative scalar; diagnose() exposes all 8 v-range slots [23..31). state_reset_registry.rs "follow-up task" descriptions updated for all 4 adaptive slots. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml/build.rs | 1 + .../src/cuda_pipeline/atoms_update_kernel.cu | 104 +++++++++++ .../ml/src/cuda_pipeline/gpu_dqn_trainer.rs | 172 +++++------------- crates/ml/src/trainers/dqn/fused_training.rs | 13 +- .../trainers/dqn/monitors/atoms_monitor.rs | 118 ++++++++++++ crates/ml/src/trainers/dqn/monitors/mod.rs | 1 + .../src/trainers/dqn/state_reset_registry.rs | 8 +- .../src/trainers/dqn/trainer/training_loop.rs | 37 +--- docs/dqn-gpu-hot-path-audit.md | 1 + docs/dqn-wire-up-audit.md | 2 + 10 files changed, 288 insertions(+), 169 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/atoms_update_kernel.cu create mode 100644 crates/ml/src/trainers/dqn/monitors/atoms_monitor.rs diff --git a/crates/ml/build.rs b/crates/ml/build.rs index 944fba8ec..dfbc3ce15 100644 --- a/crates/ml/build.rs +++ b/crates/ml/build.rs @@ -88,6 +88,7 @@ fn main() { "epsilon_update_kernel.cu", "gamma_update_kernel.cu", "kelly_cap_update_kernel.cu", + "atoms_update_kernel.cu", ]; // ALL kernels get common header (BF16 types + wrappers) diff --git a/crates/ml/src/cuda_pipeline/atoms_update_kernel.cu b/crates/ml/src/cuda_pipeline/atoms_update_kernel.cu new file mode 100644 index 000000000..b8a696b73 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/atoms_update_kernel.cu @@ -0,0 +1,104 @@ +/* atoms_update — C51 adaptive atom position recompute, GPU-driven (Plan 1 Task 9). + * + * Replaces the CPU-side `recompute_atom_positions` loop which launched the + * existing `adaptive_atom_positions` kernel once per branch. Here, all 4 + * branches are dispatched in a single launch: grid = (4, 1, 1), + * block = (256, 1, 1), shared memory = 2 * num_atoms * sizeof(float). + * blockIdx.x selects the branch (0=dir, 1=mag, 2=ord, 3=urg). + * + * Reads: + * spacing_raw[b] — learnable softmax spacing params for branch b, + * accessed via pre-resolved device pointers passed as args. + * isv_signals — ISV bus (per-branch v-range at slots 23+2*b / 24+2*b). + * + * Writes: + * positions_out — atom grid [4 * num_atoms], branch b at offset b*num_atoms. + * + * Formula (same as adaptive_atom_positions in experience_kernels.cu): + * spacing = softmax(spacing_raw) + * positions[j] = v_min + (v_max - v_min) * cumsum(spacing)[j] + * with v_min = isv[23+2*b] - isv[24+2*b], v_max = isv[23+2*b] + isv[24+2*b]. + */ + +extern "C" __global__ void atoms_update( + unsigned long long spacing_ptr_b0, /* device ptr to spacing_raw[dir] */ + unsigned long long spacing_ptr_b1, /* device ptr to spacing_raw[mag] */ + unsigned long long spacing_ptr_b2, /* device ptr to spacing_raw[ord] */ + unsigned long long spacing_ptr_b3, /* device ptr to spacing_raw[urg] */ + float* __restrict__ positions_out, /* [4 * num_atoms] */ + int num_atoms, + const float* __restrict__ isv_signals +) { + int branch = (int)blockIdx.x; + if (branch >= 4) return; + + int tid = threadIdx.x; + extern __shared__ float shmem[]; + float* s_softmax = shmem; + float* s_cumsum = shmem + num_atoms; + + /* Per-branch v-range from ISV: centre at 23+2*b, half at 24+2*b */ + float v_center = isv_signals[23 + 2 * branch]; + float v_half = isv_signals[24 + 2 * branch]; + float v_min = v_center - v_half; + float v_max = v_center + v_half; + + /* Resolve spacing_raw pointer for this branch */ + const float* spacing_raw; + switch (branch) { + case 0: spacing_raw = (const float*)spacing_ptr_b0; break; + case 1: spacing_raw = (const float*)spacing_ptr_b1; break; + case 2: spacing_raw = (const float*)spacing_ptr_b2; break; + default: spacing_raw = (const float*)spacing_ptr_b3; break; + } + + /* Load spacing_raw into shared memory */ + if (tid < num_atoms) { + s_softmax[tid] = spacing_raw[tid]; + } + __syncthreads(); + + /* Find max for numerical stability (single thread, num_atoms <= 256 is tiny) */ + if (tid == 0) { + float max_val = s_softmax[0]; + for (int j = 1; j < num_atoms; j++) + max_val = fmaxf(max_val, s_softmax[j]); + s_cumsum[0] = max_val; /* store max temporarily */ + } + __syncthreads(); + float max_val = s_cumsum[0]; + __syncthreads(); + + /* Compute softmax */ + if (tid < num_atoms) { + s_softmax[tid] = expf(s_softmax[tid] - max_val); + } + __syncthreads(); + + /* Sum for normalization */ + if (tid == 0) { + float sum = 0.0f; + for (int j = 0; j < num_atoms; j++) + sum += s_softmax[j]; + float inv_sum = 1.0f / fmaxf(sum, 1e-8f); + for (int j = 0; j < num_atoms; j++) + s_softmax[j] *= inv_sum; + } + __syncthreads(); + + /* Inclusive prefix sum (sequential, tid=0 only) */ + if (tid == 0) { + s_cumsum[0] = s_softmax[0]; + for (int j = 1; j < num_atoms; j++) + s_cumsum[j] = s_cumsum[j-1] + s_softmax[j]; + } + __syncthreads(); + + /* Write positions: positions[0] = v_min, positions[j>0] = v_min + range * cumsum[j-1] */ + float* out = positions_out + (long long)branch * num_atoms; + if (tid < num_atoms) { + float range = v_max - v_min; + float frac = (tid == 0) ? 0.0f : s_cumsum[tid - 1]; + out[tid] = v_min + range * frac; + } +} diff --git a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs index 87a2e003d..fa3985483 100644 --- a/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs +++ b/crates/ml/src/cuda_pipeline/gpu_dqn_trainer.rs @@ -85,6 +85,8 @@ static EPSILON_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/e static GAMMA_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/gamma_update_kernel.cubin")); /// Plan 1 Task 11: GPU-driven effective Kelly cap. static KELLY_CAP_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/kelly_cap_update_kernel.cubin")); +/// Plan 1 Task 9: GPU-driven C51 atom position recompute (4-branch, single launch). +static ATOMS_UPDATE_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/atoms_update_kernel.cubin")); /// Mamba2 temporal scan configuration. const MAMBA2_HISTORY_K: usize = 8; // Rolling history length @@ -228,8 +230,8 @@ pub const SHARPE_EMA_INDEX: usize = 22; /// ISV v-range unification bundle (spec 2026-04-23): per-branch Q-support /// centre + half-width, 2 slots per branch (direction=0, magnitude=1, /// order=2, urgency=3). Written by `update_eval_v_range` from the per-branch -/// Q-stats EMA state; read by `adaptive_atom_positions` (atom grid), -/// `warm_start_atom_positions` (quantile clamp) and `update_per_sample_support` +/// Q-stats EMA state; read by the `atoms_update` GPU kernel (`launch_atoms_update`, +/// Plan 1 Task 9) and `update_per_sample_support` /// (loss kernel projection range). Bootstrap values at construction + fold /// reset: `v_center = 0`, `v_half = (config.v_max - config.v_min) / 2` — at /// epoch 1 this reproduces the pre-spec `[config.v_min, config.v_max]` @@ -1416,6 +1418,9 @@ pub struct GpuDqnTrainer { /// Plan 1 Task 11: GPU-driven effective Kelly cap. Cold-path (per-epoch). /// Reads ISV[12] + portfolio_state[n_envs, PS_STRIDE]; writes ISV[KELLY_CAP_EFF_INDEX=44]. kelly_cap_update_kernel: CudaFunction, + /// Plan 1 Task 9: GPU-driven C51 atom position recompute. Cold-path (per-epoch + SGD step). + /// Reads ISV v-range slots [23..31) + spacing_raw params; writes atom_positions_buf. + atoms_update_kernel: CudaFunction, /// Host-side cached per-component norms populated at epoch boundary /// from the pinned result slot. Index order matches `grad_mag_*_ratio` /// accessors: `[iqn, cql, cql_sx, c51, c51_bs, distill, rec, pred, ens]`. @@ -2325,136 +2330,46 @@ impl GpuDqnTrainer { // NO graph invalidation — lr_dev_ptr is captured, value is read at replay time. } - /// Recompute adaptive atom positions from learned spacing parameters. - /// Call once per epoch (positions are slow-moving). + /// Plan 1 Task 9: GPU-driven atom position recompute (all 4 branches, single launch). /// - /// Per-branch v-range bounds come from the ISV signal bus (slots - /// `V_CENTER_{branch}_INDEX` and `V_HALF_{branch}_INDEX`). The bootstrap - /// written at construction reproduces `[config.v_min, config.v_max]` for - /// every branch before the first `update_eval_v_range` call, so epoch 1 - /// behaviour matches the pre-ISV path byte-for-byte. - pub(crate) fn recompute_atom_positions(&self) -> Result<(), MLError> { + /// Dispatches `atoms_update` kernel with grid=(4,1,1), block=(256,1,1). + /// Each block handles one branch; reads ISV v-range slots [23..31) and + /// spacing_raw params; writes atom_positions_buf in-place. Replaces the + /// CPU-side loop over `adaptive_atom_positions` per branch. Must be called + /// once per epoch boundary and after each `step_atom_positions` SGD update. + pub(crate) fn launch_atoms_update(&self) -> Result<(), MLError> { let na = self.config.num_atoms; let param_sizes = compute_param_sizes(&self.config); let shmem = (na * 2) * std::mem::size_of::(); - let atom_positions_buf_ptr = self.atom_positions_buf.raw_ptr(); - let isv_dev_ptr = self.isv_signals_dev_ptr; + let positions_out = self.atom_positions_buf.raw_ptr(); + let isv_dev_ptr = self.isv_signals_dev_ptr; - for branch in 0..4_usize { - let spacing_ptr = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 52 + branch); - let out_offset = (branch * na * std::mem::size_of::()) as u64; - let out_ptr = atom_positions_buf_ptr + out_offset; - let branch_i32 = branch as i32; + let sp_b0 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 52); + let sp_b1 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 53); + let sp_b2 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 54); + let sp_b3 = self.ptrs.params_ptr + padded_byte_offset(¶m_sizes, 55); + let na_i32 = na as i32; - // ISV v-range diag: log host-side ISV read used by the kernel. - // The kernel reads ISV via isv_dev_ptr (device mirror); the - // pinned mirror tracks the same writes from update_eval_v_range. - let (diag_center, diag_half) = if !self.isv_signals_pinned.is_null() { - unsafe { - let c = *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * branch); - let h = *self.isv_signals_pinned.add(V_CENTER_DIR_INDEX + 2 * branch + 1); - (c, h) - } - } else { - (f32::NAN, f32::NAN) - }; - tracing::info!( - target: "isv_vrange_diag", - branch, diag_center, diag_half, - "recompute_atom_positions ISV read" - ); - - unsafe { - self.stream.launch_builder(&self.adaptive_atom_kernel) - .arg(&spacing_ptr) - .arg(&out_ptr) - .arg(&(na as i32)) - .arg(&branch_i32) - .arg(&isv_dev_ptr) - .launch(LaunchConfig { - grid_dim: (1, 1, 1), - block_dim: (256, 1, 1), - shared_mem_bytes: shmem as u32, - }) - .map_err(|e| MLError::ModelError(format!("adaptive_atom_positions branch {branch}: {e}")))?; - } + unsafe { + self.stream + .launch_builder(&self.atoms_update_kernel) + .arg(&sp_b0) + .arg(&sp_b1) + .arg(&sp_b2) + .arg(&sp_b3) + .arg(&positions_out) + .arg(&na_i32) + .arg(&isv_dev_ptr) + .launch(LaunchConfig { + grid_dim: (4, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: shmem as u32, + }) + .map_err(|e| MLError::ModelError(format!("atoms_update launch: {e}")))?; } Ok(()) } - /// Warm-start C51 atom positions from empirical reward quantiles. - /// - /// Called once per epoch after experience collection, before training steps. - /// Shared empirical reward quantiles are clamped to each branch's - /// ISV-derived `[centre - half, centre + half]` support before being - /// tiled into `atom_positions_buf`. Prior to the first ISV write, the - /// bootstrap values (centre=0, half=(config.v_max-config.v_min)/2) make - /// the clamp equivalent to the old `config.v_{min,max}` clamp for every - /// branch — epoch 1 behaviour is preserved byte-for-byte. Per-branch - /// specialisation then flows through the existing `step_atom_positions` - /// SGD optimiser. - /// - /// The HtoD transfer is ~1us for 4 * num_atoms floats (typically 4 * 51 = 204). - pub(crate) fn warm_start_atom_positions(&mut self, quantiles: &[f32]) -> Result<(), MLError> { - let na = self.config.num_atoms; - if quantiles.len() != na { - return Ok(()); // skip on mismatch - } - - // Fallback bounds used when the ISV bus has not yet been written — - // keep the pre-spec behaviour (full config range, shared across - // branches) as a safety net. - let v_min_f = self.config.v_min; - let v_max_f = self.config.v_max; - let isv_ptr = self.isv_signals_pinned; - - let mut host_data = vec![0.0_f32; 4 * na]; - for branch in 0..4_usize { - let (b_min, b_max) = if isv_ptr.is_null() { - (v_min_f, v_max_f) - } else { - let centre = unsafe { *isv_ptr.add(V_CENTER_DIR_INDEX + 2 * branch) }; - let half = unsafe { *isv_ptr.add(V_CENTER_DIR_INDEX + 2 * branch + 1) }; - let b_min = centre - half; - let b_max = centre + half; - // Guard against degenerate (half=0) reads; fall back to full - // config range so warm-start never produces zero-range atoms. - if (b_max - b_min).abs() < 1e-6 { (v_min_f, v_max_f) } - else { (b_min, b_max) } - }; - tracing::info!( - target: "isv_vrange_diag", - branch, b_min, b_max, - isv_ptr_null = isv_ptr.is_null(), - "warm_start clamp bounds" - ); - let slot = &mut host_data[branch * na..(branch + 1) * na]; - for (atom_idx, q) in quantiles.iter().enumerate() { - slot[atom_idx] = q.clamp(b_min, b_max); - } - } - - self.stream.memcpy_htod(&host_data, &mut self.atom_positions_buf) - .map_err(|e| MLError::ModelError(format!("warm_start_atom_positions HtoD: {e}")))?; - - // ISV v-range diag: log what we just wrote per branch (min/max of - // atom_positions) — host_data is the source of truth (no need for a - // DtoH roundtrip since memcpy_htod doesn't mutate source on failure - // and succeeded above). - for branch in 0..4_usize { - let slot = &host_data[branch * na..(branch + 1) * na]; - let amin = slot.iter().cloned().fold(f32::INFINITY, f32::min); - let amax = slot.iter().cloned().fold(f32::NEG_INFINITY, f32::max); - tracing::info!( - target: "isv_vrange_diag", - branch, amin, amax, - "warm_start atom_positions written" - ); - } - - Ok(()) - } - /// SGD update on adaptive atom spacing_raw parameters. /// Maximizes atom entropy (gradient ascent) — concentrates atoms where mass exists. /// Uses scale_f32_kernel to decay spacing_raw toward zero: softmax(0,...,0) = uniform. @@ -2487,8 +2402,8 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("atom_position_sgd branch {branch}: {e}")))?; } } - // Recompute positions from updated spacing_raw - self.recompute_atom_positions()?; + // Recompute positions from updated spacing_raw via GPU atoms_update kernel. + self.launch_atoms_update()?; Ok(()) } @@ -7070,6 +6985,14 @@ impl GpuDqnTrainer { .map_err(|e| MLError::ModelError(format!("kelly_cap_update load: {e}")))? }; + // Plan 1 Task 9: load atoms_update kernel (cold-path, per-epoch + SGD step). + let atoms_update_kernel = { + let module = stream.context().load_cubin(ATOMS_UPDATE_CUBIN.to_vec()) + .map_err(|e| MLError::ModelError(format!("atoms_update cubin load: {e}")))?; + module.load_function("atoms_update") + .map_err(|e| MLError::ModelError(format!("atoms_update load: {e}")))? + }; + // eval_v_range — pinned device-mapped (zero-copy write from CPU). let eval_v_range_pinned: *mut f32 = unsafe { let flags = cudarc::driver::sys::CU_MEMHOSTALLOC_DEVICEMAP; @@ -9127,6 +9050,7 @@ impl GpuDqnTrainer { epsilon_update_kernel, gamma_update_kernel, kelly_cap_update_kernel, + atoms_update_kernel, grad_component_norms_mag: [0.0_f32; 9], grad_component_norms_dir: [0.0_f32; 9], grad_component_norms_trunk: [0.0_f32; 9], diff --git a/crates/ml/src/trainers/dqn/fused_training.rs b/crates/ml/src/trainers/dqn/fused_training.rs index 0b007154a..5266cc362 100644 --- a/crates/ml/src/trainers/dqn/fused_training.rs +++ b/crates/ml/src/trainers/dqn/fused_training.rs @@ -2960,9 +2960,10 @@ impl FusedTrainingCtx { /// Per-sample epsilon from IQL expectile gap. pub(crate) fn per_sample_epsilon_ptr(&self) -> u64 { self.gpu_iql.per_sample_epsilon_ptr() } pub(crate) fn num_atoms(&self) -> usize { self.trainer.config().num_atoms } - /// Warm-start C51 atom positions from empirical reward quantiles (once per epoch). - pub(crate) fn warm_start_atom_positions(&mut self, quantiles: &[f32]) -> Result<(), crate::MLError> { - self.trainer.warm_start_atom_positions(quantiles) + /// Plan 1 Task 9: GPU-driven atom position recompute (all 4 branches, single launch). + pub(crate) fn launch_atoms_update(&self) -> anyhow::Result<()> { + self.trainer.launch_atoms_update() + .map_err(|e| anyhow::anyhow!("launch_atoms_update: {e}")) } pub(crate) fn update_adaptive_clip(&mut self, grad_norm: f32) { self.trainer.update_adaptive_clip(grad_norm); } pub(crate) fn adaptive_clip_value(&self) -> f32 { self.trainer.adaptive_clip_value() } @@ -3055,12 +3056,6 @@ impl FusedTrainingCtx { self.trainer.set_lr(lr); } - /// Recompute adaptive atom positions from learned spacing parameters (once per epoch). - pub(crate) fn recompute_atom_positions(&self) -> anyhow::Result<()> { - self.trainer.recompute_atom_positions() - .map_err(|e| anyhow::anyhow!("recompute_atom_positions: {e}")) - } - /// Read adam_step counter. pub(crate) fn adam_step_val(&self) -> i32 { self.trainer.adam_step_val() } diff --git a/crates/ml/src/trainers/dqn/monitors/atoms_monitor.rs b/crates/ml/src/trainers/dqn/monitors/atoms_monitor.rs new file mode 100644 index 000000000..2833705c3 --- /dev/null +++ b/crates/ml/src/trainers/dqn/monitors/atoms_monitor.rs @@ -0,0 +1,118 @@ +//! CPU-side observer for the GPU-driven C51 adaptive atom grid. +//! +//! The `atoms_update` CUDA kernel (Plan 1 Task 9) reads ISV v-range slots +//! [23..31) (per-branch centre + half-width) and spacing_raw parameters to +//! compute adaptive atom positions for all 4 branches in a single GPU launch. +//! This monitor reads the representative `V_CENTER_DIR_INDEX=23` scalar for +//! fire-rate tracking, and exposes all 8 v-range ISV slots in `diagnose()`. +//! +//! The monitor's `read()` returns `isv[V_CENTER_DIR_INDEX]`. `diagnose()` +//! exposes the direction branch v-range (centre + half), all branch v-ranges, +//! and fire-rate. + +use crate::cuda_pipeline::gpu_dqn_trainer::{ + V_CENTER_DIR_INDEX, V_HALF_DIR_INDEX, + V_CENTER_MAG_INDEX, V_HALF_MAG_INDEX, + V_CENTER_ORD_INDEX, V_HALF_ORD_INDEX, + V_CENTER_URG_INDEX, V_HALF_URG_INDEX, +}; +use crate::trainers::dqn::adaptive_monitor::{ + AdaptiveMonitor, DiagSnapshot, FireRateStats, IsvBus, +}; + +pub(crate) struct AtomsMonitor { + last_center: f32, + fire_rate: FireRateStats, +} + +impl AtomsMonitor { + pub(crate) fn new() -> Self { + Self { last_center: 0.0, fire_rate: FireRateStats::default() } + } +} + +impl AdaptiveMonitor for AtomsMonitor { + fn read(&self, isv: &IsvBus<'_>) -> f32 { + isv.read(V_CENTER_DIR_INDEX) + } + + fn diagnose(&self, isv: &IsvBus<'_>) -> DiagSnapshot { + DiagSnapshot::new() + .with("v_center_dir", isv.read(V_CENTER_DIR_INDEX)) + .with("v_half_dir", isv.read(V_HALF_DIR_INDEX)) + .with("v_center_mag", isv.read(V_CENTER_MAG_INDEX)) + .with("v_half_mag", isv.read(V_HALF_MAG_INDEX)) + .with("v_center_ord", isv.read(V_CENTER_ORD_INDEX)) + .with("v_half_ord", isv.read(V_HALF_ORD_INDEX)) + .with("v_center_urg", isv.read(V_CENTER_URG_INDEX)) + .with("v_half_urg", isv.read(V_HALF_URG_INDEX)) + .with("fire_rate", self.fire_rate.fire_rate()) + } + + fn observe(&mut self, current: f32) { + let fired = (current - self.last_center).abs() > 1e-6; + self.fire_rate.record_fire(fired); + self.last_center = current; + } + + fn fire_rate(&self) -> &FireRateStats { + &self.fire_rate + } + + fn name(&self) -> &'static str { + "atoms" + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn make_isv(center_dir: f32, half_dir: f32) -> Vec { + let mut v = vec![0.0f32; 49]; + v[V_CENTER_DIR_INDEX] = center_dir; + v[V_HALF_DIR_INDEX] = half_dir; + v[V_CENTER_MAG_INDEX] = 0.1; + v[V_HALF_MAG_INDEX] = 0.5; + v + } + + #[test] + fn read_returns_v_center_dir_slot() { + let mut buf = make_isv(0.05, 2.5); + let bus = IsvBus::new(&mut buf); + let mon = AtomsMonitor::new(); + let v = mon.read(&bus); + assert!((v - 0.05).abs() < 1e-7, "expected v_center_dir=0.05, got {v}"); + } + + #[test] + fn diagnose_exposes_all_fields() { + let mut buf = make_isv(0.10, 3.0); + let bus = IsvBus::new(&mut buf); + let mon = AtomsMonitor::new(); + let snap = mon.diagnose(&bus); + assert!(snap.fields.contains_key("v_center_dir")); + assert!(snap.fields.contains_key("v_half_dir")); + assert!(snap.fields.contains_key("v_center_mag")); + assert!(snap.fields.contains_key("v_half_mag")); + assert!(snap.fields.contains_key("v_center_ord")); + assert!(snap.fields.contains_key("v_half_ord")); + assert!(snap.fields.contains_key("v_center_urg")); + assert!(snap.fields.contains_key("v_half_urg")); + assert!(snap.fields.contains_key("fire_rate")); + let center = snap.fields["v_center_dir"]; + assert!((center - 0.10).abs() < 1e-6, "expected v_center_dir=0.10, got {center}"); + } + + #[test] + fn observe_tracks_fire_rate_on_change() { + let mut mon = AtomsMonitor::new(); + mon.observe(0.0); + mon.observe(0.0); + mon.observe(0.05); + let rate = mon.fire_rate().fire_rate(); + assert!((rate - 2.0 / 3.0).abs() < 1e-6, + "expected 2/3 fire rate, got {rate}"); + } +} diff --git a/crates/ml/src/trainers/dqn/monitors/mod.rs b/crates/ml/src/trainers/dqn/monitors/mod.rs index 5c02b06f1..5527665e4 100644 --- a/crates/ml/src/trainers/dqn/monitors/mod.rs +++ b/crates/ml/src/trainers/dqn/monitors/mod.rs @@ -5,6 +5,7 @@ //! read those slots for HEALTH_DIAG emission and `controller_activity` //! smoke-test fire-rate tracking. No CPU-side adaptive computation. +pub(crate) mod atoms_monitor; pub(crate) mod epsilon_monitor; pub(crate) mod gamma_monitor; pub(crate) mod grad_balancer_monitor; diff --git a/crates/ml/src/trainers/dqn/state_reset_registry.rs b/crates/ml/src/trainers/dqn/state_reset_registry.rs index 1dc6f54a3..2c1a3af28 100644 --- a/crates/ml/src/trainers/dqn/state_reset_registry.rs +++ b/crates/ml/src/trainers/dqn/state_reset_registry.rs @@ -165,22 +165,22 @@ impl StateResetRegistry { RegistryEntry { name: "isv_epoch_idx", category: ResetCategory::FoldReset, - description: "ISV[EPOCH_IDX_INDEX=39] — current epoch index; CPU writes at epoch boundary (follow-up task)", + description: "ISV[EPOCH_IDX_INDEX=39] — current epoch index; CPU writes at epoch boundary (Plan 1 Task 13)", }, RegistryEntry { name: "isv_epsilon_eff", category: ResetCategory::FoldReset, - description: "ISV[EPSILON_EFF_INDEX=41] — effective epsilon; GPU epsilon-adaptive kernel fills (follow-up task)", + description: "ISV[EPSILON_EFF_INDEX=41] — effective epsilon; GPU epsilon_update kernel fills (Plan 1 Task 14)", }, RegistryEntry { name: "isv_tau_eff", category: ResetCategory::FoldReset, - description: "ISV[TAU_EFF_INDEX=42] — effective tau; GPU tau-adaptive kernel fills (follow-up task)", + description: "ISV[TAU_EFF_INDEX=42] — effective tau; GPU tau_update kernel fills (Plan 1 Task 13)", }, RegistryEntry { name: "isv_gamma_eff", category: ResetCategory::FoldReset, - description: "ISV[GAMMA_EFF_INDEX=43] — effective gamma; GPU gamma-adaptive kernel fills (follow-up task)", + description: "ISV[GAMMA_EFF_INDEX=43] — effective gamma; GPU gamma_update kernel fills (Plan 1 Task 10)", }, RegistryEntry { name: "isv_kelly_cap_eff", diff --git a/crates/ml/src/trainers/dqn/trainer/training_loop.rs b/crates/ml/src/trainers/dqn/trainer/training_loop.rs index 904a98f82..52c397d60 100644 --- a/crates/ml/src/trainers/dqn/trainer/training_loop.rs +++ b/crates/ml/src/trainers/dqn/trainer/training_loop.rs @@ -352,10 +352,10 @@ impl DQNTrainer { } } - // Recompute adaptive atom positions from learned spacing_raw params (slow-moving, once per epoch). + // Plan 1 Task 9: GPU-driven atom position recompute (4-branch, single launch). if let Some(ref fused) = self.fused_ctx { - if let Err(e) = fused.recompute_atom_positions() { - tracing::warn!(epoch, "recompute_atom_positions failed: {e}"); + if let Err(e) = fused.launch_atoms_update() { + tracing::warn!(epoch, "launch_atoms_update failed: {e}"); } } @@ -517,38 +517,11 @@ impl DQNTrainer { // v8: TD(λ) is now wired in collect_experiences_gpu — self-bootstrap with // rewards as Q(s') approximation, overwrites n-step when td_lambda > 0.01. - // C51 atom warm-start: extract empirical reward quantiles from the bitonic - // sort pipeline and initialize atom positions for this epoch. The existing - // step_atom_positions SGD optimizer refines from here during training steps. if let Some(ref mut collector) = self.gpu_experience_collector { let count = collector.last_experience_count(); - let num_atoms = self.hyperparams.num_atoms; - match collector.compute_reward_quantiles(count, num_atoms) { - Ok(quantiles) => { - // Per-branch clamp now lives inside - // `warm_start_atom_positions` (ISV v-range spec Phase 2c, - // commit 9deda5f65). That path reads each branch's - // `[centre - half, centre + half]` from the ISV pinned - // bus and clamps the quantile array per branch. An - // outer static clamp to `config.v_{min,max}` here would - // be redundant double-clamping to a wider bound — it - // hides which layer owns the support and never - // contributes once ISV is populated. Pass the raw - // quantiles straight through. - if let Some(ref mut fused) = self.fused_ctx { - if let Err(e) = fused.warm_start_atom_positions(&quantiles) { - warn!("C51 atom warm-start failed (non-fatal): {e}"); - } - } - } - Err(e) => { - debug!("C51 reward quantile extraction failed (non-fatal): {e}"); - } - } - // Robust PopArt: extract median/IQR from the already-sorted reward buffer. - // The bitonic sort from compute_reward_quantiles above leaves reward_sort_keys_buf - // sorted, so extract_median_iqr just reads 3 positions (tiny DtoH: 3 floats). + // Robust PopArt: extract median/IQR from the reward sort buffer. + // extract_median_iqr reads 3 positions from the sorted buffer (tiny DtoH: 3 floats). if self.hyperparams.popart_robust { match collector.extract_median_iqr(count) { Ok((median, iqr)) => { diff --git a/docs/dqn-gpu-hot-path-audit.md b/docs/dqn-gpu-hot-path-audit.md index 38e4c790e..8d629f369 100644 --- a/docs/dqn-gpu-hot-path-audit.md +++ b/docs/dqn-gpu-hot-path-audit.md @@ -71,6 +71,7 @@ | `epsilon_update_kernel.cu` (Plan 1 Task 14) | single-thread kernel: reads ISV[39,40,12]; writes ISV[41] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; zero hot-path overhead | OK | | `gamma_update_kernel.cu` (Plan 1 Task 10) | single-thread kernel: reads ISV[12]; writes ISV[43] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; ISV[43] read by fill_gamma_buf on hot path via read_isv_signal_at (pinned, zero-copy) | OK | | `kelly_cap_update_kernel.cu` (Plan 1 Task 11) | single-thread kernel: reads ISV[12] + portfolio_states[n_envs, PS_STRIDE]; aggregates mean half-Kelly × health-safety; writes ISV[44] | COLD-PATH | Per-epoch boundary call from `training_loop.rs`; 1-thread, 1-block; portfolio_states dev ptr passed from GpuExperienceCollector without CPU roundtrip | OK | +| `atoms_update_kernel.cu` (Plan 1 Task 9) | 4-block kernel (one per branch): reads ISV[23..31) + spacing_raw params[52..56); writes atom_positions_buf[4 * num_atoms] via softmax + cumsum | COLD-PATH | Per-epoch boundary + SGD step (every 50 steps); grid=(4,1,1), block=(256,1,1); replaces CPU loop that launched adaptive_atom_positions 4× separately | OK | ## Summary diff --git a/docs/dqn-wire-up-audit.md b/docs/dqn-wire-up-audit.md index 03a8519db..01de1d2f6 100644 --- a/docs/dqn-wire-up-audit.md +++ b/docs/dqn-wire-up-audit.md @@ -52,6 +52,8 @@ | `cuda_pipeline/gamma_update_kernel.cu` | `GpuDqnTrainer::launch_gamma_update` → `FusedTrainingCtx::launch_gamma_update` → `training_loop.rs` epoch boundary | Wired | Plan 1 Task 10 — cold-path (per-epoch) | — | | `trainers/dqn/monitors/kelly_cap_monitor.rs` | Read-only observer for kelly_cap_update kernel output (ISV slot 44); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 11 | — | | `cuda_pipeline/kelly_cap_update_kernel.cu` | `GpuDqnTrainer::launch_kelly_cap_update` → `FusedTrainingCtx::launch_kelly_cap_update` → `training_loop.rs` epoch boundary; takes portfolio_states dev ptr + n_envs | Wired | Plan 1 Task 11 — cold-path (per-epoch) | — | +| `trainers/dqn/monitors/atoms_monitor.rs` | Read-only observer for atoms_update kernel inputs (ISV v-range slots 23..31); consumers: HEALTH_DIAG + controller_activity smoke | Wired | Plan 1 Task 9 | — | +| `cuda_pipeline/atoms_update_kernel.cu` | `GpuDqnTrainer::launch_atoms_update` → `FusedTrainingCtx::launch_atoms_update` → `training_loop.rs` epoch boundary; replaces per-branch `recompute_atom_positions` CPU loop | Wired | Plan 1 Task 9 — cold-path (per-epoch + SGD step) | — | ## CUDA Pipeline — Rust Wrappers