EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

fix FlashMLA cudagraph config (#4691)

Browse Source 
Co-authored-by: yinfan98 <1106310035@qq.com>
This commit is contained in:
lukec
2025-03-25 12:06:58 +08:00
committed by GitHub
parent f01b092519
commit 57eec0bfbc
1 changed files with 4 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/layers/attention/flashmla_backend.py
View File

@@ -92,7 +92,7 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
if forward_batch.forward_mode.is_decode_or_idle():
if spec_info is None:
max_seqlen_pad = triton.cdiv(
forward_batch.seq_lens.max().item(), PAGE_SIZE
forward_batch.decode_seq_lens_cpu.max().item(), PAGE_SIZE
)
block_kv_indices = torch.full(
(bs, max_seqlen_pad),
@@ -206,8 +206,10 @@ class FlashMLABackend(FlashInferMLAAttnBackend):
):
if forward_mode.is_decode_or_idle():
assert seq_lens_cpu is not None
seq_lens = seq_lens[:bs]
max_seqlen_pad = triton.cdiv(seq_lens.max().item(), PAGE_SIZE)
seq_lens_cpu = seq_lens_cpu[:bs]
max_seqlen_pad = triton.cdiv(seq_lens_cpu.max().item(), PAGE_SIZE)
create_flashmla_kv_indices_triton[(bs,)](
self.req_to_token,
req_pool_indices[:bs],