feat(sgl-kernel): add InfLLM v2 attention kernels (#29383)
Co-authored-by: Size Wang <paulgeorge13hhhhh@gmail.com> Co-authored-by: lijiayi <lijiayi@modelbest.cn> Co-authored-by: suhmily10 <suhmily@gmail.com> Co-authored-by: Xiaoyue Xu <xiaoyue.xu.me@gmail.com> Co-authored-by: hansjohn <74091612+hansjohn@users.noreply.github.com> Co-authored-by: zhangyan <1762895426@qq.com>
This commit is contained in:
co-authored by
Size Wang
lijiayi
suhmily10
Xiaoyue Xu
hansjohn
zhangyan
parent
be70bfbdbb
commit
9bd02dc5b9
@@ -112,6 +112,25 @@ int64_t cutlass_mla_get_workspace_size(
|
||||
int64_t sm_count = 0,
|
||||
int64_t num_kv_splits = 1 /* Set to 1 to avoid cuda_graph issue by default. */);
|
||||
|
||||
/*
|
||||
* From csrc/infllm_v2
|
||||
*/
|
||||
void infllm_v2_max_pooling_1d_varlen(
|
||||
at::Tensor input,
|
||||
at::Tensor output,
|
||||
at::Tensor cu_seqlens_q,
|
||||
at::Tensor cu_seqlens_k,
|
||||
at::Tensor cache_lens,
|
||||
int64_t max_seqlen_q,
|
||||
int64_t max_seqlen_k,
|
||||
int64_t kernel_size,
|
||||
int64_t stride,
|
||||
int64_t padding,
|
||||
int64_t block_size,
|
||||
int64_t local_blocks,
|
||||
int64_t init_blocks,
|
||||
int64_t total_q);
|
||||
|
||||
/*
|
||||
* From csrc/elementwise
|
||||
*/
|
||||
|
||||
Reference in New Issue
Block a user