Auto Model Detection, Simplified Args, and Overflow Safeguards

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "candle-vllm"
-version = "0.2.1"
+version = "0.3.0"
 edition = "2021"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

src/backend/cache.rs

@@ -1,5 +1,3 @@
-#[cfg(feature = "cuda")]
-use crate::{openai::responses::APIError, try_api};
 #[cfg(feature = "metal")]
 use candle_core::{
     backend::BackendStorage, CpuStorage, Device, IndexOp, Layout, MetalDevice, MetalStorage,
@@ -9,7 +7,7 @@ use candle_core::{
 use candle_core::{
     cuda_backend::cudarc::driver::{CudaSlice, DevicePtr},
     cuda_backend::CudaStorageSlice,
-    Device, IndexOp, Storage, Tensor,
+    Device, IndexOp, Result, Storage, Tensor,
 };
 #[cfg(feature = "cuda")]
 use kernels::ffi::{copy_blocks_bf16, copy_blocks_f16, copy_blocks_f32};
@@ -22,26 +20,26 @@ pub unsafe fn copy_blocks(
     key_caches: Vec<&mut Tensor>,
     value_caches: Vec<&mut Tensor>,
     block_mapping: HashMap<usize, Vec<usize>>,
-) -> Result<(), APIError> {
+) -> Result<()> {
     use candle_core::DType;
 
     let cache_dev = key_caches.first().unwrap().device();
     let Device::Cuda(dev) = cache_dev else {
         panic!("Expected the key caches to be on a CUDA device.")
     };
     if !cache_dev.same_device(value_caches.first().unwrap().device()) {
-        return Err(APIError::new(format!(
+        candle_core::bail!(
             "`key` and `value` caches have different devices, got {:?} and {:?} respectively.",
             cache_dev,
             value_caches.first().unwrap().device()
-        )));
+        )
     }
     if key_caches.first().unwrap().dtype() != value_caches.first().unwrap().dtype() {
-        return Err(APIError::new(format!(
+        candle_core::bail!(
             "Key and value caches have different types, got {:?} and {:?}.",
             key_caches.first().unwrap().dtype(),
             value_caches.first().unwrap().dtype()
-        )));
+        )
     }
     let num_layers: u32 = key_caches.len().try_into().unwrap();
     if num_layers == 0 {
@@ -55,8 +53,8 @@ pub unsafe fn copy_blocks(
     let mut dtype = DType::F32;
 
     for (key_cache, value_cache) in zip(&key_caches, &value_caches) {
-        try_api!(key_cache.to_device(cache_dev));
-        try_api!(value_cache.to_device(cache_dev));
+        key_cache.to_device(cache_dev)?;
+        value_cache.to_device(cache_dev)?;
 
         let key_offset: u64 = key_cache
             .storage_and_layout()
@@ -99,9 +97,7 @@ pub unsafe fn copy_blocks(
                 (ptr_key, ptr_value)
             }
             _ => {
-                return Err(APIError::from(
-                    "only f32, f16 and bf16 input data type supported!",
-                ));
+                candle_core::bail!("only f32, f16 and bf16 input data type supported!")
             }
         };
         key_cache_ptrs.push(key_ptr + key_offset);
@@ -121,7 +117,10 @@ pub unsafe fn copy_blocks(
     let value_cache_ptr = value_cache_ptrs.as_mut_ptr() as *mut core::ffi::c_void;
     let block_mapping_ptr = block_mapping_vec.as_mut_ptr() as *const core::ffi::c_void;
 
-    let numel_per_block: u32 = try_api!(key_caches.first().unwrap().i(0))
+    let numel_per_block: u32 = key_caches
+        .first()
+        .unwrap()
+        .i(0)?
         .shape()
         .dims()
         .iter()
@@ -174,19 +173,23 @@ pub fn swap_blocks(
     src: Tensor,
     dst: &mut Tensor,
     block_mapping: HashMap<usize, usize>,
-) -> Result<(), APIError> {
+) -> Result<()> {
     let block_size_in_bytes = src.dtype().size_in_bytes() * src.dims()[0];
     match (src.device(), dst.device()) {
         (Device::Cuda(src_dev), Device::Cuda(dst_dev)) => {
             if src_dev.ordinal() != dst_dev.ordinal() {
-                return Err(APIError::new(format!("Tensors must be on the same device to copy, got ordinals {} (src) and {} (dst).", src_dev.ordinal(), dst_dev.ordinal())))
+                candle_core::bail!("Tensors must be on the same device to copy, got ordinals {} (src) and {} (dst).", src_dev.ordinal(), dst_dev.ordinal())
             }
             let (src_storage, src_layout) = src.storage_and_layout();
             let (dst_storage, dst_layout) = dst.storage_and_layout();
             assert!(matches!(&*src_storage, Storage::Cuda(_)));
             assert!(matches!(&*dst_storage, Storage::Cuda(_)));
-            let Storage::Cuda(src_storage) = &*src_storage else { unreachable!() };
-            let Storage::Cuda(dst_storage) = &*dst_storage else { unreachable!() };
+            let Storage::Cuda(src_storage) = &*src_storage else {
+                unreachable!()
+            };
+            let Storage::Cuda(dst_storage) = &*dst_storage else {
+                unreachable!()
+            };
             let (src_ptr, dst_ptr) = match (&src_storage.slice, &dst_storage.slice) {
                 (CudaStorageSlice::BF16(slice_src), CudaStorageSlice::BF16(slice_dst)) => {
                     let ptr_src = *slice_src.slice(src_layout.start_offset()..).device_ptr();
@@ -204,7 +207,7 @@ pub fn swap_blocks(
                     (ptr_src, ptr_dst)
                 }
                 _ => {
-                    return Err(APIError::from("only f32, f16 and bf16 input data type supported!"));
+                    candle_core::bail!("only f32, f16 and bf16 input data type supported!");
                 }
             };
             // let src_ptr = src_storage.as_cuda_slice::<u8>().map_err(APIError::from)?.device_ptr() + TryInto::<u64>::try_into(src_layout.start_offset()).unwrap();
@@ -214,33 +217,51 @@ pub fn swap_blocks(
                 let src_offset: u64 = (src_block_number * block_size_in_bytes).try_into().unwrap();
                 let dst_offset: u64 = (dst_block_number * block_size_in_bytes).try_into().unwrap();
                 // u8s because we copy by bytes
-                let src_slice: CudaSlice<u8> = unsafe { src_dev.upgrade_device_ptr(src_ptr+src_offset, block_size_in_bytes) };
-                let mut dst_slice = unsafe { dst_dev.upgrade_device_ptr(dst_ptr+dst_offset, block_size_in_bytes) };
-
-                try_api!(src_dev.dtod_copy(&src_slice, &mut dst_slice));
+                let src_slice: CudaSlice<u8> = unsafe {
+                    src_dev.upgrade_device_ptr(src_ptr + src_offset, block_size_in_bytes)
+                };
+                let mut dst_slice = unsafe {
+                    dst_dev.upgrade_device_ptr(dst_ptr + dst_offset, block_size_in_bytes)
+                };
+
+                src_dev
+                    .dtod_copy(&src_slice, &mut dst_slice)
+                    .map_err(candle_core::Error::wrap)?;
             }
         }
         (Device::Cpu, Device::Cuda(dst_dev)) => {
             let (src_storage, _src_layout) = src.storage_and_layout();
             let (dst_storage, dst_layout) = dst.storage_and_layout();
             assert!(matches!(&*src_storage, Storage::Cpu(_)));
             assert!(matches!(&*dst_storage, Storage::Cuda(_)));
-            let Storage::Cpu(src_storage) = &*src_storage else { unreachable!() };
-            let Storage::Cuda(dst_storage) = &*dst_storage else { unreachable!() };
-            let dst_ptr = dst_storage.as_cuda_slice::<u8>().map_err(APIError::from)?.device_ptr() + TryInto::<u64>::try_into(dst_layout.start_offset()).unwrap();
-            let src_slice = try_api!(src_storage.as_slice());
+            let Storage::Cpu(src_storage) = &*src_storage else {
+                unreachable!()
+            };
+            let Storage::Cuda(dst_storage) = &*dst_storage else {
+                unreachable!()
+            };
+            let dst_ptr = dst_storage.as_cuda_slice::<u8>()?.device_ptr()
+                + TryInto::<u64>::try_into(dst_layout.start_offset()).unwrap();
+            let src_slice = src_storage.as_slice()?;
 
             for (src_block_number, dst_block_number) in block_mapping {
                 let src_offset = src_block_number * block_size_in_bytes;
                 let dst_offset: u64 = (dst_block_number * block_size_in_bytes).try_into().unwrap();
                 // u8s because we copy by bytes
-                let mut dst_slice: CudaSlice<u8> = unsafe { dst_dev.upgrade_device_ptr(dst_ptr+dst_offset, block_size_in_bytes) };
+                let mut dst_slice: CudaSlice<u8> = unsafe {
+                    dst_dev.upgrade_device_ptr(dst_ptr + dst_offset, block_size_in_bytes)
+                };
 
-                try_api!(dst_dev.htod_sync_copy_into(&src_slice[src_offset..src_offset+block_size_in_bytes], &mut dst_slice));
+                dst_dev
+                    .htod_sync_copy_into(
+                        &src_slice[src_offset..src_offset + block_size_in_bytes],
+                        &mut dst_slice,
+                    )
+                    .map_err(candle_core::Error::wrap)?;
             }
         }
         (src, dst) => {
-            return Err(APIError::new(format!("Tensors must be on either the GPU or CPU to swap,, got {src:?} (src) and {dst:?} (dst).")))
+            candle_core::bail!("Tensors must be on either the GPU or CPU to swap,, got {src:?} (src) and {dst:?} (dst).")
         }
     }
 

src/backend/gguf.rs

@@ -618,3 +618,15 @@ impl TryFrom<Normalizer<'_>> for NormalizerWrapper {
         Ok(value)
     }
 }
+
+pub fn get_arch_and_num_of_layers(ct: gguf_file::Content) -> Result<(String, usize)> {
+    let md_get = |s: &str| match ct.metadata.get(s) {
+        None => candle_core::bail!("cannot find {s} in metadata"),
+        Some(v) => Ok(v),
+    };
+    let architecture = md_get("general.architecture")?.to_string()?;
+
+    let nlayers =
+        md_get(format!("{}.block_count", architecture.as_str()).as_str())?.to_u32()? as usize;
+    Ok((architecture.clone(), nlayers))
+}

src/backend/gptq.rs

@@ -1,15 +1,17 @@
+#[allow(unused_imports)]
 use candle::backend::BackendStorage;
 #[cfg(feature = "cuda")]
 use candle::CudaStorage;
+#[allow(unused_imports)]
 use candle::{CpuStorage, DType, Layout, Result, Shape, Storage, Tensor};
 use candle_core as candle;
-use half::{bf16, f16};
 #[cfg(feature = "cuda")]
 use kernels::ffi::{
     awq_repack, gemm_half_q_half_alt, gptq_repack, marlin_4bit_bf16, marlin_4bit_f16,
     marlin_awq_4bit_bf16, marlin_awq_4bit_f16,
 };
 
+#[allow(unused)]
 struct GPTQMatMul {
     qzeros: Option<Tensor>,
     g_idx: Option<Tensor>,
@@ -228,8 +230,10 @@ impl candle::CustomOp3 for GPTQMatMul {
         scale_l: &Layout,
     ) -> Result<(CudaStorage, Shape)> {
         match x.dtype() {
-            DType::F16 => self.cuda_fwd_t::<f16>(x, x_l, qweight, qweight_l, scale, scale_l),
-            DType::BF16 => self.cuda_fwd_t::<bf16>(x, x_l, qweight, qweight_l, scale, scale_l),
+            DType::F16 => self.cuda_fwd_t::<half::f16>(x, x_l, qweight, qweight_l, scale, scale_l),
+            DType::BF16 => {
+                self.cuda_fwd_t::<half::bf16>(x, x_l, qweight, qweight_l, scale, scale_l)
+            }
             dt => candle::bail!("GPTQMatMul is only supported for f16 and bf16 ({dt:?})"),
         }
     }

src/backend/mod.rs

@@ -6,10 +6,11 @@ mod paged_attention;
 pub fn get_or_load_func(
     ptx_file: &'static str,
     kernel_base: &str,
-    dtype: DType,
+    dtype: candle_core::DType,
     suffix: Option<&str>,
     device: &CudaDevice,
 ) -> Result<CudaFunction, APIError> {
+    use candle_core::DType;
     let spec = match dtype {
         DType::U8 => "_u8",
         DType::U32 => "_u32",
@@ -30,9 +31,9 @@ pub fn get_or_load_func(
         .map_err(APIError::from)
 }
 
+#[allow(unused_imports)]
 use crate::openai::responses::APIError;
 pub use cache::*;
-use candle_core::DType;
 #[cfg(feature = "cuda")]
 use candle_core::{cuda_backend::cudarc::driver::CudaFunction, CudaDevice};
 pub use gptq::*;

src/backend/paged_attention.rs

@@ -5,8 +5,6 @@ use candle::CudaStorage;
 use candle::MetalStorage;
 use candle::{CpuStorage, DType, Layout, Result, Shape, Storage, Tensor};
 use candle_core as candle;
-use half::{bf16, f16};
-use std::ffi::c_int;
 #[allow(dead_code)]
 struct PagedAttention {
     softmax_scale: f32,
@@ -30,6 +28,7 @@ impl PagedAttention {
     ) -> Result<(CudaStorage, Shape)> {
         use candle::cuda_backend::cudarc::driver::DevicePtr;
         use candle::cuda_backend::WrapErr;
+        use core::ffi::c_int;
         let dtype = q.dtype();
         let internal_type = match dtype {
             DType::F16 => 0,
@@ -480,8 +479,8 @@ impl candle::CustomOp1 for PagedAttention {
     fn cuda_fwd(&self, q: &CudaStorage, q_l: &Layout) -> Result<(CudaStorage, Shape)> {
         match q.dtype() {
             DType::F32 => self.cuda_fwd_t::<f32>(q, q_l),
-            DType::F16 => self.cuda_fwd_t::<f16>(q, q_l),
-            DType::BF16 => self.cuda_fwd_t::<bf16>(q, q_l),
+            DType::F16 => self.cuda_fwd_t::<half::f16>(q, q_l),
+            DType::BF16 => self.cuda_fwd_t::<half::bf16>(q, q_l),
             dt => candle::bail!("paged-attention is only supported for f32/f16/bf16 ({dt:?})"),
         }
     }
@@ -557,6 +556,7 @@ impl ReshapeCache {
         slot_mapping: &Tensor,
     ) -> Result<()> {
         use candle::cuda_backend::cudarc::driver::DevicePtr;
+        use core::ffi::c_int;
         let dtype = k.dtype();
         let dev = k.device();
         let internal_type = match dtype {
@@ -842,15 +842,15 @@ impl candle::InplaceOp1 for ReshapeCache {
                 &self.value_cache,
                 &self.slot_mapping,
             ),
-            DType::F16 => self.cuda_fwd_t::<f16>(
+            DType::F16 => self.cuda_fwd_t::<half::f16>(
                 k,
                 k_l,
                 &self.value,
                 &self.key_cache,
                 &self.value_cache,
                 &self.slot_mapping,
             ),
-            DType::BF16 => self.cuda_fwd_t::<bf16>(
+            DType::BF16 => self.cuda_fwd_t::<half::bf16>(
                 k,
                 k_l,
                 &self.value,