replicate/flash-attn2

1#pragma once2 3#include <torch/torch.h>4 5// std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>6std::vector<torch::Tensor>7mha_fwd(8    torch::Tensor &q, 9    const torch::Tensor &k, 10    const torch::Tensor &v,11    c10::optional<torch::Tensor> out_,\12    c10::optional<torch::Tensor> alibi_slopes_,13    const double p_dropout, 14    const double softmax_scale, 15    bool is_causal,16    const int64_t window_size_left, 17    const int64_t window_size_right,18    const double softcap, 19    const bool return_softmax,20    c10::optional<at::Generator> gen_);21 22std::vector<torch::Tensor>23mha_varlen_fwd(24    at::Tensor &q,  // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i25    const torch::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.26    const torch::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.27    c10::optional<torch::Tensor> out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i28    const torch::Tensor &cu_seqlens_q,  // b+129    const torch::Tensor &cu_seqlens_k,  // b+130    c10::optional<torch::Tensor> seqused_k, // b. If given, only this many elements of each batch element's keys are used.31    // c10::optional<const at::Tensor> leftpad_k_, // batch_size32    c10::optional<torch::Tensor> leftpad_k_, // batch_size33    c10::optional<torch::Tensor> block_table_, // batch_size x max_num_blocks_per_seq34    c10::optional<torch::Tensor> alibi_slopes_, // num_heads or b x num_heads35    int64_t max_seqlen_q,36    const int64_t max_seqlen_k,37    const double p_dropout,38    const double softmax_scale,39    const bool zero_tensors,40    bool is_causal,41    int64_t window_size_left,42    int64_t window_size_right,43    const double softcap,44    const bool return_softmax,45    std::optional<at::Generator> gen_); 46 47std::vector<torch::Tensor>48mha_bwd(const torch::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)49        const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size50        const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size51        const torch::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size52        const torch::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size53        const torch::Tensor &softmax_lse,                  // b x h x seqlen_q54        const c10::optional<torch::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size55        const c10::optional<torch::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size56        const c10::optional<torch::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size57        const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads58        const double p_dropout,                            // probability to drop59        const double softmax_scale,60        const bool is_causal,61        const int64_t window_size_left,62        const int64_t window_size_right,63        const double softcap,64        const bool deterministic,65        c10::optional<at::Generator> gen_,66        const c10::optional<torch::Tensor> &rng_state);67 68 69std::vector<torch::Tensor>70mha_varlen_bwd(71        const torch::Tensor &dout,                         // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)72        const torch::Tensor &q,                            // batch_size x seqlen_q x num_heads x head_size73        const torch::Tensor &k,                            // batch_size x seqlen_k x num_heads_k x head_size74        const torch::Tensor &v,                            // batch_size x seqlen_k x num_heads_k x head_size75        const torch::Tensor &out,                          // batch_size x seqlen_q x num_heads x head_size76        const torch::Tensor &softmax_lse,                  // b x h x seqlen_q77        const c10::optional<torch::Tensor> &dq_,           // batch_size x seqlen_q x num_heads x head_size78        const c10::optional<torch::Tensor> &dk_,           // batch_size x seqlen_k x num_heads_k x head_size79        const c10::optional<torch::Tensor> &dv_,           // batch_size x seqlen_k x num_heads_k x head_size80        const torch::Tensor &cu_seqlens_q,                 // batch_size + 181        const torch::Tensor &cu_seqlens_k,                 // batch_size + 182        const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or b x num_heads83        const int64_t max_seqlen_q,84        const int64_t max_seqlen_k,85        const double p_dropout,86        const double softmax_scale,87        const bool zero_tensors,88        const bool is_causal,89        const int64_t window_size_left,90        const int64_t window_size_right,91        const double softcap,92        const bool deterministic,93        c10::optional<at::Generator> gen_,94        const c10::optional<torch::Tensor> &rng_state);95 96std::vector<torch::Tensor>97mha_fwd_kvcache(98        const torch::Tensor &q,                                // batch_size x seqlen_q x num_heads x head_size99        const torch::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.100        const torch::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.101        const c10::optional<torch::Tensor> &k_,                // batch_size x seqlen_knew x num_heads_k x head_size102        const c10::optional<torch::Tensor> &v_,                // batch_size x seqlen_knew x num_heads_k x head_size103        const c10::optional<torch::Tensor> &seqlens_k_,        // batch_size104        const c10::optional<torch::Tensor> &rotary_cos_,       // seqlen_ro x (rotary_dim / 2)105        const c10::optional<torch::Tensor> &rotary_sin_,       // seqlen_ro x (rotary_dim / 2)106        const c10::optional<torch::Tensor> &cache_batch_idx_,  // indices to index into the KV cache107        const c10::optional<torch::Tensor> &leftpad_k_,        // batch_size108        const c10::optional<torch::Tensor> &block_table_,      // batch_size x max_num_blocks_per_seq109        const c10::optional<torch::Tensor> &alibi_slopes_,     // num_heads or batch_size x num_heads110        const c10::optional<torch::Tensor> &out_,              // batch_size x seqlen_q x num_heads x head_size111        const double softmax_scale,112        bool is_causal,113        const int64_t window_size_left,114        const int64_t window_size_right,115        const double softcap,116        bool is_rotary_interleaved,117        const int64_t num_splits);118