huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-20 20:09:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update for the latest cudarc.

Browse Source
This commit is contained in:
laurent
2025-04-11 14:02:41 +02:00
parent c87f0fa5d6
commit 543b5b5898
2 changed files with 15 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-core/Cargo.toml
View File

@ -42,7 +42,7 @@ clap = { workspace = true }
criterion = { workspace = true } criterion = { workspace = true }
[features] [features]
default = [] default = ["cuda"]
cuda = ["cudarc", "dep:candle-kernels", "dep:ug-cuda"] cuda = ["cudarc", "dep:candle-kernels", "dep:ug-cuda"]
cudnn = ["cuda", "cudarc/cudnn"] cudnn = ["cuda", "cudarc/cudnn"]
mkl = ["dep:libc", "dep:intel-mkl-src"] mkl = ["dep:libc", "dep:intel-mkl-src"]

41
candle-core/examples/cuda_basics.rs
View File

@ -15,7 +15,7 @@ fn cuda_graph() -> Result<()> {
Device::Cuda(dev) => dev, Device::Cuda(dev) => dev,
_ => unreachable!(), _ => unreachable!(),
}; };
let cu_stream = cu_device.cu_stream(); let cu_stream = cu_device.cuda_stream();
{ {
// load_ptx cannot be called while capturing the stream so we need this to happen // load_ptx cannot be called while capturing the stream so we need this to happen
// beforehand. // beforehand.
@ -31,14 +31,9 @@ fn cuda_graph() -> Result<()> {
device.synchronize()?; device.synchronize()?;
} }
if USE_CUDA_GRAPH { if USE_CUDA_GRAPH {
unsafe { cu_stream.begin_capture(
cudarc::driver::sys::lib() cudarc::driver::sys::CUstreamCaptureMode::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
.cuStreamBeginCapture_v2( )?;
*cu_stream,
cudarc::driver::sys::CUstreamCaptureMode_enum::CU_STREAM_CAPTURE_MODE_THREAD_LOCAL,
)
.result()?
}
} }
{ {
let u = Tensor::zeros((4096, 4096), candle_core::DType::F32, &device)? let u = Tensor::zeros((4096, 4096), candle_core::DType::F32, &device)?
@ -55,28 +50,17 @@ fn cuda_graph() -> Result<()> {
} }
} }
if USE_CUDA_GRAPH { if USE_CUDA_GRAPH {
let cu_graph: cudarc::driver::sys::CUgraph = unsafe { println!("capturing graph");
let mut cu_graph = std::mem::MaybeUninit::uninit(); let cu_graph = match cu_stream.end_capture(
cudarc::driver::sys::lib() cudarc::driver::sys::CUgraphInstantiate_flags::CUDA_GRAPH_INSTANTIATE_FLAG_USE_NODE_PRIORITY
.cuStreamEndCapture(*cu_stream, cu_graph.as_mut_ptr()) )? {
.result()?; None => anyhow::bail!("no graph captured"),
cu_graph.assume_init() Some(cu_graph) => cu_graph,
};
let cu_graph_e: cudarc::driver::sys::CUgraphExec = unsafe {
let mut cu_graph_e = std::mem::MaybeUninit::uninit();
cudarc::driver::sys::lib()
.cuGraphInstantiateWithFlags(cu_graph_e.as_mut_ptr(), cu_graph, 0)
.result()?;
cu_graph_e.assume_init()
}; };
println!("graph captured!"); println!("graph captured!");
for i in 1..100 { for i in 1..100 {
println!("graph exec {i}"); println!("graph exec {i}");
unsafe { cu_graph.launch()?;
cudarc::driver::sys::lib()
.cuGraphLaunch(cu_graph_e, *cu_stream)
.result()?
}
println!("sync"); println!("sync");
if let Err(err) = device.synchronize() { if let Err(err) = device.synchronize() {
println!("err: {err:?}") println!("err: {err:?}")
@ -92,6 +76,9 @@ fn cuda_graph() -> Result<()> {
fn main() -> Result<()> { fn main() -> Result<()> {
cuda_graph()?; cuda_graph()?;
return Ok(()); return Ok(());
}
fn _matmul() -> Result<()> {
let device = Device::new_cuda_with_stream(0)?; let device = Device::new_cuda_with_stream(0)?;
let x = Tensor::randn(0f32, 1.0, (8 * 4096, 8 * 4096), &device)? let x = Tensor::randn(0f32, 1.0, (8 * 4096, 8 * 4096), &device)?
.to_dtype(candle_core::DType::BF16)?; .to_dtype(candle_core::DType::BF16)?;