Add runtime validation on CUDA compute cap for BF16

Author: alvarobartt

backends/candle/src/lib.rs

@@ -211,12 +211,40 @@ impl CandleBackend {
         }
         .map_err(|err| BackendError::Start(err.to_string()))?;
 
-        // Get candle dtype
         let dtype = if &dtype == "float32" {
             Ok(DType::F32)
         } else if &dtype == "float16" {
             Ok(DType::F16)
         } else if &dtype == "bfloat16" {
+            match &device {
+                Device::Cpu => {
+                    return Err(BackendError::Start(
+                        "BFloat16 is not supported on CPU. Use float16 or float32 instead."
+                            .to_string(),
+                    ));
+                }
+                #[cfg(feature = "cuda")]
+                Device::Cuda(_) => {
+                    let compute_cap = get_runtime_compute_cap().map_err(|e| {
+                        BackendError::Start(format!("Failed to get CUDA compute capability: {e:?}"))
+                    })?;
+                    if compute_cap < 80 {
+                        return Err(BackendError::Start(format!(
+                            "BFloat16 requires CUDA compute capability >= 8.0 (Ampere or newer), \
+                             but found {}.{}. Use float16 or float32 instead.",
+                            compute_cap / 10,
+                            compute_cap % 10
+                        )));
+                    }
+                }
+                #[cfg(not(feature = "cuda"))]
+                Device::Cuda(_) => {
+                    return Err(BackendError::Start(
+                        "CUDA feature is not enabled".to_string(),
+                    ));
+                }
+                Device::Metal(_) => (),
+            }
             Ok(DType::BF16)
         } else {
             Err(BackendError::Start(format!(

backends/src/dtype.rs

@@ -13,10 +13,11 @@ pub enum DType {
     Float16,
     #[cfg(any(feature = "python", feature = "candle", feature = "ort"))]
     Float32,
+    // NOTE: For CUDA, BFloat16 requires Ampere (SM 80) or newer, which is validated at runtime, as
+    // there are no specific features for the different CUDA compute capabilities to filter out
+    // Turing and Volta from having `DType::Bfloat16`.
     #[cfg(any(
         feature = "python",
-        // NOTE: It needs to exclude Turing (and Volta) i.e., anything before Ampere, as it won't
-        // have support for BF16
         all(feature = "candle", any(feature = "metal", feature = "cuda"))
     ))]
     Bfloat16,

router/src/main.rs

@@ -34,7 +34,9 @@ struct Args {
     #[clap(long, env)]
     tokenization_workers: Option<usize>,
 
-    /// The dtype to be forced upon the model.
+    /// The dtype to be forced upon the model, otherwise the value set in `dtype` (or `torch_dtype`
+    /// as fallback) in the `config.json` file is used. Note that `bfloat16` is not supported on
+    /// CPU, neither for Turing on CUDA, but only from Ampere onwards.
     #[clap(long, env, value_enum)]
     dtype: Option<DType>,