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 )