diff --git a/crates/ml-alpha/src/rl/reward.rs b/crates/ml-alpha/src/rl/reward.rs index 5cb9c20eb..5c7b635ab 100644 --- a/crates/ml-alpha/src/rl/reward.rs +++ b/crates/ml-alpha/src/rl/reward.rs @@ -108,4 +108,22 @@ pub trait RlLobBackend { fn pos_book_and_market_targets_mut( &mut self, ) -> (&CudaSlice, &CudaSlice, &CudaSlice, &mut CudaSlice); + + /// 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<()>; } diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 8d76303a3..836f50845 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -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::()) 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 diff --git a/crates/ml-backtesting/src/sim/mod.rs b/crates/ml-backtesting/src/sim/mod.rs index 96c94df33..5d8a2e34a 100644 --- a/crates/ml-backtesting/src/sim/mod.rs +++ b/crates/ml-backtesting/src/sim/mod.rs @@ -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::(); + 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, &CudaSlice, &CudaSlice, &mut CudaSlice) { (&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