replicate/flash-mla

1#pragma once2 3#include <torch/torch.h>4 5std::vector<torch::Tensor> 6get_mla_metadata(7    torch::Tensor &seqlens_k,8    const int64_t num_heads_per_head_k,9    const int64_t num_heads_k10);11 12std::vector<torch::Tensor> 13mha_fwd_kvcache_mla(14    torch::Tensor &q,15    const torch::Tensor &kcache,16    const c10::optional<torch::Tensor> &vcache_,17    const int64_t head_size_v,18    const torch::Tensor &seqlens_k,19    const torch::Tensor &block_table,20    const double softmax_scale,21    bool is_causal,22    const torch::Tensor &tile_scheduler_metadata,23    const torch::Tensor &num_splits24);