chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

fix seqlen bug for trtllm_mla's draft_extend (#12295)

Browse Source
This commit is contained in:
bmac3
2025-10-28 14:47:47 -07:00
committed by GitHub
parent 84a9d0eab2
commit 8d6ab1cb88
1 changed files with 10 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/layers/attention/trtllm_mla_backend.py
View File

@@ -944,8 +944,16 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
metadata.max_seq_len_k + forward_batch.spec_info.draft_token_num
)
else:
seq_lens = forward_batch.seq_lens.to(torch.int32)
max_seq_len = metadata.max_seq_len_k
# forward_batch.seq_lens is the seq_lens of the prev_context + verified tokens.
# To account for pad_draft_extend_query, we need seq_lens = prev_context + max_draft_tokens.
# This will ensure queries align with kvs correctly when calling
# flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla.
seq_lens = (
forward_batch.seq_lens
- metadata.seq_lens_q
+ metadata.max_seq_len_q
).to(torch.int32)
max_seq_len = metadata.max_seq_len_k + metadata.max_seq_len_q
# Check if we're in CUDA graph mode (buffers are pre-allocated)
if self.padded_q_buffer is not None:
# Use pre-allocated buffer for CUDA graph compatibility