[Feature] Add zero bubble for spec v2

python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py

@@ -5,7 +5,9 @@
 from typing import TYPE_CHECKING
 
 import torch
+import torch.nn.functional as F
 
+from sglang.srt.environ import envs
 from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 
@@ -162,6 +164,16 @@ def process_prebuilt(
             hidden_states_list = [req.hidden_states_tensor for req in self.reqs]
             hidden_states = torch.stack(hidden_states_list, dim=0).to(self.device)
 
+            enable_spec_v2_zero_bubble = envs.SGLANG_SPEC_V2_ZERO_BUBBLE.get()
+
+            if enable_spec_v2_zero_bubble and server_args.speculative_num_steps > 1:
+                topk_pad_size = (
+                    server_args.speculative_num_steps * num_states - topk_p.shape[-1]
+                )
+
+                topk_p = F.pad(topk_p, (0, topk_pad_size))
+                topk_index = F.pad(topk_index, (0, topk_pad_size))
+
             # local import to avoid circular import
             from sglang.srt.speculative.eagle_info import EagleDraftInput
 

python/sglang/srt/environ.py

@@ -449,6 +449,7 @@ class Envs:
     SGLANG_SPEC_ENABLE_STRICT_FILTER_CHECK = EnvBool(True)
     SGLANG_SPEC_NAN_DETECTION = EnvBool(False)
     SGLANG_SPEC_OOB_DETECTION = EnvBool(False)
+    SGLANG_SPEC_V2_ZERO_BUBBLE = EnvBool(False)
 
     # VLM
     SGLANG_VLM_CACHE_SIZE_MB = EnvInt(100)

python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py

@@ -179,10 +179,10 @@ def replay(
         # Replay
         if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
             if forward_batch.forward_mode.is_target_verify():
-                seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
+                seq_lens_cpu = forward_batch.seq_lens_cpu + self.num_tokens_per_bs
                 seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
             else:
-                seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
+                seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
                     self.bs - self.raw_bs
                 )
             thread = threading.Thread(target=self._update_inputs, args=(seq_lens,))

python/sglang/srt/model_executor/forward_batch_info.py

@@ -63,7 +63,7 @@
     is_npu,
     support_triton,
 )
-from sglang.srt.utils.common import ceil_align
+from sglang.srt.utils.common import ceil_align, is_pin_memory_available
 
 if TYPE_CHECKING:
     from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -486,6 +486,7 @@ def init_new(
             rids=[req.rid for req in batch.reqs],
         )
         device = model_runner.device
