From 67834119fcfd961b4852458abe9426fbdfb2fd76 Mon Sep 17 00:00:00 2001
From: Laurent Mazare <laurent.mazare@gmail.com>
Date: Mon, 31 Jul 2023 10:04:39 +0100
Subject: Fix the flash-attention function names. (#282)

---
 candle-flash-attn/src/lib.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

(limited to 'candle-flash-attn')

diff --git a/candle-flash-attn/src/lib.rs b/candle-flash-attn/src/lib.rs
index 99b05229..092743f1 100644
--- a/candle-flash-attn/src/lib.rs
+++ b/candle-flash-attn/src/lib.rs
@@ -17,7 +17,7 @@ fn round_multiple(x: usize, m: usize) -> usize {
 
 impl candle::CustomOp3 for FlashAttn {
     fn name(&self) -> &'static str {
-        "flash-hdim32-sm80"
+        "flash-attn"
     }
 
     fn cpu_fwd(
@@ -192,7 +192,7 @@ struct FlashAttnVarLen {
 
 impl candle::CustomOp3 for FlashAttnVarLen {
     fn name(&self) -> &'static str {
-        "flash-hdim32-sm80"
+        "flash-attn-varlen"
     }
 
     fn cpu_fwd(
-- 
cgit v1.2.3