FlashAttention-2 release

4f285b35 · Tri Dao · 6d48e14a · 6d48e14a · 6d48e14a · 6d48e14a
Commit 4f285b35 authored Jul 17, 2023 by Tri Dao
20 changed files
--- a/csrc/flash_attn/src/fmha_block_dgrad_kernel_1xN_loop.h
+++ b/csrc/flash_attn/src/fmha_block_dgrad_kernel_1xN_loop.h
--- a/csrc/flash_attn/src/fmha_block_fprop_fp16_kernel.sm80.cu
+++ b/csrc/flash_attn/src/fmha_block_fprop_fp16_kernel.sm80.cu
--- a/csrc/flash_attn/src/fmha_block_fprop_kernel_1xN.h
+++ b/csrc/flash_attn/src/fmha_block_fprop_kernel_1xN.h
--- a/csrc/flash_attn/src/fmha_blockmask.h
+++ b/csrc/flash_attn/src/fmha_blockmask.h
--- a/csrc/flash_attn/src/fmha_bwd_hdim128.cu
+++ b/csrc/flash_attn/src/fmha_bwd_hdim128.cu
-// Copyright (c) 2022, Tri Dao.
-// Splitting the different head dimensions to different files to speed up compilation.
-#include "fmha_bwd_launch_template.h"
-void run_fmha_bwd_hdim128(FMHA_dgrad_params &params, cudaStream_t stream, const bool configure) {
-    FP16_SWITCH(params.is_bf16, ([&] {
-        using Kernel_traits = FMHA_kernel_traits<128, 128, 16, 1, 8, 0x100u, elem_type>;
-        run_fmha_bwd_loop<Kernel_traits>(params, stream, configure);
-    }));
-}
\ No newline at end of file
--- a/csrc/flash_attn/src/fmha_bwd_hdim32.cu
+++ b/csrc/flash_attn/src/fmha_bwd_hdim32.cu
--- a/csrc/flash_attn/src/fmha_bwd_hdim64.cu
+++ b/csrc/flash_attn/src/fmha_bwd_hdim64.cu
--- a/csrc/flash_attn/src/fmha_bwd_launch_template.h
+++ b/csrc/flash_attn/src/fmha_bwd_launch_template.h
--- a/csrc/flash_attn/src/fmha_dgrad_kernel_1xN_loop.h
+++ b/csrc/flash_attn/src/fmha_dgrad_kernel_1xN_loop.h
--- a/csrc/flash_attn/src/fmha_fprop_kernel_1xN.h
+++ b/csrc/flash_attn/src/fmha_fprop_kernel_1xN.h
--- a/csrc/flash_attn/src/fmha_fwd_hdim128.cu
+++ b/csrc/flash_attn/src/fmha_fwd_hdim128.cu
--- a/csrc/flash_attn/src/fmha_fwd_hdim32.cu
+++ b/csrc/flash_attn/src/fmha_fwd_hdim32.cu
--- a/csrc/flash_attn/src/fmha_fwd_hdim64.cu
+++ b/csrc/flash_attn/src/fmha_fwd_hdim64.cu
--- a/csrc/flash_attn/src/fmha_fwd_launch_template.h
+++ b/csrc/flash_attn/src/fmha_fwd_launch_template.h
--- a/csrc/flash_attn/src/fmha_kernel.h
+++ b/csrc/flash_attn/src/fmha_kernel.h
--- a/csrc/flash_attn/src/fmha_utils.h
+++ b/csrc/flash_attn/src/fmha_utils.h
--- a/csrc/flash_attn/src/kernel_traits.h
+++ b/csrc/flash_attn/src/kernel_traits.h
--- a/csrc/flash_attn/src/kernel_traits_sm90.h
+++ b/csrc/flash_attn/src/kernel_traits_sm90.h
--- a/csrc/flash_attn/src/philox.cuh
+++ b/csrc/flash_attn/src/philox.cuh
--- a/csrc/flash_attn/src/softmax.h
+++ b/csrc/flash_attn/src/softmax.h