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

Re-revert the reverted revision (bf16 gemm metal)

Browse Source
This commit is contained in:
Ivar Flakstad
2024-08-06 12:37:36 +02:00
parent bb191c25d5
commit e178aacead
2 changed files with 21 additions and 23 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(_) => true,
Self::Metal(_) | Self::Cpu => false,
Self::Cuda(_) | Self::Metal(_) => true,
Self::Cpu => false,
}
}

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

@ -1162,27 +1162,25 @@ fn gemm() {
);
// bgemm sanity test
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]
);
}
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);