sglang 0.5.4__py3-none-any.whl → 0.5.4.post2__py3-none-any.whl

sglang/srt/speculative/spec_info.py

@@ -1,46 +1,320 @@
+from __future__ import annotations
+
+import threading
 from abc import ABC, abstractmethod
+from collections import defaultdict
 from enum import IntEnum, auto
-from functools import lru_cache
-from typing import List, Tuple
+from typing import (
+    Any,
+    Callable,
+    DefaultDict,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Set,
+    Tuple,
+    Union,
+)
 
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch
 
+DraftWorkerClass = Callable[..., Any]
+DraftWorkerFactory = Callable[..., Any]
 
-class SpeculativeAlgorithm(IntEnum):
-    NONE = auto()
-    EAGLE = auto()
-    EAGLE3 = auto()
-    STANDALONE = auto()
-    NGRAM = auto()
 
-    def is_none(self):
-        return self == SpeculativeAlgorithm.NONE
+class _SpeculativeAlgorithmMeta(type):
+    def __iter__(cls) -> Iterator["SpeculativeAlgorithm"]:
+        return iter(cls._registration_order)
 
-    def is_eagle(self):
-        return self == SpeculativeAlgorithm.EAGLE or self == SpeculativeAlgorithm.EAGLE3
 
-    def is_eagle3(self):
-        return self == SpeculativeAlgorithm.EAGLE3
+class SpeculativeAlgorithm(metaclass=_SpeculativeAlgorithmMeta):
+    """Registry-backed representation of speculative decoding algorithms."""
 
-    def is_standalone(self):
-        return self == SpeculativeAlgorithm.STANDALONE
+    __slots__ = ("name", "value", "_draft_worker_factory")
 
-    def is_ngram(self):
-        return self == SpeculativeAlgorithm.NGRAM
+    _registry_by_name: Dict[str, "SpeculativeAlgorithm"] = {}
+    _registry_by_value: Dict[int, "SpeculativeAlgorithm"] = {}
+    _registration_order: List["SpeculativeAlgorithm"] = []
+    _flags: DefaultDict[str, Set[int]] = defaultdict(set)
+    _next_value: int = 0
 
