feat(SpecEagleV2): add standalone_worker_v2 (#12625)

Co-authored-by: liugaoji.lgj <liugaoji.lgj@alibaba-inc.com>

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -141,6 +141,7 @@ class EAGLEDraftCudaGraphRunner:
             cuda_graph_bs = (
                 max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
                 if self.model_runner.spec_algorithm.is_eagle()
+                or self.model_runner.spec_algorithm.is_standalone()
                 else max(forward_batch.global_num_tokens_cpu)
             )
         else:
@@ -328,6 +329,7 @@ class EAGLEDraftCudaGraphRunner:
             max_batch_size = (
                 max_num_tokens // self.num_tokens_per_bs
                 if self.model_runner.spec_algorithm.is_eagle()
+                or self.model_runner.spec_algorithm.is_standalone()
                 else max_num_tokens
             )
             index = bisect.bisect_left(self.capture_bs, max_batch_size)

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -189,6 +189,7 @@ class EAGLEDraftExtendCudaGraphRunner:
             cuda_graph_bs = (
                 max(forward_batch.global_num_tokens_cpu) // self.num_tokens_per_bs
                 if self.model_runner.spec_algorithm.is_eagle()
+                or self.model_runner.spec_algorithm.is_standalone()
                 else max(forward_batch.global_num_tokens_cpu)
             )
         else:

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -12,6 +12,7 @@ from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_r
 from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
     EAGLEDraftNpuGraphRunner,
 )
+from sglang.srt.layers.attention.triton_backend import TritonMultiStepDraftBackend
 from sglang.srt.layers.moe.utils import (
     speculative_moe_a2a_backend_context,
     speculative_moe_backend_context,
@@ -50,12 +51,14 @@ from sglang.srt.utils.common import (
     empty_context,
     fast_topk,
     get_available_gpu_memory,
+    is_cuda,
     is_npu,
     next_power_of_2,
 )
 from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
 
 _is_npu = is_npu()
+_is_cuda = is_cuda()
 
 logger = logging.getLogger(__name__)
 
@@ -251,8 +254,14 @@ class EagleDraftWorker(BaseDraftWorker):
             "cuda": EAGLEDraftExtendCudaGraphRunner,
         }
         # Capture extend
-        # FIXME cuda not support draft_extend capture
-        if self.draft_extend_attn_backend and _is_npu:
+        # TODO: support draft extend cuda graph for more attention backends
+        if self.draft_extend_attn_backend and (
+            _is_npu
+            or (
+                _is_cuda
+                and isinstance(self.draft_attn_backend, TritonMultiStepDraftBackend)
+            )
+        ):
             tic = time.perf_counter()
             before_mem = get_available_gpu_memory(self.device, self.gpu_id)
             logger.info(
@@ -506,6 +515,9 @@ class EagleDraftWorker(BaseDraftWorker):
                 self.plan_stream
             )
 