+        _pin = is_pin_memory_available(device)
 
         if batch.extend_input_logprob_token_ids is not None:
             ret.extend_input_logprob_token_ids_gpu = (
@@ -494,9 +495,9 @@ def init_new(
 
         num_tokens = len(batch.input_ids) if batch.input_ids is not None else 0
         if enable_num_token_non_padded(model_runner.server_args):
-            ret.num_token_non_padded = torch.tensor(num_tokens, dtype=torch.int32).to(
-                device, non_blocking=True
-            )
+            ret.num_token_non_padded = torch.tensor(
+                num_tokens, dtype=torch.int32, pin_memory=_pin
+            ).to(device, non_blocking=True)
         ret.num_token_non_padded_cpu = num_tokens
 
         # For MLP sync
@@ -516,15 +517,18 @@ def init_new(
             ret.original_global_num_tokens_cpu = batch.global_num_tokens
             ret.global_num_tokens_cpu = global_num_tokens
             ret.global_num_tokens_gpu = torch.tensor(
-                global_num_tokens, dtype=torch.int64
+                global_num_tokens, dtype=torch.int64, pin_memory=_pin
             ).to(device, non_blocking=True)
 
             ret.global_num_tokens_for_logprob_cpu = global_num_tokens_for_logprob
             ret.global_num_tokens_for_logprob_gpu = torch.tensor(
-                global_num_tokens_for_logprob, dtype=torch.int64
+                global_num_tokens_for_logprob, dtype=torch.int64, pin_memory=_pin
             ).to(device, non_blocking=True)
 
         if ret.forward_mode.is_idle():
+            if _is_npu:
+                # This synchronize is necessary to prevent the system from hanging on npu.
+                torch.npu.synchronize()
             ret.positions = torch.empty((0,), dtype=torch.int64, device=device)
             return ret
 
@@ -540,6 +544,7 @@ def init_new(
                     for i in range(block_offset, block_offset + block_size)
                 ],
                 dtype=positions_dtype,
+                pin_memory=_pin,
             ).to(device, non_blocking=True)
         elif (
             ret.spec_info is not None
@@ -555,10 +560,10 @@ def init_new(
             assert isinstance(batch.extend_seq_lens, list)
             assert isinstance(batch.extend_prefix_lens, list)
             ret.extend_seq_lens = torch.tensor(
-                batch.extend_seq_lens, dtype=torch.int32
+                batch.extend_seq_lens, dtype=torch.int32, pin_memory=_pin
             ).to(device, non_blocking=True)
             ret.extend_prefix_lens = torch.tensor(
-                batch.extend_prefix_lens, dtype=torch.int32
+                batch.extend_prefix_lens, dtype=torch.int32, pin_memory=_pin
             ).to(device, non_blocking=True)
             ret.extend_num_tokens = batch.extend_num_tokens
             positions, ret.extend_start_loc = compute_position(
@@ -761,6 +766,7 @@ def _compute_mrope_positions(
         # batch_size * [3 * seq_len]
         batch_size = self.seq_lens_cpu.shape[0]
         mrope_positions_list = [[]] * batch_size
+        _pin = is_pin_memory_available(model_runner.device)
         for batch_idx in range(batch_size):
             mm_input = batch.multimodal_inputs[batch_idx]
             if self.forward_mode.is_decode():
@@ -812,10 +818,20 @@ def _compute_mrope_positions(
                         )
                 mrope_positions_list[batch_idx] = mrope_positions
 
-        self.mrope_positions = torch.cat(
-            [pos for pos in mrope_positions_list],
-            dim=1,
-        ).to(dtype=torch.int64, device=model_runner.device, non_blocking=True)
+        if _pin:
+            self.mrope_positions = (
+                torch.cat(
+                    [pos for pos in mrope_positions_list],
+                    dim=1,
+                )
+                .pin_memory()
+                .to(dtype=torch.int64, device=model_runner.device, non_blocking=True)
+            )
+        else:
+            self.mrope_positions = torch.cat(
+                [pos for pos in mrope_positions_list],
+                dim=1,
+            ).to(dtype=torch.int64, device=model_runner.device, non_blocking=True)
 
     def _pad_tensor_to_size(self, tensor: torch.Tensor, size: int, *, value: int = 0):
         if value == 0:

python/sglang/srt/model_executor/model_runner.py

@@ -348,6 +348,7 @@ def __init__(
         self.init_new_workspace = False
         self.draft_model_idx = draft_model_idx
         self.enable_hisparse = server_args.enable_hisparse
+        self.enable_spec_v2_zero_bubble = envs.SGLANG_SPEC_V2_ZERO_BUBBLE.get()
 
         self.remote_instance_transfer_engine = None
         self.remote_instance_transfer_engine_session_id = ""
@@ -2920,6 +2921,7 @@ def _forward_raw(
             and forward_batch.global_num_tokens_gpu is not None
             and require_gathered_buffer(self.server_args)
             and not is_nsa_enable_prefill_cp()
+            and not self.enable_spec_v2_zero_bubble
         ):
             forward_batch.adjust_num_token_non_padded_for_attn_tp(
                 server_args=self.server_args,

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -6,6 +6,7 @@
 
 import torch
 
+from sglang.srt.environ import envs
 from sglang.srt.layers.dp_attention import DpPaddingMode, set_dp_buffer_len
 from sglang.srt.model_executor.cuda_graph_runner import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
@@ -79,6 +80,7 @@ def __init__(self, eagle_worker: EAGLEWorker):
         )
         self.enable_pdmux = False
         self.deepep_adapter = DeepEPCudaGraphRunnerAdapter()
+        self.enable_spec_v2_zero_bubble = envs.SGLANG_SPEC_V2_ZERO_BUBBLE.get()
 
         # Batch sizes to capture
         self.capture_bs, self.compile_bs = get_batch_sizes_to_capture(model_runner)
@@ -329,7 +331,13 @@ def run_once():
             output_cache_loc_backup = forward_batch.out_cache_loc
             hidden_states_backup = forward_batch.spec_info.hidden_states
 
-            ret = self.eagle_worker.draft_forward(forward_batch)
+            if self.enable_spec_v2_zero_bubble:
+                assert hasattr(
+                    self.eagle_worker, "draft_forward_zero_bubble"
+                ), "`Spec v2 zero bubble` just support when enable `overlap scheduler` and enable `eagle algorithm` now"
+                ret = self.eagle_worker.draft_forward_zero_bubble(forward_batch)
+            else:
+                ret = self.eagle_worker.draft_forward(forward_batch)
 
             forward_batch.out_cache_loc = output_cache_loc_backup
             forward_batch.spec_info.hidden_states = hidden_states_backup
@@ -348,6 +356,10 @@ def run_once():
 
     def _postprocess_output_to_raw_bs(self, out, raw_bs):
         # Keep the variables name for readability
+        if self.enable_spec_v2_zero_bubble:
+            ret_topk_p_list, ret_topk_index_list = (t[:raw_bs] for t in out)
+            return ret_topk_p_list, ret_topk_index_list
+
         parent_list, top_scores_index, draft_tokens = (t[:raw_bs] for t in out)
         return parent_list, top_scores_index, draft_tokens