-    @lru_cache(maxsize=None)
-    @staticmethod
-    def from_string(name: str):
-        name_map = {
-            "EAGLE": SpeculativeAlgorithm.EAGLE,
-            "EAGLE3": SpeculativeAlgorithm.EAGLE3,
-            "STANDALONE": SpeculativeAlgorithm.STANDALONE,
-            "NGRAM": SpeculativeAlgorithm.NGRAM,
-            None: SpeculativeAlgorithm.NONE,
-        }
-        if name is not None:
-            name = name.upper()
-        return name_map[name]
+    def __init__(
+        self,
+        name: str,
+        value: int,
+        draft_worker_factory: Optional[DraftWorkerFactory] = None,
+    ):
+        self.name = name
+        self.value = value
+        self._draft_worker_factory = draft_worker_factory
+
+    def __repr__(self) -> str:  # pragma: no cover - trivial
+        return f"SpeculativeAlgorithm.{self.name}"
+
+    def __str__(self) -> str:  # pragma: no cover - trivial
+        return self.name
+
+    def __hash__(self) -> int:
+        return hash(self.value)
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, SpeculativeAlgorithm):
+            return self.value == other.value
+        return NotImplemented
+
+    def __int__(self) -> int:
+        return self.value
+
+    @classmethod
+    def register(
+        cls,
+        name: str,
+        *,
+        aliases: Optional[Sequence[str]] = None,
+        value: Optional[int] = None,
+        draft_worker_factory: Optional[DraftWorkerFactory] = None,
+    ) -> SpeculativeAlgorithm:
+        normalized_name = name.upper()
+        if normalized_name in cls._registry_by_name:
+            raise ValueError(
+                f"SpeculativeAlgorithm '{normalized_name}' already registered"
+            )
+
+        if value is None:
+            value = cls._next_value
+        cls._next_value = max(cls._next_value, value + 1)
+
+        algorithm = cls(
+            normalized_name,
+            value,
+            draft_worker_factory=draft_worker_factory,
+        )
+
+        cls._registry_by_name[normalized_name] = algorithm
+        cls._registry_by_value[value] = algorithm
+        cls._registration_order.append(algorithm)
+        setattr(cls, normalized_name, algorithm)
+
+        if aliases:
+            cls.register_aliases(algorithm, *aliases)
+
+        return algorithm
+
+    @classmethod
+    def register_aliases(cls, algorithm: SpeculativeAlgorithm, *aliases: str) -> None:
+        for alias in aliases:
+            cls._registry_by_name[alias.upper()] = algorithm
+
+    @classmethod
+    def register_draft_worker(
+        cls,
+        algorithm: SpeculativeAlgorithm | str,
+        factory: DraftWorkerFactory,
+    ) -> None:
+        algo = cls._ensure_algorithm(algorithm)
+        algo._draft_worker_factory = factory
+
+    @classmethod
+    def _ensure_algorithm(
+        cls, algorithm: SpeculativeAlgorithm | str
+    ) -> SpeculativeAlgorithm:
+        if isinstance(algorithm, SpeculativeAlgorithm):
+            return algorithm
+        if isinstance(algorithm, str):
+            return cls.from_string(algorithm)
+        raise TypeError(f"Unsupported algorithm identifier: {algorithm!r}")
+
+    @classmethod
+    def _add_flag(
+        cls, flag: str | Sequence[str], algorithm: SpeculativeAlgorithm | str
+    ) -> None:
+        algo = cls._ensure_algorithm(algorithm)
+        if isinstance(flag, str):
+            flag_iter = (flag,)
+        else:
+            flag_iter = flag
+        for flag_name in flag_iter:
+            cls._flags[flag_name.upper()].add(algo.value)
+
+    @classmethod
+    def from_string(cls, name: Optional[str]) -> SpeculativeAlgorithm:
+        if name is None:
+            return cls.NONE
+        try:
+            return cls._registry_by_name[name.upper()]
+        except KeyError as exc:
+            raise ValueError(f"Unknown speculative algorithm '{name}'") from exc
+
+    @classmethod
+    def from_value(cls, value: int) -> SpeculativeAlgorithm:
+        try:
+            return cls._registry_by_value[value]
+        except KeyError as exc:
+            raise ValueError(f"Unknown speculative algorithm id {value}") from exc
+
+    def _has_flag(self, flag: str) -> bool:
+        return self.value in type(self)._flags.get(flag.upper(), set())
+
+    def is_none(self) -> bool:
+        return self is SpeculativeAlgorithm.NONE
+
+    def is_eagle(self) -> bool:
+        return self._has_flag("EAGLE")
+
+    def is_eagle3(self) -> bool:
+        return self._has_flag("EAGLE3")
+
+    def is_standalone(self) -> bool:
+        return self._has_flag("STANDALONE")
+
+    def is_ngram(self) -> bool:
+        return self._has_flag("NGRAM")
+
+    def create_draft_worker(self, **factory_kwargs: Any) -> Any:
+        if self._draft_worker_factory is None:
+            return None
+        return self._draft_worker_factory(self, **factory_kwargs)
+
+
+# Registry helpers backed by `SpeculativeAlgorithm`.
+_LOCK = threading.RLock()
+_REGISTERED_WORKERS: Dict[SpeculativeAlgorithm, DraftWorkerClass] = {}
+_FLAG_MARKERS: Dict[str, Callable[[Union[SpeculativeAlgorithm, str]], None]] = {
+    "EAGLE": lambda algorithm: SpeculativeAlgorithm._add_flag("EAGLE", algorithm),
+    "EAGLE3": lambda algorithm: SpeculativeAlgorithm._add_flag("EAGLE3", algorithm),
+    "STANDALONE": lambda algorithm: SpeculativeAlgorithm._add_flag(
+        "STANDALONE", algorithm
+    ),
+    "NGRAM": lambda algorithm: SpeculativeAlgorithm._add_flag("NGRAM", algorithm),
+}
+
+
+def _wrap_worker_class(worker_cls: DraftWorkerClass) -> DraftWorkerFactory:
+    def _factory(_: SpeculativeAlgorithm, **kwargs: Any) -> Any:
+        return worker_cls(**kwargs)
+
+    return _factory
+
+
+def register_speculative_algorithm(
+    name: str,
+    worker_cls: DraftWorkerClass,
+    *,
+    aliases: Optional[Sequence[str]] = None,
+    flags: Optional[Iterable[str]] = None,
+    value: Optional[int] = None,
+    override_worker: bool = False,
+) -> SpeculativeAlgorithm:
+    """Register a speculative algorithm and the associated draft worker class.
+
+    Example:
+        >>> from sglang.srt.speculative.spec_info import register_speculative_algorithm
+        >>> register_speculative_algorithm("MY_ALGO", MyDraftWorker, flags=("EAGLE",))
+    """
+
+    name_upper = name.upper()
+    with _LOCK:
+        try:
+            algorithm = SpeculativeAlgorithm.from_string(name_upper)
+            exists = True
+        except ValueError:
+            algorithm = SpeculativeAlgorithm.register(
+                name_upper,
+                aliases=aliases,
+                value=value,
+            )
+            SpeculativeAlgorithm.register_draft_worker(
+                algorithm, _wrap_worker_class(worker_cls)
+            )
+            exists = False
+
+        if exists:
+            if aliases:
+                SpeculativeAlgorithm.register_aliases(algorithm, *aliases)
+            if not override_worker and algorithm in _REGISTERED_WORKERS:
+                raise ValueError(
+                    f"Worker already registered for {algorithm!r}. "
+                    "Pass override_worker=True to replace it."
+                )
+            SpeculativeAlgorithm.register_draft_worker(
+                algorithm, _wrap_worker_class(worker_cls)
+            )
+
+        _REGISTERED_WORKERS[algorithm] = worker_cls
+
+        if flags:
+            for flag in flags:
+                marker = _FLAG_MARKERS.get(flag.upper())
+                if marker is None:
+                    raise ValueError(f"Unsupported flag '{flag}'")
+                marker(algorithm)
+
+        return algorithm
+
+
+def list_registered_workers() -> Dict[str, DraftWorkerClass]:
+    """Return a snapshot of registered speculative worker classes keyed by algorithm name."""
+    with _LOCK:
+        return {algo.name: cls for algo, cls in _REGISTERED_WORKERS.items()}
+
+
+def _create_eagle_worker(**kwargs: Any) -> Any:
+    enable_overlap = kwargs.pop("enable_overlap", False)
+    if enable_overlap:
+        from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2
+
+        return EAGLEWorkerV2(**kwargs)
+
+    from sglang.srt.speculative.eagle_worker import EAGLEWorker
+
+    return EAGLEWorker(**kwargs)
+
+
+def _create_standalone_worker(**kwargs: Any) -> Any:
+    from sglang.srt.speculative.standalone_worker import StandaloneWorker
+
+    return StandaloneWorker(**kwargs)
+
+
+def _create_ngram_worker(**kwargs: Any) -> Any:
+    from sglang.srt.speculative.ngram_worker import NGRAMWorker
+
+    return NGRAMWorker(**kwargs)
+
+
+# Register built-in algorithms.
+# Third-party integrations should import `SpeculativeAlgorithm` and either
+# call `register_speculative_algorithm` or use the helpers below to attach
+# additional draft workers.
+SpeculativeAlgorithm.register("NONE")
+
+register_speculative_algorithm(
+    "EAGLE",
+    aliases=("NEXTN",),
+    worker_cls=_create_eagle_worker,
+    flags=("EAGLE",),
+)
+
+register_speculative_algorithm(
+    "EAGLE3",
+    worker_cls=_create_eagle_worker,
+    flags=("EAGLE", "EAGLE3"),
+)
+
+register_speculative_algorithm(
+    "STANDALONE",
+    worker_cls=_create_standalone_worker,
+    flags=("STANDALONE",),
+)
+
+register_speculative_algorithm(
+    "NGRAM",
+    worker_cls=_create_ngram_worker,
+    flags=("NGRAM",),
+)
 
 
 class SpecInputType(IntEnum):

