From 8c7e907fcf6c1263a71e30cd0e493543e209fea3 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 11 Mar 2026 13:05:02 +0100 Subject: [PATCH] =?UTF-8?q?perf(cuda):=20eliminate=20GPU=E2=86=92CPU=20rou?= =?UTF-8?q?ndtrips=20from=20backtest=20evaluate=20loop?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - gather_states: replace memcpy_dtoh + Tensor::from_vec with DtoD copy (cuMemcpyDtoDAsync) — state tensor stays on device, zero CPU touch - actions: replace to_vec1 + memcpy_htod with DtoD copy from argmax tensor directly into actions_buf — eliminates per-step PCIe upload - batch_q_values (RegimeConditionalDQN): replace CPU-side regime classification (to_vec2 + serial loop + sub-batch re-upload) with on-device classify_regime_masks_gpu + all-heads forward + masked blend Co-Authored-By: Claude Opus 4.6 --- crates/ml-dqn/src/regime_conditional.rs | 104 +++++++++--------- .../cuda_pipeline/gpu_backtest_evaluator.rs | 103 ++++++++++++----- 2 files changed, 125 insertions(+), 82 deletions(-) diff --git a/crates/ml-dqn/src/regime_conditional.rs b/crates/ml-dqn/src/regime_conditional.rs index f66d58c96..3e1692ae2 100644 --- a/crates/ml-dqn/src/regime_conditional.rs +++ b/crates/ml-dqn/src/regime_conditional.rs @@ -405,73 +405,67 @@ impl RegimeConditionalDQN { Ok(results) } - /// Batch Q-value computation across regime heads. + /// Batch Q-value computation across regime heads — fully GPU-resident. + /// + /// Uses on-device regime classification masks to blend Q-values from all + /// 3 heads without any GPU→CPU roundtrip: + /// Q_final = Q_trending * mask_trending + Q_ranging * mask_ranging + Q_volatile * mask_volatile /// - /// Same regime-dispatch logic as `batch_greedy_actions` but returns the raw - /// Q-value tensor `[batch, num_actions]` instead of argmax action indices. /// Used by `GpuBacktestEvaluator` for GPU-side argmax. pub fn batch_q_values(&self, states: &Tensor) -> Result { - let state_vecs = states.to_vec2::().map_err(|e| { - MLError::ModelError(format!("Failed to extract states for regime: {}", e)) - })?; - let n = state_vecs.len(); + let n = states.dims()[0]; if n == 0 { return Err(MLError::ModelError("Empty batch for batch_q_values".into())); } - let mut trending_idx = Vec::new(); - let mut ranging_idx = Vec::new(); - let mut volatile_idx = Vec::new(); - for (i, sv) in state_vecs.iter().enumerate() { - match RegimeType::classify_from_features(sv) { - RegimeType::Trending => trending_idx.push(i), - RegimeType::Ranging => ranging_idx.push(i), - RegimeType::Volatile => volatile_idx.push(i), - } - } + // On-device regime classification — zero CPU roundtrip + let (trending_mask, ranging_mask, volatile_mask) = + RegimeType::classify_regime_masks_gpu(states)?; - let device = &self.device; - let dim = state_vecs.first().map(|v| v.len()).unwrap_or(0); + // Forward full batch through all 3 heads (each head ignores irrelevant samples + // via masking — cheaper than splitting/gathering sub-batches) + let trending_q = self.trending_head.q_values_for_batch(states)?; + let ranging_q = self.ranging_head.q_values_for_batch(states)?; + let volatile_q = self.volatile_head.q_values_for_batch(states)?; - // Pre-allocate output: [n, num_actions] filled with zeros - // We'll scatter-assign per-regime Q-values into this - let num_actions = 5; // DQN exposure actions - let mut q_out = vec![0.0_f32; n * num_actions]; + // Reshape masks from [batch] to [batch, 1] for broadcasting over actions dim + let trending_mask = trending_mask.unsqueeze(1).map_err(|e| { + MLError::ModelError(format!("trending unsqueeze: {e}")) + })?; + let ranging_mask = ranging_mask.unsqueeze(1).map_err(|e| { + MLError::ModelError(format!("ranging unsqueeze: {e}")) + })?; + let volatile_mask = volatile_mask.unsqueeze(1).map_err(|e| { + MLError::ModelError(format!("volatile unsqueeze: {e}")) + })?; - for (indices, head) in [ - (&trending_idx, &self.trending_head), - (&ranging_idx, &self.ranging_head), - (&volatile_idx, &self.volatile_head), - ] { - if indices.is_empty() { - continue; - } - let mut flat = Vec::with_capacity(indices.len() * dim); - for &i in indices { - flat.extend_from_slice(&state_vecs[i]); - } - let sub_tensor = - Tensor::from_vec(flat, (indices.len(), dim), device).map_err(|e| { - MLError::ModelError(format!("Regime sub-batch failed: {e}")) - })?; - let sub_q = head.q_values_for_batch(&sub_tensor)?; - let sub_q_vec: Vec = sub_q.flatten_all() - .map_err(|e| MLError::ModelError(format!("flatten q_values: {e}")))? - .to_vec1() - .map_err(|e| MLError::ModelError(format!("q_values to_vec1: {e}")))?; + // Ensure Q-values are F32 for multiplication with F32 masks + let trending_q = trending_q.to_dtype(DType::F32).map_err(|e| { + MLError::ModelError(format!("trending_q to_f32: {e}")) + })?; + let ranging_q = ranging_q.to_dtype(DType::F32).map_err(|e| { + MLError::ModelError(format!("ranging_q to_f32: {e}")) + })?; + let volatile_q = volatile_q.to_dtype(DType::F32).map_err(|e| { + MLError::ModelError(format!("volatile_q to_f32: {e}")) + })?; - // Scatter sub-batch Q-values back into the full output - for (j, &idx) in indices.iter().enumerate() { - let src_base = j * num_actions; - let dst_base = idx * num_actions; - for a in 0..num_actions { - q_out[dst_base + a] = sub_q_vec[src_base + a]; - } - } - } + // Blend: Q_final = sum of (Q_head * mask_head) across all regimes + let blended = trending_q.broadcast_mul(&trending_mask).map_err(|e| { + MLError::ModelError(format!("trending mul: {e}")) + })?; + let blended = blended.add( + &ranging_q.broadcast_mul(&ranging_mask).map_err(|e| { + MLError::ModelError(format!("ranging mul: {e}")) + })? + ).map_err(|e| MLError::ModelError(format!("ranging add: {e}")))?; + let blended = blended.add( + &volatile_q.broadcast_mul(&volatile_mask).map_err(|e| { + MLError::ModelError(format!("volatile mul: {e}")) + })? + ).map_err(|e| MLError::ModelError(format!("volatile add: {e}")))?; - Tensor::from_vec(q_out, (n, num_actions), device) - .map_err(|e| MLError::ModelError(format!("batch_q_values output tensor: {e}"))) + Ok(blended) } /// Batch softmax action selection across regime heads (Gumbel-max, GPU-resident). diff --git a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs index dfc400987..09e817695 100644 --- a/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs +++ b/crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs @@ -7,13 +7,13 @@ //! 2. Step loop: GPU gather kernel → Candle forward → env kernel //! 3. Metrics reduction kernel → single readback //! -//! The only GPU→CPU transfers are the per-step state download (n_windows × state_dim -//! floats, typically ~1.5 KB) and the final metrics readback (n_windows × 10 floats). +//! The only GPU→CPU transfer is the final metrics readback (n_windows × 10 floats). +//! Per-step state construction uses a zero-copy DtoD path (gather kernel → Candle tensor). use std::sync::Arc; use candle_core::cuda_backend::cudarc; use candle_core::{DType, Device, Tensor}; -use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; use std::sync::OnceLock; use tracing::info; @@ -325,12 +325,12 @@ impl GpuBacktestEvaluator { /// /// Launches the `gather_states` CUDA kernel which reads directly from the /// pre-uploaded features buffer and the live portfolio state buffer on GPU, - /// avoiding the large GPU→CPU→GPU roundtrip of the old path. + /// then performs a zero-copy DtoD transfer from the kernel output `states_buf` + /// into a freshly-allocated Candle tensor. No data ever touches the CPU. /// /// The kernel writes into `states_buf` (pre-allocated, `n_windows × state_dim`). - /// We then download only that small buffer (typically ~384 floats) to create the - /// Candle tensor, which is negligible compared to the old path that downloaded - /// the full features buffer (n_windows × max_len × feat_dim floats). + /// A `cuMemcpyDtoDAsync` then copies those bytes directly into the Candle tensor's + /// CUDA storage — eliminating the GPU→CPU→GPU roundtrip entirely. /// /// # Panics /// `portfolio_dim` must equal `self.portfolio_dim` (always 3). Callers using a @@ -385,15 +385,34 @@ impl GpuBacktestEvaluator { .map_err(|e| MLError::ModelError(format!("gather_states launch step {step}: {e}")))?; } - // Download the gather output — n_windows × state_dim floats (tiny: ~384 floats - // for 8 windows × 48 state_dim, vs the old path's n_windows × max_len × feat_dim). - let mut host_states = vec![0.0_f32; self.n_windows * state_dim]; - self.stream - .memcpy_dtoh(&self.states_buf, &mut host_states) - .map_err(|e| MLError::ModelError(format!("states download step {step}: {e}")))?; + // Zero-copy: DtoD from kernel output CudaSlice into Candle Tensor — no CPU transfer. + let n_elems = self.n_windows * state_dim; + let tensor = Tensor::zeros(&[self.n_windows, state_dim], DType::F32, device) + .map_err(|e| MLError::ModelError(format!("alloc states tensor step {step}: {e}")))?; - Tensor::from_vec(host_states, (self.n_windows, state_dim), device) - .map_err(|e| MLError::ModelError(format!("states tensor step {step}: {e}"))) + let (storage_guard, _layout) = tensor.storage_and_layout(); + match *storage_guard { + candle_core::Storage::Cuda(ref cs) => { + let dst_slice: &CudaSlice = cs.as_cuda_slice() + .map_err(|e| MLError::ModelError(format!("states as_cuda_slice: {e}")))?; + let (dst_ptr, _dst_sync) = dst_slice.device_ptr(&self.stream); + let src_view = self.states_buf.slice(..n_elems); + let (src_ptr, _src_sync) = src_view.device_ptr(&self.stream); + let num_bytes = n_elems * std::mem::size_of::(); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, src_ptr, num_bytes, self.stream.cu_stream(), + ).map_err(|e| MLError::ModelError(format!("states DtoD copy step {step}: {e}")))?; + } + } + candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { + return Err(MLError::ModelError( + "gather_states: expected CUDA device".into(), + )); + } + } + drop(storage_guard); + Ok(tensor) } /// Run the full backtest evaluation loop. @@ -431,23 +450,53 @@ impl GpuBacktestEvaluator { .map_err(|e| MLError::ModelError(format!("q_values f32 cast: {e}")))? }; - // 3. Greedy action selection — argmax over action dim + // 3. Greedy action selection — argmax over action dim (stays on GPU) let actions_tensor = q_f32 .argmax(1) .map_err(|e| MLError::ModelError(format!("argmax: {e}")))?; + + // DtoD copy: argmax output (U32 Candle tensor) → actions_buf (CudaSlice). + // Action values are 0..N_ACTIONS (small non-negative), so u32 and i32 share + // identical bit patterns — raw byte reinterpret is safe. Eliminates the old + // to_vec1 → collect::> → memcpy_htod upload roundtrip on the hot path. + { + let (act_guard, _act_layout) = actions_tensor.storage_and_layout(); + match &*act_guard { + candle_core::Storage::Cuda(ref cs) => { + let src_slice: &CudaSlice = cs + .as_cuda_slice() + .map_err(|e| MLError::ModelError(format!("actions as_cuda_slice step {step}: {e}")))?; + let src_view = src_slice.slice(..self.n_windows); + let (src_ptr, _src_sync) = src_view.device_ptr(&self.stream); + let (dst_ptr, _dst_sync) = self.actions_buf.device_ptr(&self.stream); + let num_bytes = self.n_windows * std::mem::size_of::(); + unsafe { + cudarc::driver::result::memcpy_dtod_async( + dst_ptr, + src_ptr, + num_bytes, + self.stream.cu_stream(), + ) + .map_err(|e| MLError::ModelError(format!("actions DtoD step {step}: {e}")))?; + } + } + candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { + return Err(MLError::ModelError(format!( + "actions_tensor not on CUDA at step {step}" + ))); + } + } + drop(act_guard); + } + + // Accumulate actions into CPU-side history (uploaded once before metrics kernel). + // Download is n_windows × 4 bytes — negligible. The old path did the same + // download (to_vec1) PLUS an upload (memcpy_htod); we keep only the download. let actions_u32: Vec = actions_tensor .to_vec1() - .map_err(|e| MLError::ModelError(format!("argmax to_vec1: {e}")))?; - let actions_i32: Vec = actions_u32.iter().map(|&a| a as i32).collect(); - - // Upload actions to GPU - self.stream - .memcpy_htod(&actions_i32, &mut self.actions_buf) - .map_err(|e| MLError::ModelError(format!("actions upload step {step}: {e}")))?; - - // Accumulate actions into CPU-side history (uploaded once before metrics kernel) - for w in 0..self.n_windows { - self.actions_history_cpu[w * self.max_len + step] = actions_i32[w]; + .map_err(|e| MLError::ModelError(format!("argmax to_vec1 step {step}: {e}")))?; + for (w, &a) in actions_u32.iter().enumerate() { + self.actions_history_cpu[w * self.max_len + step] = a as i32; } // 4. Launch env step kernel — one thread per window