gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf(cuda): eliminate GPU→CPU roundtrips from backtest evaluate loop

Browse Source 
- 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 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-11 13:05:02 +01:00
parent 057fefe846
commit 8c7e907fcf
2 changed files with 125 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

104
crates/ml-dqn/src/regime_conditional.rs
View File

@@ -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<Tensor, MLError> {
let state_vecs = states.to_vec2::<f32>().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<f32> = 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).

103
crates/ml/src/cuda_pipeline/gpu_backtest_evaluator.rs
View File

@@ -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<f32> = 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::<f32>();
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<i32>).
// 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::<Vec<i32>> → 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<u32> = 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::<i32>();
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<u32> = actions_tensor
.to_vec1()
.map_err(|e| MLError::ModelError(format!("argmax to_vec1: {e}")))?;
let actions_i32: Vec<i32> = 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