From f16816f0438b79fa4a3efaa817c4277acb4f09ba Mon Sep 17 00:00:00 2001 From: Yongji Wu <30348494+libertyeagle@users.noreply.github.com> Date: Thu, 28 May 2026 21:55:33 -0700 Subject: [PATCH] fix: copy seq_lens in TRTLLM MHA draft decode cuda graph capture (#26521) --- python/sglang/srt/layers/attention/trtllm_mha_backend.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index e270f2ccb..068791b02 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -324,6 +324,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[ "cache_seqlens" ][:bs] + metadata.cache_seqlens_int32.copy_( + seq_lens + self.speculative_step_id + 1 + ) metadata.max_seq_len_k = seq_lens.max().item() + ( self.speculative_step_id + 1 )