"""Internal storage backing ``SpeculativeAlgorithm.register``. Plugins should use that classmethod API; do not import from this module directly. """ from __future__ import annotations import logging from typing import TYPE_CHECKING, Callable, Dict, Optional, Type import torch if TYPE_CHECKING: from sglang.srt.managers.overlap_utils import FutureMap from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.spec_info import SpecInput WorkerFactory = Callable[["ServerArgs"], Type] ServerArgsValidator = Callable[["ServerArgs"], None] logger = logging.getLogger(__name__) class CustomSpecAlgo: """A plugin-registered speculative algorithm. Duck-types ``SpeculativeAlgorithm`` enum values (same ``is_*()`` / ``create_worker`` interface). Plugins may subclass this to override any ``is_*()`` / ``supports_*()`` / ``create_worker`` method (e.g. to integrate with builtin-specific branches like ``if spec_algorithm.is_eagle():`` in scheduler / model_runner). Pass the subclass via ``spec_class=...`` at registration. Defaults: all ``is_*()`` return ``False`` except ``is_speculative``. ``supports_overlap=False`` is deprecated: the spec V1 worker path has been removed, so such algorithms run on the V2 scheduler schema with overlap disabled (synchronous). Migrate plugin workers to the V2 schema and overlap scheduling. """ def __init__( self, name: str, factory: WorkerFactory, *, supports_overlap: bool = False, validate_server_args: Optional[ServerArgsValidator] = None, ): self.name = name self.factory = factory self.supports_overlap = supports_overlap self.validate_server_args = validate_server_args def __repr__(self) -> str: return f"CustomSpecAlgo({self.name!r})" def is_some(self) -> bool: return True def is_none(self) -> bool: return False def is_speculative(self) -> bool: return True def is_eagle(self) -> bool: return False def is_eagle3(self) -> bool: return False def is_frozen_kv_mtp(self) -> bool: return False def is_dflash(self) -> bool: return False def is_dspark(self) -> bool: return False def is_dflash_family(self) -> bool: return False def is_standalone(self) -> bool: return False def is_ngram(self) -> bool: return False def supports_target_verify_for_draft(self) -> bool: return False def supports_ragged_verify(self) -> bool: return False def has_draft_kv(self) -> bool: # Conservative default: the larger KV reserve. return True def handle_server_args(self, server_args: ServerArgs) -> None: pass def create_worker(self, server_args: ServerArgs) -> Type: if not server_args.disable_overlap_schedule and not self.supports_overlap: raise ValueError( f"Speculative algorithm {self.name} does not support overlap scheduling." ) if not self.supports_overlap: # Reached only when overlap is disabled, so the algorithm really # does run synchronously on the V2 schema below. logger.warning( "Speculative algorithm %s is registered with " "supports_overlap=False, which is deprecated: the spec V1 " "worker path has been removed, and the algorithm now runs on " "the V2 scheduler schema with overlap disabled (synchronous). " "Migrate the plugin worker to support overlap scheduling.", self.name, ) return self.factory(server_args) def get_num_tokens_per_bs_for_target_verify( self, num_draft_tokens: int, is_draft_worker: bool ) -> int: # FIXME: Remove this after the forward mode refactor. Target verify is # essentially a fixed sequence length prefill/extend with full cuda # graph support. We can use it for target verify, or we can use it for # other cases which is not target verify but fixed length prefill. # Here, we expose this interface to allow the other use cases. return num_draft_tokens def build_disagg_draft_input( self, batch: ScheduleBatch, server_args: ServerArgs, last_tokens_tensor: torch.Tensor, future_map: FutureMap, ) -> Optional[SpecInput]: return None _REGISTRY: Dict[str, CustomSpecAlgo] = {} # CLI spellings that are not ``SpeculativeAlgorithm`` members but still resolve # to a builtin (e.g. NEXTN -> EAGLE). Reserved alongside the enum members so # plugins cannot shadow them. _RESERVED_ALIASES = frozenset({"NEXTN"}) def _reserved_names() -> frozenset: """Names plugins cannot register under: every ``SpeculativeAlgorithm`` member plus ``_RESERVED_ALIASES``. Derived from the enum (lazily, to avoid a circular import — ``spec_info`` imports this module) so any new builtin is reserved automatically without editing a second list. """ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm return frozenset(algo.name for algo in SpeculativeAlgorithm) | _RESERVED_ALIASES def _assert_custom_spec_algo_conforms(spec_class: Type[CustomSpecAlgo]) -> None: """Fail fast if ``spec_class`` drifts from the ``SpeculativeAlgorithm`` duck-typing contract. ``from_string`` returns either type and callers dispatch on the shared ``is_*()`` / ``supports_*()`` interface without isinstance checks, so every such method on the enum must also exist on the registered spec class — otherwise a plugin-registered algo hits ``AttributeError`` at a call site (this is how ``is_some`` / ``is_frozen_kv_mtp`` silently went missing). New predicates are covered automatically; no second list to maintain. Called from ``register_algorithm`` rather than at import time because ``spec_info`` imports this module, so ``SpeculativeAlgorithm`` does not yet exist while this module is loading; at registration time it is fully defined. """ # NOTE: use ``vars()`` not ``dir()`` for the enum — ``EnumMeta.__dir__`` # hides instance methods, so ``dir(SpeculativeAlgorithm)`` would yield an # empty interface and turn this guard into a silent no-op. from sglang.srt.speculative.spec_info import SpeculativeAlgorithm interface = { name for name in vars(SpeculativeAlgorithm) if name.startswith(("is_", "supports_")) } missing = sorted(interface - set(dir(spec_class))) if missing: raise TypeError( f"{spec_class.__name__} is missing duck-typed methods from " f"SpeculativeAlgorithm: {missing}. Add them to {spec_class.__name__} " "so plugin-registered algorithms stay dispatchable." ) def register_algorithm( name: str, *, supports_overlap: bool = False, validate_server_args: Optional[ServerArgsValidator] = None, spec_class: Type[CustomSpecAlgo] = CustomSpecAlgo, ) -> Callable[[WorkerFactory], WorkerFactory]: """Return a decorator that registers a plugin algorithm under ``name``. Pass a ``spec_class`` subclass of ``CustomSpecAlgo`` to override any ``is_*()`` / ``supports_*()`` / ``create_worker`` method. """ upper = name.upper() if upper in _reserved_names(): raise ValueError( f"'{upper}' is a reserved speculative algorithm name; cannot be re-registered." ) if upper in _REGISTRY: raise ValueError(f"Speculative algorithm '{upper}' already registered.") _assert_custom_spec_algo_conforms(spec_class) def decorator(factory: WorkerFactory) -> WorkerFactory: _REGISTRY[upper] = spec_class( name=upper, factory=factory, supports_overlap=supports_overlap, validate_server_args=validate_server_args, ) return factory return decorator def get_spec(name: Optional[str]) -> Optional[CustomSpecAlgo]: """Return the registered spec for ``name``, or ``None`` for builtin / unknown names.""" if name is None: return None return _REGISTRY.get(name.upper())