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

Add back the bf16 flash-attn kernels. (#730)

Browse Source
This commit is contained in:
Laurent Mazare
2023-09-04 08:50:52 +02:00
committed by GitHub
parent 20512ba408
commit d0cdea95a5
4 changed files with 25 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
candle-flash-attn/src/lib.rs
View File

@ -146,6 +146,7 @@ impl candle::CustomOp3 for FlashAttn {
/* seqlen_q_rounded */ seqlen_q_rounded as u32,
/* seqlen_k_rounded */ seqlen_k_rounded as u32,
/* is_causal */ causal,
/* is_bf16 */ 0,
)
}
@ -354,6 +355,7 @@ impl candle::CustomOp3 for FlashAttnVarLen {
/* seqlen_q_rounded */ seqlen_q_rounded as u32,
/* seqlen_k_rounded */ seqlen_k_rounded as u32,
/* is_causal */ causal,
/* is_bf16 */ 0,
)
}