From 3cfa99de3f2acb190d28d1c9637ee63aab0aef23 Mon Sep 17 00:00:00 2001
From: jgrusewski <jgrusewski@localhost>
Date: Sat, 28 Feb 2026 02:48:33 +0100
Subject: [PATCH] feat(ppo): pre-upload rollout states to GPU, eliminate
 per-step tensor allocs

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
---
 crates/ml/src/trainers/ppo.rs | 57 +++++++++++++++++++++--------------
 1 file changed, 35 insertions(+), 22 deletions(-)

diff --git a/crates/ml/src/trainers/ppo.rs b/crates/ml/src/trainers/ppo.rs
index 44bb7767a..5839da6d5 100644
--- a/crates/ml/src/trainers/ppo.rs
+++ b/crates/ml/src/trainers/ppo.rs
@@ -23,6 +23,7 @@ use crate::gpu::memory_profile::{self, resolve_batch_size};
 use crate::ppo::gae::GAEConfig;
 use crate::ppo::ppo::{PPOConfig, PPO};
 use crate::ppo::trajectories::{Trajectory, TrajectoryBatch, TrajectoryStep};
+use crate::cuda_pipeline::PpoGpuData;
 use crate::MLError;
 
 /// PPO training hyperparameters (matches gRPC PpoParams)
@@ -501,15 +502,25 @@ impl PpoTrainer {
         // Track position for PnL calculation
         let mut position: i8 = 0; // -1: short, 0: neutral, 1: long
 
+        // Phase 1: Pre-upload all states to GPU (single allocation, reused per-step)
+        let gpu_data = PpoGpuData::upload(market_data, &self.device).ok();
+        if gpu_data.is_some() {
+            tracing::debug!("PPO: Pre-uploaded {} states to GPU", num_steps);
+        }
+
         for step_idx in 0..num_steps {
             let state = &market_data[step_idx];
 
-            // Select action using policy
-            let action_probs = model.actor.action_probabilities(&Tensor::from_vec(
-                state.clone(),
-                (1, state.len()),
-                &self.device,
-            )?)?;
+            // Use pre-uploaded GPU tensor or fall back to per-step allocation
+            let state_tensor = if let Some(ref gpu) = gpu_data {
+                gpu.step_state(step_idx)
+                    .map_err(|e| MLError::ModelError(format!("PPO GPU state access failed: {e}")))?
+            } else {
+                Tensor::from_vec(state.clone(), (1, state.len()), &self.device)?
+            };
+
+            // Actor forward pass — use pre-uploaded tensor
+            let action_probs = model.actor.action_probabilities(&state_tensor)?;
 
             let action_idx = self.sample_action(&action_probs)?;
             let action = TradingAction::from_int(action_idx as u8).unwrap_or(TradingAction::Hold);
@@ -522,11 +533,7 @@ impl PpoTrainer {
 
             let value = model
                 .critic
-                .forward(&Tensor::from_vec(
-                    state.clone(),
-                    (1, state.len()),
-                    &self.device,
-                )?)?
+                .forward(&state_tensor)? // Reuse same GPU tensor — eliminates second clone+alloc
                 .flatten_all()?
                 .to_vec1::<f32>()?[0]; // Flatten [1, 1] to vec, take first element
 
@@ -576,6 +583,12 @@ impl PpoTrainer {
         let model = self.model.lock().await;
         let mut all_trajectories = Vec::new();
 
+        // Phase 1: Pre-upload all states to GPU (single allocation, reused per-step)
+        let gpu_data = PpoGpuData::upload(market_data, &self.device).ok();
+        if gpu_data.is_some() {
+            tracing::debug!("PPO vectorized: Pre-uploaded {} states to GPU", num_steps);
+        }
+
         for env_id in 0..num_envs {
             let start_idx = env_id * steps_per_env;
             let end_idx = ((env_id + 1) * steps_per_env).min(num_steps);
@@ -591,12 +604,16 @@ impl PpoTrainer {
             for step_idx in start_idx..end_idx {
                 let state = &market_data[step_idx];
 
-                // Get action from policy
-                let action_probs = model.actor.action_probabilities(&Tensor::from_vec(
-                    state.clone(),
-                    (1, state.len()),
-                    &self.device,
-                )?)?;
+                // Use pre-uploaded GPU tensor or fall back to per-step allocation
+                let state_tensor = if let Some(ref gpu) = gpu_data {
+                    gpu.step_state(step_idx)
+                        .map_err(|e| MLError::ModelError(format!("PPO GPU state access failed: {e}")))?
+                } else {
+                    Tensor::from_vec(state.clone(), (1, state.len()), &self.device)?
+                };
+
+                // Get action from policy — use pre-uploaded tensor
+                let action_probs = model.actor.action_probabilities(&state_tensor)?;
 
                 let action_idx = self.sample_action(&action_probs)?;
                 let action =
@@ -608,11 +625,7 @@ impl PpoTrainer {
 
                 let value = model
                     .critic
-                    .forward(&Tensor::from_vec(
-                        state.clone(),
-                        (1, state.len()),
-                        &self.device,
-                    )?)?
+                    .forward(&state_tensor)? // Reuse same GPU tensor — eliminates second clone+alloc
                     .flatten_all()?
                     .to_vec1::<f32>()?[0];