minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Yongji Wu
2026-05-28 21:55:33 -07:00
committed by GitHub
parent b47366fbf9
commit f16816f043
1 changed files with 3 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+3
View File
@@ -324,6 +324,9 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[ metadata.cache_seqlens_int32 = self.decode_cuda_graph_metadata[
"cache_seqlens" "cache_seqlens"
][:bs] ][:bs]
metadata.cache_seqlens_int32.copy_(
seq_lens + self.speculative_step_id + 1
)
metadata.max_seq_len_k = seq_lens.max().item() + ( metadata.max_seq_len_k = seq_lens.max().item() + (
self.speculative_step_id + 1 self.speculative_step_id + 1
) )