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

perf(rl): DiagStaging snapshot_async takes raw u64 ptrs — zero device_ptr() calls

Browse Source 
Replaces 15 `&CudaSlice` parameters with pre-extracted `u64` device
pointers (`CudaSlice::raw_ptr()`). Each prior `device_ptr()` call
triggered cudarc's inter-stream event tracking (`cuStreamWaitEvent` +
`cuEventRecord`), adding sync overhead on every step. The raw pointers
are stable for the trainer's lifetime (all buffers pre-allocated at
init), so the event tracking was unnecessary.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 12:18:37 +02:00
parent 148454d37d
commit 241f8b347b
2 changed files with 68 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -599,21 +599,21 @@ fn main() -> Result<()> {
// sync_and_swap completes the copy).
diag_staging
.snapshot_async(
&trainer.isv_d,
&trainer.rewards_d,
&trainer.dones_d,
&trainer.actions_d,
&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,
trainer.isv_d.raw_ptr(),
trainer.rewards_d.raw_ptr(),
trainer.dones_d.raw_ptr(),
trainer.actions_d.raw_ptr(),
trainer.raw_rewards_d.raw_ptr(),
trainer.trade_duration_emit_d.raw_ptr(),
trainer.outcome_ema_d.raw_ptr(),
trainer.prev_position_lots_d.raw_ptr(),
trainer.pyramid_units_count_d.raw_ptr(),
trainer.unit_entry_price_d.raw_ptr(),
trainer.unit_entry_step_d.raw_ptr(),
trainer.unit_lots_d.raw_ptr(),
trainer.unit_trail_distance_d.raw_ptr(),
trainer.close_unit_index_d.raw_ptr(),
trainer.frd_logits_d.raw_ptr(),
)
.context("diag snapshot_async")?;

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

@@ -23,7 +23,7 @@
use std::sync::Arc;
use anyhow::{Context, Result};
use cudarc::driver::{CudaSlice, CudaStream, DevicePtr};
use cudarc::driver::CudaStream;
use ml_core::device::MlDevice;
use crate::pinned_mem::MappedF32Buffer;
@@ -133,30 +133,38 @@ impl DiagStaging {
}
/// Launch async DtoD copies from device buffers into the CURRENT staging
/// buffer. Non-blocking on the training stream — runs on the diag_stream.
/// buffer. Non-blocking on the training stream -- runs on the diag_stream.
///
/// The cudarc v0.19 `device_ptr()` call on a `CudaSlice` with a different
/// 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.
/// Accepts pre-extracted raw `u64` device pointers (`CudaSlice::raw_ptr()`)
/// instead of `&CudaSlice` references. This bypasses cudarc's
/// `device_ptr()` event-tracking path which internally does
/// `cuStreamWaitEvent` + `cuEventRecord` per call -- eliminating 15
/// inter-stream sync events per step.
///
/// Safety contract: callers must pass pointers from `CudaSlice` fields
/// whose device addresses are stable for the trainer's lifetime (all
/// pre-allocated at init). The training stream must have completed
/// writing to these buffers before the diag_stream reads them; this is
/// guaranteed by CUDA's same-context stream ordering when the training
/// stream's prior kernel launches have been enqueued before this call.
#[allow(clippy::too_many_arguments)]
pub fn snapshot_async(
&mut self,
isv_d: &CudaSlice<f32>,
rewards_d: &CudaSlice<f32>,
dones_d: &CudaSlice<f32>,
actions_d: &CudaSlice<i32>,
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>,
isv_ptr: u64,
rewards_ptr: u64,
dones_ptr: u64,
actions_ptr: u64,
raw_rewards_ptr: u64,
trade_duration_ptr: u64,
outcome_ema_ptr: u64,
position_lots_ptr: u64,
pyramid_count_ptr: u64,
unit_entry_price_ptr: u64,
unit_entry_step_ptr: u64,
unit_lots_ptr: u64,
unit_trail_ptr: u64,
close_unit_index_ptr: u64,
frd_logits_ptr: u64,
) -> Result<()> {
let buf = &self.bufs[self.current];
let s = self.diag_stream.cu_stream();
@@ -164,150 +172,135 @@ impl DiagStaging {
unsafe {
// ISV
let (src, _g) = isv_d.device_ptr(&self.diag_stream);
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr,
src,
isv_ptr,
self.isv_len * 4,
s,
)
.context("diag snapshot isv")?;
// rewards
let (src, _g) = rewards_d.device_ptr(&self.diag_stream);
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.rewards_offset * 4) as u64,
src,
rewards_ptr,
b * 4,
s,
)
.context("diag snapshot rewards")?;
// dones
let (src, _g) = dones_d.device_ptr(&self.diag_stream);
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.dones_offset * 4) as u64,
src,
dones_ptr,
b * 4,
s,
)
.context("diag snapshot dones")?;
// actions (i32 → same byte width as f32, reinterpret on host)
let (src, _g) = actions_d.device_ptr(&self.diag_stream);
// actions (i32 -- same byte width as f32, reinterpret on host)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.actions_offset * 4) as u64,
src,
actions_ptr,
b * 4,
s,
)
.context("diag snapshot actions")?;
// raw_rewards
let (src, _g) = raw_rewards_d.device_ptr(&self.diag_stream);
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.raw_rewards_offset * 4) as u64,
src,
raw_rewards_ptr,
b * 4,
s,
)
.context("diag snapshot raw_rewards")?;
// trade_duration
let (src, _g) = trade_duration_d.device_ptr(&self.diag_stream);
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.trade_duration_offset * 4) as u64,
src,
trade_duration_ptr,
b * 4,
s,
)
.context("diag snapshot trade_duration")?;
// outcome_ema
let (src, _g) = outcome_ema_d.device_ptr(&self.diag_stream);
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.outcome_ema_offset * 4) as u64,
src,
outcome_ema_ptr,
b * 4,
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);
// position_lots (i32 -- raw bytes into f32 slot)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.position_lots_offset * 4) as u64,
src,
position_lots_ptr,
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);
// pyramid_count (i32 -- raw bytes into f32 slot)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.pyramid_count_offset * 4) as u64,
src,
pyramid_count_ptr,
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);
// unit_entry_price (f32, B x 4)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.unit_entry_price_offset * 4) as u64,
src,
unit_entry_price_ptr,
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);
// unit_entry_step (i32, B x 4 -- raw bytes)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.unit_entry_step_offset * 4) as u64,
src,
unit_entry_step_ptr,
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);
// unit_lots (i32, B x 4 -- raw bytes)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.unit_lots_offset * 4) as u64,
src,
unit_lots_ptr,
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);
// unit_trail (f32, B x 4)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.unit_trail_offset * 4) as u64,
src,
unit_trail_ptr,
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);
// close_unit_index (i32, B -- raw bytes)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.close_unit_index_offset * 4) as u64,
src,
close_unit_index_ptr,
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);
// frd_logits (f32, B x FRD_OUT_DIM)
cudarc::driver::result::memcpy_dtod_async(
buf.dev_ptr + (self.frd_logits_offset * 4) as u64,
src,
frd_logits_ptr,
b * FRD_OUT_DIM * 4,
s,
)