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

[perf] simplify _apply_cuda_graph_metadata for draft extend in trtllm_mla backend (#29077)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-06-24 15:45:26 -07:00
committed by GitHub
parent 85a3f1f75c
commit e26bceb81e
1 changed files with 4 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+4 -15
View File
@@ -345,7 +345,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
if forward_mode.is_target_verify():
metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device)
elif forward_mode.is_draft_extend_v2():
num_tokens_per_bs = num_tokens // bs
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 = torch.arange(
@@ -385,24 +385,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
if forward_mode.is_target_verify():
seq_lens = seq_lens[:bs] + self.num_draft_tokens
metadata.seq_lens_k.copy_(seq_lens.to(dtype=torch.int32))
metadata.seq_lens_k.copy_(seq_lens)
elif forward_mode.is_draft_extend_v2():
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[:bs].fill_(num_tokens_per_bs)
# see NOTE(draft_extend seq_len handling)
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))
seq_lens = seq_lens[:bs]
metadata.seq_lens_k.copy_(seq_lens)
# Update block indices for new sequences.
create_flashmla_kv_indices_triton[