| 1 | #include "core/registration.h" |
| 2 | #include "flash_kda.h" |
| 3 | |
| 4 | #include <torch/csrc/stable/library.h> |
| 5 | |
| 6 | STABLE_TORCH_LIBRARY(_flashkda_C, m) { |
| 7 | m.def("get_workspace_size(int T_total, int H, int N=1) -> int"); |
| 8 | m.def( |
| 9 | "fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, " |
| 10 | "Tensor(a!) out, Tensor(c!) workspace, Tensor A_log, Tensor dt_bias, " |
| 11 | "float lower_bound, " |
| 12 | "Tensor? initial_state=None, Tensor(b!)? final_state=None, " |
| 13 | "Tensor? cu_seqlens=None, Tensor(d!)? checkpoint_state=None, " |
| 14 | "Tensor? checkpoint_offsets=None) -> ()"); |
| 15 | } |
| 16 | |
| 17 | STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CompositeExplicitAutograd, m) { |
| 18 | m.impl("get_workspace_size", TORCH_BOX(&get_workspace_size)); |
| 19 | } |
| 20 | |
| 21 | STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { |
| 22 | m.impl("fwd", TORCH_BOX(&fwd)); |
| 23 | } |
| 24 | |
| 25 | REGISTER_EXTENSION(_flashkda_C) |