File size: 6,989 Bytes
a7165c8 876ac68 ab4cc6a 876ac68 ab4cc6a 876ac68 ab4cc6a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 |
#pragma once
#include <torch/torch.h>
// std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
std::vector<torch::Tensor>
mha_fwd(
torch::Tensor &q,
const torch::Tensor &k,
const torch::Tensor &v,
c10::optional<torch::Tensor> out_,\
c10::optional<torch::Tensor> alibi_slopes_,
const double p_dropout,
const double softmax_scale,
bool is_causal,
const int64_t window_size_left,
const int64_t window_size_right,
const double softcap,
const bool return_softmax,
c10::optional<at::Generator> gen_);
std::vector<torch::Tensor>
mha_varlen_fwd(
at::Tensor &q, // total_q x num_heads x head_size, total_q := \sum_{i=0}^{b} s_i
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.
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.
c10::optional<torch::Tensor> out_, // total_q x num_heads x head_size, total_k := \sum_{i=0}^{b} s_i
const torch::Tensor &cu_seqlens_q, // b+1
const torch::Tensor &cu_seqlens_k, // b+1
c10::optional<torch::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<torch::Tensor> leftpad_k_, // batch_size
c10::optional<torch::Tensor> block_table_, // batch_size x max_num_blocks_per_seq
c10::optional<torch::Tensor> alibi_slopes_, // num_heads or b x num_heads
int64_t max_seqlen_q,
const int64_t max_seqlen_k,
const double p_dropout,
const double softmax_scale,
const bool zero_tensors,
bool is_causal,
int64_t window_size_left,
int64_t window_size_right,
const double softcap,
const bool return_softmax,
std::optional<at::Generator> gen_);
std::vector<torch::Tensor>
mha_bwd(const torch::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
const torch::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const torch::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size
const torch::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size
const torch::Tensor &out, // batch_size x seqlen_q x num_heads x head_size
const torch::Tensor &softmax_lse, // b x h x seqlen_q
const c10::optional<torch::Tensor> &dq_, // batch_size x seqlen_q x num_heads x head_size
const c10::optional<torch::Tensor> &dk_, // batch_size x seqlen_k x num_heads_k x head_size
const c10::optional<torch::Tensor> &dv_, // batch_size x seqlen_k x num_heads_k x head_size
const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
const double p_dropout, // probability to drop
const double softmax_scale,
const bool is_causal,
const int64_t window_size_left,
const int64_t window_size_right,
const double softcap,
const bool deterministic,
c10::optional<at::Generator> gen_,
const c10::optional<torch::Tensor> &rng_state);
std::vector<torch::Tensor>
mha_varlen_bwd(
const torch::Tensor &dout, // batch_size x seqlen_q x num_heads, x multiple_of(head_size_og, 8)
const torch::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
const torch::Tensor &k, // batch_size x seqlen_k x num_heads_k x head_size
const torch::Tensor &v, // batch_size x seqlen_k x num_heads_k x head_size
const torch::Tensor &out, // batch_size x seqlen_q x num_heads x head_size
const torch::Tensor &softmax_lse, // b x h x seqlen_q
const c10::optional<torch::Tensor> &dq_, // batch_size x seqlen_q x num_heads x head_size
const c10::optional<torch::Tensor> &dk_, // batch_size x seqlen_k x num_heads_k x head_size
const c10::optional<torch::Tensor> &dv_, // batch_size x seqlen_k x num_heads_k x head_size
const torch::Tensor &cu_seqlens_q, // batch_size + 1
const torch::Tensor &cu_seqlens_k, // batch_size + 1
const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or b x num_heads
const int64_t max_seqlen_q,
const int64_t max_seqlen_k,
const double p_dropout,
const double softmax_scale,
const bool zero_tensors,
const bool is_causal,
const int64_t window_size_left,
const int64_t window_size_right,
const double softcap,
const bool deterministic,
c10::optional<at::Generator> gen_,
const c10::optional<torch::Tensor> &rng_state);
std::vector<torch::Tensor>
mha_fwd_kvcache(
const torch::Tensor &q, // batch_size x seqlen_q x num_heads x head_size
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.
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.
const c10::optional<torch::Tensor> &k_, // batch_size x seqlen_knew x num_heads_k x head_size
const c10::optional<torch::Tensor> &v_, // batch_size x seqlen_knew x num_heads_k x head_size
const c10::optional<torch::Tensor> &seqlens_k_, // batch_size
const c10::optional<torch::Tensor> &rotary_cos_, // seqlen_ro x (rotary_dim / 2)
const c10::optional<torch::Tensor> &rotary_sin_, // seqlen_ro x (rotary_dim / 2)
const c10::optional<torch::Tensor> &cache_batch_idx_, // indices to index into the KV cache
const c10::optional<torch::Tensor> &leftpad_k_, // batch_size
const c10::optional<torch::Tensor> &block_table_, // batch_size x max_num_blocks_per_seq
const c10::optional<torch::Tensor> &alibi_slopes_, // num_heads or batch_size x num_heads
const c10::optional<torch::Tensor> &out_, // batch_size x seqlen_q x num_heads x head_size
const double softmax_scale,
bool is_causal,
const int64_t window_size_left,
const int64_t window_size_right,
const double softcap,
bool is_rotary_interleaved,
const int64_t num_splits);
|