From dc4f54182374e9a52cefdd0b96eb468b40f04529 Mon Sep 17 00:00:00 2001 From: yinghui <32845984+cicirori@users.noreply.github.com> Date: Wed, 5 Nov 2025 01:35:02 -0800 Subject: [PATCH] fix trtllm_mla attention backend when disabling cuda graph. (#12687) --- python/sglang/srt/layers/attention/trtllm_mla_backend.py | 4 ++-- 1 file changed, 2 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 730ed7bc0..65ae9c4dc 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -585,7 +585,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): if forward_batch.forward_mode.is_target_verify(): max_seq = max_seq + self.num_draft_tokens seq_lens = seq_lens + self.num_draft_tokens - self.forward_decode_metadata.seq_lens_k = seq_lens + self.forward_decode_metadata.seq_lens_k = seq_lens.to(torch.int32) elif forward_batch.forward_mode.is_draft_extend(include_v2=True): max_seq = forward_batch.seq_lens_cpu.max().item() @@ -604,7 +604,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend): self.forward_decode_metadata.sum_seq_lens_q = sum_seq_lens_q self.forward_decode_metadata.cu_seqlens_q = cu_seqlens_q self.forward_decode_metadata.seq_lens_q = forward_batch.extend_seq_lens - self.forward_decode_metadata.seq_lens_k = seq_lens + self.forward_decode_metadata.seq_lens_k = seq_lens.to(torch.int32) max_seqlen_pad = self._calc_padded_blocks(max_seq) block_kv_indices = self._create_block_kv_indices(