mirror of
https://github.com/vllm-project/vllm.git
synced 2026-08-20 20:50:15 +00:00
11f88260a2
Signed-off-by: Thien Tran <[email protected]> Co-authored-by: Shengqi Chen <[email protected]>
25 lines
762 B
C++
25 lines
762 B
C++
#include "core/registration.h"
|
|
#include "flash_kda.h"
|
|
|
|
#include <torch/csrc/stable/library.h>
|
|
|
|
STABLE_TORCH_LIBRARY(_flashkda_C, m) {
|
|
m.def("get_workspace_size(int T_total, int H, int N=1) -> int");
|
|
m.def(
|
|
"fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, "
|
|
"Tensor(a!) out, Tensor(c!) workspace, Tensor A_log, Tensor dt_bias, "
|
|
"float lower_bound, "
|
|
"Tensor? initial_state=None, Tensor(b!)? final_state=None, "
|
|
"Tensor? cu_seqlens=None) -> ()");
|
|
}
|
|
|
|
STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CompositeExplicitAutograd, m) {
|
|
m.impl("get_workspace_size", TORCH_BOX(&get_workspace_size));
|
|
}
|
|
|
|
STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) {
|
|
m.impl("fwd", TORCH_BOX(&fwd));
|
|
}
|
|
|
|
REGISTER_EXTENSION(_flashkda_C)
|