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

feat(rl): wire lobsim book data from GPU SoA — proper rewards in GPU loader path

Browse Source 
Add apply_snapshot_from_device to RlLobBackend: DtoD copy from
perception SoA buffers into lobsim book arrays + book_update kernel.
No host roundtrip, no sync (stream-ordered).

step_with_lobsim_gpu now calls apply_snapshot_from_device with batch 0's
last-snapshot (K-1) book data. Produces non-zero dones, pnl, wr.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-05-26 21:29:05 +02:00
parent 3552d08501
commit 6f645df11f
3 changed files with 119 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
crates/ml-alpha/src/rl/reward.rs
View File

@@ -108,4 +108,22 @@ pub trait RlLobBackend {
fn pos_book_and_market_targets_mut(
&mut self,
) -> (&CudaSlice<u8>, &CudaSlice<f32>, &CudaSlice<f32>, &mut CudaSlice<i32>);
/// Device-to-device variant of `apply_snapshot`. Copies 4 ×
/// `BOOK_LEVELS` floats from the given source device pointers into
/// the lobsim's internal bid/ask staging buffers, then launches the
/// `book_update_apply_snapshot` kernel to broadcast into per-backtest
/// book state — identical to `apply_snapshot` except the source is a
/// device buffer (e.g. the perception trainer's SoA) rather than
/// host memory.
///
/// Each source pointer addresses `BOOK_LEVELS` contiguous `f32`
/// values on the device. The copy is stream-ordered (no sync).
fn apply_snapshot_from_device(
&mut self,
bid_px_src: u64,
bid_sz_src: u64,
ask_px_src: u64,
ask_sz_src: u64,
) -> Result<()>;
}

39
crates/ml-alpha/src/trainer/integrated.rs
View File

