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

Revert the bf16 gemm metal changes for now. (#2386)

Browse Source
This commit is contained in:
Laurent Mazare
2024-08-01 22:08:47 +01:00
committed by GitHub
parent 6991a37b94
commit 0fcb40b229
3 changed files with 23 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
candle-core/src/device.rs
View File

@ -173,8 +173,8 @@ impl Device {
pub fn supports_bf16(&self) -> bool {
match self {
Self::Cuda(_) | Self::Metal(_) => true,
Self::Cpu => false,
Self::Cuda(_) => true,
Self::Metal(_) | Self::Cpu => false,
}
}

BIN
candle-metal-kernels/src/libMetalFlashAttention.metallib
View File

Binary file not shown.

40
candle-metal-kernels/src/tests.rs
View File

@ -1162,25 +1162,27 @@ fn gemm() {
);
// bgemm sanity test
let (b, m, n, k) = (1, 2, 4, 3);
let lhs_stride = vec![m * k, k, 1];
let lhs: Vec<bf16> = (0..b * m * k).map(|f| bf16::from_f32(f as f32)).collect();
let rhs_stride = vec![n * k, n, 1];
let rhs: Vec<bf16> = (0..b * n * k).map(|f| bf16::from_f32(f as f32)).collect();
let results = run_gemm(
"bgemm",
(b, m, n, k),
&lhs,
lhs_stride,
0,
&rhs,
rhs_stride,
0,
);
assert_eq!(
approx_bf16(results, 4),
vec![20.0, 23.0, 26.0, 29.0, 56.0, 68.0, 80.0, 92.0]
);
if false {
let (b, m, n, k) = (1, 2, 4, 3);
let lhs_stride = vec![m * k, k, 1];
let lhs: Vec<bf16> = (0..b * m * k).map(|f| bf16::from_f32(f as f32)).collect();
let rhs_stride = vec![n * k, n, 1];
let rhs: Vec<bf16> = (0..b * n * k).map(|f| bf16::from_f32(f as f32)).collect();
let results = run_gemm(
"bgemm",
(b, m, n, k),
&lhs,
lhs_stride,
0,
&rhs,
rhs_stride,
0,
);
assert_eq!(
approx_bf16(results, 4),
vec![20.0, 23.0, 26.0, 29.0, 56.0, 68.0, 80.0, 92.0]
);
}
// hgemm sanity test
let (b, m, n, k) = (1, 2, 4, 3);