Add support for cuda streams. (#2532)

huggingface · Oct 2, 2024 · 7b60bda · 7b60bda
1 parent 9363006
commit 7b60bda
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 0 deletions.
diff --git a/candle-core/src/cuda_backend/device.rs b/candle-core/src/cuda_backend/device.rs
@@ -144,6 +144,20 @@ impl CudaDevice {
  }
 }
 
+impl CudaDevice {
+ pub fn new_with_stream(ordinal: usize) -> Result<Self> {
+ let device = cudarc::driver::CudaDevice::new_with_stream(ordinal).w()?;
+ let blas = cudarc::cublas::CudaBlas::new(device.clone()).w()?;
+ let curand = cudarc::curand::CudaRng::new(299792458, device.clone()).w()?;
+ Ok(Self {
+ id: DeviceId::new(),
+ device,
+ blas: Arc::new(blas),
+ curand: Arc::new(Mutex::new(CudaRng(curand))),
+ })
+ }
+}
+
 impl BackendDevice for CudaDevice {
  type Storage = CudaStorage;
 

diff --git a/candle-core/src/device.rs b/candle-core/src/device.rs
@@ -130,6 +130,10 @@ impl Device {
  Ok(Self::Cuda(crate::CudaDevice::new(ordinal)?))
  }
 
+ pub fn new_cuda_with_stream(ordinal: usize) -> Result<Self> {
+ Ok(Self::Cuda(crate::CudaDevice::new_with_stream(ordinal)?))
+ }
+
  pub fn new_metal(ordinal: usize) -> Result<Self> {
  Ok(Self::Metal(crate::MetalDevice::new(ordinal)?))
  }

diff --git a/candle-core/src/dummy_cuda_backend.rs b/candle-core/src/dummy_cuda_backend.rs
@@ -14,6 +14,12 @@ macro_rules! fail {
  };
 }
 
+impl CudaDevice {
+ pub fn new_with_stream(_: usize) -> Result<Self> {
+ Err(Error::NotCompiledWithCudaSupport)
+ }
+}
+
 impl crate::backend::BackendStorage for CudaStorage {
  type Device = CudaDevice;