From 8d6ab1cb88b5d00ea67a67df01b4077e23a489fe Mon Sep 17 00:00:00 2001 From: bmac3 <55830617+bmac3@users.noreply.github.com> Date: Tue, 28 Oct 2025 14:47:47 -0700 Subject: [PATCH] fix seqlen bug for trtllm_mla's draft_extend (#12295) --- .../srt/layers/attention/trtllm_mla_backend.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index f6d31dbc0..1882881e5 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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