WaveCut's picture
Add ordered and approximate INT8 attention decode layouts
99323da verified
Raw History Blame Contribute Delete
316 Bytes
#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);