#include <torch/extension.h>

#include <vector>

std::vector<torch::Tensor> launch_vea_b_numeric(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale,
    bool capture);
torch::Tensor launch_vea_b_numeric_into(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale,
    torch::Tensor output);
std::vector<torch::Tensor> launch_sage_checkpoint(
    torch::Tensor q, torch::Tensor k, torch::Tensor v,
    torch::Tensor q_scale, torch::Tensor k_scale, torch::Tensor v_scale);
torch::Tensor numeric_kernel_attributes();

PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
  module.def("vea_b", &launch_vea_b_numeric, "Aligned VEA-B numerical prototype");
  module.def("vea_b_into", &launch_vea_b_numeric_into, "Aligned VEA-B numerical prototype into output");
  module.def("sage_checkpoint", &launch_sage_checkpoint, "Pinned Sage-layout checkpoint oracle");
  module.def("attributes", &numeric_kernel_attributes, "Numerical kernel attributes");
}