From 151a31482968ff5e60e6a4ff88a7eaa273d17e24 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 13 Aug 2026 16:10:54 -0700 Subject: [PATCH] [Fix] Make the DSpark draft num_token_non_padded host-to-device copy non-blocking (#34782) --- .../sglang/srt/speculative/dspark_components/dspark_draft.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index ff5d3396d..ca7a6a76f 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -355,8 +355,8 @@ class DraftBlockProposer: spec_algorithm=SpeculativeAlgorithm.DSPARK, spec_info=self._draft_block_spec_info, capture_hidden_mode=CaptureHiddenMode.NULL, - num_token_non_padded=torch.tensor( - draft_num_tokens, dtype=torch.int32, device=device + num_token_non_padded=torch.tensor(draft_num_tokens, dtype=torch.int32).to( + device, non_blocking=True ), num_token_non_padded_cpu=draft_num_tokens, )