huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 02:38:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a forward_via_f16 method to the qmatmul op. (#2138)

Browse Source
This commit is contained in:
Laurent Mazare
2024-04-28 20:35:01 +02:00
committed by GitHub
parent eb26e2467e
commit 287013ef28
1 changed files with 19 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
candle-core/src/quantized/mod.rs
View File

@ -439,6 +439,25 @@ impl QMatMul {
pub fn from_qtensor(qtensor: QTensor) -> Result<Self> {
Self::from_arc(std::sync::Arc::new(qtensor))
}
pub fn dequantize_f16(&self) -> Result<Tensor> {
match self {
Self::QTensor(t) => t.dequantize_f16(&t.device()),
Self::Tensor(t) => t.to_dtype(DType::F16),
Self::TensorF16(t) => Ok(t.clone()),
}
}
pub fn forward_via_f16(&self, xs: &Tensor) -> Result<Tensor> {
let w = self.dequantize_f16()?;
let in_dtype = xs.dtype();
let w = match *xs.dims() {
[b1, b2, _, _] => w.broadcast_left((b1, b2))?.t()?,
[bsize, _, _] => w.broadcast_left(bsize)?.t()?,
_ => w.t()?,
};
xs.to_dtype(DType::F16)?.matmul(&w)?.to_dtype(in_dtype)
}
}
impl crate::CustomOp1 for QTensor {