fix: copy seq_lens in TRTLLM MHA draft decode cuda graph capture (#26521)

2026-05-28 21:55:33 -07:00
parent b47366fbf9
commit f16816f043
1 changed files with 3 additions and 0 deletions
@@ -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
                )