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

[Bugfix] Fix CUDA graph replay issues in trtllm_mla draft_extend (#21987)

Browse Source 
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Khoa Pham
2026-04-03 01:45:13 -07:00
committed by GitHub
co-authored by Claude Opus 4.6
parent 4f84ce5807
commit cd75d54fc5
2 changed files with 15 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+13 -13
View File
@@ -543,21 +543,21 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
metadata.seq_lens_k.copy_(seq_lens.to(dtype=torch.int32))
del seq_lens_sum # not handle "num_draft_tokens" but we do not need it
elif forward_mode.is_draft_extend(include_v2=True):
accept_length = spec_info.accept_length[:bs]
if spec_info.accept_length_cpu:
metadata.max_seq_len_q = max(spec_info.accept_length_cpu[:bs]) + 1
metadata.sum_seq_lens_q = sum(spec_info.accept_length_cpu[:bs]) + bs
else:
metadata.max_seq_len_q = 1
metadata.sum_seq_lens_q = bs
# draft_extend uses (accept_length + 1) query tokens per sequence
extend_seq_lens = accept_length + 1
metadata.cu_seqlens_q[1:].copy_(
torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32)
num_tokens_per_bs = self.num_draft_tokens
metadata.max_seq_len_q = num_tokens_per_bs
metadata.sum_seq_lens_q = num_tokens_per_bs * bs
metadata.cu_seqlens_q[: bs + 1].copy_(
torch.arange(
0,
bs * num_tokens_per_bs + 1,
step=num_tokens_per_bs,
dtype=torch.int32,
device=seq_lens.device,
)
)
metadata.seq_lens_q.copy_(extend_seq_lens)
metadata.seq_lens_q[:bs].fill_(num_tokens_per_bs)
# see NOTE(draft_extend seq_len handling)
seq_lens = seq_lens[:bs] - metadata.seq_lens_q + metadata.max_seq_len_q
seq_lens = seq_lens[:bs] - metadata.seq_lens_q[:bs] + metadata.max_seq_len_q
metadata.seq_lens_k.copy_(seq_lens.to(torch.int32))
# Update block indices for new sequences.