replicate/flash-mla

1#include <torch/library.h>2 3#include "registration.h"4#include "torch_binding.h"5 6TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {7  ops.def("get_mla_metadata(Tensor! seqlens_k, int num_heads_per_head_k, int num_heads_k) -> Tensor[]");8  ops.impl("get_mla_metadata", torch::kCUDA, &get_mla_metadata);9 10  // TOOD: remove last unknown_param when resolved11  ops.def("mha_fwd_kvcache_mla(Tensor! q, Tensor! kcache, Tensor? vcache_, int head_size_v, Tensor! seqlens_k, Tensor! block_table, float softmax_scale, bool is_causal_, Tensor! tile_scheduler_metadata, Tensor! num_splits) -> Tensor[]");12  ops.impl("mha_fwd_kvcache_mla", torch::kCUDA, &mha_fwd_kvcache_mla);13}14 15REGISTER_EXTENSION(TORCH_EXTENSION_NAME)16