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

perf(diag): move ALL remaining device reads to async staging — zero syncs

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 01:11:59 +02:00
parent 5dd5a08963
commit 5f3e3578a0
2 changed files with 276 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

106
crates/ml-alpha/examples/alpha_rl_train.rs
View File

@@ -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<u8> = 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 {

246
crates/ml-alpha/src/trainer/diag_staging.rs
View File

@@ -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<CudaStream>,
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<f32>,
@@ -107,6 +149,14 @@ impl DiagStaging {
raw_rewards_d: &CudaSlice<f32>,
trade_duration_d: &CudaSlice<f32>,
outcome_ema_d: &CudaSlice<f32>,
position_lots_d: &CudaSlice<i32>,
pyramid_count_d: &CudaSlice<i32>,
unit_entry_price_d: &CudaSlice<f32>,
unit_entry_step_d: &CudaSlice<i32>,
unit_lots_d: &CudaSlice<i32>,
unit_trail_d: &CudaSlice<f32>,
close_unit_index_d: &CudaSlice<i32>,
frd_logits_d: &CudaSlice<f32>,
) -> 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<i32> {
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<i32> {
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<i32> {
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<i32> {
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<i32> {
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,
)
}
}
}