From da2aff98c654431fb001b6fb9ed36cdd0fd51ec3 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 17 Mar 2026 14:35:52 +0100 Subject: [PATCH] feat(cuda): add CudaPolicyNetwork and CudaValueNetwork to ml-ppo cuda_nn GPU-native MLP networks that use CudaLinear layers + cuBLAS sgemm internally. These are drop-in replacements for the Candle-based PolicyNetwork and ValueNetwork, providing the same forward/softmax/ log_softmax API surface but with zero Candle overhead. Co-Authored-By: Claude Opus 4.6 (1M context) --- crates/ml-ppo/src/cuda_nn/mod.rs | 2 + crates/ml-ppo/src/cuda_nn/networks.rs | 225 ++++++++++++++++++++++++++ 2 files changed, 227 insertions(+) create mode 100644 crates/ml-ppo/src/cuda_nn/networks.rs diff --git a/crates/ml-ppo/src/cuda_nn/mod.rs b/crates/ml-ppo/src/cuda_nn/mod.rs index a71585f2d..3ce3c9e5f 100644 --- a/crates/ml-ppo/src/cuda_nn/mod.rs +++ b/crates/ml-ppo/src/cuda_nn/mod.rs @@ -11,6 +11,7 @@ pub mod linear; pub mod activations; pub mod adam; pub mod lstm; +pub mod networks; pub mod softmax; pub mod tensor_util; @@ -18,6 +19,7 @@ pub use linear::CudaLinear; pub use activations::{cuda_relu, cuda_relu_backward, cuda_tanh, cuda_sigmoid}; pub use adam::CudaAdam; pub use lstm::CudaLSTM; +pub use networks::{CudaPolicyNetwork, CudaValueNetwork}; pub use softmax::{cuda_softmax, cuda_log_softmax}; pub use tensor_util::{CudaVec, cuda_zeros, cuda_ones, cuda_from_slice}; diff --git a/crates/ml-ppo/src/cuda_nn/networks.rs b/crates/ml-ppo/src/cuda_nn/networks.rs new file mode 100644 index 000000000..aa2986fa9 --- /dev/null +++ b/crates/ml-ppo/src/cuda_nn/networks.rs @@ -0,0 +1,225 @@ +//! GPU-native policy and value networks using cuBLAS-backed linear layers. +//! +//! Drop-in replacements for the Candle-based `PolicyNetwork` and `ValueNetwork` +//! in `ppo.rs`. All computation stays on GPU via cuBLAS sgemm + CUDA kernels. + +use candle_core::cuda_backend::cudarc; +use cudarc::driver::CudaSlice; + +use ml_core::MLError; + +use super::linear::CudaLinear; +use super::activations::cuda_relu; +use super::softmax::{cuda_softmax, cuda_log_softmax}; +use super::tensor_util::{CudaVec, cuda_from_slice}; +use super::GpuContext; + +// --------------------------------------------------------------------------- +// CudaPolicyNetwork +// --------------------------------------------------------------------------- + +/// GPU-native policy network: `state -> action logits` +/// +/// Architecture: `Linear -> ReLU -> Linear -> ReLU -> ... -> Linear (logits)` +/// All layers use cuBLAS sgemm internally. No Candle tensors. +pub struct CudaPolicyNetwork { + layers: Vec, + ctx: GpuContext, +} + +impl std::fmt::Debug for CudaPolicyNetwork { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaPolicyNetwork") + .field("num_layers", &self.layers.len()) + .finish() + } +} + +impl CudaPolicyNetwork { + /// Create a new policy network. + /// + /// # Arguments + /// * `input_dim` - State dimension + /// * `hidden_dims` - Hidden layer dimensions + /// * `output_dim` - Number of actions + /// * `ctx` - Shared GPU context + pub fn new( + input_dim: usize, + hidden_dims: &[usize], + output_dim: usize, + ctx: GpuContext, + ) -> Result { + let mut layers = Vec::new(); + let mut current_dim = input_dim; + + for &hidden_dim in hidden_dims { + layers.push(CudaLinear::new(ctx.clone(), current_dim, hidden_dim)?); + current_dim = hidden_dim; + } + + // Output layer (logits, no activation) + layers.push(CudaLinear::new(ctx.clone(), current_dim, output_dim)?); + + Ok(Self { layers, ctx }) + } + + /// Forward pass: state -> logits + /// + /// Input: `[batch_size * input_dim]` flattened row-major + /// Output: `CudaVec` of `[batch_size * num_actions]` + pub fn forward(&self, input: &CudaSlice, batch_size: usize) -> Result { + let mut x = self.layers.first() + .ok_or_else(|| MLError::ConfigError("Policy network has no layers".to_owned()))? + .forward(input, batch_size)?; + + for layer in self.layers.iter().skip(1) { + // ReLU activation (not on last layer) + let is_last = std::ptr::eq(layer, self.layers.last().expect("checked above")); + if !is_last { + // Apply ReLU to previous layer output, then run next linear + let relu_out = cuda_relu(&self.ctx.stream, &x.data, x.len)?; + x = layer.forward(&relu_out.data, batch_size)?; + } else { + // Apply ReLU to the penultimate output, then pass to final layer + let relu_out = cuda_relu(&self.ctx.stream, &x.data, x.len)?; + x = layer.forward(&relu_out.data, batch_size)?; + } + } + + Ok(x) + } + + /// Compute action probabilities: softmax(logits) + pub fn action_probabilities(&self, input: &CudaSlice, batch_size: usize, num_actions: usize) -> Result { + let logits = self.forward(input, batch_size)?; + cuda_softmax(&self.ctx.stream, &logits.data, batch_size, num_actions) + } + + /// Compute log-softmax of logits (for log-probability computation) + pub fn log_softmax(&self, input: &CudaSlice, batch_size: usize, num_actions: usize) -> Result { + let logits = self.forward(input, batch_size)?; + cuda_log_softmax(&self.ctx.stream, &logits.data, batch_size, num_actions) + } + + /// Get all layers for optimizer registration. + pub fn layers_mut(&mut self) -> &mut [CudaLinear] { + &mut self.layers + } + + /// Get all layers (immutable). + pub fn layers(&self) -> &[CudaLinear] { + &self.layers + } + + /// Get the GPU context. + pub fn ctx(&self) -> &GpuContext { + &self.ctx + } +} + +// --------------------------------------------------------------------------- +// CudaValueNetwork +// --------------------------------------------------------------------------- + +/// GPU-native value network: `state -> scalar value` +/// +/// Architecture: `Linear -> ReLU -> Linear -> ReLU -> ... -> Linear (1)` +pub struct CudaValueNetwork { + layers: Vec, + ctx: GpuContext, +} + +impl std::fmt::Debug for CudaValueNetwork { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaValueNetwork") + .field("num_layers", &self.layers.len()) + .finish() + } +} + +impl CudaValueNetwork { + /// Create a new value network. + /// + /// # Arguments + /// * `input_dim` - State dimension + /// * `hidden_dims` - Hidden layer dimensions + /// * `ctx` - Shared GPU context + pub fn new( + input_dim: usize, + hidden_dims: &[usize], + ctx: GpuContext, + ) -> Result { + let mut layers = Vec::new(); + let mut current_dim = input_dim; + + for &hidden_dim in hidden_dims { + layers.push(CudaLinear::new(ctx.clone(), current_dim, hidden_dim)?); + current_dim = hidden_dim; + } + + // Output layer: single value + layers.push(CudaLinear::new(ctx.clone(), current_dim, 1)?); + + Ok(Self { layers, ctx }) + } + + /// Forward pass: state -> value `[batch_size]` + /// + /// Input: `[batch_size * input_dim]` flattened row-major + /// Output: `CudaVec` of `[batch_size]` (one value per sample) + pub fn forward(&self, input: &CudaSlice, batch_size: usize) -> Result { + let mut x = self.layers.first() + .ok_or_else(|| MLError::ConfigError("Value network has no layers".to_owned()))? + .forward(input, batch_size)?; + + for layer in self.layers.iter().skip(1) { + let relu_out = cuda_relu(&self.ctx.stream, &x.data, x.len)?; + x = layer.forward(&relu_out.data, batch_size)?; + } + + // Output is [batch_size, 1], the CudaVec has batch_size elements + // since CudaLinear outputs [batch * 1] + Ok(x) + } + + /// Get all layers for optimizer registration. + pub fn layers_mut(&mut self) -> &mut [CudaLinear] { + &mut self.layers + } + + /// Get all layers (immutable). + pub fn layers(&self) -> &[CudaLinear] { + &self.layers + } + + /// Get the GPU context. + pub fn ctx(&self) -> &GpuContext { + &self.ctx + } +} + +// --------------------------------------------------------------------------- +// Conversion utilities: Candle Tensor <-> CudaSlice +// --------------------------------------------------------------------------- + +/// Upload a flat f32 slice to GPU as CudaVec. +/// +/// This is the bridge between host-side trajectory data (Vec) and +/// GPU-native network inputs. +pub fn states_to_gpu( + ctx: &GpuContext, + states_flat: &[f32], +) -> Result { + cuda_from_slice(&ctx.stream, states_flat) +} + +/// Download a CudaVec to host Vec. +/// +/// Used at API boundaries where callers need host-side results +/// (e.g., action selection during trajectory collection). +pub fn gpu_to_host( + ctx: &GpuContext, + gpu_vec: &CudaVec, +) -> Result, MLError> { + gpu_vec.to_vec(&ctx.stream) +}