[ROCm] Add DeepSeek-V4 DSpark speculative decoding support for AMD GPU(MI350X / MI355X, gfx950)

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

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -390,37 +390,45 @@ 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."
             )
         if batch.out_cache_loc is None:
             raise RuntimeError("DSpark prefill expected out_cache_loc, but got None.")
 
-        device = next_token_ids.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
+        device = self.model_runner.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,
             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,