+        if forward_batch.spec_info.accept_length is None:
+            forward_batch.spec_info.accept_length = batch_result.accept_lens
+
         # Run draft extend batch in the main compute stream
         can_cuda_graph = (
             self.cuda_graph_runner_for_draft_extend

python/sglang/srt/speculative/spec_info.py

@@ -161,6 +161,9 @@ class SpeculativeAlgorithm(metaclass=_SpeculativeAlgorithmMeta):
     def is_none(self) -> bool:
         return self is SpeculativeAlgorithm.NONE
 
+    def supports_spec_v2(self) -> bool:
+        return self.is_eagle() or self.is_eagle3() or self.is_standalone()
+
     def is_eagle(self) -> bool:
         return self._has_flag("EAGLE")
 
@@ -274,6 +277,12 @@ def _create_eagle_worker(**kwargs: Any) -> Any:
 
 
 def _create_standalone_worker(**kwargs: Any) -> Any:
+    enable_overlap = kwargs.pop("enable_overlap", False)
+    if enable_overlap:
+        from sglang.srt.speculative.standalone_worker_v2 import StandaloneWorkerV2
+
+        return StandaloneWorkerV2(**kwargs)
+
     from sglang.srt.speculative.standalone_worker import StandaloneWorker
 
     return StandaloneWorker(**kwargs)

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -0,0 +1,169 @@
+import contextlib
+import logging
+from typing import Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.layers.moe.utils import speculative_moe_backend_context
+from sglang.srt.managers.tp_worker import TpModelWorker
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.speculative.eagle_utils import TreeMaskMode
+from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker, EAGLEWorkerV2
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_utils import draft_tp_context
+from sglang.srt.utils import empty_context, get_bool_env_var, is_cuda
+
+if is_cuda():
+    from sgl_kernel import segment_packbits  # noqa: F401
+
+logger = logging.getLogger(__name__)
+SGLANG_RETURN_ORIGINAL_LOGPROB = get_bool_env_var("SGLANG_RETURN_ORIGINAL_LOGPROB")
+
+
+def _get_plan_stream(
+    device: str,
+) -> Tuple[any, contextlib.AbstractContextManager]:
+    if envs.SGLANG_ENABLE_OVERLAP_PLAN_STREAM.get():
+        plan_stream = torch.get_device_module(device).Stream()
+        plan_stream_ctx = torch.get_device_module(device).stream(plan_stream)
+        return plan_stream, plan_stream_ctx
+    else:
+        return None, contextlib.nullcontext()
+
+
+class StandaloneDraftWorker(EagleDraftWorker):
+    """Custom EagleDraftWorker that doesn't share embeddings/lm_head with target model."""
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        dp_rank: int,
+        moe_ep_rank: int,
+        nccl_port: int,
+        target_worker: TpModelWorker,
+    ):
+        # copy args
+        self.server_args = server_args
+        self.gpu_id = gpu_id
+        self.tp_rank = tp_rank
+        self.dp_rank = dp_rank
+        self.moe_ep_rank = moe_ep_rank
+        self.nccl_port = nccl_port
+        self.target_worker = target_worker
+
+        # Args for easy access
+        self.device = server_args.device
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        # Set constant
+        from sglang.srt.speculative.eagle_info import EagleDraftInput
+
+        EagleDraftInput.ALLOC_LEN_PER_DECODE = max(
+            self.speculative_num_steps * self.topk, self.speculative_num_draft_tokens
+        )
+
+        # Do not capture cuda graph in `TpModelWorker` init,
+        # will capture later with init_cuda_graphs()
+        backup_disable_cuda_graph = server_args.disable_cuda_graph
+        server_args.disable_cuda_graph = True
+
+        # Share the allocator with a target worker.
+        # Draft and target worker own their own KV cache pools.
+        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
+            target_worker.get_memory_pool()
+        )
+        with empty_context():
+            # Init draft worker
+            self.draft_worker = TpModelWorker(
+                server_args=server_args,
+                gpu_id=gpu_id,
+                tp_rank=tp_rank,
+                pp_rank=0,  # FIXME
+                dp_rank=dp_rank,
+                moe_ep_rank=moe_ep_rank,
+                nccl_port=nccl_port,
+                is_draft_worker=True,
+                req_to_token_pool=self.req_to_token_pool,
+                token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
+            )
+
+        # Alias for better readability
+        self.draft_runner = self.draft_worker.model_runner
+
+        self.init_token_map()
+        self.init_lm_head()
+
+        # Init attention backend and cuda graphs
+        self.draft_runner.server_args.disable_cuda_graph = backup_disable_cuda_graph
+        self.draft_tp_context = (
+            draft_tp_context if server_args.enable_dp_attention else empty_context
+        )
+        with self.draft_tp_context(
+            self.draft_runner.tp_group
+        ), speculative_moe_backend_context():
+            self.init_attention_backend()
+            self.init_cuda_graphs()
+        self.tree_mask_mode = TreeMaskMode.FULL_MASK
+
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+
+    def init_lm_head(self):
+        """Override to prevent sharing embeddings and lm_head with target model."""
+        # For standalone worker, we don't share embeddings and lm_head
+        # The draft model uses its own embeddings and lm_head
+        pass
+
+
+class StandaloneWorkerV2(EAGLEWorkerV2):
+
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        gpu_id: int,
+        tp_rank: int,
+        dp_rank: Optional[int],
+        moe_ep_rank: int,
+        nccl_port: int,
+        target_worker: TpModelWorker,
+    ):
+        # Parse arguments
+        self.server_args = server_args
+        self.topk = server_args.speculative_eagle_topk
+        self.speculative_num_steps = server_args.speculative_num_steps
+        self.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens
+        self.enable_nan_detection = server_args.enable_nan_detection
+        self.gpu_id = gpu_id
+        self.device = server_args.device
+        self._target_worker = target_worker
+        self.page_size = server_args.page_size
+        self.speculative_algorithm = SpeculativeAlgorithm.from_string(
+            server_args.speculative_algorithm
+        )
+
+        self.req_to_token_pool, self.token_to_kv_pool_allocator = (
+            target_worker.get_memory_pool()
+        )
+
+        # Override the context length of the draft model to be the same as the target model.
+        server_args.context_length = target_worker.model_runner.model_config.context_len
+
+        # Create our custom draft worker that doesn't share embeddings/lm_head
+        self._draft_worker = StandaloneDraftWorker(
+            server_args, gpu_id, tp_rank, dp_rank, moe_ep_rank, nccl_port, target_worker
+        )
+
+        # Some dummy tensors
+        self.num_new_pages_per_topk = torch.empty(
+            (), dtype=torch.int64, device=self.device
+        )
+        self.extend_lens = torch.empty((), dtype=torch.int64, device=self.device)
+
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)