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

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<Arc<CudaStream>>,
+}
+
+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<Self, MLError> {
+        #[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<CudaStream>> {
+        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());
+    }
+}

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