From ed91f12dd206fd1a77bc70cf11209593245e253b Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 10 Mar 2026 14:33:33 +0100 Subject: [PATCH] feat(cuda): add batch_route_exposure_to_factored kernel for GPU-resident action routing Adds a standalone CUDA kernel that converts DQN exposure indices (0-4) to factored action indices (0-44) entirely on GPU, reusing the existing route_order() device function from common_device_functions.cuh. Wired into GpuActionSelector as route_exposure_to_factored() method, following the same DtoD copy pattern as the existing select_actions methods. This eliminates a GPU->CPU->GPU roundtrip when post-hoc routing is needed after epsilon_greedy_select. Co-Authored-By: Claude Opus 4.6 --- .../cuda_pipeline/epsilon_greedy_kernel.cu | 32 +++++ .../src/cuda_pipeline/gpu_action_selector.rs | 134 ++++++++++++++++++ 2 files changed, 166 insertions(+) diff --git a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu index 921caf9f3..883c1180f 100644 --- a/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu +++ b/crates/ml/src/cuda_pipeline/epsilon_greedy_kernel.cu @@ -219,3 +219,35 @@ extern "C" __global__ void branching_action_select( rng_states[idx] = rng; actions_out[idx] = factored; } + +/** + * Batch convert exposure indices (0-4) to factored action indices (0-44). + * + * Applies OrderRouter::route() logic on GPU via route_order() device function: + * factored = exposure * 9 + order_type * 3 + urgency + * + * Same arithmetic as the fused epsilon_greedy_routed kernel but as a standalone + * kernel for post-hoc routing of epsilon_greedy_select output. + * + * Launch: grid=(ceil(N/256),1,1), block=(256,1,1). + */ +extern "C" __global__ void batch_route_exposure_to_factored( + const unsigned int* __restrict__ exposure_actions, /* [N] 0-4 */ + unsigned int* __restrict__ factored_actions, /* [N] 0-44 output */ + float spread, + float median_spread, + float volatility, + float median_volatility, + int N +) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= N) return; + + /* Route order: determine order_type and urgency from market conditions */ + int order_type, urgency; + route_order(spread, median_spread, volatility, median_volatility, + &order_type, &urgency); + + unsigned int exp = exposure_actions[i]; + factored_actions[i] = exp * 9u + (unsigned int)order_type * 3u + (unsigned int)urgency; +} diff --git a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs index 34f906546..b308c0a78 100644 --- a/crates/ml/src/cuda_pipeline/gpu_action_selector.rs +++ b/crates/ml/src/cuda_pipeline/gpu_action_selector.rs @@ -44,6 +44,7 @@ pub struct GpuActionSelector { kernel_func: CudaFunction, routed_kernel_func: CudaFunction, branching_kernel_func: CudaFunction, + route_func: CudaFunction, rng_states: CudaSlice, actions_buf: CudaSlice, fill_mask_buf: CudaSlice, @@ -98,6 +99,13 @@ impl GpuActionSelector { .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| { @@ -130,6 +138,7 @@ impl GpuActionSelector { kernel_func, routed_kernel_func, branching_kernel_func, + route_func, rng_states, actions_buf, fill_mask_buf, @@ -603,6 +612,131 @@ impl GpuActionSelector { Ok(out_tensor) } + + /// 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, + }; + 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}")) + })?; + } + + drop(exp_guard); + + // 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(), + )); + } + } + drop(out_guard); + + Ok(out_tensor) + } } impl std::fmt::Debug for GpuActionSelector {