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

JySongWithZhangCe · Aug 1, 2024 · 0fcb40b · 0fcb40b
1 parent 6991a37
commit 0fcb40b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 21 deletions.
diff --git a/candle-core/src/device.rs b/candle-core/src/device.rs
@@ -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,
         }
     }
 

diff --git a/candle-metal-kernels/src/libMetalFlashAttention.metallib b/candle-metal-kernels/src/libMetalFlashAttention.metallib
diff --git a/candle-metal-kernels/src/tests.rs b/candle-metal-kernels/src/tests.rs
@@ -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);