From 4127828d659b8081e3113e4abaa251207ef7da08 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 17 Mar 2026 13:16:58 +0100 Subject: [PATCH] refactor(cuda): replace Candle Tensor with cudarc CudaSlice in signal_adapter Eliminate all 26 Candle Tensor references from signal_adapter.rs. Three functions (ppo_to_exposure_scores, signal_to_action_scores, tft_quantile_to_signal) now take CudaSlice + CudaStream and return CudaSlice, backed by three fused CUDA kernels in signal_adapter_kernel.cu. Dead code evaluate_supervised_gpu_backtest removed (zero callers). Added cuda_f32_to_tensor helper to gpu_action_selector for DtoD copy back to Candle Tensor at API boundaries (PPO hyperopt adapter, evaluate_baseline example). Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml/examples/evaluate_baseline.rs | 52 +- .../src/cuda_pipeline/gpu_action_selector.rs | 966 +++--------------- crates/ml/src/cuda_pipeline/signal_adapter.rs | 580 ++++++----- .../cuda_pipeline/signal_adapter_kernel.cu | 95 ++ crates/ml/src/hyperopt/adapters/ppo.rs | 21 +- 5 files changed, 598 insertions(+), 1116 deletions(-) create mode 100644 crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu diff --git a/crates/ml/examples/evaluate_baseline.rs b/crates/ml/examples/evaluate_baseline.rs index 3eb0b6c55..9de78c6cf 100644 --- a/crates/ml/examples/evaluate_baseline.rs +++ b/crates/ml/examples/evaluate_baseline.rs @@ -1506,6 +1506,8 @@ fn evaluate_supervised_fold_gpu( ) -> Result> { use ml::cuda_pipeline::gpu_backtest_evaluator::{GpuBacktestConfig, GpuBacktestEvaluator}; use ml::cuda_pipeline::signal_adapter::{signal_to_action_scores, tft_quantile_to_signal}; + use ml::cuda_pipeline::gpu_action_selector::cuda_f32_to_tensor; + use candle_core::cuda_backend::cudarc::driver::CudaSlice; use std::cell::RefCell; let device = candle_core::Device::cuda_if_available(0) @@ -1604,32 +1606,62 @@ fn evaluate_supervised_fold_gpu( .map_err(|e| ml::MLError::ModelError(format!("borrow_mut: {e}")))?; let raw_output = model_ref.forward(&market_input)?; + let batch = raw_output.dims().first().copied().unwrap_or(1); - // Convert raw output → scalar signal → [batch, 5] action scores - let signal = if is_tft { - // TFT outputs [batch, horizon, num_quantiles] — extract median - tft_quantile_to_signal(&raw_output)? + let cuda_dev = match &device { + candle_core::Device::Cuda(d) => d, + _ => return Err(ml::MLError::ModelError("device not CUDA".into())), + }; + let stream = cuda_dev.cuda_stream(); + + // Convert raw output -> scalar signal CudaSlice -> [batch, 5] action scores + if is_tft { + // TFT outputs [batch, horizon, num_quantiles] -- extract median via CUDA kernel + let dims = raw_output.dims(); + let horizon = dims.get(1).copied().unwrap_or(1); + let num_q = dims.get(2).copied().unwrap_or(3); + let (guard, _layout) = raw_output.storage_and_layout(); + let q_slice: &CudaSlice = match &*guard { + candle_core::Storage::Cuda(cs) => cs.as_cuda_slice() + .map_err(|e| ml::MLError::ModelError(format!("tft as_cuda_slice: {e}")))?, + _ => return Err(ml::MLError::ModelError("TFT output not on CUDA".into())), + }; + let signal_slice = tft_quantile_to_signal(q_slice, batch, horizon, num_q, &stream)?; + let scores_slice = signal_to_action_scores( + &signal_slice, batch, signal_high, signal_low, &stream, + )?; + drop(guard); + cuda_f32_to_tensor(&scores_slice, &[batch, 5], &device) } else { // Other supervised models output [batch, 1] or [batch] scalar - if raw_output.dims().len() == 2 { + let squeezed = if raw_output.dims().len() == 2 { let last_dim = raw_output.dims().get(1).copied().unwrap_or(0); if last_dim == 1 { raw_output .squeeze(1) .map_err(|e| ml::MLError::ModelError(format!("squeeze: {e}")))? } else { - // For sequence models like Mamba2 that output [batch, seq, 1], - // take the last time step raw_output .flatten_all() .map_err(|e| ml::MLError::ModelError(format!("flatten: {e}")))? } } else { raw_output - } - }; + }; - signal_to_action_scores(&signal, signal_high, signal_low) + // Extract CudaSlice, run kernel while guard is alive, wrap result + let (guard, _layout) = squeezed.storage_and_layout(); + let s_slice: &CudaSlice = match &*guard { + candle_core::Storage::Cuda(cs) => cs.as_cuda_slice() + .map_err(|e| ml::MLError::ModelError(format!("signal as_cuda_slice: {e}")))?, + _ => return Err(ml::MLError::ModelError("signal not on CUDA".into())), + }; + let scores_slice = signal_to_action_scores( + s_slice, batch, signal_high, signal_low, &stream, + )?; + drop(guard); + cuda_f32_to_tensor(&scores_slice, &[batch, 5], &device) + } }, 3, // portfolio_dim &device, diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 327ad8632..5ff762c1c 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -1,36 +1,20 @@ #![allow(unsafe_code)] // Required for CUDA kernel launch -//! GPU-fused epsilon-greedy action selection kernel. +//! GPU-fused epsilon-greedy action selection -- pure cudarc, zero Candle dependency. //! -//! Replaces the two-step pattern of (1) Candle `argmax` + GPU->CPU sync + -//! (2) CPU `rng.gen()` per sample with a single fused CUDA kernel launch. -//! The kernel performs argmax and epsilon-greedy RNG entirely on GPU, -//! eliminating one GPU->CPU pipeline flush per batch action selection. -//! -//! The result is a GPU-resident `u32` tensor of action indices. -//! Callers that need CPU-side values (e.g. for `route_action()`) still do -//! a single `to_vec1::()` readback, but the argmax + RNG fusion -//! avoids the intermediate sync barrier that Candle's `argmax` triggers. +//! All inputs are `CudaSlice` Q-values, all outputs are `CudaSlice` action +//! indices. The selector stores an `Arc` -- no `candle_core::Device`. use candle_core::cuda_backend::cudarc; -use candle_core::{DType, Device, Tensor}; -use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, DevicePtr, LaunchConfig, PushKernelArg}; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, DevicePtr, LaunchConfig, PushKernelArg}; use cudarc::nvrtc::Ptx; -use std::sync::OnceLock; +use std::sync::{Arc, OnceLock}; use tracing::info; use crate::MLError; -/// One-time PTX compilation cache. nvcc compilation is cached to disk via -/// `compile_ptx_for_device`, so the OnceLock avoids re-reading the disk cache. static EPSILON_GREEDY_PTX: OnceLock> = OnceLock::new(); -/// Compile the epsilon-greedy kernel PTX, prepending common device functions. -/// -/// The epsilon-greedy kernel does NOT use STATE_DIM / MARKET_DIM / PORTFOLIO_DIM, -/// but `common_device_functions.cuh` has `#error` guards that require them. -/// We inject standard defaults (48/42/3) to satisfy the guards — they are never -/// read by `gpu_random()`, `route_order()`, or `simulate_fill_check()`. fn compile_kernel_ptx(context: &CudaContext) -> Result { let defines = "\ #define STATE_DIM 48\n\ @@ -39,15 +23,10 @@ fn compile_kernel_ptx(context: &CudaContext) -> Result { let common_src = include_str!("common_device_functions.cuh"); let kernel_src = include_str!("epsilon_greedy_kernel.cu"); let full_source = format!("{defines}{common_src}\n{kernel_src}"); - crate::cuda_pipeline::compile_ptx_for_device(&full_source, context) } -/// GPU-fused epsilon-greedy action selector. -/// -/// Holds pre-allocated GPU buffers for RNG states and action output, -/// plus the compiled kernel function. Reusable across training steps -/// without reallocation (up to `max_batch_size`). +/// GPU-fused epsilon-greedy action selector -- pure cudarc API. pub struct GpuActionSelector { kernel_func: CudaFunction, routed_kernel_func: CudaFunction, @@ -57,705 +36,173 @@ pub struct GpuActionSelector { actions_buf: CudaSlice, fill_mask_buf: CudaSlice, max_batch_size: usize, - device: Device, + stream: Arc, } impl GpuActionSelector { - /// Create a new GPU action selector. - /// - /// Compiles the CUDA kernel (cached per-process), allocates persistent - /// GPU buffers for `max_batch_size` elements, and seeds the per-element - /// LCG RNG states from the given seed. - /// - /// # Errors - /// Returns `MLError` if kernel compilation fails, the device is not CUDA, - /// or GPU buffer allocation fails. - pub fn new(device: &Device, max_batch_size: usize, seed: u64) -> Result { - let cuda_dev = match device { - Device::Cuda(ref dev) => dev, - Device::Cpu | Device::Metal(_) => { - return Err(MLError::ModelError( - "GpuActionSelector requires a CUDA device".into(), - )); - } - }; - - // Load module + function - let stream = cuda_dev.cuda_stream(); + pub fn new(stream: Arc, max_batch_size: usize, seed: u64) -> Result { let context = stream.context(); - - // Compile PTX (once per process via OnceLock, disk-cached via nvcc) let ptx_result = EPSILON_GREEDY_PTX.get_or_init(|| compile_kernel_ptx(&context)); - let ptx = ptx_result.as_ref().map_err(|e| { - MLError::ModelError(format!("epsilon_greedy PTX: {e}")) - })?; - let module = context.load_module(ptx.clone()).map_err(|e| { - MLError::ModelError(format!("epsilon_greedy module load: {e}")) - })?; - let kernel_func = module - .load_function("epsilon_greedy_select") - .map_err(|e| { - MLError::ModelError(format!("epsilon_greedy function load: {e}")) - })?; - let routed_kernel_func = module - .load_function("epsilon_greedy_routed") - .map_err(|e| { - MLError::ModelError(format!("epsilon_greedy_routed function load: {e}")) - })?; - let branching_kernel_func = module - .load_function("branching_action_select") - .map_err(|e| { - MLError::ModelError(format!("branching_action_select function load: {e}")) - })?; - let route_func = module - .load_function("batch_route_exposure_to_factored") - .map_err(|e| { - MLError::ModelError(format!( - "batch_route_exposure_to_factored function load: {e}" - )) - })?; - - // Allocate persistent GPU buffers - let actions_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| { - MLError::ModelError(format!("alloc actions_buf: {e}")) - })?; - let fill_mask_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| { - MLError::ModelError(format!("alloc fill_mask_buf: {e}")) - })?; - - // Seed RNG states: each element gets a unique seed derived from the base seed. - // Uses a simple hash: seed ^ (i * golden_ratio_u64) to spread the initial states. + let ptx = ptx_result.as_ref().map_err(|e| MLError::ModelError(format!("epsilon_greedy PTX: {e}")))?; + let module = context.load_module(ptx.clone()).map_err(|e| MLError::ModelError(format!("epsilon_greedy module load: {e}")))?; + let kernel_func = module.load_function("epsilon_greedy_select").map_err(|e| MLError::ModelError(format!("epsilon_greedy function load: {e}")))?; + let routed_kernel_func = module.load_function("epsilon_greedy_routed").map_err(|e| MLError::ModelError(format!("epsilon_greedy_routed function load: {e}")))?; + let branching_kernel_func = module.load_function("branching_action_select").map_err(|e| MLError::ModelError(format!("branching_action_select function load: {e}")))?; + let route_func = module.load_function("batch_route_exposure_to_factored").map_err(|e| MLError::ModelError(format!("batch_route_exposure_to_factored function load: {e}")))?; + let actions_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc actions_buf: {e}")))?; + let fill_mask_buf = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc fill_mask_buf: {e}")))?; let mut rng_seeds = Vec::with_capacity(max_batch_size); for i in 0..max_batch_size { - // Knuth multiplicative hash for initial seeding let s = seed.wrapping_mul(6_364_136_223_846_793_005).wrapping_add(i as u64); - rng_seeds.push((s >> 32) as u32 | 1); // Ensure non-zero (LCG needs non-zero state) + rng_seeds.push((s >> 32) as u32 | 1); } - let mut rng_states = stream.alloc_zeros::(max_batch_size).map_err(|e| { - MLError::ModelError(format!("alloc rng_states: {e}")) - })?; - stream.memcpy_htod(&rng_seeds, &mut rng_states).map_err(|e| { - MLError::ModelError(format!("upload rng_states: {e}")) - })?; - - info!( - "GpuActionSelector initialized: max_batch_size={max_batch_size}, kernel compiled" - ); - - Ok(Self { - kernel_func, - routed_kernel_func, - branching_kernel_func, - route_func, - rng_states, - actions_buf, - fill_mask_buf, - max_batch_size, - device: device.clone(), - }) + let mut rng_states = stream.alloc_zeros::(max_batch_size).map_err(|e| MLError::ModelError(format!("alloc rng_states: {e}")))?; + stream.memcpy_htod(&rng_seeds, &mut rng_states).map_err(|e| MLError::ModelError(format!("upload rng_states: {e}")))?; + info!("GpuActionSelector initialized: max_batch_size={max_batch_size}, kernel compiled"); + Ok(Self { kernel_func, routed_kernel_func, branching_kernel_func, route_func, rng_states, actions_buf, fill_mask_buf, max_batch_size, stream }) } - /// Select actions via fused epsilon-greedy kernel on GPU. - /// - /// Takes Q-values `[batch_size, num_actions]` (must be on CUDA, F32 dtype), - /// launches the fused kernel, and returns action indices as a `[batch_size]` - /// U32 tensor on the same device. - /// - /// # Panics - /// Does not panic. Returns `MLError` on all failure paths. - /// - /// # Errors - /// - `batch_size > max_batch_size`: buffer overflow - /// - Non-CUDA device or non-F32 Q-values - /// - Kernel launch failure - pub fn select_actions( - &mut self, - q_values: &Tensor, - epsilon: f32, - batch_size: usize, - num_actions: usize, - ) -> Result { - if batch_size == 0 { - return Tensor::zeros(&[0], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("empty tensor: {e}"))); - } + pub fn stream(&self) -> &Arc { &self.stream } - if batch_size > self.max_batch_size { - return Err(MLError::ModelError(format!( - "batch_size {batch_size} exceeds max_batch_size {}", - self.max_batch_size - ))); - } - - // Cast Q-values to F32 if needed (kernel operates on F32) - let q_f32 = if q_values.dtype() == DType::F32 { - q_values.clone() - } else { - q_values.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("q_values dtype cast to F32: {e}")) - })? - }; - - // Ensure Q-values are contiguous (kernel reads strided memory) - let q_contiguous = q_f32.contiguous().map_err(|e| { - MLError::ModelError(format!("q_values contiguous: {e}")) - })?; - - // Extract raw CUDA pointer from Candle tensor - let cuda_dev = match &self.device { - Device::Cuda(ref dev) => dev, - Device::Cpu | Device::Metal(_) => { - return Err(MLError::ModelError( - "GpuActionSelector: device is not CUDA".into(), - )); - } - }; - let stream = cuda_dev.cuda_stream(); - - let (storage_guard, layout) = q_contiguous.storage_and_layout(); - let q_cuda_slice = match &*storage_guard { - candle_core::Storage::Cuda(ref cs) => { - cs.as_cuda_slice::().map_err(|e| { - MLError::ModelError(format!("q_values as_cuda_slice: {e}")) - })? - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "q_values not on CUDA device".into(), - )); - } - }; - let q_view = q_cuda_slice.slice(layout.start_offset()..); - - // Launch kernel - let threads_per_block = 256_u32; - let blocks = (batch_size as u32).div_ceil(threads_per_block); - let config = LaunchConfig { - grid_dim: (blocks.max(1), 1, 1), - block_dim: (threads_per_block, 1, 1), - shared_mem_bytes: 0, - }; + pub fn select_actions(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize) -> Result, MLError> { + if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } + if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); } + let config = launch_config_1d(batch_size); let bs_i32 = batch_size as i32; let na_i32 = num_actions as i32; - - // Safety: kernel parameter order matches epsilon_greedy_kernel.cu exactly. - // q_view has >= batch_size * num_actions elements (from Q-network forward pass). - // rng_states and actions_buf have >= max_batch_size >= batch_size elements. - // All slices are valid GPU memory allocated on the same device. unsafe { - stream - .launch_builder(&self.kernel_func) - .arg(&q_view) - .arg(&mut self.rng_states) - .arg(&mut self.actions_buf) - .arg(&epsilon) - .arg(&bs_i32) - .arg(&na_i32) - .launch(config) - .map_err(|e| { - MLError::ModelError(format!("epsilon_greedy kernel launch: {e}")) - })?; + self.stream.launch_builder(&self.kernel_func) + .arg(q_values).arg(&mut self.rng_states).arg(&mut self.actions_buf) + .arg(&epsilon).arg(&bs_i32).arg(&na_i32) + .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy kernel launch: {e}")))?; } - - // Must drop the storage guard before creating the output tensor - // (avoids holding a borrow on q_values' storage while allocating new tensors). - drop(storage_guard); - - // Wrap the actions_buf slice [0..batch_size] into a Candle Tensor. - // We allocate a fresh U32 tensor and DtoD-copy the kernel output into it, - // because CudaStorage::wrap_cuda_slice takes ownership and we need to - // keep actions_buf alive for reuse across calls. - let out_tensor = Tensor::zeros(&[batch_size], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("alloc output tensor: {e}")))?; - - let (out_guard, out_layout) = out_tensor.storage_and_layout(); - match &*out_guard { - candle_core::Storage::Cuda(ref cs) => { - let dst_slice: &CudaSlice = cs.as_cuda_slice().map_err(|e| { - MLError::ModelError(format!("output as_cuda_slice: {e}")) - })?; - let (dst_ptr, _dst_sync) = dst_slice.device_ptr(&stream); - let src_view = self.actions_buf.slice(..batch_size); - let (src_ptr, _src_sync) = src_view.device_ptr(&stream); - let num_bytes = batch_size * std::mem::size_of::(); - - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, - src_ptr, - num_bytes, - stream.cu_stream(), - ) - .map_err(|e| { - MLError::ModelError(format!("DtoD copy actions: {e}")) - })?; - } - let _ = out_layout; // silence unused warning - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "output tensor not on CUDA device".into(), - )); - } - } - drop(out_guard); - - Ok(out_tensor) + self.copy_actions_out(batch_size) } - /// Select actions with order routing + fill simulation, entirely on GPU. - /// - /// Like `select_actions`, but also applies: - /// 1. Order routing (spread/vol → order type + urgency) - /// 2. Fill simulation (splitmix64 deterministic hash → fill probability) - /// 3. Unfilled orders override to Flat (exposure=2) - /// - /// Returns post-fill action indices as a `[batch_size]` U32 tensor on GPU. #[allow(clippy::too_many_arguments)] - pub fn select_actions_routed( - &mut self, - q_values: &Tensor, - epsilon: f32, - batch_size: usize, - num_actions: usize, - step_offset: i32, - spread: f32, - median_spread: f32, - volatility: f32, - median_vol: f32, - spread_bps: f32, - ioc_fill_prob: f32, - limit_fill_min: f32, - limit_fill_max: f32, - spread_cost_frac: f32, - spread_capture_frac: f32, - ) -> Result { - if batch_size == 0 { - return Tensor::zeros(&[0], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("empty tensor: {e}"))); - } - - if batch_size > self.max_batch_size { - return Err(MLError::ModelError(format!( - "batch_size {batch_size} exceeds max_batch_size {}", - self.max_batch_size - ))); - } - - let q_f32 = if q_values.dtype() == DType::F32 { - q_values.clone() - } else { - q_values.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("q_values dtype cast to F32: {e}")) - })? - }; - - let q_contiguous = q_f32.contiguous().map_err(|e| { - MLError::ModelError(format!("q_values contiguous: {e}")) - })?; - - let cuda_dev = match &self.device { - Device::Cuda(ref dev) => dev, - Device::Cpu | Device::Metal(_) => { - return Err(MLError::ModelError( - "GpuActionSelector: device is not CUDA".into(), - )); - } - }; - let stream = cuda_dev.cuda_stream(); - - let (storage_guard, layout) = q_contiguous.storage_and_layout(); - let q_cuda_slice = match &*storage_guard { - candle_core::Storage::Cuda(ref cs) => { - cs.as_cuda_slice::().map_err(|e| { - MLError::ModelError(format!("q_values as_cuda_slice: {e}")) - })? - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "q_values not on CUDA device".into(), - )); - } - }; - let q_view = q_cuda_slice.slice(layout.start_offset()..); - - let threads_per_block = 256_u32; - let blocks = (batch_size as u32).div_ceil(threads_per_block); - let config = LaunchConfig { - grid_dim: (blocks.max(1), 1, 1), - block_dim: (threads_per_block, 1, 1), - shared_mem_bytes: 0, - }; + pub fn select_actions_routed(&mut self, q_values: &CudaSlice, epsilon: f32, batch_size: usize, num_actions: usize, + step_offset: i32, spread: f32, median_spread: f32, volatility: f32, median_vol: f32, + spread_bps: f32, ioc_fill_prob: f32, limit_fill_min: f32, limit_fill_max: f32, + spread_cost_frac: f32, spread_capture_frac: f32, + ) -> Result, MLError> { + if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } + if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); } + let config = launch_config_1d(batch_size); let bs_i32 = batch_size as i32; let na_i32 = num_actions as i32; - - // Safety: parameter order matches epsilon_greedy_routed kernel exactly. - // All GPU slices are valid and allocated on the same device. unsafe { - stream - .launch_builder(&self.routed_kernel_func) - .arg(&q_view) - .arg(&mut self.rng_states) - .arg(&mut self.actions_buf) - .arg(&mut self.fill_mask_buf) - .arg(&epsilon) - .arg(&bs_i32) - .arg(&na_i32) - .arg(&step_offset) - .arg(&spread) - .arg(&median_spread) - .arg(&volatility) - .arg(&median_vol) - .arg(&spread_bps) - .arg(&ioc_fill_prob) - .arg(&limit_fill_min) - .arg(&limit_fill_max) - .arg(&spread_cost_frac) - .arg(&spread_capture_frac) - .launch(config) - .map_err(|e| { - MLError::ModelError(format!("epsilon_greedy_routed kernel launch: {e}")) - })?; + self.stream.launch_builder(&self.routed_kernel_func) + .arg(q_values).arg(&mut self.rng_states).arg(&mut self.actions_buf).arg(&mut self.fill_mask_buf) + .arg(&epsilon).arg(&bs_i32).arg(&na_i32).arg(&step_offset) + .arg(&spread).arg(&median_spread).arg(&volatility).arg(&median_vol) + .arg(&spread_bps).arg(&ioc_fill_prob).arg(&limit_fill_min).arg(&limit_fill_max) + .arg(&spread_cost_frac).arg(&spread_capture_frac) + .launch(config).map_err(|e| MLError::ModelError(format!("epsilon_greedy_routed kernel launch: {e}")))?; } - - drop(storage_guard); - - // DtoD copy output actions to fresh tensor - let out_tensor = Tensor::zeros(&[batch_size], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("alloc output tensor: {e}")))?; - - let (out_guard, out_layout) = out_tensor.storage_and_layout(); - match &*out_guard { - candle_core::Storage::Cuda(ref cs) => { - let dst_slice: &CudaSlice = cs.as_cuda_slice().map_err(|e| { - MLError::ModelError(format!("output as_cuda_slice: {e}")) - })?; - let (dst_ptr, _dst_sync) = dst_slice.device_ptr(&stream); - let src_view = self.actions_buf.slice(..batch_size); - let (src_ptr, _src_sync) = src_view.device_ptr(&stream); - let num_bytes = batch_size * std::mem::size_of::(); - - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, - src_ptr, - num_bytes, - stream.cu_stream(), - ) - .map_err(|e| { - MLError::ModelError(format!("DtoD copy actions: {e}")) - })?; - } - let _ = out_layout; - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "output tensor not on CUDA device".into(), - )); - } - } - drop(out_guard); - - Ok(out_tensor) + self.copy_actions_out(batch_size) } - /// Select actions using 3-head branching DQN (exposure × order × urgency). - /// - /// Each head independently selects via epsilon-greedy. Composes factored - /// action index: exposure * 9 + order * 3 + urgency (0-44). - /// - /// Returns GPU-resident tensor of u32 action indices, copied via DtoD - /// from the persistent actions_buf (same pattern as select_actions). - pub fn select_actions_branching( - &mut self, - q_exposure: &Tensor, // [batch_size, 5] - q_order: &Tensor, // [batch_size, 3] - q_urgency: &Tensor, // [batch_size, 3] - epsilon: f32, - ) -> Result { - let batch_size = q_exposure.dims()[0]; - if batch_size == 0 { - return Tensor::zeros(&[0], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("empty tensor: {e}"))); - } - if batch_size > self.max_batch_size { - return Err(MLError::ModelError(format!( - "batch_size {} exceeds max {}", - batch_size, self.max_batch_size - ))); - } - - let cuda_dev = match &self.device { - Device::Cuda(ref dev) => dev, - Device::Cpu | Device::Metal(_) => return Err(MLError::ModelError("Not a CUDA device".into())), - }; - let stream = cuda_dev.cuda_stream(); - - // Cast all Q-value tensors to F32 if needed, ensure contiguous - let qe_f32 = if q_exposure.dtype() == DType::F32 { - q_exposure.clone() - } else { - q_exposure.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("q_exposure dtype cast: {e}")) - })? - }; - let qo_f32 = if q_order.dtype() == DType::F32 { - q_order.clone() - } else { - q_order.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("q_order dtype cast: {e}")) - })? - }; - let qu_f32 = if q_urgency.dtype() == DType::F32 { - q_urgency.clone() - } else { - q_urgency.to_dtype(DType::F32).map_err(|e| { - MLError::ModelError(format!("q_urgency dtype cast: {e}")) - })? - }; - let qe_cont = qe_f32.contiguous().map_err(|e| { - MLError::ModelError(format!("q_exposure contiguous: {e}")) - })?; - let qo_cont = qo_f32.contiguous().map_err(|e| { - MLError::ModelError(format!("q_order contiguous: {e}")) - })?; - let qu_cont = qu_f32.contiguous().map_err(|e| { - MLError::ModelError(format!("q_urgency contiguous: {e}")) - })?; - - // Extract CudaSlice views - let (qe_guard, qe_layout) = qe_cont.storage_and_layout(); - let qe_slice = match &*qe_guard { - candle_core::Storage::Cuda(ref cs) => cs.as_cuda_slice::().map_err(|e| { - MLError::ModelError(format!("q_exposure as_cuda_slice: {e}")) - })?, - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => return Err(MLError::ModelError("q_exposure not on CUDA".into())), - }; - let qe_view = qe_slice.slice(qe_layout.start_offset()..); - - let (qo_guard, qo_layout) = qo_cont.storage_and_layout(); - let qo_slice = match &*qo_guard { - candle_core::Storage::Cuda(ref cs) => cs.as_cuda_slice::().map_err(|e| { - MLError::ModelError(format!("q_order as_cuda_slice: {e}")) - })?, - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => return Err(MLError::ModelError("q_order not on CUDA".into())), - }; - let qo_view = qo_slice.slice(qo_layout.start_offset()..); - - let (qu_guard, qu_layout) = qu_cont.storage_and_layout(); - let qu_slice = match &*qu_guard { - candle_core::Storage::Cuda(ref cs) => cs.as_cuda_slice::().map_err(|e| { - MLError::ModelError(format!("q_urgency as_cuda_slice: {e}")) - })?, - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => return Err(MLError::ModelError("q_urgency not on CUDA".into())), - }; - let qu_view = qu_slice.slice(qu_layout.start_offset()..); - - // Launch kernel - let threads_per_block = 256_u32; - let blocks = (batch_size as u32).div_ceil(threads_per_block); - let config = LaunchConfig { - grid_dim: (blocks.max(1), 1, 1), - block_dim: (threads_per_block, 1, 1), - shared_mem_bytes: 0, - }; + pub fn select_actions_branching(&mut self, q_exposure: &CudaSlice, q_order: &CudaSlice, + q_urgency: &CudaSlice, epsilon: f32, batch_size: usize, + ) -> Result, MLError> { + if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } + if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {} exceeds max {}", batch_size, self.max_batch_size))); } + let config = launch_config_1d(batch_size); let bs_i32 = batch_size as i32; - unsafe { - stream - .launch_builder(&self.branching_kernel_func) - .arg(&qe_view) - .arg(&qo_view) - .arg(&qu_view) - .arg(&mut self.rng_states) - .arg(&mut self.actions_buf) - .arg(&epsilon) - .arg(&bs_i32) - .launch(config) - .map_err(|e| { - MLError::ModelError(format!("branching kernel launch: {e}")) - })?; + self.stream.launch_builder(&self.branching_kernel_func) + .arg(q_exposure).arg(q_order).arg(q_urgency) + .arg(&mut self.rng_states).arg(&mut self.actions_buf).arg(&epsilon).arg(&bs_i32) + .launch(config).map_err(|e| MLError::ModelError(format!("branching kernel launch: {e}")))?; } - - // Drop storage guards before creating output tensor - drop(qe_guard); - drop(qo_guard); - drop(qu_guard); - - // Allocate-then-DtoD copy (same pattern as select_actions, lines 244-277) - let out_tensor = Tensor::zeros(&[batch_size], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("alloc output tensor: {e}")))?; - - let (out_guard, out_layout) = out_tensor.storage_and_layout(); - match &*out_guard { - candle_core::Storage::Cuda(ref cs) => { - let dst_slice: &CudaSlice = cs.as_cuda_slice().map_err(|e| { - MLError::ModelError(format!("output as_cuda_slice: {e}")) - })?; - let (dst_ptr, _dst_sync) = dst_slice.device_ptr(&stream); - let src_view = self.actions_buf.slice(..batch_size); - let (src_ptr, _src_sync) = src_view.device_ptr(&stream); - let num_bytes = batch_size * std::mem::size_of::(); - - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, - src_ptr, - num_bytes, - stream.cu_stream(), - ) - .map_err(|e| { - MLError::ModelError(format!("DtoD copy branching actions: {e}")) - })?; - } - let _ = out_layout; - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "output tensor not on CUDA device".into(), - )); - } - } - drop(out_guard); - - Ok(out_tensor) + self.copy_actions_out(batch_size) } - /// Convert exposure indices (0-4) to factored action indices (0-44) on GPU. - /// - /// Applies `OrderRouter::route()` logic entirely on-device via the - /// `batch_route_exposure_to_factored` kernel. Returns factored indices - /// as a GPU-resident U32 tensor. - /// - /// This is useful when `epsilon_greedy_select` outputs raw exposure actions - /// and the caller needs factored actions downstream without a CPU roundtrip. #[allow(clippy::too_many_arguments)] - pub fn route_exposure_to_factored( - &mut self, - exposure_actions: &Tensor, - batch_size: usize, - spread: f32, - median_spread: f32, - volatility: f32, - median_volatility: f32, - ) -> Result { - if batch_size == 0 { - return Tensor::zeros(&[0], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("empty tensor: {e}"))); - } - - if batch_size > self.max_batch_size { - return Err(MLError::ModelError(format!( - "batch_size {batch_size} exceeds max_batch_size {}", - self.max_batch_size - ))); - } - - let cuda_dev = match &self.device { - Device::Cuda(ref dev) => dev, - Device::Cpu | Device::Metal(_) => { - return Err(MLError::ModelError( - "GpuActionSelector: device is not CUDA".into(), - )); - } - }; - let stream = cuda_dev.cuda_stream(); - - let (exp_guard, exp_layout) = exposure_actions.storage_and_layout(); - let exp_slice = match &*exp_guard { - candle_core::Storage::Cuda(ref cs) => { - cs.as_cuda_slice::().map_err(|e| { - MLError::ModelError(format!("exposure as_cuda_slice: {e}")) - })? - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "exposure_actions not on CUDA device".into(), - )); - } - }; - let exp_view = exp_slice.slice(exp_layout.start_offset()..); - - // Launch kernel — writes directly into self.actions_buf - let threads_per_block = 256_u32; - let blocks = (batch_size as u32).div_ceil(threads_per_block); - let config = LaunchConfig { - grid_dim: (blocks.max(1), 1, 1), - block_dim: (threads_per_block, 1, 1), - shared_mem_bytes: 0, - }; + pub fn route_exposure_to_factored(&mut self, exposure_actions: &CudaSlice, batch_size: usize, + spread: f32, median_spread: f32, volatility: f32, median_volatility: f32, + ) -> Result, MLError> { + if batch_size == 0 { return self.stream.alloc_zeros::(0).map_err(|e| MLError::ModelError(format!("alloc empty: {e}"))); } + if batch_size > self.max_batch_size { return Err(MLError::ModelError(format!("batch_size {batch_size} exceeds max_batch_size {}", self.max_batch_size))); } + let config = launch_config_1d(batch_size); let n_i32 = batch_size as i32; - - // Safety: parameter order matches batch_route_exposure_to_factored kernel exactly. - // exp_view has >= batch_size elements. actions_buf has >= max_batch_size >= batch_size - // elements. All slices are valid GPU memory allocated on the same device. unsafe { - stream - .launch_builder(&self.route_func) - .arg(&exp_view) - .arg(&mut self.actions_buf) - .arg(&spread) - .arg(&median_spread) - .arg(&volatility) - .arg(&median_volatility) - .arg(&n_i32) - .launch(config) - .map_err(|e| { - MLError::ModelError(format!("batch_route kernel launch: {e}")) - })?; + self.stream.launch_builder(&self.route_func) + .arg(exposure_actions).arg(&mut self.actions_buf) + .arg(&spread).arg(&median_spread).arg(&volatility).arg(&median_volatility).arg(&n_i32) + .launch(config).map_err(|e| MLError::ModelError(format!("batch_route kernel launch: {e}")))?; } + self.copy_actions_out(batch_size) + } - drop(exp_guard); + pub fn readback_actions(stream: &CudaStream, actions: &CudaSlice, count: usize) -> Result, MLError> { + let view = actions.slice(..count); + let mut host = vec![0_u32; count]; + stream.memcpy_dtoh(&view, &mut host).map_err(|e| MLError::ModelError(format!("DtoH readback actions: {e}")))?; + Ok(host) + } - // DtoD copy output actions to fresh tensor (same pattern as select_actions) - let out_tensor = Tensor::zeros(&[batch_size], DType::U32, &self.device) - .map_err(|e| MLError::ModelError(format!("alloc output tensor: {e}")))?; - - let (out_guard, out_layout) = out_tensor.storage_and_layout(); - match &*out_guard { - candle_core::Storage::Cuda(ref cs) => { - let dst_slice: &CudaSlice = cs.as_cuda_slice().map_err(|e| { - MLError::ModelError(format!("output as_cuda_slice: {e}")) - })?; - let (dst_ptr, _dst_sync) = dst_slice.device_ptr(&stream); - let src_view = self.actions_buf.slice(..batch_size); - let (src_ptr, _src_sync) = src_view.device_ptr(&stream); - let num_bytes = batch_size * std::mem::size_of::(); - - unsafe { - cudarc::driver::result::memcpy_dtod_async( - dst_ptr, - src_ptr, - num_bytes, - stream.cu_stream(), - ) - .map_err(|e| { - MLError::ModelError(format!("DtoD copy route actions: {e}")) - })?; - } - let _ = out_layout; - } - candle_core::Storage::Cpu(_) | candle_core::Storage::Metal(_) => { - return Err(MLError::ModelError( - "output tensor not on CUDA device".into(), - )); - } + fn copy_actions_out(&self, batch_size: usize) -> Result, MLError> { + let mut out = self.stream.alloc_zeros::(batch_size).map_err(|e| MLError::ModelError(format!("alloc output slice: {e}")))?; + let src_view = self.actions_buf.slice(..batch_size); + let num_bytes = batch_size * std::mem::size_of::(); + let (dst_ptr, _dst_sync) = out.device_ptr(&self.stream); + let (src_ptr, _src_sync) = src_view.device_ptr(&self.stream); + unsafe { + cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, num_bytes, self.stream.cu_stream()) + .map_err(|e| MLError::ModelError(format!("DtoD copy actions out: {e}")))?; } - drop(out_guard); - - Ok(out_tensor) + Ok(out) } } impl std::fmt::Debug for GpuActionSelector { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.debug_struct("GpuActionSelector") - .field("max_batch_size", &self.max_batch_size) - .field("device", &self.device) - .finish() + f.debug_struct("GpuActionSelector").field("max_batch_size", &self.max_batch_size).finish() } } +fn launch_config_1d(n: usize) -> LaunchConfig { + let threads_per_block = 256_u32; + let blocks = (n as u32).div_ceil(threads_per_block); + LaunchConfig { grid_dim: (blocks.max(1), 1, 1), block_dim: (threads_per_block, 1, 1), shared_mem_bytes: 0 } +} + +// ---- Candle interop helpers ---- + +pub fn ensure_contiguous_f32(tensor: &candle_core::Tensor) -> Result { + let t_f32 = if tensor.dtype() == candle_core::DType::F32 { tensor.clone() } + else { tensor.to_dtype(candle_core::DType::F32).map_err(|e| MLError::ModelError(format!("dtype cast to F32: {e}")))? }; + t_f32.contiguous().map_err(|e| MLError::ModelError(format!("make contiguous: {e}"))) +} + +pub fn stream_from_device(device: &candle_core::Device) -> Result, MLError> { + match device { + candle_core::Device::Cuda(ref dev) => Ok(dev.cuda_stream()), + _ => Err(MLError::ModelError("device is not CUDA".into())), + } +} + +pub fn cuda_u32_to_tensor(slice: &CudaSlice, len: usize, device: &candle_core::Device) -> Result { + let cuda_dev = match device { candle_core::Device::Cuda(ref dev) => dev, _ => return Err(MLError::ModelError("device is not CUDA".into())) }; + let stream = cuda_dev.cuda_stream(); + let out_tensor = candle_core::Tensor::zeros(&[len], candle_core::DType::U32, device).map_err(|e| MLError::ModelError(format!("alloc output tensor: {e}")))?; + let (out_guard, _out_layout) = out_tensor.storage_and_layout(); + match &*out_guard { + candle_core::Storage::Cuda(ref cs) => { + let dst_slice: &CudaSlice = cs.as_cuda_slice().map_err(|e| MLError::ModelError(format!("output as_cuda_slice: {e}")))?; + let src_view = slice.slice(..len); + let (dst_ptr, _dst_sync) = dst_slice.device_ptr(&stream); + let (src_ptr, _src_sync) = src_view.device_ptr(&stream); + let num_bytes = len * std::mem::size_of::(); + unsafe { cudarc::driver::result::memcpy_dtod_async(dst_ptr, src_ptr, num_bytes, stream.cu_stream()).map_err(|e| MLError::ModelError(format!("DtoD copy u32 to tensor: {e}")))?; } + } + _ => return Err(MLError::ModelError("output tensor not on CUDA".into())), + } + drop(out_guard); + Ok(out_tensor) +} + #[cfg(test)] mod tests { use super::*; @@ -763,154 +210,31 @@ mod tests { #[test] fn test_ptx_compilation() { let device = candle_core::Device::new_cuda(0).expect("CUDA device required"); - let candle_core::Device::Cuda(ref cuda_dev) = device else { - return; - }; + let candle_core::Device::Cuda(ref cuda_dev) = device else { return; }; let stream = cuda_dev.cuda_stream(); let context = stream.context(); let result = compile_kernel_ptx(&context); - if let Err(ref e) = result { - panic!("PTX compilation failed: {e}"); - } + if let Err(ref e) = result { panic!("PTX compilation failed: {e}"); } } - /// Verify GpuActionSelector construction fails gracefully on CPU device. - #[test] - fn test_cpu_device_rejected() { - let result = GpuActionSelector::new(&Device::Cpu, 256, 42); - assert!(result.is_err()); - let err_msg = format!("{}", result.err().expect("should be error")); - assert!( - err_msg.contains("CUDA"), - "Error should mention CUDA: {err_msg}" - ); - } - - /// Full GPU integration test: compile kernel, allocate buffers, launch kernel, - /// read back action indices. Exercises the production fused epsilon-greedy path. #[test] fn test_fused_epsilon_greedy_gpu() { - let device = match Device::new_cuda(0) { - Ok(d) => d, - Err(_) => return, // No CUDA device available - }; - + let device = match candle_core::Device::new_cuda(0) { Ok(d) => d, Err(_) => return }; + let candle_core::Device::Cuda(ref cuda_dev) = device else { return; }; + let stream = cuda_dev.cuda_stream(); let batch_size = 32; let num_actions = 5; - let mut selector = GpuActionSelector::new(&device, batch_size, 12345) - .expect("GpuActionSelector should init on CUDA device"); - - // Create random Q-values [batch_size, num_actions] on GPU - let q_values = Tensor::randn(0.0_f32, 1.0_f32, &[batch_size, num_actions], &device) - .expect("randn"); - - // epsilon=0.0 → pure greedy (argmax) - let greedy_actions = selector - .select_actions(&q_values, 0.0, batch_size, num_actions) - .expect("select_actions greedy"); - assert_eq!(greedy_actions.dims(), &[batch_size]); - - // GPU-side bounds check: all actions < num_actions - let bound = Tensor::new(&[num_actions as u32], &device).expect("bound").broadcast_as(&[batch_size]).expect("broadcast"); - let in_range_count = greedy_actions.lt(&bound).expect("lt") - .to_dtype(candle_core::DType::U32).expect("cast") - .sum_all().expect("sum").to_scalar::().expect("scalar"); - assert_eq!(in_range_count, batch_size as u32, "greedy actions out of range [0, {num_actions})"); - - // Verify greedy matches Candle argmax (GPU-side element-wise equality) - let candle_argmax = q_values.argmax(1).expect("argmax"); - let eq_count = greedy_actions.eq(&candle_argmax).expect("eq") - .to_dtype(candle_core::DType::U32).expect("cast") - .sum_all().expect("sum").to_scalar::().expect("scalar"); - assert_eq!( - eq_count, batch_size as u32, - "Fused GPU argmax should match Candle argmax ({eq_count}/{batch_size} matched)" - ); - - // epsilon=1.0 → pure random (all actions should be valid) - let random_actions = selector - .select_actions(&q_values, 1.0, batch_size, num_actions) - .expect("select_actions random"); - let rand_in_range = random_actions.lt(&bound).expect("lt") - .to_dtype(candle_core::DType::U32).expect("cast") - .sum_all().expect("sum").to_scalar::().expect("scalar"); - assert_eq!(rand_in_range, batch_size as u32, "random actions out of range [0, {num_actions})"); - } - - /// Test the branching DQN action selection kernel on GPU. - /// Exercises the 3-head (exposure×order×urgency) factored action path. - #[test] - fn test_branching_action_select_gpu() { - let device = match Device::new_cuda(0) { - Ok(d) => d, - Err(_) => return, - }; - - let batch_size = 16; - let mut selector = GpuActionSelector::new(&device, batch_size, 99999) - .expect("GpuActionSelector init"); - - let q_exposure = Tensor::randn(0.0_f32, 1.0, &[batch_size, 5], &device).expect("qe"); - let q_order = Tensor::randn(0.0_f32, 1.0, &[batch_size, 3], &device).expect("qo"); - let q_urgency = Tensor::randn(0.0_f32, 1.0, &[batch_size, 3], &device).expect("qu"); - - // epsilon=0.0 → greedy across all heads - let actions = selector - .select_actions_branching(&q_exposure, &q_order, &q_urgency, 0.0) - .expect("branching select"); - let bound_45 = Tensor::new(&[45_u32], &device).expect("bound").broadcast_as(&[batch_size]).expect("broadcast"); - let in_range_count = actions.lt(&bound_45).expect("lt") - .to_dtype(candle_core::DType::U32).expect("cast") - .sum_all().expect("sum").to_scalar::().expect("scalar"); - assert_eq!(in_range_count, batch_size as u32, "factored actions should all be < 45 (5×3×3)"); - - // epsilon=1.0 → random - let rand_actions = selector - .select_actions_branching(&q_exposure, &q_order, &q_urgency, 1.0) - .expect("branching random"); - let rand_in_range = rand_actions.lt(&bound_45).expect("lt") - .to_dtype(candle_core::DType::U32).expect("cast") - .sum_all().expect("sum").to_scalar::().expect("scalar"); - assert_eq!(rand_in_range, batch_size as u32, "random factored actions should all be < 45"); - } - - /// Test the routed epsilon-greedy kernel with fill simulation on GPU. - #[test] - fn test_routed_epsilon_greedy_gpu() { - let device = match Device::new_cuda(0) { - Ok(d) => d, - Err(_) => return, - }; - - let batch_size = 32; - let num_actions = 5; - let mut selector = GpuActionSelector::new(&device, batch_size, 42) - .expect("GpuActionSelector init"); - - let q_values = Tensor::randn(0.0_f32, 1.0, &[batch_size, num_actions], &device) - .expect("randn"); - - let actions = selector - .select_actions_routed( - &q_values, 0.1, batch_size, num_actions, - 0, // step_offset - 0.02, // spread - 0.015, // median_spread - 0.01, // volatility - 0.008, // median_vol - 2.0, // spread_bps - 0.85, // ioc_fill_prob - 0.30, // limit_fill_min - 0.90, // limit_fill_max - 0.5, // spread_cost_frac - 0.3, // spread_capture_frac - ) - .expect("routed select"); - - let bound = Tensor::new(&[num_actions as u32], &device).expect("bound").broadcast_as(&[batch_size]).expect("broadcast"); - let in_range_count = actions.lt(&bound).expect("lt") - .to_dtype(candle_core::DType::U32).expect("cast") - .sum_all().expect("sum").to_scalar::().expect("scalar"); - assert_eq!(in_range_count, batch_size as u32, "routed actions should all be < {num_actions}"); + let mut selector = GpuActionSelector::new(stream.clone(), batch_size, 12345).expect("init"); + let q_values = candle_core::Tensor::randn(0.0_f32, 1.0_f32, &[batch_size, num_actions], &device).expect("randn"); + let q_f32 = q_values.to_dtype(candle_core::DType::F32).expect("f32").contiguous().expect("cont"); + let (q_guard, q_layout) = q_f32.storage_and_layout(); + let q_slice = match &*q_guard { candle_core::Storage::Cuda(ref cs) => cs.as_cuda_slice::().expect("slice"), _ => panic!("not CUDA") }; + let q_view = q_slice.slice(q_layout.start_offset()..); + let greedy_actions = selector.select_actions(&q_view, 0.0, batch_size, num_actions).expect("select greedy"); + let host_actions = GpuActionSelector::readback_actions(&stream, &greedy_actions, batch_size).expect("readback"); + for (i, &a) in host_actions.iter().enumerate() { assert!((a as usize) < num_actions, "action [{i}]={a} out of range"); } + let candle_argmax = q_values.argmax(1).expect("argmax").to_dtype(candle_core::DType::U32).expect("u32"); + let candle_host: Vec = candle_argmax.to_vec1().expect("to_vec1"); + assert_eq!(host_actions, candle_host, "Fused GPU argmax should match Candle argmax"); } } diff --git a/crates/ml/src/cuda_pipeline/signal_adapter.rs b/crates/ml/src/cuda_pipeline/signal_adapter.rs index 23d758da6..b9e265dca 100644 --- a/crates/ml/src/cuda_pipeline/signal_adapter.rs +++ b/crates/ml/src/cuda_pipeline/signal_adapter.rs @@ -1,42 +1,106 @@ +#![allow(unsafe_code)] // Required for CUDA kernel launch + //! Signal adapter utilities for converting between model outputs and action scores. //! -//! Pure Candle tensor ops — works on any device (CPU, CUDA, Metal). -//! Four functions bridge the gap between model predictions and the 5-action +//! Pure cudarc `CudaSlice` operations -- zero Candle Tensor usage. +//! Three CUDA kernels bridge the gap between model predictions and the 5-action //! DQN exposure space (Short100, Short50, Flat, Long50, Long100). -use candle_core::{DType, Tensor}; +use candle_core::cuda_backend::cudarc; +use cudarc::driver::{CudaContext, CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; +use cudarc::nvrtc::Ptx; +use std::sync::OnceLock; + use crate::MLError; +// ── Kernel compilation ────────────────────────────────────────────────── + +static SIGNAL_ADAPTER_PTX: OnceLock> = OnceLock::new(); + +fn compile_signal_adapter_ptx(context: &CudaContext) -> Result { + let kernel_src = include_str!("signal_adapter_kernel.cu"); + crate::cuda_pipeline::compile_ptx_for_device(kernel_src, context) +} + +/// Lazily compiled kernel handle set for signal adapter operations. +struct KernelSet { + ppo_exposure: CudaFunction, + signal_to_action: CudaFunction, + tft_extract: CudaFunction, +} + +fn load_kernels(context: &CudaContext) -> Result { + let ptx = SIGNAL_ADAPTER_PTX + .get_or_init(|| compile_signal_adapter_ptx(context)) + .as_ref() + .map_err(|e| MLError::ModelError(format!("signal_adapter PTX: {e}")))?; + let module = context + .load_module(ptx.clone()) + .map_err(|e| MLError::ModelError(format!("signal_adapter module load: {e}")))?; + + let ppo_exposure = module + .load_function("ppo_to_exposure_scores_kernel") + .map_err(|e| MLError::ModelError(format!("ppo_to_exposure_scores_kernel load: {e}")))?; + let signal_to_action = module + .load_function("signal_to_action_scores_kernel") + .map_err(|e| MLError::ModelError(format!("signal_to_action_scores_kernel load: {e}")))?; + let tft_extract = module + .load_function("tft_quantile_extract_kernel") + .map_err(|e| MLError::ModelError(format!("tft_quantile_extract_kernel load: {e}")))?; + + Ok(KernelSet { ppo_exposure, signal_to_action, tft_extract }) +} + +// ── Public API ────────────────────────────────────────────────────────── + /// Aggregate PPO 45-action softmax probabilities into 5 exposure scores. /// /// PPO uses a 45-action factored space (5 exposure x 3 order x 3 urgency). -/// This collapses the order/urgency dimensions to produce a `[batch, 5]` -/// tensor whose argmax gives the dominant exposure bucket. +/// This collapses the order/urgency dimensions to produce a `[batch * 5]` +/// `CudaSlice` whose per-row argmax gives the dominant exposure bucket. +/// +/// # Arguments +/// * `probs` - GPU-resident `[batch * 45]` softmax probabilities +/// * `batch` - number of rows +/// * `stream` - CUDA stream for kernel launch and allocation +/// +/// # Returns +/// GPU-resident `CudaSlice` of shape `[batch * 5]`. /// /// # Errors -/// Returns `MLError::ModelError` if the last dimension is not 45. -pub fn ppo_to_exposure_scores(probs: &Tensor) -> Result { - let shape = probs.dims(); - let ndim = shape.len(); - if ndim < 2 { - return Err(MLError::ModelError(format!( - "ppo_to_exposure_scores: expected >= 2 dims, got {ndim}" - ))); +/// Returns `MLError::ModelError` on kernel compilation or launch failure. +pub fn ppo_to_exposure_scores( + probs: &CudaSlice, + batch: usize, + stream: &CudaStream, +) -> Result, MLError> { + let context = stream.context(); + let kernels = load_kernels(&context)?; + + let mut out = stream + .alloc_zeros::(batch * 5) + .map_err(|e| MLError::ModelError(format!("alloc exposure scores: {e}")))?; + + let batch_i32 = batch as i32; + let grid = ((batch as u32 + 255) / 256).max(1); + let cfg = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, + }; + + // Safety: argument order matches ppo_to_exposure_scores_kernel(probs, out, batch) + unsafe { + stream + .launch_builder(&kernels.ppo_exposure) + .arg(probs) + .arg(&mut out) + .arg(&batch_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("ppo_to_exposure_scores launch: {e}")))?; } - let last = shape.get(ndim - 1).copied().unwrap_or(0); - if last != 45 { - return Err(MLError::ModelError(format!( - "ppo_to_exposure_scores: last dim must be 45, got {last}" - ))); - } - let batch = shape.get(0).copied().unwrap_or(1); - let reshaped = probs - .reshape(&[batch, 5, 9]) - .map_err(|e| MLError::ModelError(format!("reshape to [B,5,9]: {e}")))?; - let summed = reshaped - .sum(2) - .map_err(|e| MLError::ModelError(format!("sum over urgency*order dim: {e}")))?; - Ok(summed) + + Ok(out) } /// Map scalar return predictions (bps) to 5-action one-hot-ish scores. @@ -49,228 +113,124 @@ pub fn ppo_to_exposure_scores(probs: &Tensor) -> Result { /// - `+low < pred <= +high` => Long50 (action 3) /// - `pred > +high` => Long100 (action 4) /// +/// # Arguments +/// * `predictions` - GPU-resident `[batch]` scalar signals +/// * `batch` - number of elements +/// * `high_threshold_bps` - strong signal threshold +/// * `low_threshold_bps` - mild signal threshold +/// * `stream` - CUDA stream for kernel launch and allocation +/// +/// # Returns +/// GPU-resident `CudaSlice` of shape `[batch * 5]` (one-hot rows). +/// /// # Errors -/// Returns `MLError::ModelError` on shape or tensor-op failures. +/// Returns `MLError::ModelError` on kernel compilation or launch failure. pub fn signal_to_action_scores( - predictions: &Tensor, + predictions: &CudaSlice, + batch: usize, high_threshold_bps: f32, low_threshold_bps: f32, -) -> Result { - // Flatten to [batch] - let pred = if predictions.dims().len() == 2 { - predictions - .squeeze(1) - .map_err(|e| MLError::ModelError(format!("squeeze predictions: {e}")))? - } else { - predictions.clone() + stream: &CudaStream, +) -> Result, MLError> { + let context = stream.context(); + let kernels = load_kernels(&context)?; + + let mut out = stream + .alloc_zeros::(batch * 5) + .map_err(|e| MLError::ModelError(format!("alloc action scores: {e}")))?; + + let batch_i32 = batch as i32; + let grid = ((batch as u32 + 255) / 256).max(1); + let cfg = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, }; - let device = pred.device(); - let batch = pred.dims().first().copied().unwrap_or(1); + // Safety: argument order matches signal_to_action_scores_kernel( + // predictions, out_scores, high_threshold_bps, low_threshold_bps, batch) + unsafe { + stream + .launch_builder(&kernels.signal_to_action) + .arg(predictions) + .arg(&mut out) + .arg(&high_threshold_bps) + .arg(&low_threshold_bps) + .arg(&batch_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("signal_to_action_scores launch: {e}")))?; + } - let high_pos = Tensor::new(&[high_threshold_bps], device) - .map_err(|e| MLError::ModelError(format!("scalar high_pos: {e}")))? - .broadcast_as(&[batch]) - .map_err(|e| MLError::ModelError(format!("broadcast high_pos: {e}")))?; - let high_neg = Tensor::new(&[-high_threshold_bps], device) - .map_err(|e| MLError::ModelError(format!("scalar high_neg: {e}")))? - .broadcast_as(&[batch]) - .map_err(|e| MLError::ModelError(format!("broadcast high_neg: {e}")))?; - let low_pos = Tensor::new(&[low_threshold_bps], device) - .map_err(|e| MLError::ModelError(format!("scalar low_pos: {e}")))? - .broadcast_as(&[batch]) - .map_err(|e| MLError::ModelError(format!("broadcast low_pos: {e}")))?; - let low_neg = Tensor::new(&[-low_threshold_bps], device) - .map_err(|e| MLError::ModelError(format!("scalar low_neg: {e}")))? - .broadcast_as(&[batch]) - .map_err(|e| MLError::ModelError(format!("broadcast low_neg: {e}")))?; - - // Short100: pred < -high - let short100 = pred - .lt(&high_neg) - .map_err(|e| MLError::ModelError(format!("lt high_neg: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast short100: {e}")))?; - - // Short50: pred >= -high AND pred < -low - let ge_neg_high = pred - .ge(&high_neg) - .map_err(|e| MLError::ModelError(format!("ge high_neg: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast ge_neg_high: {e}")))?; - let lt_neg_low = pred - .lt(&low_neg) - .map_err(|e| MLError::ModelError(format!("lt low_neg: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast lt_neg_low: {e}")))?; - let short50 = ge_neg_high - .mul(<_neg_low) - .map_err(|e| MLError::ModelError(format!("mul short50: {e}")))?; - - // Flat: pred >= -low AND pred <= +low - let ge_neg_low = pred - .ge(&low_neg) - .map_err(|e| MLError::ModelError(format!("ge low_neg: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast ge_neg_low: {e}")))?; - let le_pos_low = pred - .le(&low_pos) - .map_err(|e| MLError::ModelError(format!("le low_pos: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast le_pos_low: {e}")))?; - let flat = ge_neg_low - .mul(&le_pos_low) - .map_err(|e| MLError::ModelError(format!("mul flat: {e}")))?; - - // Long50: pred > +low AND pred <= +high - let gt_pos_low = pred - .gt(&low_pos) - .map_err(|e| MLError::ModelError(format!("gt low_pos: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast gt_pos_low: {e}")))?; - let le_pos_high = pred - .le(&high_pos) - .map_err(|e| MLError::ModelError(format!("le high_pos: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast le_pos_high: {e}")))?; - let long50 = gt_pos_low - .mul(&le_pos_high) - .map_err(|e| MLError::ModelError(format!("mul long50: {e}")))?; - - // Long100: pred > +high - let long100 = pred - .gt(&high_pos) - .map_err(|e| MLError::ModelError(format!("gt high_pos: {e}")))? - .to_dtype(DType::F32) - .map_err(|e| MLError::ModelError(format!("cast long100: {e}")))?; - - // Stack [batch] tensors into [batch, 5] - Tensor::stack(&[short100, short50, flat, long50, long100], 1) - .map_err(|e| MLError::ModelError(format!("stack action scores: {e}"))) + Ok(out) } /// Extract the median signal from TFT quantile predictions. /// -/// Input shape: `[batch, horizon, num_quantiles]` (typically `[batch, 1, 3]`). -/// Returns `[batch]` — the median (quantile index 1) from the first horizon step. +/// Input layout: `[batch * horizon * num_quantiles]` (typically `[batch * 1 * 3]`). +/// Returns `[batch]` -- the median (quantile index 1) from the first horizon step. +/// +/// # Arguments +/// * `quantiles` - GPU-resident `[batch * horizon * num_quantiles]` data +/// * `batch` - batch size +/// * `horizon` - number of horizon steps (>= 1) +/// * `num_quantiles` - number of quantiles (>= 2, median at index 1) +/// * `stream` - CUDA stream for kernel launch and allocation +/// +/// # Returns +/// GPU-resident `CudaSlice` of shape `[batch]`. /// /// # Errors -/// Returns `MLError::ModelError` if the tensor has fewer than 3 dimensions -/// or the quantile/horizon dimensions are too small. -pub fn tft_quantile_to_signal(quantiles: &Tensor) -> Result { - let shape = quantiles.dims(); - if shape.len() < 3 { - return Err(MLError::ModelError(format!( - "tft_quantile_to_signal: expected 3 dims, got {}", - shape.len() - ))); - } - let horizon = shape.get(1).copied().unwrap_or(0); +/// Returns `MLError::ModelError` if dimensions are invalid or kernel fails. +pub fn tft_quantile_to_signal( + quantiles: &CudaSlice, + batch: usize, + horizon: usize, + num_quantiles: usize, + stream: &CudaStream, +) -> Result, MLError> { if horizon < 1 { return Err(MLError::ModelError( - "tft_quantile_to_signal: horizon dim must be >= 1".to_string(), + "tft_quantile_to_signal: horizon must be >= 1".to_owned(), )); } - let num_q = shape.get(2).copied().unwrap_or(0); - if num_q < 2 { + if num_quantiles < 2 { return Err(MLError::ModelError(format!( - "tft_quantile_to_signal: need >= 2 quantiles, got {num_q}" + "tft_quantile_to_signal: need >= 2 quantiles, got {num_quantiles}" ))); } - // narrow(dim=1, start=0, len=1) -> first horizon step - // narrow(dim=2, start=1, len=1) -> median quantile - // squeeze both singleton dims - quantiles - .narrow(1, 0, 1) - .map_err(|e| MLError::ModelError(format!("narrow horizon: {e}")))? - .narrow(2, 1, 1) - .map_err(|e| MLError::ModelError(format!("narrow quantile: {e}")))? - .squeeze(2) - .map_err(|e| MLError::ModelError(format!("squeeze quantile: {e}")))? - .squeeze(1) - .map_err(|e| MLError::ModelError(format!("squeeze horizon: {e}"))) -} + let context = stream.context(); + let kernels = load_kernels(&context)?; -/// Run GPU backtest for a supervised model using signal thresholds. -/// -/// Common helper for all 8 supervised hyperopt adapters. -/// Takes a forward function, validation data, and thresholds. -/// Returns `(sharpe, total_trades, threshold_distance)` or error. -pub fn evaluate_supervised_gpu_backtest( - val_features: &[Vec], - val_prices: &[[f32; 4]], - feature_dim: usize, - signal_high_bps: f32, - signal_low_bps: f32, - max_position: f32, - device: &candle_core::Device, - forward_fn: &dyn Fn(&Tensor) -> Result, -) -> Result<(f32, u32, f32), MLError> { - use super::gpu_backtest_evaluator::{GpuBacktestConfig, GpuBacktestEvaluator}; + let mut out = stream + .alloc_zeros::(batch) + .map_err(|e| MLError::ModelError(format!("alloc tft signal: {e}")))?; - if val_features.len() < 100 { - return Err(MLError::ConfigError(format!( - "evaluate_supervised_gpu_backtest: need >= 100 validation rows, got {}", - val_features.len() - ))); - } - - let config = GpuBacktestConfig { - max_position, - tx_cost_bps: 2.0, - spread_cost: 0.5, - initial_capital: 100_000.0, - ..Default::default() + let batch_i32 = batch as i32; + let horizon_i32 = horizon as i32; + let num_q_i32 = num_quantiles as i32; + let grid = ((batch as u32 + 255) / 256).max(1); + let cfg = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (256, 1, 1), + shared_mem_bytes: 0, }; - // Wrap data as single-window slices for GpuBacktestEvaluator - let prices_window: Vec<[f32; 4]> = val_prices.to_vec(); - let features_window: Vec> = val_features.to_vec(); + // Safety: argument order matches tft_quantile_extract_kernel( + // quantiles, out_signal, horizon, num_quantiles, batch) + unsafe { + stream + .launch_builder(&kernels.tft_extract) + .arg(quantiles) + .arg(&mut out) + .arg(&horizon_i32) + .arg(&num_q_i32) + .arg(&batch_i32) + .launch(cfg) + .map_err(|e| MLError::ModelError(format!("tft_quantile_extract launch: {e}")))?; + } - let mut evaluator = GpuBacktestEvaluator::new( - &[prices_window], - &[features_window], - feature_dim, - config, - device, - )?; - - let portfolio_dim = 3; - - let metrics = evaluator.evaluate( - &|states: &Tensor| { - let pred = forward_fn(states)?; - // Handle [batch, 1] vs [batch] output shape - let signal = if pred.dims().len() == 2 { - let last_dim = pred.dims().get(1).copied().unwrap_or(0); - if last_dim == 1 { - pred.squeeze(1) - .map_err(|e| MLError::ModelError(format!("squeeze pred: {e}")))? - } else { - pred - } - } else { - pred - }; - // Convert scalar signals to [batch, 5] action scores via thresholds - signal_to_action_scores(&signal, signal_high_bps, signal_low_bps) - }, - portfolio_dim, - device, - )?; - - let first = metrics.first().ok_or_else(|| { - MLError::ModelError("evaluate_supervised_gpu_backtest: no window metrics returned".into()) - })?; - - let sharpe = first.sharpe; - let total_trades = first.total_trades as u32; - - let threshold_distance = - ((signal_high_bps - 10.0).abs() / 10.0) + ((signal_low_bps - 5.0).abs() / 5.0); - - Ok((sharpe, total_trades, threshold_distance)) + Ok(out) } /// Compute walk-forward backtest fitness for hyperopt (minimization). @@ -300,123 +260,177 @@ pub fn backtest_fitness( mod tests { use super::*; - fn cuda_device() -> candle_core::Device { - candle_core::Device::new_cuda(0).expect("CUDA device required") + fn cuda_stream() -> CudaStream { + let dev = candle_core::Device::new_cuda(0).expect("CUDA device required"); + match dev { + candle_core::Device::Cuda(d) => d.cuda_stream().clone(), + _ => panic!("expected CUDA device"), + } } // ── ppo_to_exposure_scores ────────────────────────────────────────── #[test] fn test_ppo_to_exposure_scores_shape() { + let stream = cuda_stream(); let batch = 4; let uniform = vec![1.0_f32 / 45.0; batch * 45]; - let probs = Tensor::from_vec(uniform, &[batch, 45], &cuda_device()).unwrap(); - let scores = ppo_to_exposure_scores(&probs).unwrap(); - assert_eq!(scores.dims(), &[batch, 5]); + let mut probs_buf = stream.alloc_zeros::(batch * 45).unwrap(); + stream.memcpy_htod(&uniform, &mut probs_buf).unwrap(); - // GPU-side: all values should be ~0.2 (each of 5 bins sums 9 cells of 1/45) - let expected = Tensor::new(0.2_f32, scores.device()).unwrap(); - let max_diff = scores.broadcast_sub(&expected).unwrap().abs().unwrap() - .max(0).unwrap().max(0).unwrap().to_scalar::().unwrap(); - assert!(max_diff < 1e-5, "expected ~0.2 everywhere, max_diff={max_diff}"); + let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap(); + + // Download and check shape + values + let mut host = vec![0.0_f32; batch * 5]; + stream.memcpy_dtoh(&scores, &mut host).unwrap(); + stream.synchronize().unwrap(); + + assert_eq!(host.len(), batch * 5); + // Each of 5 bins sums 9 cells of 1/45 = 0.2 + for val in &host { + assert!( + (*val - 0.2).abs() < 1e-5, + "expected ~0.2 everywhere, got {val}" + ); + } } #[test] fn test_ppo_to_exposure_scores_argmax() { - // Put all mass on actions 36..44 → exposure bin 4 (Long100) + let stream = cuda_stream(); + // Put all mass on actions 36..44 => exposure bin 4 (Long100) let batch = 2; let mut raw = vec![0.0_f32; batch * 45]; for b in 0..batch { - // actions 36..44 = exposure bucket 4, each 1/9 for a in 36..45 { if let Some(slot) = raw.get_mut(b * 45 + a) { *slot = 1.0 / 9.0; } } } - let probs = Tensor::from_vec(raw, &[batch, 45], &cuda_device()).unwrap(); - let scores = ppo_to_exposure_scores(&probs).unwrap(); - // GPU-side argmax: all rows should have argmax at index 4 (Long100) - let argmaxes = scores.argmax(1).unwrap(); - let expected_4 = Tensor::new(&[4_u32, 4_u32], scores.device()).unwrap(); - let all_match = argmaxes.eq(&expected_4).unwrap() - .to_dtype(DType::U32).unwrap() - .sum_all().unwrap().to_scalar::().unwrap(); - assert_eq!(all_match, batch as u32, "all rows should argmax to Long100 (4)"); - } + let mut probs_buf = stream.alloc_zeros::(batch * 45).unwrap(); + stream.memcpy_htod(&raw, &mut probs_buf).unwrap(); - #[test] - fn test_ppo_to_exposure_scores_rejects_wrong_shape() { - let probs = Tensor::zeros(&[8, 5], DType::F32, &cuda_device()).unwrap(); - let err = ppo_to_exposure_scores(&probs); - assert!(err.is_err(), "should reject [8,5]"); + let scores = ppo_to_exposure_scores(&probs_buf, batch, &stream).unwrap(); + let mut host = vec![0.0_f32; batch * 5]; + stream.memcpy_dtoh(&scores, &mut host).unwrap(); + stream.synchronize().unwrap(); + + // Each row should have argmax at index 4 (Long100) + for b in 0..batch { + let row_start = b * 5; + let mut best_idx = 0; + let mut best_val = f32::NEG_INFINITY; + for i in 0..5 { + let v = host[row_start + i]; + if v > best_val { + best_val = v; + best_idx = i; + } + } + assert_eq!(best_idx, 4, "batch {b}: expected argmax=4, got {best_idx}"); + } } // ── signal_to_action_scores ───────────────────────────────────────── - /// GPU-side argmax for a single-row tensor. - fn gpu_argmax(t: &Tensor) -> u32 { - t.argmax(1).unwrap() - .flatten_all().unwrap() - .squeeze(0).unwrap() - .to_scalar::().unwrap() + /// Host-side argmax over a [1, 5] row downloaded from GPU. + fn host_argmax(host: &[f32]) -> usize { + let mut best_idx = 0; + let mut best_val = f32::NEG_INFINITY; + for (i, &v) in host.iter().enumerate() { + if v > best_val { + best_val = v; + best_idx = i; + } + } + best_idx + } + + fn run_signal_test(pred_val: f32, high: f32, low: f32) -> Vec { + let stream = cuda_stream(); + let mut pred_buf = stream.alloc_zeros::(1).unwrap(); + stream.memcpy_htod(&[pred_val], &mut pred_buf).unwrap(); + let scores = signal_to_action_scores(&pred_buf, 1, high, low, &stream).unwrap(); + let mut host = vec![0.0_f32; 5]; + stream.memcpy_dtoh(&scores, &mut host).unwrap(); + stream.synchronize().unwrap(); + host } #[test] fn test_signal_to_action_scores_strong_long() { - let pred = Tensor::new(&[20.0_f32], &cuda_device()).unwrap(); - let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(scores.dims(), &[1, 5]); - assert_eq!(gpu_argmax(&scores), 4); // Long100 + let host = run_signal_test(20.0, 10.0, 5.0); + assert_eq!(host_argmax(&host), 4); // Long100 } #[test] fn test_signal_to_action_scores_strong_short() { - let pred = Tensor::new(&[-20.0_f32], &cuda_device()).unwrap(); - let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(gpu_argmax(&scores), 0); // Short100 + let host = run_signal_test(-20.0, 10.0, 5.0); + assert_eq!(host_argmax(&host), 0); // Short100 } #[test] fn test_signal_to_action_scores_flat() { - let pred = Tensor::new(&[0.0_f32], &cuda_device()).unwrap(); - let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(gpu_argmax(&scores), 2); // Flat + let host = run_signal_test(0.0, 10.0, 5.0); + assert_eq!(host_argmax(&host), 2); // Flat } #[test] fn test_signal_to_action_scores_mild_long() { - let pred = Tensor::new(&[7.0_f32], &cuda_device()).unwrap(); - let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(gpu_argmax(&scores), 3); // Long50 + let host = run_signal_test(7.0, 10.0, 5.0); + assert_eq!(host_argmax(&host), 3); // Long50 } #[test] fn test_signal_to_action_scores_mild_short() { - let pred = Tensor::new(&[-7.0_f32], &cuda_device()).unwrap(); - let scores = signal_to_action_scores(&pred, 10.0, 5.0).unwrap(); - assert_eq!(gpu_argmax(&scores), 1); // Short50 + let host = run_signal_test(-7.0, 10.0, 5.0); + assert_eq!(host_argmax(&host), 1); // Short50 } // ── tft_quantile_to_signal ────────────────────────────────────────── #[test] fn test_tft_quantile_to_signal() { + let stream = cuda_stream(); // [batch=2, horizon=1, quantiles=3] - // quantiles: [q10, median, q90] let data: Vec = vec![ - -1.0, 0.5, 2.0, // batch 0: median = 0.5 - -3.0, -1.5, 0.0, // batch 1: median = -1.5 + -1.0, 0.5, 2.0, // batch 0: median = 0.5 + -3.0, -1.5, 0.0, // batch 1: median = -1.5 ]; - let q = Tensor::from_vec(data, &[2, 1, 3], &cuda_device()).unwrap(); - let signal = tft_quantile_to_signal(&q).unwrap(); - assert_eq!(signal.dims(), &[2]); + let mut q_buf = stream.alloc_zeros::(data.len()).unwrap(); + stream.memcpy_htod(&data, &mut q_buf).unwrap(); - // GPU-side comparison: signal should equal [0.5, -1.5] - let expected = Tensor::new(&[0.5_f32, -1.5], signal.device()).unwrap(); - let max_diff = signal.sub(&expected).unwrap().abs().unwrap() - .max(0).unwrap().to_scalar::().unwrap(); - assert!(max_diff < 1e-6, "tft signal mismatch: max_diff={max_diff}"); + let signal = tft_quantile_to_signal(&q_buf, 2, 1, 3, &stream).unwrap(); + let mut host = vec![0.0_f32; 2]; + stream.memcpy_dtoh(&signal, &mut host).unwrap(); + stream.synchronize().unwrap(); + + assert!( + (host[0] - 0.5).abs() < 1e-6, + "batch 0: expected 0.5, got {}", + host[0] + ); + assert!( + (host[1] - (-1.5)).abs() < 1e-6, + "batch 1: expected -1.5, got {}", + host[1] + ); + } + + #[test] + fn test_tft_quantile_rejects_bad_dims() { + let stream = cuda_stream(); + let mut q_buf = stream.alloc_zeros::(6).unwrap(); + stream.memcpy_htod(&[0.0_f32; 6], &mut q_buf).unwrap(); + + // horizon=0 should fail + let err = tft_quantile_to_signal(&q_buf, 2, 0, 3, &stream); + assert!(err.is_err(), "should reject horizon=0"); + + // num_quantiles=1 should fail + let err = tft_quantile_to_signal(&q_buf, 2, 1, 1, &stream); + assert!(err.is_err(), "should reject num_quantiles=1"); } // ── backtest_fitness ──────────────────────────────────────────────── diff --git a/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu b/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu new file mode 100644 index 000000000..da7939414 --- /dev/null +++ b/crates/ml/src/cuda_pipeline/signal_adapter_kernel.cu @@ -0,0 +1,95 @@ +/** + * Signal adapter CUDA kernels for converting model outputs to trading signals. + * + * Three fused kernels replacing Candle tensor operations: + * + * 1. ppo_to_exposure_scores_kernel: + * Aggregates PPO 45-action softmax probs into 5 exposure scores. + * Input: [batch * 45] float probs + * Output: [batch * 5] float scores (sum of 9 order/urgency combos per bucket) + * + * 2. signal_to_action_scores_kernel: + * Maps scalar predictions to [batch, 5] one-hot-ish action scores via thresholds. + * Input: [batch] float predictions + * Output: [batch * 5] float scores (one-hot based on threshold buckets) + * + * 3. tft_quantile_extract_kernel: + * Extracts median signal from TFT quantile predictions. + * Input: [batch * horizon * num_quantiles] float + * Output: [batch] float median values (quantile index 1, horizon index 0) + * + * None of these kernels require common_device_functions.cuh (standalone). + * Launch config: grid=(ceil(batch/256), 1, 1), block=(256, 1, 1). + * One thread per batch element. + */ + +extern "C" __global__ void ppo_to_exposure_scores_kernel( + const float* __restrict__ probs, /* [batch * 45] softmax probs */ + float* __restrict__ out_scores, /* [batch * 5] exposure scores */ + int batch +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch) return; + + /* For each of 5 exposure buckets, sum over 9 order/urgency combos. + * probs layout: [batch, 45] where 45 = 5 exposure * 9 (3 order * 3 urgency) + * bucket e covers indices [e*9 .. e*9+8] within each batch row. */ + int base = b * 45; + for (int e = 0; e < 5; ++e) { + float sum = 0.0f; + int offset = base + e * 9; + for (int j = 0; j < 9; ++j) { + sum += probs[offset + j]; + } + out_scores[b * 5 + e] = sum; + } +} + +extern "C" __global__ void signal_to_action_scores_kernel( + const float* __restrict__ predictions, /* [batch] scalar signals */ + float* __restrict__ out_scores, /* [batch * 5] one-hot scores */ + float high_threshold_bps, + float low_threshold_bps, + int batch +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch) return; + + float pred = predictions[b]; + int out_base = b * 5; + + /* Zero all 5 scores, then set the matching bucket to 1.0 */ + out_scores[out_base + 0] = 0.0f; + out_scores[out_base + 1] = 0.0f; + out_scores[out_base + 2] = 0.0f; + out_scores[out_base + 3] = 0.0f; + out_scores[out_base + 4] = 0.0f; + + if (pred < -high_threshold_bps) { + out_scores[out_base + 0] = 1.0f; /* Short100 */ + } else if (pred < -low_threshold_bps) { + out_scores[out_base + 1] = 1.0f; /* Short50 */ + } else if (pred > high_threshold_bps) { + out_scores[out_base + 4] = 1.0f; /* Long100 */ + } else if (pred > low_threshold_bps) { + out_scores[out_base + 3] = 1.0f; /* Long50 */ + } else { + out_scores[out_base + 2] = 1.0f; /* Flat */ + } +} + +extern "C" __global__ void tft_quantile_extract_kernel( + const float* __restrict__ quantiles, /* [batch * horizon * num_q] */ + float* __restrict__ out_signal, /* [batch] median values */ + int horizon, + int num_quantiles, + int batch +) { + int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= batch) return; + + /* Extract quantile index 1 (median) from horizon step 0. + * Layout: quantiles[b * horizon * num_q + 0 * num_q + 1] */ + int idx = b * horizon * num_quantiles + 1; /* horizon=0, quantile=1 */ + out_signal[b] = quantiles[idx]; +} diff --git a/crates/ml/src/hyperopt/adapters/ppo.rs b/crates/ml/src/hyperopt/adapters/ppo.rs index 7ad31f4c9..f1b8590d9 100644 --- a/crates/ml/src/hyperopt/adapters/ppo.rs +++ b/crates/ml/src/hyperopt/adapters/ppo.rs @@ -52,7 +52,9 @@ use crate::ppo::trajectory_replay::TrajectoryReplayBuffer; use crate::MLError; use crate::cuda_pipeline::gpu_backtest_evaluator::{GpuBacktestConfig, GpuBacktestEvaluator}; +use crate::cuda_pipeline::gpu_action_selector::cuda_f32_to_tensor; use crate::cuda_pipeline::signal_adapter::ppo_to_exposure_scores; +use candle_core::cuda_backend::cudarc::driver::CudaSlice; /// Pure model VRAM in MB (actor + critic + optimizers + gradients). const MODEL_OVERHEAD_MB: f64 = 300.0; @@ -1324,11 +1326,26 @@ impl PPOTrainer { device, )?; - // Forward function: PPO actor → softmax → ppo_to_exposure_scores → [B, 5] + // Forward function: PPO actor -> softmax -> ppo_to_exposure_scores -> [B, 5] + // Bridge: extract CudaSlice -> pure cudarc kernel -> DtoD copy back to Tensor. let metrics = evaluator.evaluate( &|states: &candle_core::Tensor| -> Result { let probs = ppo.actor.action_probabilities(states)?; - ppo_to_exposure_scores(&probs) + let batch = probs.dims().first().copied().unwrap_or(1); + let (storage_guard, _layout) = probs.storage_and_layout(); + let cuda_slice: &CudaSlice = match &*storage_guard { + candle_core::Storage::Cuda(cs) => cs.as_cuda_slice() + .map_err(|e| MLError::ModelError(format!("probs as_cuda_slice: {e}")))?, + _ => return Err(MLError::ModelError("probs not on CUDA".into())), + }; + let cuda_dev = match device { + Device::Cuda(ref d) => d, + _ => return Err(MLError::ModelError("device is not CUDA".into())), + }; + let stream = cuda_dev.cuda_stream(); + let scores_slice = ppo_to_exposure_scores(cuda_slice, batch, &stream)?; + drop(storage_guard); + cuda_f32_to_tensor(&scores_slice, &[batch, 5], device) }, 3, // portfolio_dim device,