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
@@ -71,6 +71,10 @@ else:
|
||||
shuffle_rows,
|
||||
)
|
||||
from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda
|
||||
from sgl_kernel.infllm_v2 import (
|
||||
infllmv2_attn_stage1,
|
||||
max_pooling_1d_varlen,
|
||||
)
|
||||
from sgl_kernel.kvcacheio import (
|
||||
transfer_kv_all_layer,
|
||||
transfer_kv_all_layer_mla,
|
||||
|
||||
Reference in New Issue
Block a user