diff --git a/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs b/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs index 16b5d5c26..7395aa986 100644 --- a/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_ppo_collector.rs @@ -172,6 +172,12 @@ pub struct GpuPpoExperienceCollector { readback_f32_host: Vec, readback_i32_staging: CudaSlice, readback_i32_host: Vec, + + // Pre-allocated host-side staging for reset_episodes() to avoid per-call Vec allocations. + // portfolio_init_staging is re-filled with structured data each reset; rng_seed_staging + // is re-filled with random seeds. Both avoid hitting the allocator on every epoch boundary. + portfolio_init_staging: Vec, + rng_seed_staging: Vec, } impl GpuPpoExperienceCollector { @@ -351,6 +357,11 @@ impl GpuPpoExperienceCollector { })?; let readback_i32_host = vec![0_i32; i32_staging_len]; + // ---- Step 6c: Pre-allocate host staging for reset_episodes() ---- + // Avoids per-epoch Vec allocations for portfolio init and RNG seeds. + let portfolio_init_staging = vec![0.0_f32; MAX_EPISODES * PORTFOLIO_STATE_SIZE]; + let rng_seed_staging = vec![0_u32; MAX_EPISODES]; + // ---- Step 7: Log buffer sizes ---- let per_episode_bytes = (PORTFOLIO_STATE_SIZE + BARRIER_STATE_SIZE) * 4 + DIVERSITY_WINDOW * 4 @@ -388,6 +399,8 @@ impl GpuPpoExperienceCollector { readback_f32_host, readback_i32_staging, readback_i32_host, + portfolio_init_staging, + rng_seed_staging, }) } @@ -672,56 +685,74 @@ impl GpuPpoExperienceCollector { /// /// Re-initializes portfolio states to starting capital, zeros barrier /// states and diversity windows, and regenerates RNG seeds. + /// + /// Zero-fill buffers use `memset_zeros` (async GPU memset) instead of + /// allocating CPU-side `Vec`s and uploading via HtoD. Structured data + /// (portfolio init, RNG seeds) reuses pre-allocated host staging vectors + /// to avoid per-epoch allocations. pub fn reset_episodes( &mut self, initial_capital: f32, avg_spread: f32, cash_reserve_pct: f32, ) -> Result<(), MLError> { - // Reset portfolio states - let mut portfolio_init = vec![0.0_f32; MAX_EPISODES * PORTFOLIO_STATE_SIZE]; + // Reset portfolio states using pre-allocated staging buffer + // (structured init -- cannot use memset_zeros) for i in 0..MAX_EPISODES { let off = i * PORTFOLIO_STATE_SIZE; - portfolio_init[off] = initial_capital; - portfolio_init[off + 3] = initial_capital; - portfolio_init[off + 4] = avg_spread; - portfolio_init[off + 6] = cash_reserve_pct; + // Zero all fields first + for j in 0..PORTFOLIO_STATE_SIZE { + if let Some(slot) = self.portfolio_init_staging.get_mut(off + j) { + *slot = 0.0; + } + } + // Set non-zero fields + if let Some(slot) = self.portfolio_init_staging.get_mut(off) { + *slot = initial_capital; // cash + } + if let Some(slot) = self.portfolio_init_staging.get_mut(off + 3) { + *slot = initial_capital; // initial_cap + } + if let Some(slot) = self.portfolio_init_staging.get_mut(off + 4) { + *slot = avg_spread; // spread + } + if let Some(slot) = self.portfolio_init_staging.get_mut(off + 6) { + *slot = cash_reserve_pct; // reserve_pct + } } self.stream - .memcpy_htod(&portfolio_init, &mut self.portfolio_states) + .memcpy_htod(&self.portfolio_init_staging, &mut self.portfolio_states) .map_err(|e| { MLError::ModelError(format!("Failed to reset portfolio_states: {e}")) })?; - // Zero barrier states - let barrier_zeros = vec![0.0_f32; MAX_EPISODES * BARRIER_STATE_SIZE]; + // Zero barrier states via async GPU memset (no CPU allocation) self.stream - .memcpy_htod(&barrier_zeros, &mut self.barrier_states) + .memset_zeros(&mut self.barrier_states) .map_err(|e| { MLError::ModelError(format!("Failed to reset barrier_states: {e}")) })?; - // Zero diversity windows and metas - let div_window_zeros = vec![0_i32; MAX_EPISODES * DIVERSITY_WINDOW]; + // Zero diversity windows via async GPU memset (no CPU allocation) self.stream - .memcpy_htod(&div_window_zeros, &mut self.diversity_windows) + .memset_zeros(&mut self.diversity_windows) .map_err(|e| { MLError::ModelError(format!("Failed to reset diversity_windows: {e}")) })?; - let div_meta_zeros = vec![0_i32; MAX_EPISODES * 2]; + // Zero diversity metas via async GPU memset (no CPU allocation) self.stream - .memcpy_htod(&div_meta_zeros, &mut self.diversity_metas) + .memset_zeros(&mut self.diversity_metas) .map_err(|e| { MLError::ModelError(format!("Failed to reset diversity_metas: {e}")) })?; - // Fresh RNG seeds - let rng_seeds: Vec = (0..MAX_EPISODES) - .map(|_| fastrand::u32(..)) - .collect(); + // Fresh RNG seeds using pre-allocated staging buffer + for slot in &mut self.rng_seed_staging { + *slot = fastrand::u32(..); + } self.stream - .memcpy_htod(&rng_seeds, &mut self.rng_states) + .memcpy_htod(&self.rng_seed_staging, &mut self.rng_states) .map_err(|e| { MLError::ModelError(format!("Failed to reset rng_states: {e}")) })?; @@ -730,7 +761,7 @@ impl GpuPpoExperienceCollector { initial_capital, avg_spread, cash_reserve_pct, - "PPO episode state reset complete" + "PPO episode state reset complete (memset_zeros path)" ); Ok(())