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(): if forward_mode.is_target_verify():
metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device) metadata.seq_lens_k = torch.zeros((bs,), dtype=torch.int32, device=device)
elif forward_mode.is_draft_extend_v2(): 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.max_seq_len_q = num_tokens_per_bs
metadata.sum_seq_lens_q = num_tokens_per_bs * bs metadata.sum_seq_lens_q = num_tokens_per_bs * bs
metadata.cu_seqlens_q = torch.arange( metadata.cu_seqlens_q = torch.arange(
@@ -385,24 +385,13 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
if forward_mode.is_target_verify(): if forward_mode.is_target_verify():
seq_lens = seq_lens[:bs] + self.num_draft_tokens 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(): elif forward_mode.is_draft_extend_v2():
num_tokens_per_bs = self.num_draft_tokens num_tokens_per_bs = self.num_draft_tokens
metadata.max_seq_len_q = num_tokens_per_bs metadata.max_seq_len_q = num_tokens_per_bs
metadata.sum_seq_lens_q = num_tokens_per_bs * bs metadata.sum_seq_lens_q = num_tokens_per_bs * bs
metadata.cu_seqlens_q[: bs + 1].copy_( seq_lens = seq_lens[:bs]
torch.arange( metadata.seq_lens_k.copy_(seq_lens)
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))
# Update block indices for new sequences. # Update block indices for new sequences.
create_flashmla_kv_indices_triton[ create_flashmla_kv_indices_triton[