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." )