sglang/srt/speculative/spec_utils.py

@@ -19,16 +19,22 @@ from sglang.srt.distributed.parallel_state import (
 from sglang.srt.environ import envs
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import Req
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda, is_hip, is_npu, next_power_of_2
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+_is_npu = is_npu()
 
 if TYPE_CHECKING:
     from sglang.srt.speculative.eagle_info import EagleVerifyInput
 
 
-if is_cuda():
+if _is_cuda:
     from sgl_kernel import fast_topk
-elif is_hip():
+elif _is_hip:
     from sgl_kernel import fast_topk
+else:
+    from sglang.srt.utils.common import fast_topk
 
 
 logger = logging.getLogger(__name__)
@@ -39,7 +45,7 @@ SIMULATE_ACC_LEN = envs.SGLANG_SIMULATE_ACC_LEN.get()  # turn off if < 0
 SIMULATE_ACC_METHOD = envs.SGLANG_SIMULATE_ACC_METHOD.get()
 
 TREE_TRAVERSE_TIME_THRESHOLD = 1  # TODO: set this properly
-TREE_SPEC_KERNEL_AVAILABLE = is_cuda()  # This kernel is only available for CUDA now
+TREE_SPEC_KERNEL_AVAILABLE = _is_cuda  # This kernel is only available for CUDA now
 
 
 @triton.jit
@@ -103,6 +109,36 @@ def assign_req_to_token_pool(
         load_offset += BLOCK_SIZE
 
 
+def assign_req_to_token_pool_func(
+    req_pool_indices: torch.Tensor,
+    req_to_token: torch.Tensor,
+    start_offset: torch.Tensor,
+    end_offset: torch.Tensor,
+    out_cache_loc: torch.Tensor,
+    batch_size: int,
+):
+    if _is_cuda or _is_hip:
+        assign_req_to_token_pool[(batch_size,)](
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+            req_to_token.shape[1],
+            next_power_of_2(batch_size),
+        )
+    elif _is_npu:
+        import sgl_kernel_npu  # noqa: F401
+
+        torch.ops.npu.cache_loc_assign(
+            req_pool_indices,
+            req_to_token,
+            start_offset,
+            end_offset,
+            out_cache_loc,
+        )
+
+
 @triton.jit
 def assign_draft_cache_locs(
     req_pool_indices,
@@ -331,7 +367,7 @@ def get_target_cache_loc(
     )
 
 
-@torch.compile(dynamic=True)
+@torch.compile(dynamic=True, disable=_is_npu)
 def get_src_tgt_cache_loc(
     seq_lens: torch.Tensor,
     out_cache_loc: torch.Tensor,
@@ -381,7 +417,7 @@ def filter_finished_cache_loc_kernel(
     )
 
 
-@torch.compile(dynamic=True)
+@torch.compile(dynamic=True, disable=_is_npu)
 def create_accept_length_filter(
     accept_length: torch.Tensor,
     unfinished_index_device: torch.Tensor,
@@ -395,7 +431,7 @@ def create_accept_length_filter(
     return accept_length_filter
 
 
-@torch.compile(dynamic=True)
+@torch.compile(dynamic=True, disable=_is_npu)
 def select_top_k_tokens(
     i: int,
     topk_p: torch.Tensor,
@@ -413,7 +449,7 @@ def select_top_k_tokens(
         tree_info = (
             topk_p.unsqueeze(1),  # shape: (b, 1, topk)
             topk_index,  # shape: (b, topk)
-            torch.arange(-1, topk, dtype=torch.long, device="cuda")
+            torch.arange(-1, topk, dtype=torch.long, device=hidden_states.device)
             .unsqueeze(0)
             .repeat(topk_p.shape[0], 1),  # shape: (b, topk + 1)
         )