From 792c5a722e9ccb0aa62b14529ddb088b5aeb546b Mon Sep 17 00:00:00 2001 From: Hiroshi Shinaoka Date: Sat, 18 Jul 2026 13:25:49 +0900 Subject: [PATCH] fix(cuda): gate CUDA 12.8 features at runtime --- crates/cubecl-cuda/src/compute/server.rs | 24 ++++++ crates/cubecl-cuda/src/runtime.rs | 79 ++++++++++++++++++- .../tests/cudarc_feature_contract.rs | 18 +++++ 3 files changed, 117 insertions(+), 4 deletions(-) diff --git a/crates/cubecl-cuda/src/compute/server.rs b/crates/cubecl-cuda/src/compute/server.rs index b80cc32a75..2d7ee9361f 100644 --- a/crates/cubecl-cuda/src/compute/server.rs +++ b/crates/cubecl-cuda/src/compute/server.rs @@ -8,6 +8,7 @@ use crate::{ stream::CudaStreamBackend, sync::Fence, }, + runtime::CudaApiVersions, }; use cubecl_common::{ backtrace::BackTrace, bytes::Bytes, profile::ProfileDuration, stream_id::StreamId, @@ -50,6 +51,7 @@ pub(crate) const MB: usize = 1024 * 1024; #[derive(Debug)] pub struct CudaServer { ctx: CudaContext, + cuda_api_versions: CudaApiVersions, device_id: DeviceId, streams: MultiStream, utilities: Arc>, @@ -555,6 +557,7 @@ impl CudaServer { /// Create a new cuda server. pub(crate) fn new( ctx: CudaContext, + cuda_api_versions: CudaApiVersions, mem_props: MemoryDeviceProperties, mem_config: MemoryConfiguration, mem_alignment: usize, @@ -573,6 +576,7 @@ impl CudaServer { Self { ctx, + cuda_api_versions, device_id, streams: MultiStream::new( utilities.logger.clone(), @@ -663,6 +667,7 @@ impl CudaServer { .compilation_options .supports_features .grid_constants; + let supports_cuda_12_8 = self.cuda_api_versions.supports_cuda_12_8(); let mut command = self.command( stream_id, bindings.buffers.iter(), @@ -746,6 +751,25 @@ impl CudaServer { .collect(); let elem_stride: Vec<_> = map.elem_stride.iter().rev().map(|s| *s as u32).collect(); + #[cfg(cuda_12080)] + if !supports_cuda_12_8 + && (matches!(&map.format, TensorMapFormat::Im2colWide(_)) + || matches!( + map.swizzle, + TensorMapSwizzle::B128Atom32B + | TensorMapSwizzle::B128Atom32BFlip8B + | TensorMapSwizzle::B128Atom64B + ) + || matches!(map.storage_ty, StorageType::Packed(ty, 2) if ty.size_bits() == 4)) + { + return Err(LaunchError::Unknown { + reason: "CUDA driver and NVRTC 12.8 or newer are required for this tensor map" + .into(), + backtrace: BackTrace::capture(), + } + .into()); + } + match &map.format { // SAFETY: `map_ptr` is a zeroed `MaybeUninit`. `device_ptr` is a // valid device pointer. Shape, strides, tile_size, and elem_stride vectors diff --git a/crates/cubecl-cuda/src/runtime.rs b/crates/cubecl-cuda/src/runtime.rs index a69742d748..a5339d9df4 100644 --- a/crates/cubecl-cuda/src/runtime.rs +++ b/crates/cubecl-cuda/src/runtime.rs @@ -31,9 +31,55 @@ use cubecl_cpp::{ use cubecl_runtime::{ allocator::PitchedMemoryLayoutPolicy, client::ComputeClient, logging::ServerLogger, }; -use cudarc::driver::sys::{CUDA_VERSION, cuDeviceTotalMem_v2}; +use cudarc::driver::sys::{cuDeviceTotalMem_v2, cuDriverGetVersion}; use std::{mem::MaybeUninit, sync::Arc}; +const CUDA_12_8: i32 = 12080; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct CudaApiVersions { + driver: i32, + nvrtc: i32, +} + +impl CudaApiVersions { + const fn new(driver: i32, nvrtc: i32) -> Self { + Self { driver, nvrtc } + } + + const fn effective(self) -> i32 { + if self.driver < self.nvrtc { + self.driver + } else { + self.nvrtc + } + } + + pub(crate) const fn supports_cuda_12_8(self) -> bool { + self.effective() >= CUDA_12_8 + } + + fn detect() -> Self { + let mut driver = 0; + let mut nvrtc_major = 0; + let mut nvrtc_minor = 0; + + // SAFETY: Both functions only write integer version components to the + // valid output pointers supplied here. The CUDA driver is initialized + // immediately before this query in `CudaServer::init`. + unsafe { + cuDriverGetVersion(&mut driver) + .result() + .expect("failed to query CUDA driver version"); + cudarc::nvrtc::sys::nvrtcVersion(&mut nvrtc_major, &mut nvrtc_minor) + .result() + .expect("failed to query NVRTC version"); + } + + Self::new(driver, nvrtc_major * 1000 + nvrtc_minor * 10) + } +} + /// Options configuring the CUDA runtime. #[derive(Default)] pub struct RuntimeOptions { @@ -51,6 +97,7 @@ impl DeviceService for CudaServer { // To get the supported WMMA features, and memory properties, we have to initialize the server immediately. cudarc::driver::result::init().unwrap(); + let cuda_api_versions = CudaApiVersions::detect(); let device_index = device.index as i32; let device_ptr = cudarc::driver::result::device::get(device_index).unwrap(); let arch_major; @@ -217,7 +264,7 @@ impl DeviceService for CudaServer { .matmul .ldmatrix .insert(ElemType::Float(FloatKind::BF16).into()); - comp_opts.supports_features.fast_tanh = CUDA_VERSION >= 12080; + comp_opts.supports_features.fast_tanh = cuda_api_versions.supports_cuda_12_8(); } if arch_version >= 80 { @@ -265,7 +312,7 @@ impl DeviceService for CudaServer { ); } - if arch_version >= 100 { + if arch_version >= 100 && cuda_api_versions.supports_cuda_12_8() { device_props.features.tma.insert(Tma::Im2colWide); // Breaks swizzle so disable for now and fix in a PR specifically for this // if CUDA_VERSION >= 12090 { @@ -296,7 +343,7 @@ impl DeviceService for CudaServer { TypeUsage::Conversion | TypeUsage::Buffer, ); - if CUDA_VERSION >= 12080 { + if cuda_api_versions.supports_cuda_12_8() { device_props.features.tma.insert(Tma::SwizzleAtomicity); } } @@ -320,6 +367,7 @@ impl DeviceService for CudaServer { CudaServer::new( cuda_ctx, + cuda_api_versions, mem_properties, options.memory_config, mem_alignment, @@ -398,3 +446,26 @@ impl Runtime for CudaRuntime { .collect() } } + +#[cfg(test)] +mod tests { + use super::CudaApiVersions; + + #[test] + fn cuda_12_8_capabilities_require_both_driver_and_nvrtc() { + let cases = [ + (12040, 12040, 12040, false), + (12060, 12060, 12060, false), + (12080, 12080, 12080, true), + (13000, 12080, 12080, true), + (12080, 12040, 12040, false), + (12040, 12080, 12040, false), + ]; + + for (driver, nvrtc, effective, supports_cuda_12_8) in cases { + let versions = CudaApiVersions::new(driver, nvrtc); + assert_eq!(versions.effective(), effective); + assert_eq!(versions.supports_cuda_12_8(), supports_cuda_12_8); + } + } +} diff --git a/crates/cubecl-cuda/tests/cudarc_feature_contract.rs b/crates/cubecl-cuda/tests/cudarc_feature_contract.rs index fcdd21e395..ae37eb5f59 100644 --- a/crates/cubecl-cuda/tests/cudarc_feature_contract.rs +++ b/crates/cubecl-cuda/tests/cudarc_feature_contract.rs @@ -20,3 +20,21 @@ fn cudarc_build_dependency_only_enables_build_script_requirements() { r#"[build-dependencies] cudarc = { version = "0.19.0", default-features = false, features = ["std", "driver", "dynamic-loading", "cuda-12080"] }"# )); } + +#[test] +fn cuda_12_8_tensor_map_symbol_is_runtime_guarded() { + let runtime = include_str!("../src/runtime.rs"); + let server = include_str!("../src/compute/server.rs"); + + assert!(runtime.contains("cuDriverGetVersion")); + assert!(runtime.contains("nvrtcVersion")); + + let tensor_map_loop = server + .find("for TensorMapBinding") + .expect("tensor map encoding loop must exist"); + let symbol_call = server[tensor_map_loop..] + .find("cuTensorMapEncodeIm2colWide(") + .expect("Im2colWide symbol call must exist"); + let guarded_prefix = &server[tensor_map_loop..tensor_map_loop + symbol_call]; + assert!(guarded_prefix.contains("supports_cuda_12_8")); +}