#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);