From 9efff5b8601658d4c3ff8322db570ebc0046157f Mon Sep 17 00:00:00 2001 From: inkcherry Date: Thu, 17 Sep 2026 16:42:32 +0000 Subject: [PATCH] avoid host sync in DSpark prefill slot expansion --- .../speculative/dspark_components/dspark_worker_v2.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 9b4a74a14058..d80439dbfa02 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -584,11 +584,16 @@ def _forward_prefill( state_slot = final_pos = None if is_unified_kv_triton(): repeats = ctx_lens.to(torch.int64) + num_tokens = sum(batch.extend_lens) state_slot = torch.repeat_interleave( - batch.req_pool_indices.to(device=device, dtype=torch.int64), repeats + batch.req_pool_indices.to(device=device, dtype=torch.int64), + repeats, + output_size=num_tokens, ) final_pos = torch.repeat_interleave( - (draft_seq_lens + ctx_lens - 1).to(torch.int64), repeats + (draft_seq_lens + ctx_lens - 1).to(torch.int64), + repeats, + output_size=num_tokens, ) self._kv_injector.inject_target_hidden( target_hidden=logits_output.hidden_states,