File size: 316 Bytes
99323da
 
 
 
 
1
2
3
4
5
6
#pragma once
#include <torch/torch.h>
void split_k_attention(torch::Tensor &out, torch::Tensor &partial, torch::Tensor &stats,
    torch::Tensor const &q, torch::Tensor const &k, torch::Tensor const &v,
    torch::Tensor const &kscale, torch::Tensor const &vscale, torch::Tensor const &lengths, int64_t partitions);