23 lines
1.1 KiB
C++
23 lines
1.1 KiB
C++
#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");
|
|
}
|