From 550944ebf2ee14450bcf159c97e1e9c33674c814 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Mon, 2 Mar 2026 17:52:45 +0100 Subject: [PATCH] feat(ml): add CudaStreamPool for multi-stream ensemble inference CUDA stream pool with CPU no-op fallback. Foundation for StreamAwareEnsemble that runs models on separate CUDA streams. Co-Authored-By: Claude Opus 4.6 --- crates/ml/src/ensemble/cuda_streams.rs | 200 +++++++++++++++++++++++++ crates/ml/src/ensemble/mod.rs | 1 + 2 files changed, 201 insertions(+) create mode 100644 crates/ml/src/ensemble/cuda_streams.rs diff --git a/crates/ml/src/ensemble/cuda_streams.rs b/crates/ml/src/ensemble/cuda_streams.rs new file mode 100644 index 000000000..2fa814ed2 --- /dev/null +++ b/crates/ml/src/ensemble/cuda_streams.rs @@ -0,0 +1,200 @@ +//! CUDA stream pool for multi-stream ensemble inference. +//! +//! Provides parallel execution of multiple models on different CUDA streams. +//! Falls back to no-op on CPU devices. +//! +//! On CUDA devices, each stream can execute kernels concurrently with other streams, +//! enabling multiple model forward passes to overlap. This is the foundation for +//! `StreamAwareEnsemble` (Task 14) which assigns each model to a separate stream. + +use candle_core::Device; +use crate::MLError; + +#[cfg(feature = "cuda")] +use std::sync::Arc; + +#[cfg(feature = "cuda")] +type CudaStream = candle_core::cuda_backend::cudarc::driver::CudaStream; + +/// Pool of CUDA streams for parallel model inference. +/// +/// On CUDA devices, creates multiple streams via `CudaStream::fork()`. +/// On CPU, acts as a no-op (all methods succeed without doing anything). +/// +/// # Usage +/// +/// ```ignore +/// let pool = CudaStreamPool::new(&device, 4)?; +/// // Each model gets assigned stream index = model_index % pool.count() +/// // Sync all streams after inference: +/// pool.sync_all()?; +/// ``` +pub struct CudaStreamPool { + /// Number of streams (0 for CPU) + count: usize, + /// Whether this pool is on a CUDA device + is_cuda: bool, + /// The device this pool is associated with + device: Device, + /// Forked CUDA streams (empty on CPU) + #[cfg(feature = "cuda")] + streams: Vec>, +} + +impl std::fmt::Debug for CudaStreamPool { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CudaStreamPool") + .field("count", &self.count) + .field("is_cuda", &self.is_cuda) + .field("device", &self.device) + .finish() + } +} + +impl CudaStreamPool { + /// Create a new stream pool. + /// + /// On CUDA: creates `count` streams by forking the device's default stream. + /// Each forked stream automatically waits for the default stream's current + /// work to complete before starting. + /// + /// On CPU: creates a no-op pool that passes through all operations. + pub fn new(device: &Device, count: usize) -> Result { + #[cfg(feature = "cuda")] + { + if let Device::Cuda(cuda_dev) = device { + let default_stream = cuda_dev.cuda_stream(); + let mut streams = Vec::with_capacity(count); + for i in 0..count { + let stream = default_stream.fork().map_err(|e| { + MLError::DeviceError(format!( + "Failed to fork CUDA stream {i}/{count}: {e}" + )) + })?; + streams.push(stream); + } + + return Ok(Self { + count, + is_cuda: true, + device: device.clone(), + streams, + }); + } + } + + // CPU or non-cuda build: no-op pool + Ok(Self { + count: 0, + is_cuda: false, + device: device.clone(), + #[cfg(feature = "cuda")] + streams: Vec::new(), + }) + } + + /// Number of streams in the pool. + pub fn count(&self) -> usize { + self.count + } + + /// Whether this pool has CUDA streams. + pub fn is_cuda(&self) -> bool { + self.is_cuda + } + + /// Get the device associated with this pool. + pub fn device(&self) -> &Device { + &self.device + } + + /// Get a CUDA stream by index. + /// + /// Returns `None` on CPU or if index is out of bounds. + /// Typical usage: `pool.get_stream(model_index % pool.count())`. + #[cfg(feature = "cuda")] + pub fn get_stream(&self, index: usize) -> Option<&Arc> { + self.streams.get(index) + } + + /// Synchronize all streams (wait for completion). + /// + /// On CPU: no-op. + /// On CUDA: synchronizes each forked stream individually, ensuring all + /// concurrent model inference is complete before returning. + pub fn sync_all(&self) -> Result<(), MLError> { + #[cfg(feature = "cuda")] + { + for (i, stream) in self.streams.iter().enumerate() { + stream.synchronize().map_err(|e| { + MLError::DeviceError(format!( + "CUDA stream {i} sync failed: {e}" + )) + })?; + } + } + Ok(()) + } + + /// Synchronize a single stream by index. + /// + /// On CPU: no-op. + /// On CUDA: synchronizes the specified stream. + /// Returns `Ok(())` if index is out of bounds (no-op for safety). + pub fn sync_stream(&self, index: usize) -> Result<(), MLError> { + #[cfg(feature = "cuda")] + { + if let Some(stream) = self.streams.get(index) { + stream.synchronize().map_err(|e| { + MLError::DeviceError(format!( + "CUDA stream {index} sync failed: {e}" + )) + })?; + } + } + let _ = index; // suppress unused warning on non-cuda builds + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_stream_pool_cpu_noop() { + let pool = CudaStreamPool::new(&Device::Cpu, 4) + .expect("CPU pool should always succeed"); + assert_eq!(pool.count(), 0); + assert!(!pool.is_cuda()); + assert!(matches!(pool.device(), Device::Cpu)); + pool.sync_all().expect("CPU sync should be noop"); + } + + #[test] + fn test_stream_pool_cpu_sync_stream_noop() { + let pool = CudaStreamPool::new(&Device::Cpu, 4) + .expect("CPU pool should always succeed"); + // Sync on CPU is always a no-op, even for out-of-bounds indices + pool.sync_stream(0).expect("sync_stream on CPU should be noop"); + pool.sync_stream(100).expect("out-of-bounds sync_stream should be noop"); + } + + #[test] + fn test_stream_pool_cpu_zero_count() { + let pool = CudaStreamPool::new(&Device::Cpu, 0) + .expect("Zero-count CPU pool should succeed"); + assert_eq!(pool.count(), 0); + assert!(!pool.is_cuda()); + pool.sync_all().expect("sync_all on empty pool should be noop"); + } + + #[cfg(feature = "cuda")] + #[test] + fn test_stream_pool_cpu_no_streams() { + let pool = CudaStreamPool::new(&Device::Cpu, 4) + .expect("CPU pool should always succeed"); + assert!(pool.get_stream(0).is_none()); + assert!(pool.get_stream(3).is_none()); + } +} diff --git a/crates/ml/src/ensemble/mod.rs b/crates/ml/src/ensemble/mod.rs index 90912d3ec..37c64d559 100644 --- a/crates/ml/src/ensemble/mod.rs +++ b/crates/ml/src/ensemble/mod.rs @@ -25,6 +25,7 @@ pub mod conviction_gates; pub mod weight_optimizer; pub mod gate_optimizer; pub mod model_adapter; +pub mod cuda_streams; // Re-export key types that are used across ensemble modules pub use model_adapter::{EnsembleModelAdapter, build_production_strategy};