chore: import upstream snapshot with attribution

python/sglang/srt/speculative/standalone_worker_v2.py

@@ -0,0 +1,248 @@
+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.adaptive_runtime_state import (
+    AdaptiveController,
+)
+from sglang.srt.speculative.eagle_utils import default_tree_mask_mode
+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,
+        attn_cp_rank: int,
+        moe_dp_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
+        self.attn_cp_rank = attn_cp_rank
+        self.moe_dp_rank = moe_dp_rank
+
+        # 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
+        )
+
+        # Pre-allocated constants for the topk=1 chain fast path in draft_forward.
+        self._topk1_parents_prealloc = None
+        self._topk1_score_indices_prealloc = None
+        self._rebuild_topk1_chain_buffers()
+
+        # 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
+        )
+
+        # Load draft model weights only.
+        with empty_context():
+            self.draft_worker = TpModelWorker(
+                server_args=server_args,
+                gpu_id=gpu_id,
+                tp_rank=tp_rank,
+                pp_rank=0,  # spec workers don't support pipeline parallelism
+                dp_rank=dp_rank,
+                moe_ep_rank=moe_ep_rank,
+                attn_cp_rank=attn_cp_rank,
+                moe_dp_rank=moe_dp_rank,
+                nccl_port=nccl_port,
+                is_draft_worker=True,
+            )
+
+        # Alias for better readability
+        self.draft_runner = self.draft_worker.model_runner
+        self.draft_tp_context = (
+            draft_tp_context if server_args.enable_dp_attention else empty_context
+        )
+        self.tree_mask_mode = default_tree_mask_mode()
+        self.plan_stream, self.plan_stream_ctx = _get_plan_stream(self.device)
+        # draft_forward reads this (set in EagleDraftWorker.__init__, skipped here).
+        self.index_share_for_mtp_iteration = (
+            getattr(
+                self.draft_runner.model_config.hf_config,
+                "index_share_for_mtp_iteration",
+                False,
+            )
+            and self.topk == 1
+        )
+        self.dsa_index_topk = None
+        self.seed_dsa_topk_from_draft_extend = False
+        self.dsa_extend_topk_buf = None
+
+    def alloc_memory_pool(
+        self,
+        memory_pool_config=None,
+        req_to_token_pool=None,
+        token_to_kv_pool_allocator=None,
+    ):
+        """Standalone: allocate pools without sharing embeddings."""
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.draft_worker.alloc_memory_pool(
+            memory_pool_config=memory_pool_config,
+            req_to_token_pool=req_to_token_pool,
+            token_to_kv_pool_allocator=token_to_kv_pool_allocator,
+        )
+        self.init_token_map()
+        self.init_lm_head()
+
+    def init_attention_backends(self):
+        with self.draft_tp_context(
+            self.draft_runner.tp_group
+        ), speculative_moe_backend_context():
+            super().init_attention_backends()
+
+    def init_cuda_graphs(self):
+        with self.draft_tp_context(
+            self.draft_runner.tp_group
+        ), speculative_moe_backend_context():
+            super().init_cuda_graphs()
+
+    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,
+        attn_cp_rank: int,
+        moe_dp_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.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
+        )
+
+        # Override the context length of the draft model to be the same as the target model.
+        server_args.override(
+            "spec_worker.match_target_context_length",
+            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,
+            attn_cp_rank,
+            moe_dp_rank,
+            nccl_port,
+            target_worker,
+        )
+
+        self._validate_vocab_compatibility(
+            target_vocab_size=target_worker.model_runner.model_config.vocab_size,
+            target_tokenizer=target_worker.tokenizer,
+        )
+
+        # 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)
+
+        # TODO: Adaptive speculative
+        self.adaptive_controller: Optional[AdaptiveController] = None
+
+    def _validate_vocab_compatibility(
+        self,
+        target_vocab_size: int,
+        target_tokenizer,
+    ) -> None:
+        """Raise ValueError if the draft and target vocabularies are incompatible."""
+        draft_vocab_size = self._draft_worker.draft_runner.model_config.vocab_size
+        draft_tokenizer = self._draft_worker.draft_worker.tokenizer
+        if target_vocab_size != draft_vocab_size:
+            raise ValueError(
+                f"STANDALONE speculative decoding requires the draft model to share the "
+                f"same vocabulary as the target model, but got "
+                f"target vocab_size={target_vocab_size} and "
+                f"draft vocab_size={draft_vocab_size}. "
+                f"Use a draft model with a matching vocabulary, or a speculative "
+                f"algorithm that supports heterogeneous vocabularies."
+            )
+        if (
+            target_tokenizer is not None
+            and draft_tokenizer is not None
+            and hasattr(target_tokenizer, "get_vocab")
+            and hasattr(draft_tokenizer, "get_vocab")
+            and target_tokenizer.get_vocab() != draft_tokenizer.get_vocab()
+        ):
+            raise ValueError(
+                "STANDALONE speculative decoding requires the draft model to share the "
+                "same vocabulary as the target model, but the two tokenizers have "
+                "different token-to-id mappings even though their vocab sizes match. "
+                "Use a draft model with a matching vocabulary, or a speculative "
+                "algorithm that supports heterogeneous vocabularies."
+            )