Support FlashMLA backend cuda graph (#4514)

Co-authored-by: yinfan98 <1106310035@qq.com>
Co-authored-by: Hongbosherlock <hongbosherlock@gmail.com>
Co-authored-by: ispobock <ispobaoke@163.com>
This commit is contained in:
lukec
2025-03-19 08:25:34 -07:00
committed by GitHub
co-authored by yinfan98 Hongbosherlock ispobock
parent f44db16c8e
commit b6944f97a6
3 changed files with 188 additions and 32 deletions
@@ -49,7 +49,6 @@ def create_flashmla_kv_indices_triton(
kv_indices_ptr,
req_to_token_ptr_stride: tl.constexpr,
kv_indices_ptr_stride: tl.constexpr,
max_pagesize: tl.constexpr,
):
PAGED_SIZE: tl.constexpr = 64
BLOCK_SIZE: tl.constexpr = 4096