Add more QMMV cuda kernels. (#2077)

* Add more QMMV cuda kernels.

* Enable the new kernels.

* Adapt the testing.
This commit is contained in:
Laurent Mazare
2024-04-18 08:36:43 +02:00
committed by GitHub
parent ce6d08df94
commit 8de0ce6cba
3 changed files with 349 additions and 15 deletions

View File

@ -193,17 +193,25 @@ fn qmm_batch(dev: &Device) -> Result<()> {
let mm3 = rhs.forward(&lhs3)?;
assert_eq!(mm3.shape().dims(), [6, 6]);
let diff3 = (mm3.i(2..4)? - &mm)?.abs()?.sum_all()?.to_vec0::<f32>()?;
if dev.is_cuda() {
assert!(diff3 < 1e-4)
} else {
assert_eq!(diff3, 0.0)
};
assert_eq!(diff3, 0.0);
let diff3 = (mm3.i(4..)? - &mm)?.abs()?.sum_all()?.to_vec0::<f32>()?;
assert_eq!(diff3, 0.0);
let lhs4 = Tensor::cat(&[&lhs3, &lhs3], 0)?;
let mm4 = rhs.forward(&lhs4)?;
assert_eq!(mm4.shape().dims(), [12, 6]);
let diff4 = (mm4.i(..6)? - &mm3)?.abs()?.sum_all()?.to_vec0::<f32>()?;
if dev.is_cuda() {
assert!(diff3 < 1e-4)
// We use a different kernel for sizes from 1 to 8 on cuda which explains
// the difference here.
assert!(0. < diff4 && diff4 < 1e-4)
} else {
assert_eq!(diff3, 0.0)
assert_eq!(diff4, 0.0)
};
let diff4 = (mm4.i(6..)? - &mm4.i(..6)?)?
.abs()?
.sum_all()?
.to_vec0::<f32>()?;
assert_eq!(diff4, 0.0);
Ok(())
}