forked from Karylab-cklius/vllm
+13
7c6729b769
7c6729b769
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Bugen Zhao <i@bugenzhao.com> Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg> Co-authored-by: Bugen Zhao <i@bugenzhao.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Ziming Huang <zelda.huanghuang@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@inferact.ai> Co-authored-by: aoshen02 <aoshen02@users.noreply.github.com> Co-authored-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Summer Yang <girasoleyang@gmail.com> Co-authored-by: Kevin H. Luu <khluu000@gmail.com> Co-authored-by: Bowen Wang <abmfy@icloud.com> Co-authored-by: gnovack <novackgm@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: xiaozhoupy <peiyuanzhou1994@gmail.com> Co-authored-by: Roy Wang <yasong.wang@inferact.ai> Co-authored-by: Jeff (Junze) Ma <93145857+majunze2001@users.noreply.github.com> Co-authored-by: Codex <codex@openai.com>
18 lines
582 B
C++
18 lines
582 B
C++
#include "core/registration.h"
|
|
#include "flash_kda.h"
|
|
|
|
TORCH_LIBRARY(_flashkda_C, m) {
|
|
m.def("get_workspace_size(int T_total, int H, int N=1) -> int",
|
|
&get_workspace_size);
|
|
m.def(
|
|
"fwd(Tensor q, Tensor k, Tensor v, Tensor g, Tensor beta, float scale, "
|
|
"Tensor(a!) out, Tensor workspace, Tensor A_log, Tensor dt_bias, "
|
|
"float lower_bound, "
|
|
"Tensor? initial_state=None, Tensor(b!)? final_state=None, "
|
|
"Tensor? cu_seqlens=None) -> ()");
|
|
}
|
|
|
|
TORCH_LIBRARY_IMPL(_flashkda_C, CUDA, m) { m.impl("fwd", &fwd); }
|
|
|
|
REGISTER_EXTENSION(_flashkda_C)
|