From 5f3e3578a03daf759f0611342c5637fcf440988f Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 26 May 2026 01:11:59 +0200 Subject: [PATCH] =?UTF-8?q?perf(diag):=20move=20ALL=20remaining=20device?= =?UTF-8?q?=20reads=20to=20async=20staging=20=E2=80=94=20zero=20syncs?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Expanded DiagStaging with 8 more fields (position_lots, pyramid_count, unit_entry_price, unit_entry_step, unit_lots, unit_trail, close_unit_index, frd_logits). Eliminates 36 stream syncs per step that consumed 90.9% of CUDA API time per nsys profile. unit_active derived from unit_lots != 0 on host (avoids u8-to-f32 DtoD mismatch). Co-Authored-By: Claude Opus 4.7 --- crates/ml-alpha/examples/alpha_rl_train.rs | 106 ++++----- crates/ml-alpha/src/trainer/diag_staging.rs | 246 +++++++++++++++++++- 2 files changed, 276 insertions(+), 76 deletions(-) diff --git a/crates/ml-alpha/examples/alpha_rl_train.rs b/crates/ml-alpha/examples/alpha_rl_train.rs index fdac44db2..3ae6ee51a 100644 --- a/crates/ml-alpha/examples/alpha_rl_train.rs +++ b/crates/ml-alpha/examples/alpha_rl_train.rs @@ -89,7 +89,7 @@ use ml_alpha::rl::isv_slots::{ }; use ml_alpha::trainer::diag_staging::DiagStaging; use ml_alpha::trainer::integrated::{ - read_slice_d_pub, read_slice_i32_d_pub, read_slice_u8_d_pub, write_slice_i32_d_pub, + write_slice_i32_d_pub, IntegratedStepStats, IntegratedTrainer, IntegratedTrainerConfig, }; use ml_alpha::trainer::perception::PerceptionTrainerConfig; @@ -442,11 +442,12 @@ fn main() -> Result<()> { eprintln!("per-step diag JSONL: {}", diag_path.display()); // Async non-blocking diagnostic staging: separate CUDA stream + - // double-buffered mapped-pinned memory. Replaces 7 blocking - // read_slice_d calls (ISV, rewards, dones, actions, raw_rewards, - // trade_duration, outcome_ema) with async DtoD copies. The host - // reads the PREVIOUS step's buffer while the current step's copies - // run concurrently on the diag stream. One-step delay on diag data. + // double-buffered mapped-pinned memory. ALL per-step device reads + // (ISV, rewards, dones, actions, raw_rewards, trade_duration, + // outcome_ema, position_lots, pyramid_count, unit state x4, + // close_unit_index, frd_logits) use async DtoD copies — zero + // stream syncs on the training stream. The host reads the PREVIOUS + // step's buffer while current step copies run on the diag stream. let mut diag_staging = DiagStaging::new(&dev, cli.n_backtests) .context("DiagStaging::new")?; @@ -577,11 +578,11 @@ fn main() -> Result<()> { process::exit(2); } - // Launch async DtoD copies of the 7 high-frequency diag buffers - // into the current staging buffer. Non-blocking on the training - // stream — the copies run on diag_staging's separate stream. - // The host will read these values at the START of the NEXT step - // iteration (after sync_and_swap completes the copy). + // Launch async DtoD copies of ALL diag buffers into the current + // staging buffer. Non-blocking on the training stream — the copies + // run on diag_staging's separate stream. The host will read these + // values at the START of the NEXT step iteration (after + // sync_and_swap completes the copy). diag_staging .snapshot_async( &trainer.isv_d, @@ -591,20 +592,22 @@ fn main() -> Result<()> { &trainer.raw_rewards_d, &trainer.trade_duration_emit_d, &trainer.outcome_ema_d, + &trainer.prev_position_lots_d, + &trainer.pyramid_units_count_d, + &trainer.unit_entry_price_d, + &trainer.unit_entry_step_d, + &trainer.unit_lots_d, + &trainer.unit_trail_distance_d, + &trainer.close_unit_index_d, + &trainer.frd_logits_d, ) .context("diag snapshot_async")?; // ── Per-step diag dump (async staging reads). ──────────────── - // The 7 high-frequency reads (ISV, rewards, dones, actions, - // raw_rewards, trade_duration, outcome_ema) now come from the - // DiagStaging double-buffer — zero GPU stalls on the training - // stream. On step 0 these report zeros (no prior snapshot); - // from step 1+ they report step N-1's data (one-step delay, - // acceptable for diagnostics). - // - // Remaining reads (position_lots, unit state, pyramid_count, - // close_unit_index, frd_logits) still use blocking helpers — - // they're less frequent or structurally complex. + // ALL diag reads come from the DiagStaging double-buffer — zero + // GPU stalls on the training stream. On step 0 these report + // zeros (no prior snapshot); from step 1+ they report step N-1's + // data (one-step delay, acceptable for diagnostics). let rewards_host = diag_staging.read_rewards(); let dones_host = diag_staging.read_dones(); let actions_raw = diag_staging.read_actions_raw(); @@ -615,43 +618,19 @@ fn main() -> Result<()> { let raw_rewards_host = diag_staging.read_raw_rewards(); let trade_duration_host = diag_staging.read_trade_duration(); let outcome_ema_host = diag_staging.read_outcome_ema(); - - // Remaining blocking reads — structurally complex or low-frequency. - let dev_stream = dev.cuda_stream().context("cuda_stream for diag")?; - // audit — position state per batch (signed lot count from - // `prev_position_lots_d`; positive=long, negative=short, 0=flat). - // Helps post-hoc analysis understand whether the agent is - // actually trading or stuck in some attractor. - let position_lots_host = read_slice_i32_d_pub( - dev_stream, - &trainer.prev_position_lots_d, - cli.n_backtests, - ) - .context("diag: read prev_position_lots_d")?; - let pyramid_count_host = read_slice_i32_d_pub( - dev_stream, - &trainer.pyramid_units_count_d, - cli.n_backtests, - ) - .context("diag: read pyramid_units_count_d")?; - let unit_entry_price_host = read_slice_d_pub( - dev_stream, &trainer.unit_entry_price_d, cli.n_backtests * 4, - ).context("diag: read unit_entry_price_d")?; - let unit_entry_step_host = read_slice_i32_d_pub( - dev_stream, &trainer.unit_entry_step_d, cli.n_backtests * 4, - ).context("diag: read unit_entry_step_d")?; - let unit_lots_host = read_slice_i32_d_pub( - dev_stream, &trainer.unit_lots_d, cli.n_backtests * 4, - ).context("diag: read unit_lots_d")?; - let unit_trail_host = read_slice_d_pub( - dev_stream, &trainer.unit_trail_distance_d, cli.n_backtests * 4, - ).context("diag: read unit_trail_distance_d")?; - let unit_active_host = read_slice_u8_d_pub( - dev_stream, &trainer.unit_active_d, cli.n_backtests * 4, - ).context("diag: read unit_active_d")?; - let close_unit_index_host = read_slice_i32_d_pub( - dev_stream, &trainer.close_unit_index_d, cli.n_backtests, - ).context("diag: read close_unit_index_d")?; + let position_lots_host = diag_staging.read_position_lots(); + let pyramid_count_host = diag_staging.read_pyramid_count(); + let unit_entry_price_host = diag_staging.read_unit_entry_price(); + let unit_entry_step_host = diag_staging.read_unit_entry_step(); + let unit_lots_host = diag_staging.read_unit_lots(); + let unit_trail_host = diag_staging.read_unit_trail(); + // Derive active mask from unit_lots (lots != 0 → active). + // Avoids a u8-to-f32 type mismatch in DtoD staging. + let unit_active_host: Vec = unit_lots_host + .iter() + .map(|&l| if l != 0 { 1u8 } else { 0u8 }) + .collect(); + let close_unit_index_host = diag_staging.read_close_unit_index(); let mut act_hist = [0u32; N_ACTIONS]; for &a in &actions_host { @@ -765,15 +744,10 @@ fn main() -> Result<()> { }; // FRD diag — per-horizon softmax entropy + argmax index, - // averaged across the batch. Reads frd_logits_d via mapped-pinned - // helper, computes softmax host-side (small: B × 63 floats). + // averaged across the batch. Reads from DiagStaging (async, + // zero GPU stalls). Computes softmax host-side (small: B x 63). let frd_diag = { - let frd_logits = read_slice_d_pub( - dev_stream, - &trainer.frd_logits_d, - cli.n_backtests * FRD_OUT_DIM, - ) - .context("diag: read frd_logits_d")?; + let frd_logits = diag_staging.read_frd_logits(); let mut per_h_entropy = [0.0_f32; FRD_N_HORIZONS]; let mut per_h_argmax_sum = [0.0_f32; FRD_N_HORIZONS]; for b in 0..cli.n_backtests { diff --git a/crates/ml-alpha/src/trainer/diag_staging.rs b/crates/ml-alpha/src/trainer/diag_staging.rs index 187f4e520..522af0575 100644 --- a/crates/ml-alpha/src/trainer/diag_staging.rs +++ b/crates/ml-alpha/src/trainer/diag_staging.rs @@ -27,14 +27,21 @@ use cudarc::driver::{CudaSlice, CudaStream, DevicePtr}; use ml_core::device::MlDevice; use crate::pinned_mem::MappedF32Buffer; +use crate::rl::frd::FRD_OUT_DIM; use crate::rl::isv_slots::RL_SLOTS_END; /// Double-buffered mapped-pinned staging for async diagnostic reads. /// /// Layout within each buffer (all f32, contiguous): -/// `[ISV (RL_SLOTS_END) | rewards (B) | dones (B) | actions (B) | raw_rewards (B) | trade_duration (B) | outcome_ema (B)]` +/// `[ISV (RL_SLOTS_END) | rewards (B) | dones (B) | actions (B) | raw_rewards (B) +/// | trade_duration (B) | outcome_ema (B) | position_lots (B) | pyramid_count (B) +/// | unit_entry_price (B×4) | unit_entry_step (B×4) | unit_lots (B×4) +/// | unit_trail (B×4) | close_unit_index (B) | frd_logits (B×FRD_OUT_DIM)]` /// -/// `actions` are stored as raw i32 bit patterns reinterpreted as f32 bytes. +/// i32 fields (`actions`, `position_lots`, `pyramid_count`, `unit_entry_step`, +/// `unit_lots`, `close_unit_index`) are stored as raw bit patterns in f32 slots. +/// DtoD copies raw bytes (4 bytes per element, same width); host reinterprets +/// via `f32::to_bits() as i32`. pub struct DiagStaging { diag_stream: Arc, bufs: [MappedF32Buffer; 2], @@ -50,6 +57,14 @@ pub struct DiagStaging { pub raw_rewards_offset: usize, pub trade_duration_offset: usize, pub outcome_ema_offset: usize, + pub position_lots_offset: usize, + pub pyramid_count_offset: usize, + pub unit_entry_price_offset: usize, + pub unit_entry_step_offset: usize, + pub unit_lots_offset: usize, + pub unit_trail_offset: usize, + pub close_unit_index_offset: usize, + pub frd_logits_offset: usize, pub b_size: usize, } @@ -65,9 +80,27 @@ impl DiagStaging { .new_stream() .map_err(|e| anyhow::anyhow!("DiagStaging new_stream: {e}"))?; - // Layout: ISV | rewards | dones | actions | raw_rewards | trade_duration | outcome_ema + // Layout (f32 slots): + // ISV | rewards | dones | actions | raw_rewards | trade_duration | outcome_ema + // | position_lots | pyramid_count | unit_entry_price(B×4) | unit_entry_step(B×4) + // | unit_lots(B×4) | unit_trail(B×4) | close_unit_index | frd_logits(B×FRD_OUT_DIM) let isv_len = RL_SLOTS_END; - let n_floats = isv_len + 6 * b_size; + let mut off = isv_len; + let rewards_offset = off; off += b_size; + let dones_offset = off; off += b_size; + let actions_offset = off; off += b_size; + let raw_rewards_offset = off; off += b_size; + let trade_duration_offset = off; off += b_size; + let outcome_ema_offset = off; off += b_size; + let position_lots_offset = off; off += b_size; + let pyramid_count_offset = off; off += b_size; + let unit_entry_price_offset = off; off += b_size * 4; + let unit_entry_step_offset = off; off += b_size * 4; + let unit_lots_offset = off; off += b_size * 4; + let unit_trail_offset = off; off += b_size * 4; + let close_unit_index_offset = off; off += b_size; + let frd_logits_offset = off; off += b_size * FRD_OUT_DIM; + let n_floats = off; let buf0 = unsafe { MappedF32Buffer::new(n_floats) } .map_err(|e| anyhow::anyhow!("DiagStaging buf0: {e}"))?; @@ -81,12 +114,20 @@ impl DiagStaging { n_floats, isv_offset: 0, isv_len, - rewards_offset: isv_len, - dones_offset: isv_len + b_size, - actions_offset: isv_len + 2 * b_size, - raw_rewards_offset: isv_len + 3 * b_size, - trade_duration_offset: isv_len + 4 * b_size, - outcome_ema_offset: isv_len + 5 * b_size, + rewards_offset, + dones_offset, + actions_offset, + raw_rewards_offset, + trade_duration_offset, + outcome_ema_offset, + position_lots_offset, + pyramid_count_offset, + unit_entry_price_offset, + unit_entry_step_offset, + unit_lots_offset, + unit_trail_offset, + close_unit_index_offset, + frd_logits_offset, b_size, }) } @@ -98,6 +139,7 @@ impl DiagStaging { /// stream automatically inserts inter-stream event waits: the diag_stream /// will wait for any prior writes to the source slice before the copy /// begins. This ensures correctness without explicit event management. + #[allow(clippy::too_many_arguments)] pub fn snapshot_async( &mut self, isv_d: &CudaSlice, @@ -107,6 +149,14 @@ impl DiagStaging { raw_rewards_d: &CudaSlice, trade_duration_d: &CudaSlice, outcome_ema_d: &CudaSlice, + position_lots_d: &CudaSlice, + pyramid_count_d: &CudaSlice, + unit_entry_price_d: &CudaSlice, + unit_entry_step_d: &CudaSlice, + unit_lots_d: &CudaSlice, + unit_trail_d: &CudaSlice, + close_unit_index_d: &CudaSlice, + frd_logits_d: &CudaSlice, ) -> Result<()> { let buf = &self.bufs[self.current]; let s = self.diag_stream.cu_stream(); @@ -182,6 +232,86 @@ impl DiagStaging { s, ) .context("diag snapshot outcome_ema")?; + + // position_lots (i32 → raw bytes into f32 slot) + let (src, _g) = position_lots_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.position_lots_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot position_lots")?; + + // pyramid_count (i32 → raw bytes into f32 slot) + let (src, _g) = pyramid_count_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.pyramid_count_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot pyramid_count")?; + + // unit_entry_price (f32, B×4) + let (src, _g) = unit_entry_price_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.unit_entry_price_offset * 4) as u64, + src, + b * 4 * 4, + s, + ) + .context("diag snapshot unit_entry_price")?; + + // unit_entry_step (i32, B×4 → raw bytes) + let (src, _g) = unit_entry_step_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.unit_entry_step_offset * 4) as u64, + src, + b * 4 * 4, + s, + ) + .context("diag snapshot unit_entry_step")?; + + // unit_lots (i32, B×4 → raw bytes) + let (src, _g) = unit_lots_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.unit_lots_offset * 4) as u64, + src, + b * 4 * 4, + s, + ) + .context("diag snapshot unit_lots")?; + + // unit_trail (f32, B×4) + let (src, _g) = unit_trail_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.unit_trail_offset * 4) as u64, + src, + b * 4 * 4, + s, + ) + .context("diag snapshot unit_trail")?; + + // close_unit_index (i32, B → raw bytes) + let (src, _g) = close_unit_index_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.close_unit_index_offset * 4) as u64, + src, + b * 4, + s, + ) + .context("diag snapshot close_unit_index")?; + + // frd_logits (f32, B×FRD_OUT_DIM) + let (src, _g) = frd_logits_d.device_ptr(&self.diag_stream); + cudarc::driver::result::memcpy_dtod_async( + buf.dev_ptr + (self.frd_logits_offset * 4) as u64, + src, + b * FRD_OUT_DIM * 4, + s, + ) + .context("diag snapshot frd_logits")?; } Ok(()) @@ -255,4 +385,100 @@ impl DiagStaging { std::slice::from_raw_parts(buf.host_ptr.add(self.outcome_ema_offset), self.b_size) } } + + /// Read the PREVIOUS step's per-batch position lots as i32. + /// The underlying data is i32 stored as raw bit patterns in f32 slots. + pub fn read_position_lots(&self) -> Vec { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + let raw = unsafe { + std::slice::from_raw_parts(buf.host_ptr.add(self.position_lots_offset), self.b_size) + }; + raw.iter().map(|f| f.to_bits() as i32).collect() + } + + /// Read the PREVIOUS step's per-batch pyramid unit count as i32. + pub fn read_pyramid_count(&self) -> Vec { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + let raw = unsafe { + std::slice::from_raw_parts(buf.host_ptr.add(self.pyramid_count_offset), self.b_size) + }; + raw.iter().map(|f| f.to_bits() as i32).collect() + } + + /// Read the PREVIOUS step's per-batch per-unit entry prices (B x 4 f32). + pub fn read_unit_entry_price(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { + std::slice::from_raw_parts( + buf.host_ptr.add(self.unit_entry_price_offset), + self.b_size * 4, + ) + } + } + + /// Read the PREVIOUS step's per-batch per-unit entry steps as i32 (B x 4). + pub fn read_unit_entry_step(&self) -> Vec { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + let raw = unsafe { + std::slice::from_raw_parts( + buf.host_ptr.add(self.unit_entry_step_offset), + self.b_size * 4, + ) + }; + raw.iter().map(|f| f.to_bits() as i32).collect() + } + + /// Read the PREVIOUS step's per-batch per-unit lot counts as i32 (B x 4). + pub fn read_unit_lots(&self) -> Vec { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + let raw = unsafe { + std::slice::from_raw_parts( + buf.host_ptr.add(self.unit_lots_offset), + self.b_size * 4, + ) + }; + raw.iter().map(|f| f.to_bits() as i32).collect() + } + + /// Read the PREVIOUS step's per-batch per-unit trail distances (B x 4 f32). + pub fn read_unit_trail(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { + std::slice::from_raw_parts( + buf.host_ptr.add(self.unit_trail_offset), + self.b_size * 4, + ) + } + } + + /// Read the PREVIOUS step's per-batch close unit index as i32. + pub fn read_close_unit_index(&self) -> Vec { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + let raw = unsafe { + std::slice::from_raw_parts( + buf.host_ptr.add(self.close_unit_index_offset), + self.b_size, + ) + }; + raw.iter().map(|f| f.to_bits() as i32).collect() + } + + /// Read the PREVIOUS step's per-batch FRD logits (B x FRD_OUT_DIM f32). + pub fn read_frd_logits(&self) -> &[f32] { + let prev = 1 - self.current; + let buf = &self.bufs[prev]; + unsafe { + std::slice::from_raw_parts( + buf.host_ptr.add(self.frd_logits_offset), + self.b_size * FRD_OUT_DIM, + ) + } + } }