@@ -6650,10 +6650,10 @@ impl IntegratedTrainer {
/// * Encoder forwards read from SoA buffers pre-filled by the GPU
/// gather kernels (`forward_encoder_from_device` replaces
/// `forward_encoder(snapshots)`).
/// * `lobsim.apply_snapshot` is skipped — the lobsim's book state
/// carries forward from the previous step (one-tick stale). This
/// is acceptable for RL training: fill pricing is approximate by
/// design (the agent learns on noisy fills).
/// * `lobsim.apply_snapshot_from_device` replaces the HtoD
/// `apply_snapshot`: batch 0's last-snapshot book data is DtoD'd
/// from the SoA buffers into the lobsim's staging arrays, then the
/// `book_update_apply_snapshot` kernel broadcasts to all backtests.
/// * FRD labels are gathered by `gpu_gather_frd_labels` directly
/// into `self.frd_labels_d` (no mapped-pinned staging).
/// * `ts_ns` for `step_fill_from_market_targets` is DtoD'd from
@@ -6757,11 +6757,10 @@ impl IntegratedTrainer {
self.step_counter = self.step_counter.wrapping_add(1);
// ── From here, the kernel pipeline is identical to
// step_with_lobsim. The only remaining difference is that
// apply_snapshot + ts_ns staging are handled differently:
// step_with_lobsim. The remaining differences are:
//
// * apply_snapshot is SKIPPED (lobsim book carries forward from
// previous step — one-tick-stale pricing, acceptable for RL).
// * apply_snapshot_from_device DtoD's batch 0's last-snapshot
// book data from the SoA, replacing the HtoD apply_snapshot.
// * ts_ns for step_fill_from_market_targets is DtoD'd from the
// SoA's last-snapshot position into self.ts_ns_d.
//
@@ -7045,9 +7044,27 @@ impl IntegratedTrainer {
} // end else (warmup / capture dispatch)
// ── Step 5: GPU-pure env step (GPU data path). ─────────────────
// apply_snapshot SKIPPED: the lobsim's book state carries forward
// from the previous step. For RL training the one-tick-stale book
// pricing is acceptable — the agent learns on approximate fills.
//
// apply_snapshot from device: copy batch 0's last-snapshot book
// data from the SoA buffers into lobsim's staging arrays, then
// launch the broadcast kernel. This replaces the HtoD path used
// by step_with_lobsim and the previous "SKIPPED" placeholder that
// left the lobsim book stale (producing dones=0 / pnl=$0).
//
// Layout: SoA arrays are [B * K * BOOK_LEVELS] f32. Batch 0's
// last snapshot is at position (K-1), byte offset =
// (seq_len - 1) * BOOK_LEVELS * sizeof(f32).
{
let book_levels = crate::cfc::snap_features::BOOK_LEVELS;
let snap_byte_offset =
((seq_len - 1) * book_levels * std::mem::size_of::<f32>()) as u64;
lobsim.apply_snapshot_from_device(
soa.bid_px + snap_byte_offset,
soa.bid_sz + snap_byte_offset,
soa.ask_px + snap_byte_offset,
soa.ask_sz + snap_byte_offset,
).context("step_with_lobsim_gpu: apply_snapshot_from_device")?;
}
//
// ts_ns: DtoD from the SoA's last-snapshot ts_ns into self.ts_ns_d
// for the multires kernel. The lobsim fill kernel receives 0 as

73
crates/ml-backtesting/src/sim/mod.rs
View File

@@ -1147,6 +1147,63 @@ impl LobSimCuda {
Ok(())
}
/// Device-to-device variant of [`apply_snapshot`]. Copies 4 ×
/// `BOOK_LEVELS` floats from source device pointers into the lobsim's
/// staging buffers, then launches `book_update_apply_snapshot` to
/// broadcast the 10-level book to every backtest.
///
/// Each `*_src` argument is a raw device pointer (`CUdeviceptr`)
/// addressing `BOOK_LEVELS` contiguous `f32` values. The copies and
/// the kernel launch are stream-ordered on `self.stream` — no sync.
pub fn apply_snapshot_from_device(
&mut self,
bid_px_src: u64,
bid_sz_src: u64,
ask_px_src: u64,
ask_sz_src: u64,
) -> Result<()> {
use ml_alpha::trainer::raw_launch::{RawArgs, raw_launch, raw_memcpy_dtod_async};
let bytes = BOOK_LEVELS * std::mem::size_of::<f32>();
let raw_stream = self.stream.cu_stream();
unsafe {
raw_memcpy_dtod_async(self.bid_px_d.raw_ptr(), bid_px_src, bytes, raw_stream)
.map_err(|e| anyhow::anyhow!("apply_snapshot_from_device: bid_px DtoD: {:?}", e))?;
raw_memcpy_dtod_async(self.bid_sz_d.raw_ptr(), bid_sz_src, bytes, raw_stream)
.map_err(|e| anyhow::anyhow!("apply_snapshot_from_device: bid_sz DtoD: {:?}", e))?;
raw_memcpy_dtod_async(self.ask_px_d.raw_ptr(), ask_px_src, bytes, raw_stream)
.map_err(|e| anyhow::anyhow!("apply_snapshot_from_device: ask_px DtoD: {:?}", e))?;
raw_memcpy_dtod_async(self.ask_sz_d.raw_ptr(), ask_sz_src, bytes, raw_stream)
.map_err(|e| anyhow::anyhow!("apply_snapshot_from_device: ask_sz DtoD: {:?}", e))?;
}
let n = self.n_backtests as i32;
let mut args = RawArgs::new();
args.push_ptr(self.bid_px_d.raw_ptr());
args.push_ptr(self.bid_sz_d.raw_ptr());
args.push_ptr(self.ask_px_d.raw_ptr());
args.push_ptr(self.ask_sz_d.raw_ptr());
args.push_ptr(self.books_d.raw_ptr());
args.push_ptr(self.prev_mid_d.raw_ptr());
args.push_ptr(self.atr_mid_ema_d.raw_ptr());
args.push_ptr(self.snapshots_skipped_d.raw_ptr());
args.push_ptr(self.min_reasonable_px_d.raw_ptr());
args.push_ptr(self.max_reasonable_px_d.raw_ptr());
args.push_i32(n);
let mut ptrs = args.build_arg_ptrs();
unsafe {
raw_launch(
self.book_update_fn.cu_function(),
(self.n_backtests as u32, 1, 1),
(32, 1, 1),
0,
raw_stream,
&mut ptrs[..args.len()],
).map_err(|e| anyhow::anyhow!("apply_snapshot_from_device: book_update kernel: {:?}", e))?;
}
Ok(())
}
/// Submit a market order to one specific backtest. side: 0=buy, 1=sell.
/// Walks the book + updates Pos, then runs pnl_track to detect
/// segment_complete and emit TradeRecord if the position closed.
@@ -1357,6 +1414,22 @@ impl ml_alpha::rl::reward::RlLobBackend for LobSimCuda {
) -> (&CudaSlice<u8>, &CudaSlice<f32>, &CudaSlice<f32>, &mut CudaSlice<i32>) {
(&self.pos_d, &self.bid_px_d, &self.ask_px_d, &mut self.market_targets_d)
}
fn apply_snapshot_from_device(
&mut self,
bid_px_src: u64,
bid_sz_src: u64,
ask_px_src: u64,
ask_sz_src: u64,
) -> Result<()> {
LobSimCuda::apply_snapshot_from_device(
self,
bid_px_src,
bid_sz_src,
ask_px_src,
ask_sz_src,
)
}
}
// Re-open the impl block for `LobSimCuda` so subsequent methods (if any