chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Fix issue 19717 by making qo_indptr uniform strided instead of packed (#19807)

Browse Source
This commit is contained in:
kpham-sgl
2026-03-04 15:27:10 -08:00
committed by GitHub
parent 28c931e1a5
commit 376dfb03f7
1 changed files with 8 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
python/sglang/srt/layers/attention/triton_backend.py
View File

@@ -761,9 +761,15 @@ class TritonAttnBackend(AttentionBackend):
mask_indptr[1 : bs + 1] = torch.cumsum(seq_mask_len, dim=0)
elif forward_mode.is_draft_extend(include_v2=True):
seq_lens = seq_lens[:bs]
accept_lens = spec_info.accept_length[:bs]
num_tokens_per_bs = self.speculative_num_steps + 1
qo_indptr = self.qo_indptr[: bs + 1]
qo_indptr[1 : bs + 1] = torch.cumsum(accept_lens, dim=0)
qo_indptr[: bs + 1] = torch.arange(
0,
bs * num_tokens_per_bs + 1,
step=num_tokens_per_bs,
dtype=torch.int32,
device=self.device,
)
kv_indptr = self.kv_indptr[: bs + 1]
kv_indptr[1 : bs + 1] = torch.cumsum(seq_lens, dim=0)
kv_indices = self.cuda_graph_kv_indices