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
.github/workflows/nightly-test-nvidia.yml
+2 -1
View File
@@ -25,6 +25,7 @@ on:
- 'nightly-test-multimodal-server-2-gpu'
- 'nightly-test-perf-4-gpu-b200'
- 'nightly-test-perf-8-gpu-b200'
- 'nightly-test-specialized-8-gpu-b200'
- 'nightly-test-kernel-1-gpu-h100'
- 'nightly-test-diffusion-comparison'
- 'nightly-test-kernel-8-gpu-h200'
@@ -626,7 +627,7 @@ jobs:
# Specialized B200 tests - 8 GPU, for specific backends and configs
nightly-test-specialized-8-gpu-b200:
if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-perf-8-gpu-b200')
if: github.repository == 'sgl-project/sglang' && (inputs.job_filter == '' || inputs.job_filter == 'all' || inputs.job_filter == 'nightly-test-perf-8-gpu-b200' || inputs.job_filter == 'nightly-test-specialized-8-gpu-b200')
runs-on: 8-gpu-b200
env:
RUNNER_LABELS: 8-gpu-b200
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.