From 87521e8d9889a180a0c3c7d454784ff9879b8031 Mon Sep 17 00:00:00 2001 From: haic0 Date: Fri, 10 Jul 2026 11:46:11 -0400 Subject: [PATCH 1/2] Fix DSpark draft token metadata handling Use the active DSpark draft token count when building DeepSeek V4 HIP target-verify metadata and prepare KV injection inputs before the target prefill forward. --- .../deepseek_v4_backend_hip_radix.py | 17 ++++++++--- .../dspark_components/dspark_worker_v2.py | 29 ++++++++++--------- 2 files changed, 28 insertions(+), 18 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 28199de8065f..2f0c2f87867e 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -590,6 +590,7 @@ def init_forward_metadata_target_verify( extend_seq_lens: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, seq_lens_cpu: Optional[List[int]] = None, + num_draft_tokens: Optional[int] = None, ) -> Union[DSV4Metadata, DSV4RawVerifyMetadata]: # HIP path: build target-verify metadata eagerly even when # SGLANG_PREP_IN_CUDA_GRAPH is enabled. The raw/lazy-upgrade route can @@ -603,6 +604,7 @@ def init_forward_metadata_target_verify( seq_lens_cpu=seq_lens_cpu, out_cache_loc=out_cache_loc, use_prefill_cuda_graph=use_prefill_cuda_graph, + num_draft_tokens=num_draft_tokens, ) def init_forward_metadata_target_verify_old( @@ -613,13 +615,15 @@ def init_forward_metadata_target_verify_old( seq_lens_cpu: Optional[List[int]] = None, out_cache_loc: Optional[torch.Tensor] = None, use_prefill_cuda_graph: bool = False, + num_draft_tokens: Optional[int] = None, ) -> DSV4Metadata: batch_size = len(seq_lens) - seq_lens = seq_lens + self.speculative_num_draft_tokens - seq_lens_cpu = [x + self.speculative_num_draft_tokens for x in seq_lens_cpu] - extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * batch_size + num_draft_tokens = int(num_draft_tokens or self.speculative_num_draft_tokens) + seq_lens = seq_lens + num_draft_tokens + seq_lens_cpu = [x + num_draft_tokens for x in seq_lens_cpu] + extend_seq_lens_cpu = [num_draft_tokens] * batch_size extend_seq_lens = self._move_to_device(extend_seq_lens_cpu) - num_tokens = self.speculative_num_draft_tokens * batch_size + num_tokens = num_draft_tokens * batch_size if out_cache_loc is None: out_cache_loc = seq_lens.new_zeros(num_tokens) return self.init_forward_metadata_prefill( @@ -972,6 +976,11 @@ def init_forward_metadata(self, forward_batch: ForwardBatch) -> None: seq_lens_cpu=( seq_lens_cpu.tolist() if seq_lens_cpu is not None else None ), + num_draft_tokens=getattr( + getattr(forward_batch, "spec_info", None), + "draft_token_num", + None, + ), ) elif forward_batch.forward_mode.is_prefill(include_draft_extend_v2=True): extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu 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 6e3f8da867c5..6ad2e2cbd83d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -390,19 +390,6 @@ def _forward_prefill( self.target_worker.forward_batch_generation(batch) return self._decode_idle_result(on_publish=on_publish) - batch.capture_hidden_mode = CaptureHiddenMode.FULL - batch_output = self.target_worker.forward_batch_generation(batch) - logits_output = batch_output.logits_output - next_token_ids = batch_output.next_token_ids - batch_output.new_seq_lens = batch.seq_lens - if on_publish is not None: - on_publish(batch_output.new_seq_lens) - - if logits_output.hidden_states is None: - raise RuntimeError( - "DSpark requires target aux hidden capture for prefill, but got None. " - "Make sure the target model has DFlash layers-to-capture configured." - ) if batch.extend_lens is None or batch.prefix_lens is None: raise RuntimeError( "DSpark expected extend_lens / prefix_lens in extend mode, got None." @@ -410,7 +397,7 @@ def _forward_prefill( if batch.out_cache_loc is None: raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.") - device = next_token_ids.device + device = self.model_runner.device ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device) draft_seq_lens = torch.tensor( batch.prefix_lens, dtype=torch.int32, device=device @@ -421,6 +408,20 @@ def _forward_prefill( ctx_lens, int(sum(batch.extend_lens)), ) + + batch.capture_hidden_mode = CaptureHiddenMode.FULL + batch_output = self.target_worker.forward_batch_generation(batch) + logits_output = batch_output.logits_output + next_token_ids = batch_output.next_token_ids + batch_output.new_seq_lens = batch.seq_lens + if on_publish is not None: + on_publish(batch_output.new_seq_lens) + + if logits_output.hidden_states is None: + raise RuntimeError( + "DSpark requires target aux hidden capture for prefill, but got None. " + "Make sure the target model has DFlash layers-to-capture configured." + ) self._kv_injector.inject_target_hidden( target_hidden=logits_output.hidden_states, cache_loc=batch.out_cache_loc, From 1484632a9f1dab24b84c0953ffd41a4f8cedaf55 Mon Sep 17 00:00:00 2001 From: haic0 Date: Sat, 11 Jul 2026 04:26:00 -0400 Subject: [PATCH 2/2] Fix DSpark prefill length tensor transfers Stage prefill length tensors in pinned CPU memory before non-blocking device transfer to avoid extra synchronization in the DSpark prefill path. --- .../dspark_components/dspark_worker_v2.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 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 6ad2e2cbd83d..1c095868eac1 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -398,10 +398,17 @@ def _forward_prefill( raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.") device = self.model_runner.device - ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device) - draft_seq_lens = torch.tensor( - batch.prefix_lens, dtype=torch.int32, device=device + pin_memory = is_cuda() + ctx_lens_cpu = torch.empty( + (len(batch.extend_lens),), dtype=torch.int32, pin_memory=pin_memory ) + draft_seq_lens_cpu = torch.empty( + (len(batch.prefix_lens),), dtype=torch.int32, pin_memory=pin_memory + ) + ctx_lens_cpu.copy_(torch.tensor(batch.extend_lens, dtype=torch.int32)) + draft_seq_lens_cpu.copy_(torch.tensor(batch.prefix_lens, dtype=torch.int32)) + ctx_lens = ctx_lens_cpu.to(device, non_blocking=True) + draft_seq_lens = draft_seq_lens_cpu.to(device, non_blocking=True) positions, _ = compute_position( self.model_runner.server_args.attention_backend, draft_seq_lens,