e2182cc21d
* Integrate ck branch of ck_tile/fa_bwd_opt * Assume dq and q share the same stride * update ck * Integrate more stride of dq_acc * Revert fwd dropout * Fix paremeter order * Integrate ck with more stride * update the limit of hdim of bwd * Check argument * Add test_flash_attn_causal * Support unpad lse * Add test_flash_attn_varlen_causal, test_flash_attn_race_condition, test_flash_attn_bwd_overflow, test_flash_attn_bwd_transpose, test_flash_attn_bwd_varlen_overflow, test_flash_attn_deterministic, test_flash_attn_varlen_deterministic * Fix stride and Kn0 * Fix CK sync issue * Fix typo * Update CK for changing of fmha_fwd_args * Add kvcache tmp * Add kvcache * Fix comment * Sync behavior with ck * Update CK to develop * remove large test case * Add kvcache test * Fix page_block_size in arg * Minor fix * Fix stride error * Update seqlen of kvcache before splitkv * Fix compile error * Fix bug of hdim is not 8x * Fit ck arg * support adaptive num_splits * add more tests * Refine test tolerance * update CK * Move override_num_splits_if_necessary into cpp * update ck * Update ck * Support different flag for different version of hip * remove coerce-illegal, becasue this is not required in FA * Update ck to fix xcratch memory * Add coerce-illegal in some version * Add compile flag for rtn rounding * remove redundant init * Using env var to switch rounding mode * update ck
123 lines
7.6 KiB
C++
123 lines
7.6 KiB
C++
/******************************************************************************
|
|
* Copyright (c) 2024, Tri Dao.
|
|
******************************************************************************/
|
|
|
|
#include "flash_common.hpp"
|
|
|
|
std::vector<at::Tensor>
|
|
mha_fwd(at::Tensor &q,
|
|
const at::Tensor &k,
|
|
const at::Tensor &v,
|
|
c10::optional<at::Tensor> &out_,
|
|
c10::optional<at::Tensor> &alibi_slopes_,
|
|
const float p_dropout,
|
|
const float softmax_scale,
|
|
bool is_causal,
|
|
int window_size_left,
|
|
int window_size_right,
|
|
const float softcap,
|
|
const bool return_softmax,
|
|
c10::optional<at::Generator> gen_);
|
|
|
|
std::vector<at::Tensor>
|
|
mha_varlen_fwd(at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
|
|
const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
|
|
const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
|
|
c10::optional<at::Tensor> &out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
|
|
const at::Tensor &cu_seqlens_q, // b+1
|
|
const at::Tensor &cu_seqlens_k, // b+1
|
|
c10::optional<at::Tensor> &seqused_k, // b. If given, only this many elements of each batch element's keys are used.
|
|
c10::optional<const at::Tensor> &leftpad_k_, // batch_size
|
|
c10::optional<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
|
|
c10::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
|
|
int max_seqlen_q,
|
|
const int max_seqlen_k,
|
|
const float p_dropout,
|
|
const float softmax_scale,
|
|
const bool zero_tensors,
|
|
bool is_causal,
|
|
int window_size_left,
|
|
int window_size_right,
|
|
const float softcap,
|
|
const bool return_softmax,
|
|
c10::optional<at::Generator> gen_);
|
|
|
|
std::vector<at::Tensor>
|
|
mha_bwd(const at::Tensor &dout, // batch_size x seqlen_q x num_heads, x head_size_og
|
|
const at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
|
|
const at::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size
|
|
const at::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size
|
|
const at::Tensor &out, // batch_size x seqlen_q x num_heads x head_size
|
|
const at::Tensor &softmax_lse, // b x h x seqlen_q
|
|
c10::optional<at::Tensor> &dq_, // batch_size x seqlen_q x num_heads x head_size
|
|
c10::optional<at::Tensor> &dk_, // batch_size x seqlen_k x num_heads_k x head_size
|
|
c10::optional<at::Tensor> &dv_, // batch_size x seqlen_k x num_heads_k x head_size
|
|
c10::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
|
|
const float p_dropout, // probability to drop
|
|
const float softmax_scale,
|
|
const bool is_causal,
|
|
int window_size_left,
|
|
int window_size_right,
|
|
const float softcap,
|
|
const bool deterministic,
|
|
c10::optional<at::Generator> gen_,
|
|
c10::optional<at::Tensor> &rng_state);
|
|
|
|
std::vector<at::Tensor>
|
|
mha_varlen_bwd(const at::Tensor &dout, // total_q x num_heads x head_size
|
|
const at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
|
|
const at::Tensor &k, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
|
|
const at::Tensor &v, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
|
|
const at::Tensor &out, // total_q x num_heads x head_size
|
|
const at::Tensor &softmax_lse, // b x h x s softmax logsumexp
|
|
c10::optional<at::Tensor> &dq_, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
|
|
c10::optional<at::Tensor> &dk_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
|
|
c10::optional<at::Tensor> &dv_, // total_k x num_heads_k x head_size, total_k := \sum_{i=0}^{b} s_i
|
|
const at::Tensor &cu_seqlens_q, // b+1
|
|
const at::Tensor &cu_seqlens_k, // b+1
|
|
c10::optional<at::Tensor> &alibi_slopes_, // num_heads or b x num_heads
|
|
const int max_seqlen_q,
|
|
const int max_seqlen_k, // max sequence length to choose the kernel
|
|
const float p_dropout, // probability to drop
|
|
const float softmax_scale,
|
|
const bool zero_tensors,
|
|
const bool is_causal,
|
|
int window_size_left,
|
|
int window_size_right,
|
|
const float softcap,
|
|
const bool deterministic,
|
|
c10::optional<at::Generator> gen_,
|
|
c10::optional<at::Tensor> &rng_state);
|
|
|
|
std::vector<at::Tensor>
|
|
mha_fwd_kvcache(at::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
|
|
const at::Tensor &kcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
|
|
const at::Tensor &vcache, // batch_size_c x seqlen_k x num_heads_k x head_size or num_blocks x page_block_size x num_heads_k x head_size if there's a block_table.
|
|
c10::optional<const at::Tensor> &k_, // batch_size x seqlen_knew x num_heads_k x head_size
|
|
c10::optional<const at::Tensor> &v_, // batch_size x seqlen_knew x num_heads_k x head_size
|
|
c10::optional<const at::Tensor> &seqlens_k_, // batch_size
|
|
c10::optional<const at::Tensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
|
|
c10::optional<const at::Tensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
|
|
c10::optional<const at::Tensor> &cache_batch_idx_, // indices to index into the KV cache
|
|
c10::optional<const at::Tensor> &leftpad_k_, // batch_size
|
|
c10::optional<at::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
|
|
c10::optional<at::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
|
|
c10::optional<at::Tensor> &out_, // batch_size x seqlen_q x num_heads x head_size
|
|
const float softmax_scale,
|
|
bool is_causal,
|
|
int window_size_left,
|
|
int window_size_right,
|
|
const float softcap,
|
|
bool is_rotary_interleaved, // if true, rotary combines indices 0 & 1, else indices 0 & rotary_dim / 2
|
|
int num_splits);
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
m.doc() = "FlashAttention";
|
|
m.def("fwd", &mha_fwd, "Forward pass");
|
|
m.def("varlen_fwd", &mha_varlen_fwd, "Forward pass (variable length)");
|
|
m.def("bwd", &mha_bwd, "Backward pass");
|
|
m.def("varlen_bwd", &mha_varlen_bwd, "Backward pass (variable length)");
|
|
m.def("fwd_kvcache", &mha_fwd_kvcache, "Forward pass, with KV-cache");
|
|
}
|