candle_core/quantized/
mod.rs

1use crate::{
2    backend::BackendStorage, CpuStorage, DType, Device, Result, Shape, Storage, Tensor, D,
3};
4use k_quants::*;
5use std::borrow::Cow;
6
7#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
8pub mod avx;
9mod dummy_cuda;
10mod dummy_metal;
11pub mod ggml_file;
12pub mod gguf_file;
13pub mod imatrix_file;
14pub mod k_quants;
15#[cfg(feature = "metal")]
16pub mod metal;
17pub mod repack;
18#[cfg(target_arch = "x86_64")]
19pub(crate) mod repack_x86;
20#[cfg(not(target_arch = "wasm32"))]
21pub mod tokenizer;
22#[cfg(not(feature = "metal"))]
23mod metal {
24    pub use super::dummy_metal::*;
25}
26#[cfg(feature = "cuda")]
27pub mod cuda;
28#[cfg(feature = "cuda")]
29pub mod fast_mmq;
30#[cfg(feature = "cuda")]
31pub mod fast_mmvq;
32#[cfg(not(feature = "cuda"))]
33mod cuda {
34    pub use super::dummy_cuda::*;
35}
36
37#[cfg(any(
38    target_arch = "aarch64",
39    all(target_arch = "arm", target_feature = "neon")
40))]
41pub mod neon;
42#[cfg(target_feature = "simd128")]
43pub mod simd128;
44pub mod utils;
45use half::{bf16, f16};
46
47pub use k_quants::GgmlType;
48
49fn as_t_slice<T>(data: &[u8]) -> &[T] {
50    let size = std::mem::size_of::<T>();
51    assert_eq!(
52        data.len() % size,
53        0,
54        "Data length must be a multiple of T's size"
55    );
56    let ptr = data.as_ptr();
57    assert_eq!(
58        (ptr as usize) % std::mem::align_of::<T>(),
59        0,
60        "Data pointer must be aligned to T's alignment"
61    );
62    unsafe { std::slice::from_raw_parts(ptr as *const T, data.len() / size) }
63}
64
65pub struct QTensor {
66    storage: QStorage,
67    shape: Shape,
68    #[allow(dead_code)]
69    repacked_qs: repack::PackedCache,
70}
71
72impl Device {
73    fn qzeros(&self, elem_count: usize, dtype: GgmlDType) -> Result<QStorage> {
74        match self {
75            Device::Cpu => {
76                let storage = dtype.cpu_zeros(elem_count);
77                Ok(QStorage::Cpu(storage))
78            }
79            Device::Metal(metal) => {
80                let storage = metal::QMetalStorage::zeros(metal, elem_count, dtype)?;
81                Ok(QStorage::Metal(storage))
82            }
83            Device::Cuda(cuda) => {
84                let storage = cuda::QCudaStorage::zeros(cuda, elem_count, dtype)?;
85                Ok(QStorage::Cuda(storage))
86            }
87        }
88    }
89}
90
91pub enum QStorage {
92    Cpu(Box<dyn QuantizedType>),
93    Metal(metal::QMetalStorage),
94    Cuda(cuda::QCudaStorage),
95}
96
97impl QStorage {
98    pub fn from_data(data: Cow<'_, [u8]>, device: &Device, dtype: GgmlDType) -> Result<Self> {
99        let data: &[u8] = &data;
100        match device {
101            Device::Cpu => Ok(Self::Cpu(dtype.from_data(Cow::Borrowed(data)))),
102            Device::Metal(d) => match dtype {
103                GgmlDType::F32 => metal::load_quantized(d, as_t_slice::<f32>(data)),
104                GgmlDType::F16 => metal::load_quantized(d, as_t_slice::<f16>(data)),
105                GgmlDType::Q4_0 => metal::load_quantized(d, as_t_slice::<BlockQ4_0>(data)),
106                GgmlDType::Q4_1 => metal::load_quantized(d, as_t_slice::<BlockQ4_1>(data)),
107                GgmlDType::Q5_0 => metal::load_quantized(d, as_t_slice::<BlockQ5_0>(data)),
108                GgmlDType::Q5_1 => metal::load_quantized(d, as_t_slice::<BlockQ5_1>(data)),
109                GgmlDType::Q8_0 => metal::load_quantized(d, as_t_slice::<BlockQ8_0>(data)),
110                GgmlDType::Q8_1 => metal::load_quantized(d, as_t_slice::<BlockQ8_1>(data)),
111                GgmlDType::Q2K => metal::load_quantized(d, as_t_slice::<BlockQ2K>(data)),
112                GgmlDType::Q3K => metal::load_quantized(d, as_t_slice::<BlockQ3K>(data)),
113                GgmlDType::Q4K => metal::load_quantized(d, as_t_slice::<BlockQ4K>(data)),
114                GgmlDType::Q5K => metal::load_quantized(d, as_t_slice::<BlockQ5K>(data)),
115                GgmlDType::Q6K => metal::load_quantized(d, as_t_slice::<BlockQ6K>(data)),
116                GgmlDType::Q8K => metal::load_quantized(d, as_t_slice::<BlockQ8K>(data)),
117                GgmlDType::BF16 => metal::load_quantized(d, as_t_slice::<bf16>(data)),
118            },
119            Device::Cuda(d) => match dtype {
120                GgmlDType::F32 => cuda::load_quantized(d, as_t_slice::<f32>(data)),
121                GgmlDType::F16 => cuda::load_quantized(d, as_t_slice::<f16>(data)),
122                GgmlDType::Q4_0 => cuda::load_quantized(d, as_t_slice::<BlockQ4_0>(data)),
123                GgmlDType::Q4_1 => cuda::load_quantized(d, as_t_slice::<BlockQ4_1>(data)),
124                GgmlDType::Q5_0 => cuda::load_quantized(d, as_t_slice::<BlockQ5_0>(data)),
125                GgmlDType::Q5_1 => cuda::load_quantized(d, as_t_slice::<BlockQ5_1>(data)),
126                GgmlDType::Q8_0 => cuda::load_quantized(d, as_t_slice::<BlockQ8_0>(data)),
127                GgmlDType::Q8_1 => cuda::load_quantized(d, as_t_slice::<BlockQ8_1>(data)),
128                GgmlDType::Q2K => cuda::load_quantized(d, as_t_slice::<BlockQ2K>(data)),
129                GgmlDType::Q3K => cuda::load_quantized(d, as_t_slice::<BlockQ3K>(data)),
130                GgmlDType::Q4K => cuda::load_quantized(d, as_t_slice::<BlockQ4K>(data)),
131                GgmlDType::Q5K => cuda::load_quantized(d, as_t_slice::<BlockQ5K>(data)),
132                GgmlDType::Q6K => cuda::load_quantized(d, as_t_slice::<BlockQ6K>(data)),
133                GgmlDType::Q8K => cuda::load_quantized(d, as_t_slice::<BlockQ8K>(data)),
134                GgmlDType::BF16 => cuda::load_quantized(d, as_t_slice::<bf16>(data)),
135            },
136        }
137    }
138
139    fn block_size(&self) -> usize {
140        match self {
141            QStorage::Cpu(storage) => storage.block_size(),
142            QStorage::Metal(storage) => storage.dtype().block_size(),
143            QStorage::Cuda(storage) => storage.dtype().block_size(),
144        }
145    }
146
147    fn dtype(&self) -> GgmlDType {
148        match self {
149            QStorage::Cpu(storage) => storage.dtype(),
150            QStorage::Metal(storage) => storage.dtype(),
151            QStorage::Cuda(storage) => storage.dtype(),
152        }
153    }
154
155    fn device(&self) -> Device {
156        match self {
157            QStorage::Cpu(_storage) => Device::Cpu,
158            QStorage::Metal(storage) => Device::Metal(storage.device().clone()),
159            QStorage::Cuda(storage) => Device::Cuda(storage.device().clone()),
160        }
161    }
162
163    fn size_in_bytes(&self) -> usize {
164        match self {
165            QStorage::Cpu(storage) => storage.storage_size_in_bytes(),
166            QStorage::Metal(storage) => storage.storage_size_in_bytes(),
167            QStorage::Cuda(storage) => storage.storage_size_in_bytes(),
168        }
169    }
170
171    fn quantize(&mut self, src: &Storage) -> Result<()> {
172        match (self, src) {
173            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
174                storage.from_float(src.as_slice::<f32>()?);
175            }
176            (QStorage::Metal(storage), Storage::Metal(src)) => storage.quantize(src)?,
177            (QStorage::Cuda(storage), Storage::Cuda(src)) => storage.quantize(src)?,
178            _ => crate::bail!("Invalid quantize storage locations do not match"),
179        }
180        Ok(())
181    }
182
183    fn quantize_imatrix(
184        &mut self,
185        src: &Storage,
186        imatrix_weights: &[f32],
187        n_per_row: usize,
188    ) -> Result<()> {
189        match (self, src) {
190            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
191                storage.from_float_imatrix(src.as_slice::<f32>()?, imatrix_weights, n_per_row);
192            }
193            (QStorage::Metal(storage), Storage::Metal(src)) => {
194                storage.quantize_imatrix(src, imatrix_weights, n_per_row)?
195            }
196            (QStorage::Cuda(storage), Storage::Cuda(src)) => {
197                storage.quantize_imatrix(src, imatrix_weights, n_per_row)?
198            }
199            _ => crate::bail!("Invalid quantize storage locations do not match"),
200        }
201        Ok(())
202    }
203
204    fn quantize_onto(&mut self, src: &Storage) -> Result<()> {
205        match (self, src) {
206            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
207                storage.from_float(src.as_slice::<f32>()?);
208            }
209            (QStorage::Metal(storage), Storage::Cpu(src)) => storage.quantize_onto(src)?,
210            (QStorage::Cuda(storage), Storage::Cpu(src)) => storage.quantize_onto(src)?,
211            _ => crate::bail!("Invalid quantize source storage locations: not on cpu"),
212        }
213        Ok(())
214    }
215
216    fn quantize_imatrix_onto(
217        &mut self,
218        src: &Storage,
219        imatrix_weights: &[f32],
220        n_per_row: usize,
221    ) -> Result<()> {
222        match (self, src) {
223            (QStorage::Cpu(storage), Storage::Cpu(src)) => {
224                storage.from_float_imatrix(src.as_slice::<f32>()?, imatrix_weights, n_per_row);
225            }
226            (QStorage::Metal(storage), Storage::Cpu(src)) => {
227                storage.quantize_imatrix_onto(src, imatrix_weights, n_per_row)?
228            }
229            (QStorage::Cuda(storage), Storage::Cpu(src)) => {
230                storage.quantize_imatrix_onto(src, imatrix_weights, n_per_row)?
231            }
232            _ => crate::bail!("Invalid quantize storage locations do not match"),
233        }
234        Ok(())
235    }
236
237    fn dequantize(&self, elem_count: usize) -> Result<Storage> {
238        match self {
239            QStorage::Cpu(storage) => Ok(Storage::Cpu(storage.dequantize(elem_count)?)),
240            QStorage::Metal(storage) => Ok(Storage::Metal(storage.dequantize(elem_count)?)),
241            QStorage::Cuda(storage) => Ok(Storage::Cuda(storage.dequantize(elem_count)?)),
242        }
243    }
244
245    fn data(&self) -> Result<Cow<'_, [u8]>> {
246        match self {
247            QStorage::Cpu(storage) => {
248                let data_ptr = storage.as_ptr();
249                let size_in_bytes = storage.storage_size_in_bytes();
250                let data = unsafe { std::slice::from_raw_parts(data_ptr, size_in_bytes) };
251                Ok(Cow::from(data))
252            }
253            QStorage::Cuda(storage) => Ok(Cow::from(storage.data()?)),
254            QStorage::Metal(storage) => Ok(Cow::from(storage.data()?)),
255        }
256    }
257
258    pub fn device_ptr(&self) -> Result<*const u8> {
259        match self {
260            QStorage::Cuda(storage) => storage.device_ptr(),
261            QStorage::Metal(_) | QStorage::Cpu(_) => {
262                crate::bail!("not implemented");
263            }
264        }
265    }
266
267    #[cfg(feature = "cuda")]
268    pub fn device_ptr_with_guard<'a>(
269        &'a self,
270        stream: &'a crate::cuda_backend::cudarc::driver::CudaStream,
271    ) -> Result<(
272        *const u8,
273        crate::cuda_backend::cudarc::driver::SyncOnDrop<'a>,
274    )> {
275        match self {
276            QStorage::Cuda(storage) => storage.device_ptr_with_guard(stream),
277            QStorage::Metal(_) | QStorage::Cpu(_) => {
278                crate::bail!("not implemented");
279            }
280        }
281    }
282}
283
284#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
285pub enum GgmlDType {
286    F32,
287    F16,
288    BF16,
289    Q4_0,
290    Q4_1,
291    Q5_0,
292    Q5_1,
293    Q8_0,
294    Q8_1,
295    Q2K,
296    Q3K,
297    Q4K,
298    Q5K,
299    Q6K,
300    Q8K,
301}
302
303impl GgmlDType {
304    pub(crate) fn from_u32(u: u32) -> Result<Self> {
305        let dtype = match u {
306            0 => Self::F32,
307            1 => Self::F16,
308            2 => Self::Q4_0,
309            3 => Self::Q4_1,
310            6 => Self::Q5_0,
311            7 => Self::Q5_1,
312            8 => Self::Q8_0,
313            9 => Self::Q8_1,
314            10 => Self::Q2K,
315            11 => Self::Q3K,
316            12 => Self::Q4K,
317            13 => Self::Q5K,
318            14 => Self::Q6K,
319            15 => Self::Q8K,
320            // https://github.com/ggerganov/ggml/blob/29d87fc6676e7ed0cdfdec0804b06001d9c2bb44/include/ggml.h#L389
321            30 => Self::BF16,
322            _ => crate::bail!("unknown dtype for tensor {u}"),
323        };
324        Ok(dtype)
325    }
326
327    pub(crate) fn to_u32(self) -> u32 {
328        match self {
329            Self::F32 => 0,
330            Self::F16 => 1,
331            Self::Q4_0 => 2,
332            Self::Q4_1 => 3,
333            Self::Q5_0 => 6,
334            Self::Q5_1 => 7,
335            Self::Q8_0 => 8,
336            Self::Q8_1 => 9,
337            Self::Q2K => 10,
338            Self::Q3K => 11,
339            Self::Q4K => 12,
340            Self::Q5K => 13,
341            Self::Q6K => 14,
342            Self::Q8K => 15,
343            // https://github.com/ggerganov/ggml/blob/29d87fc6676e7ed0cdfdec0804b06001d9c2bb44/include/ggml.h#L389
344            Self::BF16 => 30,
345        }
346    }
347
348    /// The block dtype
349    pub fn cpu_zeros(&self, elem_count: usize) -> Box<dyn QuantizedType> {
350        match self {
351            Self::F32 => Box::new(vec![f32::zeros(); elem_count]),
352            Self::F16 => Box::new(vec![f16::zeros(); elem_count]),
353            Self::Q4_0 => Box::new(vec![BlockQ4_0::zeros(); elem_count / BlockQ4_0::BLCK_SIZE]),
354            Self::Q4_1 => Box::new(vec![BlockQ4_1::zeros(); elem_count / BlockQ4_1::BLCK_SIZE]),
355            Self::Q5_0 => Box::new(vec![BlockQ5_0::zeros(); elem_count / BlockQ5_0::BLCK_SIZE]),
356            Self::Q5_1 => Box::new(vec![BlockQ5_1::zeros(); elem_count / BlockQ5_1::BLCK_SIZE]),
357            Self::Q8_0 => Box::new(vec![BlockQ8_0::zeros(); elem_count / BlockQ8_0::BLCK_SIZE]),
358            Self::Q8_1 => Box::new(vec![BlockQ8_1::zeros(); elem_count / BlockQ8_1::BLCK_SIZE]),
359            Self::Q2K => Box::new(vec![BlockQ2K::zeros(); elem_count / BlockQ2K::BLCK_SIZE]),
360            Self::Q3K => Box::new(vec![BlockQ3K::zeros(); elem_count / BlockQ3K::BLCK_SIZE]),
361            Self::Q4K => Box::new(vec![BlockQ4K::zeros(); elem_count / BlockQ4K::BLCK_SIZE]),
362            Self::Q5K => Box::new(vec![BlockQ5K::zeros(); elem_count / BlockQ5K::BLCK_SIZE]),
363            Self::Q6K => Box::new(vec![BlockQ6K::zeros(); elem_count / BlockQ6K::BLCK_SIZE]),
364            Self::Q8K => Box::new(vec![BlockQ8K::zeros(); elem_count / BlockQ8K::BLCK_SIZE]),
365            Self::BF16 => Box::new(vec![bf16::zeros(); elem_count]),
366        }
367    }
368
369    pub fn from_data(&self, data: Cow<'_, [u8]>) -> Box<dyn QuantizedType> {
370        let data: &[u8] = &data;
371        match self {
372            Self::F32 => Box::new(as_t_slice::<f32>(data).to_vec()),
373            Self::F16 => Box::new(as_t_slice::<f16>(data).to_vec()),
374            Self::Q4_0 => Box::new(as_t_slice::<BlockQ4_0>(data).to_vec()),
375            Self::Q4_1 => Box::new(as_t_slice::<BlockQ4_1>(data).to_vec()),
376            Self::Q5_0 => Box::new(as_t_slice::<BlockQ5_0>(data).to_vec()),
377            Self::Q5_1 => Box::new(as_t_slice::<BlockQ5_1>(data).to_vec()),
378            Self::Q8_0 => Box::new(as_t_slice::<BlockQ8_0>(data).to_vec()),
379            Self::Q8_1 => Box::new(as_t_slice::<BlockQ8_1>(data).to_vec()),
380            Self::Q2K => Box::new(as_t_slice::<BlockQ2K>(data).to_vec()),
381            Self::Q3K => Box::new(as_t_slice::<BlockQ3K>(data).to_vec()),
382            Self::Q4K => Box::new(as_t_slice::<BlockQ4K>(data).to_vec()),
383            Self::Q5K => Box::new(as_t_slice::<BlockQ5K>(data).to_vec()),
384            Self::Q6K => Box::new(as_t_slice::<BlockQ6K>(data).to_vec()),
385            Self::Q8K => Box::new(as_t_slice::<BlockQ8K>(data).to_vec()),
386            Self::BF16 => Box::new(as_t_slice::<bf16>(data).to_vec()),
387        }
388    }
389
390    /// The type size for blocks in bytes.
391    pub fn type_size(&self) -> usize {
392        use k_quants::*;
393        match self {
394            Self::F32 => 4,
395            Self::F16 | Self::BF16 => 2,
396            Self::Q4_0 => std::mem::size_of::<BlockQ4_0>(),
397            Self::Q4_1 => std::mem::size_of::<BlockQ4_1>(),
398            Self::Q5_0 => std::mem::size_of::<BlockQ5_0>(),
399            Self::Q5_1 => std::mem::size_of::<BlockQ5_1>(),
400            // https://github.com/ggerganov/llama.cpp/blob/468ea24fb4633a0d681f7ac84089566c1c6190cb/ggml.c#L932
401            Self::Q8_0 => std::mem::size_of::<BlockQ8_0>(),
402            Self::Q8_1 => std::mem::size_of::<BlockQ8_1>(),
403            Self::Q2K => std::mem::size_of::<BlockQ2K>(),
404            Self::Q3K => std::mem::size_of::<BlockQ3K>(),
405            Self::Q4K => std::mem::size_of::<BlockQ4K>(),
406            Self::Q5K => std::mem::size_of::<BlockQ5K>(),
407            Self::Q6K => std::mem::size_of::<BlockQ6K>(),
408            Self::Q8K => std::mem::size_of::<BlockQ8K>(),
409        }
410    }
411
412    /// The block size, i.e. the number of elements stored in each block.
413    pub fn block_size(&self) -> usize {
414        match self {
415            Self::F32 => 1,
416            Self::F16 | Self::BF16 => 1,
417            Self::Q4_0 => k_quants::QK4_0,
418            Self::Q4_1 => k_quants::QK4_1,
419            Self::Q5_0 => k_quants::QK5_0,
420            Self::Q5_1 => k_quants::QK5_1,
421            Self::Q8_0 => k_quants::QK8_0,
422            Self::Q8_1 => k_quants::QK8_1,
423            Self::Q2K | Self::Q3K | Self::Q4K | Self::Q5K | Self::Q6K | Self::Q8K => k_quants::QK_K,
424        }
425    }
426}
427
428// A version of GgmlType without `vec_dot` so that it can be dyn boxed.
429pub trait QuantizedType: Send + Sync {
430    fn dtype(&self) -> GgmlDType;
431    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()>;
432    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()>;
433    fn embedding(&self, ids: &[u32], rows: usize, hidden: usize) -> Result<CpuStorage>;
434    fn dequantize(&self, elem_count: usize) -> Result<CpuStorage>;
435    fn storage_size_in_bytes(&self) -> usize;
436    fn as_ptr(&self) -> *const u8;
437    fn block_size(&self) -> usize;
438    #[allow(clippy::wrong_self_convention)]
439    fn from_float(&mut self, xs: &[f32]);
440    #[allow(clippy::wrong_self_convention)]
441    fn from_float_imatrix(&mut self, xs: &[f32], imatrix_weights: &[f32], n_per_row: usize);
442    fn size(&self) -> usize;
443}
444
445impl<T: k_quants::GgmlType + Send + Sync> QuantizedType for Vec<T> {
446    fn matmul_t(&self, mkn: (usize, usize, usize), lhs: &[f32], dst: &mut [f32]) -> Result<()> {
447        k_quants::matmul(mkn, lhs, self.as_slice(), dst)
448    }
449    fn matmul_t_f16(&self, mkn: (usize, usize, usize), lhs: &[f16], dst: &mut [f16]) -> Result<()> {
450        k_quants::matmul_f16(mkn, lhs, self.as_slice(), dst)
451    }
452
453    fn embedding(&self, ids: &[u32], rows: usize, hidden: usize) -> Result<CpuStorage> {
454        if !hidden.is_multiple_of(T::BLCK_SIZE) {
455            crate::bail!(
456                "quantized embedding hidden size {hidden} is not divisible by block size {}",
457                T::BLCK_SIZE
458            )
459        }
460        let row_blocks = hidden / T::BLCK_SIZE;
461        if self.len() != rows * row_blocks {
462            crate::bail!(
463                "quantized tensor has {} blocks, expected {}",
464                self.len(),
465                rows * row_blocks
466            )
467        }
468        let mut out = vec![0f32; ids.len() * hidden];
469        for (out_row, &row_id) in ids.iter().enumerate() {
470            let row = row_id as usize;
471            if row >= rows {
472                crate::bail!("embedding id {row} is out of range for {rows} rows")
473            }
474            let src = &self[row * row_blocks..(row + 1) * row_blocks];
475            let dst = &mut out[out_row * hidden..(out_row + 1) * hidden];
476            T::to_float(src, dst);
477        }
478        Ok(CpuStorage::F32(out))
479    }
480
481    fn size(&self) -> usize {
482        self.len() * core::mem::size_of::<T>()
483    }
484
485    fn from_float(&mut self, xs: &[f32]) {
486        T::from_float(xs, self)
487    }
488
489    fn from_float_imatrix(&mut self, xs: &[f32], imatrix_weights: &[f32], n_per_row: usize) {
490        T::from_float_imatrix(xs, self, imatrix_weights, n_per_row)
491    }
492
493    fn dtype(&self) -> GgmlDType {
494        T::DTYPE
495    }
496
497    fn block_size(&self) -> usize {
498        T::BLCK_SIZE
499    }
500
501    fn dequantize(&self, elem_count: usize) -> Result<CpuStorage> {
502        let mut ys = vec![0.0f32; elem_count];
503        T::to_float(self.as_slice(), &mut ys);
504        Ok(CpuStorage::F32(ys))
505    }
506
507    fn storage_size_in_bytes(&self) -> usize {
508        self.len() * std::mem::size_of::<T>()
509    }
510
511    fn as_ptr(&self) -> *const u8 {
512        self.as_ptr() as *const u8
513    }
514}
515
516impl std::fmt::Debug for QTensor {
517    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
518        write!(f, "QTensor[{:?}; {:?}]", self.shape, self.dtype())
519    }
520}
521
522fn check_shape(shape: &Shape, block_size: usize) -> Result<()> {
523    let dims = shape.dims();
524    if dims.is_empty() {
525        crate::bail!("scalar tensor cannot be quantized {shape:?}")
526    }
527    if !dims[dims.len() - 1].is_multiple_of(block_size) {
528        crate::bail!(
529            "quantized tensor must have their last dim divisible by block size {shape:?} {}",
530            block_size
531        )
532    }
533    Ok(())
534}
535
536impl QTensor {
537    pub fn new<S: Into<Shape>>(storage: QStorage, shape: S) -> Result<Self> {
538        let shape = shape.into();
539        check_shape(&shape, storage.block_size())?;
540        Ok(Self {
541            storage,
542            shape,
543            repacked_qs: repack::PackedCache::new(),
544        })
545    }
546
547    pub fn quantize(src: &Tensor, dtype: GgmlDType) -> Result<Self> {
548        let shape = src.shape();
549        let block_size = dtype.block_size();
550        check_shape(shape, block_size)?;
551        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
552        let elem_count = shape.elem_count();
553        if !elem_count.is_multiple_of(block_size) {
554            crate::bail!(
555                "tensor size ({shape:?}) is not divisible by block size {}",
556                block_size
557            )
558        }
559        let mut storage = src.device().qzeros(elem_count, dtype)?;
560        storage.quantize(&src.storage())?;
561        Ok(Self {
562            storage,
563            shape: shape.clone(),
564            repacked_qs: repack::PackedCache::new(),
565        })
566    }
567
568    pub fn quantize_imatrix(
569        src: &Tensor,
570        imatrix_weights: &[f32],
571        dtype: GgmlDType,
572    ) -> Result<Self> {
573        // (n_per_row/QK_K-1)*QK_K+(QK_K/32-1)*32+32=n_per_row
574        // Size of imatrix == last dim of tensor
575        let n_per_row = src.dim(D::Minus1)?;
576        if imatrix_weights.len() != n_per_row {
577            crate::bail!(
578                "imatrix weights must have the same length {} as the last dim of src {}",
579                imatrix_weights.len(),
580                src.dim(D::Minus1)?
581            );
582        }
583
584        let shape = src.shape();
585        let block_size = dtype.block_size();
586        check_shape(shape, block_size)?;
587        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
588        let elem_count = shape.elem_count();
589        if !elem_count.is_multiple_of(block_size) {
590            crate::bail!(
591                "tensor size ({shape:?}) is not divisible by block size {}",
592                block_size
593            );
594        }
595        let mut storage = src.device().qzeros(elem_count, dtype)?;
596        storage.quantize_imatrix(&src.storage(), imatrix_weights, n_per_row)?;
597        Ok(Self {
598            storage,
599            shape: shape.clone(),
600            repacked_qs: repack::PackedCache::new(),
601        })
602    }
603
604    /// Quantize `src` (currently on the CPU) to a QTensor on `dev`
605    pub fn quantize_imatrix_onto(
606        src: &Tensor,
607        imatrix_weights: &[f32],
608        dtype: GgmlDType,
609        dev: &Device,
610    ) -> Result<Self> {
611        if !src.device().is_cpu() {
612            crate::bail!(
613                "`quantize_onto` expects a `src` to be on the cpu, got {:?}.",
614                src.device()
615            )
616        }
617        // (n_per_row/QK_K-1)*QK_K+(QK_K/32-1)*32+32=n_per_row
618        // Size of imatrix == last dim of tensor
619        let n_per_row = src.dim(D::Minus1)?;
620        if imatrix_weights.len() != n_per_row {
621            crate::bail!(
622                "imatrix weights must have the same length {} as the last dim of src {}",
623                imatrix_weights.len(),
624                src.dim(D::Minus1)?
625            );
626        }
627        let shape = src.shape();
628        let block_size = dtype.block_size();
629        check_shape(shape, block_size)?;
630        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
631        let elem_count = shape.elem_count();
632        if !elem_count.is_multiple_of(block_size) {
633            crate::bail!(
634                "tensor size ({shape:?}) is not divisible by block size {}",
635                block_size
636            )
637        }
638        // storage is on the `dev`, src is on `cpu`
639        let mut storage = dev.qzeros(elem_count, dtype)?;
640        storage.quantize_imatrix_onto(&src.storage(), imatrix_weights, n_per_row)?;
641        Ok(Self {
642            storage,
643            shape: shape.clone(),
644            repacked_qs: repack::PackedCache::new(),
645        })
646    }
647
648    /// Quantize `src` (currently on the CPU) to a QTensor on `dev`
649    pub fn quantize_onto(src: &Tensor, dtype: GgmlDType, dev: &Device) -> Result<Self> {
650        if !src.device().is_cpu() {
651            crate::bail!(
652                "`quantize_onto` expects a `src` to be on the cpu, got {:?}.",
653                src.device()
654            )
655        }
656        let shape = src.shape();
657        let block_size = dtype.block_size();
658        check_shape(shape, block_size)?;
659        let src = src.to_dtype(crate::DType::F32)?.flatten_all()?;
660        let elem_count = shape.elem_count();
661        if !elem_count.is_multiple_of(block_size) {
662            crate::bail!(
663                "tensor size ({shape:?}) is not divisible by block size {}",
664                block_size
665            )
666        }
667        // storage is on the `dev`, src is on `cpu`
668        let mut storage = dev.qzeros(elem_count, dtype)?;
669        storage.quantize_onto(&src.storage())?;
670        Ok(Self {
671            storage,
672            shape: shape.clone(),
673            repacked_qs: repack::PackedCache::new(),
674        })
675    }
676
677    pub fn dtype(&self) -> GgmlDType {
678        self.storage.dtype()
679    }
680
681    pub fn device(&self) -> Device {
682        self.storage.device()
683    }
684
685    pub fn rank(&self) -> usize {
686        self.shape.rank()
687    }
688
689    pub fn shape(&self) -> &Shape {
690        &self.shape
691    }
692
693    pub fn dequantize(&self, device: &Device) -> Result<Tensor> {
694        let storage = self.storage.dequantize(self.shape.elem_count())?;
695        let none = crate::op::BackpropOp::none();
696        crate::tensor::from_storage(storage, self.shape.clone(), none, false).to_device(device)
697    }
698
699    pub fn dequantize_f16(&self, device: &Device) -> Result<Tensor> {
700        // In the CUDA case, we have a specialized kernel as this can be useful for volta
701        // architectures. https://github.com/huggingface/candle/issues/2136
702        match &self.storage {
703            QStorage::Cuda(s) => {
704                let s = s.dequantize_f16(self.shape.elem_count())?;
705                let none = crate::op::BackpropOp::none();
706                crate::tensor::from_storage(Storage::Cuda(s), self.shape.clone(), none, false)
707                    .to_device(device)
708            }
709            _ => {
710                let s = self.dequantize(device)?.to_dtype(crate::DType::F16)?;
711                Ok(s)
712            }
713        }
714    }
715
716    pub fn embedding(&self, ids: &Tensor) -> Result<Tensor> {
717        let (rows, hidden) = self.shape.dims2()?;
718        if !hidden.is_multiple_of(self.dtype().block_size()) {
719            crate::bail!(
720                "quantized embedding hidden size {hidden} is not divisible by block size {}",
721                self.dtype().block_size()
722            )
723        }
724        let mut out_shape = ids.dims().to_vec();
725        out_shape.push(hidden);
726        let device = self.device();
727        let ids = ids
728            .to_device(&device)?
729            .to_dtype(DType::U32)?
730            .flatten_all()?
731            .contiguous()?;
732        let storage = match &self.storage {
733            QStorage::Cpu(storage) => {
734                let ids = ids.to_vec1::<u32>()?;
735                Storage::Cpu(storage.embedding(&ids, rows, hidden)?)
736            }
737            QStorage::Metal(storage) => match &*ids.storage() {
738                Storage::Metal(ids_storage) => {
739                    Storage::Metal(storage.embedding(rows, hidden, ids_storage, ids.layout())?)
740                }
741                _ => unreachable!("ids were moved to the QTensor device"),
742            },
743            QStorage::Cuda(storage) => match &*ids.storage() {
744                Storage::Cuda(ids_storage) => {
745                    Storage::Cuda(storage.embedding(rows, hidden, ids_storage, ids.layout())?)
746                }
747                _ => unreachable!("ids were moved to the QTensor device"),
748            },
749        };
750        let none = crate::op::BackpropOp::none();
751        Ok(crate::tensor::from_storage(storage, out_shape, none, false))
752    }
753
754    pub fn storage_size_in_bytes(&self) -> usize {
755        self.storage.size_in_bytes()
756    }
757
758    pub fn data(&self) -> Result<Cow<'_, [u8]>> {
759        self.storage.data()
760    }
761
762    pub fn indexed_moe_forward(&self, x: &Tensor, ids: &Tensor) -> Result<Tensor> {
763        match &self.storage {
764            QStorage::Cuda(s) => match (&*x.storage(), &*ids.storage()) {
765                (Storage::Cuda(x_storage), Storage::Cuda(ids_storage)) => {
766                    let (storage, out_shape) = s.indexed_moe_forward(
767                        self.shape(),
768                        x_storage,
769                        x.layout(),
770                        ids_storage,
771                        ids.layout(),
772                    )?;
773                    Ok(crate::tensor::from_storage(
774                        Storage::Cuda(storage),
775                        out_shape,
776                        crate::op::BackpropOp::none(),
777                        false,
778                    ))
779                }
780                _ => {
781                    panic!("Non-cuda indexed_moe_forward is not implemented!");
782                }
783            },
784            _ => {
785                panic!("indexed_moe_forward is not implemented in this platform!");
786            }
787        }
788    }
789
790    pub fn device_ptr(&self) -> Result<*const u8> {
791        match &self.storage {
792            QStorage::Cuda(storage) => storage.device_ptr(),
793            QStorage::Metal(_) | QStorage::Cpu(_) => {
794                crate::bail!("not implemented");
795            }
796        }
797    }
798
799    #[cfg(feature = "cuda")]
800    pub fn device_ptr_with_guard<'a>(
801        &'a self,
802        stream: &'a crate::cuda_backend::cudarc::driver::CudaStream,
803    ) -> Result<(
804        *const u8,
805        crate::cuda_backend::cudarc::driver::SyncOnDrop<'a>,
806    )> {
807        self.storage.device_ptr_with_guard(stream)
808    }
809}
810
811#[derive(Clone, Debug)]
812pub enum QMatMul {
813    QTensor(std::sync::Arc<QTensor>),
814    Tensor(Tensor),
815    TensorF16(Tensor),
816}
817
818thread_local! {
819    static DEQUANTIZE_ALL: bool = {
820        match std::env::var("CANDLE_DEQUANTIZE_ALL") {
821            Ok(s) => {
822                !s.is_empty() && s != "0"
823            },
824            Err(_) => false,
825        }
826    }
827}
828
829thread_local! {
830    static DEQUANTIZE_ALL_F16: bool = {
831        match std::env::var("CANDLE_DEQUANTIZE_ALL_F16") {
832            Ok(s) => {
833                !s.is_empty() && s != "0"
834            },
835            Err(_) => false,
836        }
837    }
838}
839
840impl QMatMul {
841    pub fn from_arc(qtensor: std::sync::Arc<QTensor>) -> Result<Self> {
842        let dequantize = match qtensor.dtype() {
843            GgmlDType::F32 | GgmlDType::F16 | GgmlDType::BF16 => true,
844            _ => DEQUANTIZE_ALL.with(|b| *b),
845        };
846        let t = if dequantize {
847            let tensor = qtensor.dequantize(&qtensor.device())?;
848            Self::Tensor(tensor)
849        } else if DEQUANTIZE_ALL_F16.with(|b| *b) {
850            let tensor = qtensor.dequantize_f16(&qtensor.device())?;
851            Self::TensorF16(tensor)
852        } else {
853            Self::QTensor(qtensor)
854        };
855        Ok(t)
856    }
857
858    pub fn from_qtensor(qtensor: QTensor) -> Result<Self> {
859        Self::from_arc(std::sync::Arc::new(qtensor))
860    }
861
862    pub fn dequantize_f16(&self) -> Result<Tensor> {
863        match self {
864            Self::QTensor(t) => t.dequantize_f16(&t.device()),
865            Self::Tensor(t) => t.to_dtype(DType::F16),
866            Self::TensorF16(t) => Ok(t.clone()),
867        }
868    }
869
870    pub fn forward_via_f16(&self, xs: &Tensor) -> Result<Tensor> {
871        let w = self.dequantize_f16()?;
872        let in_dtype = xs.dtype();
873        let w = match *xs.dims() {
874            [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
875            [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
876            _ => w.t()?,
877        };
878        xs.to_dtype(DType::F16)?.matmul(&w)?.to_dtype(in_dtype)
879    }
880
881    pub fn indexed_moe_forward(&self, x: &Tensor, ids: &Tensor) -> Result<Tensor> {
882        match self {
883            Self::QTensor(t) => t.indexed_moe_forward(x, ids),
884            _ => {
885                panic!("Not implemented!")
886            }
887        }
888    }
889
890    pub fn embedding(&self, ids: &Tensor) -> Result<Tensor> {
891        match self {
892            Self::QTensor(t) => t.embedding(ids),
893            Self::Tensor(w) | Self::TensorF16(w) => {
894                let mut final_dims = ids.dims().to_vec();
895                final_dims.push(w.dim(D::Minus1)?);
896                let ids = ids.to_device(w.device())?.flatten_all()?;
897                w.index_select(&ids, 0)?.reshape(final_dims)
898            }
899        }
900    }
901}
902
903impl QTensor {
904    /// Fused m==1 matmul over same-dtype tensors sharing one lhs (e.g. qkv or gate+up in
905    /// decode): one lhs quantization and one parallel region. None when unsupported.
906    pub fn gemv_fused_shared_lhs(ts: &[&Self], lhs: &Tensor) -> Result<Option<Vec<Tensor>>> {
907        #[cfg(target_arch = "aarch64")]
908        {
909            if ts.is_empty() || !lhs.device().is_cpu() || lhs.dtype() != crate::DType::F32 {
910                return Ok(None);
911            }
912            if !lhs.is_contiguous() {
913                return Ok(None);
914            }
915            let dims = lhs.dims();
916            let Some((&k, batch)) = dims.split_last() else {
917                return Ok(None);
918            };
919            if batch.iter().product::<usize>() != 1 {
920                return Ok(None);
921            }
922            let mut ns = Vec::with_capacity(ts.len());
923            let mut parts: Vec<(&dyn QuantizedType, &repack::PackedCache)> =
924                Vec::with_capacity(ts.len());
925            for t in ts {
926                let (n, tk) = t.shape.dims2()?;
927                if tk != k {
928                    return Ok(None);
929                }
930                let QStorage::Cpu(s) = &t.storage else {
931                    return Ok(None);
932                };
933                ns.push(n);
934                parts.push((s.as_ref(), &t.repacked_qs));
935            }
936            let storage = lhs.storage();
937            let crate::Storage::Cpu(cpu) = &*storage else {
938                return Ok(None);
939            };
940            let slice = cpu.as_slice::<f32>()?;
941            let offset = lhs.layout().start_offset();
942            let slice = &slice[offset..offset + k];
943            let mut dsts: Vec<Vec<f32>> = ns.iter().map(|&n| vec![0f32; n]).collect();
944            if !repack::try_gemv_fused(k, slice, &parts, &mut dsts)? {
945                return Ok(None);
946            }
947            drop(storage);
948            let outs = dsts
949                .into_iter()
950                .zip(&ns)
951                .map(|(d, &n)| {
952                    let mut shape = dims.to_vec();
953                    *shape.last_mut().unwrap() = n;
954                    Tensor::from_vec(d, shape, &crate::Device::Cpu)
955                })
956                .collect::<Result<Vec<_>>>()?;
957            Ok(Some(outs))
958        }
959        #[cfg(target_arch = "x86_64")]
960        {
961            // the fused kernel is 512-bit only; downlevel tiers fall through to
962            // per-matmul calls which dispatch to the 256-bit kernels
963            if repack_x86::level() != Some(repack_x86::X86Level::Avx512Vnni) {
964                return Ok(None);
965            }
966            if ts.is_empty() || !lhs.device().is_cpu() || lhs.dtype() != crate::DType::F32 {
967                return Ok(None);
968            }
969            if !lhs.is_contiguous() {
970                return Ok(None);
971            }
972            let dims = lhs.dims();
973            let Some((&k, batch)) = dims.split_last() else {
974                return Ok(None);
975            };
976            let m: usize = batch.iter().product();
977            if m == 0 {
978                return Ok(None);
979            }
980            let dtype = ts[0].dtype();
981            let mut shapes = Vec::with_capacity(ts.len());
982            for t in ts {
983                let Ok((n, tk)) = t.shape.dims2() else {
984                    return Ok(None);
985                };
986                if tk != k
987                    || t.dtype() != dtype
988                    || !matches!(&t.storage, QStorage::Cpu(_))
989                    || !repack_x86::select(dtype, m, n, k)
990                {
991                    return Ok(None);
992                }
993                shapes.push(n);
994            }
995            let guard = lhs.storage();
996            let crate::Storage::Cpu(cpu) = &*guard else {
997                return Ok(None);
998            };
999            let slice = cpu.as_slice::<f32>()?;
1000            let offset = lhs.layout().start_offset();
1001            let lhs_data = &slice[offset..offset + m * k];
1002            let lhs_q = repack_x86::quantize_lhs(lhs_data, m, k);
1003            let packs: Vec<&repack_x86::PackedX86> = ts
1004                .iter()
1005                .enumerate()
1006                .map(|(i, t)| {
1007                    let n = shapes[i];
1008                    t.repacked_qs
1009                        .x86_get_or_init(|| repack_x86::pack(dtype, t.storage_ref(), n, k))
1010                })
1011                .collect();
1012            let parts: Vec<(&repack_x86::PackedX86, usize)> = packs
1013                .iter()
1014                .zip(shapes.iter())
1015                .map(|(p, &n)| (*p, n))
1016                .collect();
1017            let mut dsts: Vec<Vec<f32>> = shapes.iter().map(|&n| vec![0f32; m * n]).collect();
1018            repack_x86::gemv_fused(&parts, &lhs_q, m, k, &mut dsts);
1019            drop(guard);
1020            let mut out = Vec::with_capacity(ts.len());
1021            let mut out_dims = dims.to_vec();
1022            for (dst, &n) in dsts.into_iter().zip(shapes.iter()) {
1023                *out_dims.last_mut().unwrap() = n;
1024                out.push(Tensor::from_vec(
1025                    dst,
1026                    out_dims.clone(),
1027                    &crate::Device::Cpu,
1028                )?);
1029            }
1030            Ok(Some(out))
1031        }
1032
1033        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
1034        {
1035            let _ = (ts, lhs);
1036            Ok(None)
1037        }
1038    }
1039}
1040
1041impl QTensor {
1042    #[cfg(target_arch = "x86_64")]
1043    pub(crate) fn storage_ref(&self) -> &dyn QuantizedType {
1044        match &self.storage {
1045            QStorage::Cpu(s) => s.as_ref(),
1046            _ => unreachable!("cpu-only path"),
1047        }
1048    }
1049
1050    /// Indexed (MoE) matmul over stacked expert weights [n_experts, n_out, k]: each
1051    /// (token, expert-id) pair gemvs against its expert's rows via the repacked cache.
1052    /// Returns None when the layout, dtype, or device is unsupported.
1053    pub fn indexed_gemv(&self, x: &Tensor, ids: &Tensor) -> Result<Option<Tensor>> {
1054        #[cfg(target_arch = "aarch64")]
1055        {
1056            if !x.device().is_cpu()
1057                || !matches!(x.dtype(), crate::DType::F32 | crate::DType::BF16)
1058                || !x.is_contiguous()
1059            {
1060                return Ok(None);
1061            }
1062            let Ok((n_experts, n_out, k)) = self.shape.dims3() else {
1063                return Ok(None);
1064            };
1065            let Ok((batch, x_t, xk)) = x.dims3() else {
1066                return Ok(None);
1067            };
1068            let Ok((ids_b, topk)) = ids.dims2() else {
1069                return Ok(None);
1070            };
1071            if xk != k || ids_b != batch || (x_t != 1 && x_t != topk) {
1072                return Ok(None);
1073            }
1074            let QStorage::Cpu(storage) = &self.storage else {
1075                return Ok(None);
1076            };
1077            let ids_v: Vec<u32> = ids
1078                .to_dtype(crate::DType::U32)?
1079                .flatten_all()?
1080                .to_vec1::<u32>()?;
1081            if ids_v.iter().any(|&e| e as usize >= n_experts) {
1082                crate::bail!("expert index out of range");
1083            }
1084            let guard = x.storage();
1085            let crate::Storage::Cpu(cpu) = &*guard else {
1086                return Ok(None);
1087            };
1088            let offset = x.layout().start_offset();
1089            let n_rows = batch * x_t;
1090            let widened;
1091            let lhs: &[f32] = if x.dtype() == crate::DType::BF16 {
1092                let bslice = cpu.as_slice::<half::bf16>()?;
1093                widened = widen_bf16(&bslice[offset..offset + n_rows * k]);
1094                &widened
1095            } else {
1096                let slice = cpu.as_slice::<f32>()?;
1097                &slice[offset..offset + n_rows * k]
1098            };
1099            let mut dst = vec![0f32; batch * topk * n_out];
1100            let ok = repack::try_indexed_gemv(
1101                storage.as_ref(),
1102                &self.repacked_qs,
1103                n_experts,
1104                n_out,
1105                k,
1106                lhs,
1107                n_rows,
1108                &ids_v,
1109                topk,
1110                &mut dst,
1111            )?;
1112            drop(guard);
1113            if !ok {
1114                return Ok(None);
1115            }
1116            let out = Tensor::from_vec(dst, (batch, topk, n_out), &crate::Device::Cpu)?;
1117            if x.dtype() == crate::DType::BF16 {
1118                Ok(Some(out.to_dtype(crate::DType::BF16)?))
1119            } else {
1120                Ok(Some(out))
1121            }
1122        }
1123        #[cfg(not(target_arch = "aarch64"))]
1124        {
1125            let _ = (x, ids);
1126            Ok(None)
1127        }
1128    }
1129}
1130
1131fn widen_bf16(src: &[half::bf16]) -> Vec<f32> {
1132    let mut out: Vec<f32> = Vec::with_capacity(src.len());
1133    let out_ptr = out.as_mut_ptr() as usize;
1134    let n_units = src.len().div_ceil(WIDEN_CHUNK);
1135    crate::utils::barrier_pool().execute_chunked(n_units, |range| {
1136        let out_ptr = out_ptr as *mut f32;
1137        for unit in range {
1138            let lo = unit * WIDEN_CHUNK;
1139            let hi = src.len().min(lo + WIDEN_CHUNK);
1140            for (i, v) in src[lo..hi].iter().enumerate() {
1141                unsafe { *out_ptr.add(lo + i) = v.to_f32() };
1142            }
1143        }
1144    });
1145    // SAFETY: every element written by exactly one unit.
1146    unsafe { out.set_len(src.len()) };
1147    out
1148}
1149
1150fn narrow_bf16(src: &[f32]) -> Vec<half::bf16> {
1151    let mut out: Vec<half::bf16> = Vec::with_capacity(src.len());
1152    let out_ptr = out.as_mut_ptr() as usize;
1153    let n_units = src.len().div_ceil(WIDEN_CHUNK);
1154    crate::utils::barrier_pool().execute_chunked(n_units, |range| {
1155        let out_ptr = out_ptr as *mut half::bf16;
1156        for unit in range {
1157            let lo = unit * WIDEN_CHUNK;
1158            let hi = src.len().min(lo + WIDEN_CHUNK);
1159            for (i, v) in src[lo..hi].iter().enumerate() {
1160                unsafe { *out_ptr.add(lo + i) = half::bf16::from_f32(*v) };
1161            }
1162        }
1163    });
1164    // SAFETY: every element written by exactly one unit.
1165    unsafe { out.set_len(src.len()) };
1166    out
1167}
1168
1169const WIDEN_CHUNK: usize = 32 * 1024;
1170
1171impl crate::CustomOp1 for QTensor {
1172    fn name(&self) -> &'static str {
1173        "qmatmul"
1174    }
1175
1176    fn cpu_fwd(
1177        &self,
1178        storage: &crate::CpuStorage,
1179        layout: &crate::Layout,
1180    ) -> Result<(crate::CpuStorage, Shape)> {
1181        if !layout.is_contiguous() {
1182            crate::bail!("input tensor is not contiguous {layout:?}")
1183        }
1184        let src_shape = layout.shape();
1185        // self is transposed so n is first then k.
1186        let (n, k) = self.shape.dims2()?;
1187        if src_shape.rank() < 2 {
1188            crate::bail!("input tensor has only one dimension {layout:?}")
1189        }
1190        let mut dst_shape = src_shape.dims().to_vec();
1191        let last_k = dst_shape.pop().unwrap();
1192        if last_k != k {
1193            crate::bail!("input tensor {layout:?} incompatible with {:?}", self.shape)
1194        }
1195        dst_shape.push(n);
1196        let dst_shape = Shape::from(dst_shape);
1197        #[allow(clippy::infallible_destructuring_match)]
1198        let self_storage = match &self.storage {
1199            QStorage::Cpu(storage) => storage,
1200            QStorage::Metal(_) | QStorage::Cuda(_) => crate::bail!("Invalid storage"),
1201        };
1202        match storage.dtype() {
1203            DType::F32 => {
1204                let slice = storage.as_slice::<f32>()?;
1205                let slice =
1206                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
1207                let mut dst_storage = vec![0f32; dst_shape.elem_count()];
1208
1209                let mkn = (dst_shape.elem_count() / n, k, n);
1210                let used_packed = repack::try_matmul_f32(
1211                    self_storage.as_ref(),
1212                    &self.repacked_qs,
1213                    mkn,
1214                    slice,
1215                    &mut dst_storage,
1216                )?;
1217                if used_packed {
1218                    return Ok((crate::CpuStorage::F32(dst_storage), dst_shape));
1219                }
1220
1221                self_storage.matmul_t(mkn, slice, &mut dst_storage)?;
1222                Ok((crate::CpuStorage::F32(dst_storage), dst_shape))
1223            }
1224            DType::F16 => {
1225                let slice = storage.as_slice::<f16>()?;
1226                let slice =
1227                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
1228                let mut dst_storage = vec![f16::ZERO; dst_shape.elem_count()];
1229                self_storage.matmul_t_f16(
1230                    (dst_shape.elem_count() / n, k, n),
1231                    slice,
1232                    &mut dst_storage,
1233                )?;
1234                Ok((crate::CpuStorage::F16(dst_storage), dst_shape))
1235            }
1236            DType::BF16 => {
1237                // widen to f32 once and take the repacked path; output stays bf16
1238                let slice = storage.as_slice::<half::bf16>()?;
1239                let slice =
1240                    &slice[layout.start_offset()..layout.start_offset() + src_shape.elem_count()];
1241                let lhs = widen_bf16(slice);
1242                let mut dst_storage = vec![0f32; dst_shape.elem_count()];
1243
1244                let mkn = (dst_shape.elem_count() / n, k, n);
1245                let used_packed = repack::try_matmul_f32(
1246                    self_storage.as_ref(),
1247                    &self.repacked_qs,
1248                    mkn,
1249                    &lhs,
1250                    &mut dst_storage,
1251                )?;
1252                if !used_packed {
1253                    self_storage.matmul_t(mkn, &lhs, &mut dst_storage)?;
1254                }
1255                let dst: Vec<half::bf16> = narrow_bf16(&dst_storage);
1256                Ok((crate::CpuStorage::BF16(dst), dst_shape))
1257            }
1258            _ => crate::bail!("Expected f32/f16/bf16"),
1259        }
1260    }
1261
1262    fn metal_fwd(
1263        &self,
1264        storage: &crate::MetalStorage,
1265        layout: &crate::Layout,
1266    ) -> Result<(crate::MetalStorage, Shape)> {
1267        let self_storage = match &self.storage {
1268            QStorage::Metal(metal) => metal,
1269            _ => unreachable!("Cannot call metal matmul on non metal QTensor"),
1270        };
1271        self_storage.fwd(&self.shape, storage, layout)
1272    }
1273
1274    fn cuda_fwd(
1275        &self,
1276        storage: &crate::CudaStorage,
1277        layout: &crate::Layout,
1278    ) -> Result<(crate::CudaStorage, Shape)> {
1279        let self_storage = match &self.storage {
1280            QStorage::Cuda(cuda) => cuda,
1281            _ => unreachable!("Cannot call cuda matmul on non cuda QTensor"),
1282        };
1283        self_storage.fwd(&self.shape, storage, layout)
1284    }
1285}
1286
1287impl crate::Module for QMatMul {
1288    fn forward(&self, xs: &Tensor) -> Result<Tensor> {
1289        match self {
1290            Self::QTensor(t) => xs.apply_op1_no_bwd(t.as_ref()),
1291            Self::Tensor(w) => {
1292                let w = match *xs.dims() {
1293                    [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
1294                    [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
1295                    _ => w.t()?,
1296                };
1297                xs.matmul(&w)
1298            }
1299            Self::TensorF16(w) => {
1300                let in_dtype = xs.dtype();
1301                let w = match *xs.dims() {
1302                    [b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
1303                    [bsize, _, _] => w.broadcast_left(bsize)?.t()?,
1304                    _ => w.t()?,
1305                };
1306                xs.to_dtype(DType::F16)?.matmul(&w)?.to_dtype(in_dtype)
1307            }
1308        }
1309    }
1310}
candle_core/quantized/mod.rs

candle_core/quantized/
mod.rs
