diff options
Diffstat (limited to 'candle-metal-kernels/src/lib.rs')
| -rw-r--r-- | candle-metal-kernels/src/lib.rs | 2 |
1 files changed, 2 insertions, 0 deletions
diff --git a/candle-metal-kernels/src/lib.rs b/candle-metal-kernels/src/lib.rs index e0c97962..743b9fe2 100644 --- a/candle-metal-kernels/src/lib.rs +++ b/candle-metal-kernels/src/lib.rs @@ -19,6 +19,7 @@ const CAST: &str = include_str!("cast.metal"); const CONV: &str = include_str!("conv.metal"); const REDUCE: &str = include_str!("reduce.metal"); const RANDOM: &str = include_str!("random.metal"); +// Current source: https://github.com/ivarflakstad/metal-flash-attention/tree/candle const MFA: &[u8] = include_bytes!("libMetalFlashAttention.metallib"); const QUANTIZED: &str = include_str!("quantized.metal"); const SORT: &str = include_str!("sort.metal"); @@ -1564,6 +1565,7 @@ pub fn call_gemm( let bytes = match name { "sgemm" => 4, "hgemm" => 2, + "bgemm" => 2, other => { return Err(MetalKernelError::LoadLibraryError(format!( "{other} is not a valid kernel for gemm" |