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

Fix for the RWKV models. (#1955)

Browse Source 
* Fix for the RWKV models.

* More general fix + revert the rwkv hack.

* Remove the old hack.
This commit is contained in:
Laurent Mazare
2024-03-28 10:17:38 +01:00
committed by GitHub
parent ada5d7c096
commit b3484e7a5e
2 changed files with 18 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
candle-metal-kernels/src/lib.rs
View File

@ -1454,9 +1454,9 @@ pub fn call_gemm(
// lhs has shape b, m, k
// We also allow for the case where the stride on the minor dimension is not as expected but
// there is a single element.
let a_trans = if lhs_m1 == 1 && (lhs_m2 == k || b * m == 1) {
let a_trans = if (lhs_m1 == 1 || k == 1) && (lhs_m2 == k || m == 1) {
false
} else if (lhs_m1 == m || b * k == 1) && lhs_m2 == 1 {
} else if (lhs_m1 == m || k == 1) && (lhs_m2 == 1 || m == 1) {
true
} else {
return Err(MetalKernelError::MatMulNonContiguous {
@ -1466,9 +1466,9 @@ pub fn call_gemm(
})?;
};
// rhs has shape b, k, n
let b_trans = if rhs_m1 == 1 && (rhs_m2 == n || b * k == 1) {
let b_trans = if (rhs_m1 == 1 || n == 1) && (rhs_m2 == n || k == 1) {
false
} else if (rhs_m1 == k || b * n == 1) && rhs_m2 == 1 {
} else if (rhs_m1 == k || n == 1) && (rhs_m2 == 1 || k == 1) {
true
} else {
return Err(MetalKernelError::MatMulNonContiguous {