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

Handle the batch dimension in quantized MMV on metal. (#2022)

Browse Source
This commit is contained in:
Laurent Mazare
2024-04-06 20:02:24 +02:00
committed by GitHub
parent e662431acf
commit 9fd52b3b71
1 changed files with 4 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
candle-core/src/quantized/metal.rs
View File

@ -149,8 +149,11 @@ impl QMetalStorage {
let (n, k) = self_shape.dims2()?; let (n, k) = self_shape.dims2()?;
let mut dst_shape = src_shape.dims().to_vec(); let mut dst_shape = src_shape.dims().to_vec();
// We always use a single batch dimension and stack all the tensors in the batch on the
// second dimension as the implementation in candle-metal-kernels doesn't handle batch
// properly.
let (b, m) = match dst_shape.len() { let (b, m) = match dst_shape.len() {
3 => (dst_shape[0], dst_shape[1]), 3 => (1, dst_shape[0] * dst_shape[1]),
2 => (1, dst_shape[0]), 2 => (1, dst_shape[0]),
n => crate::bail!("Invalid rank {n} for quantized matmul metal"), n => crate::bail!("Invalid rank {n} for quantized matmul metal"),
}; };