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) <noreply@anthropic.com>

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};
 

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<CudaLinear>,
+    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<Self, MLError> {
+        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<f32>, batch_size: usize) -> Result<CudaVec, MLError> {
+        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<f32>, batch_size: usize, num_actions: usize) -> Result<CudaVec, MLError> {
+        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<f32>, batch_size: usize, num_actions: usize) -> Result<CudaVec, MLError> {
+        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<CudaLinear>,
+    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<Self, MLError> {
+        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<f32>, batch_size: usize) -> Result<CudaVec, MLError> {
+        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<f32>) and
+/// GPU-native network inputs.
+pub fn states_to_gpu(
+    ctx: &GpuContext,
+    states_flat: &[f32],
+) -> Result<CudaVec, MLError> {
+    cuda_from_slice(&ctx.stream, states_flat)
+}
+
+/// Download a CudaVec to host Vec<f32>.
+///
+/// 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<Vec<f32>, MLError> {
+    gpu_vec.to_vec(&ctx.stream)
+}