diff --git a/crates/ml-alpha/src/trainer/integrated.rs b/crates/ml-alpha/src/trainer/integrated.rs index 819f18a78..4c30b10a2 100644 --- a/crates/ml-alpha/src/trainer/integrated.rs +++ b/crates/ml-alpha/src/trainer/integrated.rs @@ -364,9 +364,9 @@ struct HotPathPtrs { /// Cached raw device pointers for lobsim accessors. Eliminates per-call /// vtable dispatch + CudaSlice borrow overhead on `pos_d()`, `bid_px_d()`, -/// `ask_px_d()` which are called 5-10× per step but always return the same -/// pre-allocated device buffer. Extracted once at step entry, then threaded -/// through all raw_launch sites as plain `u64`. +/// `ask_px_d()` which are called many times per step but always return the +/// same pre-allocated device buffer. Extracted once at step entry, then +/// threaded through all raw_launch sites as plain `u64`. struct LobPtrs { pos: u64, bid_px: u64, @@ -4073,6 +4073,28 @@ impl IntegratedTrainer { // element, and buffers were alloc_zeros'd at construction. The // memsets replayed on every graph launch, burning ~36 × memset // bandwidth per step for zero correctness benefit. + // + // EXCEPTION: the 3 loss scalar accumulators (ss_pi_loss, + // ss_pi_loss_entropy, ss_q_loss) are written via `atomicAdd` by + // the PPO surrogate and DQN distributional Q kernels — see + // ppo_clipped_surrogate.cu:270 and dqn_distributional_q.cu:285. + // These MUST be zeroed each step or the per-step loss values + // accumulate monotonically across all steps (root cause of + // 90f178ae9 perf-regression: l_pi=1384 at step 499 was the sum + // of all per-step l_pi values, not the current step's loss). + // v_loss uses reduce_axis0_mapped (single-writer overwrite) and + // does not need zeroing. + { + let s = self.raw_stream; + unsafe { + raw_memset_d8_zero(self.ss_q_loss_dev_ptr, std::mem::size_of::(), s) + .map_err(|e| anyhow::anyhow!("zero ss_q_loss: {:?}", e))?; + raw_memset_d8_zero(self.ss_pi_loss_dev_ptr, std::mem::size_of::(), s) + .map_err(|e| anyhow::anyhow!("zero ss_pi_loss: {:?}", e))?; + raw_memset_d8_zero(self.ss_pi_loss_entropy_dev_ptr, std::mem::size_of::(), s) + .map_err(|e| anyhow::anyhow!("zero ss_pi_loss_entropy: {:?}", e))?; + } + } // ── Step 4 (Phase R7a): inputs are trainer-owned device buffers, // populated by step_with_lobsim's GPU pipeline before this call. @@ -5415,6 +5437,7 @@ impl IntegratedTrainer { next_snapshots: &[Mbp10RawInput], lobsim: &mut dyn RlLobBackend, ) -> Result { + let lob = LobPtrs::new(lobsim); let b_size = self.cfg.perception.n_batch; if b_size == 0 { anyhow::bail!("step_with_lobsim: empty batch (n_batch = 0)"); @@ -5753,17 +5776,15 @@ impl IntegratedTrainer { // ── Gates + log_pi: OUTSIDE graph capture so they fire every // step (not captured as no-ops during warmup). Device-side kernels. { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.q_logits_d.raw_ptr()); args.push_ptr(self.atom_supports_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -5775,16 +5796,14 @@ impl IntegratedTrainer { } } { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.frd_logits_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -5862,7 +5881,7 @@ impl IntegratedTrainer { // Reads current position_lots from lobsim.pos_d for Flat-from-* // actions; reads actions from self.actions_d (just written by // rl_action_kernel above). - let pos_bytes_i = lobsim.pos_bytes() as i32; + let pos_bytes_i = lob.pos_bytes; let b_size_i = b_size as i32; // Stage ts_ns to device buffer for graph-captured multires kernel. @@ -5929,16 +5948,15 @@ impl IntegratedTrainer { // Hold time = current_step - oldest_active_unit_entry_step // (NOT steps_since_done, which counts flat steps too). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_step_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -5954,16 +5972,14 @@ impl IntegratedTrainer { // Runs BEFORE trail_mutate so the structural decay is applied // first, then agent's a7/a8 fine-tuning on top. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.unit_trail_distance_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_initial_r_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); let mut ptrs = args.build_arg_ptrs(); @@ -6003,12 +6019,10 @@ impl IntegratedTrainer { // FlatFromLong/Short on per-unit breach. Reads shared lobsim // best book for current mid. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); @@ -6033,13 +6047,12 @@ impl IntegratedTrainer { // Reads pos_state for aggregate position; writes diag fired-count // to ISV slot 505. { - let pos_d_ref_heat: &CudaSlice = lobsim.pos_d(); let block = (b_size as u32).min(256); let grid = ((b_size as u32) + block - 1) / block; let smem = 256 * std::mem::size_of::() as u32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(pos_d_ref_heat.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); args.push_i32(pos_bytes_i); @@ -6119,8 +6132,9 @@ impl IntegratedTrainer { snapshots: &[Mbp10RawInput], b_size: usize, ) -> Result { + let lob = LobPtrs::new(lobsim); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; + let pos_bytes_i = lob.pos_bytes; // ── Graph B: post-fill reward/EMA/controller pipeline ───────── // ~20 kernels from extract_realized_pnl_delta through @@ -6146,10 +6160,9 @@ impl IntegratedTrainer { // reward_shaping + raw_rewards snapshot + recent_outcome_update // in a single per-batch kernel (7→1 launch). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.prev_realized_pnl_d.raw_ptr()); args.push_ptr(self.prev_position_lots_d.raw_ptr()); args.push_ptr(self.rewards_d.raw_ptr()); @@ -6244,11 +6257,7 @@ impl IntegratedTrainer { // Trade context features — derived from unit state + current mid. { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); args.push_ptr(self.trade_context_d.raw_ptr()); @@ -6258,12 +6267,12 @@ impl IntegratedTrainer { args.push_ptr(self.unit_initial_r_d.raw_ptr()); args.push_ptr(self.unit_peak_unrealized_r_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.pos); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -6277,8 +6286,6 @@ impl IntegratedTrainer { // Multi-resolution streaming features — time-weighted EMA at 3 horizons. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let b_size_i = b_size as i32; let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); @@ -6286,10 +6293,10 @@ impl IntegratedTrainer { args.push_ptr(self.multires_output_d.raw_ptr()); args.push_ptr(self.multires_prev_mid_d.raw_ptr()); args.push_ptr(self.multires_prev_ts_ns_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); // bid_sz uses bid_px - args.push_ptr(ask_px_d.raw_ptr()); // ask_sz uses ask_px + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); + args.push_ptr(lob.bid_px); // bid_sz uses bid_px + args.push_ptr(lob.ask_px); // ask_sz uses ask_px args.push_ptr(self.isv_dev_ptr); args.push_ptr(self.ts_ns_d.raw_ptr()); args.push_i32(b_size_i); @@ -6617,23 +6624,19 @@ impl IntegratedTrainer { // mid-price into ring buffer and update peak (max for long, // min for short) while a position is open. Flat resets ring. { - let pos_d_ref = lobsim.pos_d(); - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let grid_x = ((b_size as u32) + 31) / 32; let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); - args.push_ptr(pos_d_ref.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.pos); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.hindsight.mid_ring_d.raw_ptr()); args.push_ptr(self.hindsight.ring_write_idx_d.raw_ptr()); args.push_ptr(self.hindsight.peak_mid_d.raw_ptr()); args.push_ptr(self.hindsight.entry_mid_d.raw_ptr()); args.push_ptr(self.hindsight.position_dir_d.raw_ptr()); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -6804,13 +6807,11 @@ impl IntegratedTrainer { // (ISV[552] steps). If holding would have been better, inject // a synthetic "should-have-held-longer" transition. { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let smem = (2 * 256 * std::mem::size_of::()) as u32; let cap_i = self.gpu_replay.capacity as i32; let mut args = RawArgs::new(); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_ptr(self.hindsight.closed_ring_d.raw_ptr()); args.push_ptr(self.hindsight.closed_h_t_d.raw_ptr()); @@ -7386,6 +7387,7 @@ impl IntegratedTrainer { b_size: usize, seq_len: usize, ) -> Result { + let lob = LobPtrs::new(lobsim); // ── Graph A: pre-snapshot kernel pipeline ────────────────────── if self.prefill_graph.is_some() && !self.mega_graph_enabled { unsafe { @@ -7584,17 +7586,15 @@ impl IntegratedTrainer { // ── Gates + log_pi: OUTSIDE graph capture (GPU data path). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.q_logits_d.raw_ptr()); args.push_ptr(self.atom_supports_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -7606,16 +7606,14 @@ impl IntegratedTrainer { } } { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let b_size_i = b_size as i32; - let pos_bytes_i = lobsim.pos_bytes() as i32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.frd_logits_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); - args.push_i32(pos_bytes_i); + args.push_i32(lob.pos_bytes); let mut ptrs = args.build_arg_ptrs(); unsafe { raw_launch( @@ -7737,7 +7735,7 @@ impl IntegratedTrainer { } } - let pos_bytes_i = lobsim.pos_bytes() as i32; + let pos_bytes_i = lob.pos_bytes; let b_size_i = b_size as i32; // ── Graph A2: post-snapshot / pre-fill kernel pipeline ───────── @@ -7779,13 +7777,12 @@ impl IntegratedTrainer { // Min hold check — true hold time from oldest active unit // (see eval-side launch comment above for rationale). { - let pos_d_ref: &CudaSlice = lobsim.pos_d(); let grid_x = ((b_size as u32) + 31) / 32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_step_d.raw_ptr()); - args.push_ptr(pos_d_ref.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); args.push_i32(pos_bytes_i); @@ -7802,16 +7799,14 @@ impl IntegratedTrainer { // Asymmetric trail decay { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.unit_trail_distance_d.raw_ptr()); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_initial_r_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); let mut ptrs = args.build_arg_ptrs(); @@ -7846,12 +7841,10 @@ impl IntegratedTrainer { // Trail stop check { - let bid_px_d = lobsim.bid_px_d(); - let ask_px_d = lobsim.ask_px_d(); let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(bid_px_d.raw_ptr()); - args.push_ptr(ask_px_d.raw_ptr()); + args.push_ptr(lob.bid_px); + args.push_ptr(lob.ask_px); args.push_ptr(self.unit_active_d.raw_ptr()); args.push_ptr(self.unit_entry_price_d.raw_ptr()); args.push_ptr(self.unit_lots_d.raw_ptr()); @@ -7872,13 +7865,12 @@ impl IntegratedTrainer { // Position heat cap { - let pos_d_ref_heat: &CudaSlice = lobsim.pos_d(); let block = (b_size as u32).min(256); let grid = ((b_size as u32) + block - 1) / block; let smem = 256 * std::mem::size_of::() as u32; let mut args = RawArgs::new(); args.push_ptr(self.actions_d.raw_ptr()); - args.push_ptr(pos_d_ref_heat.raw_ptr()); + args.push_ptr(lob.pos); args.push_ptr(self.isv_dev_ptr); args.push_i32(b_size_i); args.push_i32(pos_bytes_i);