Add registry for custom speculative algorithms (#23991)
This commit is contained in:
Liangsheng Yin
@@ -3321,12 +3321,28 @@ class ServerArgs:
|
||||
self.speculative_moe_runner_backend
|
||||
).is_flashinfer_trtllm(), "Currently speculative MoE runner backend doesn't support flashinfer_trtllm, please use triton or auto backend for speculative moe runner instead."
|
||||
|
||||
if self.speculative_algorithm is not None:
|
||||
self.speculative_algorithm = self.speculative_algorithm.upper()
|
||||
|
||||
self.speculative_algorithm = _resolve_speculative_algorithm_alias(
|
||||
self.speculative_algorithm,
|
||||
self.speculative_draft_model_path,
|
||||
trust_remote_code=self.trust_remote_code,
|
||||
)
|
||||
|
||||
if self.speculative_algorithm is not None:
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.speculative.spec_registry import CustomSpecAlgo
|
||||
|
||||
algo = SpeculativeAlgorithm.from_string(self.speculative_algorithm)
|
||||
|
||||
# TODO: move the per-algorithm validation below into spec module hooks.
|
||||
if (
|
||||
isinstance(algo, CustomSpecAlgo)
|
||||
and algo.validate_server_args is not None
|
||||
):
|
||||
algo.validate_server_args(self)
|
||||
|
||||
if self.speculative_skip_dp_mlp_sync:
|
||||
assert self.speculative_algorithm == "EAGLE", (
|
||||
"--speculative-skip-dp-mlp-sync is only supported with "
|
||||
@@ -5421,8 +5437,11 @@ class ServerArgs:
|
||||
parser.add_argument(
|
||||
"--speculative-algorithm",
|
||||
type=str,
|
||||
choices=["DFLASH", "EAGLE", "EAGLE3", "NEXTN", "STANDALONE", "NGRAM"],
|
||||
help="Speculative algorithm.",
|
||||
help=(
|
||||
"Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, "
|
||||
"NGRAM, DFLASH. Or any name registered via "
|
||||
"`SpeculativeAlgorithm.register`."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--speculative-draft-model-path",
|
||||
|
||||
@@ -2,7 +2,17 @@ from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum, IntEnum, auto
|
||||
from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
|
||||
from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union
|
||||
|
||||
from sglang.srt.speculative.spec_registry import (
|
||||
CustomSpecAlgo,
|
||||
ServerArgsValidator,
|
||||
WorkerFactory,
|
||||
)
|
||||
from sglang.srt.speculative.spec_registry import get_spec as _get_registered_spec
|
||||
from sglang.srt.speculative.spec_registry import (
|
||||
register_algorithm as _register_algorithm,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
|
||||
@@ -13,7 +23,11 @@ if TYPE_CHECKING:
|
||||
|
||||
|
||||
class SpeculativeAlgorithm(Enum):
|
||||
"""Speculative decoding algorithms."""
|
||||
"""Builtin speculative decoding algorithms. Plugin-registered ones are
|
||||
``CustomSpecAlgo`` instances; ``from_string`` returns either type, and
|
||||
both expose the same ``is_*()`` / ``create_worker`` interface so callers
|
||||
dispatch uniformly without isinstance checks.
|
||||
"""
|
||||
|
||||
DFLASH = auto()
|
||||
EAGLE = auto()
|
||||
@@ -24,13 +38,46 @@ class SpeculativeAlgorithm(Enum):
|
||||
NONE = auto()
|
||||
|
||||
@classmethod
|
||||
def from_string(cls, name: Optional[str]) -> SpeculativeAlgorithm:
|
||||
def from_string(
|
||||
cls, name: Optional[str]
|
||||
) -> Union[SpeculativeAlgorithm, CustomSpecAlgo]:
|
||||
if name is None:
|
||||
return cls.NONE
|
||||
upper = name.upper()
|
||||
try:
|
||||
return cls[name.upper()]
|
||||
return cls[upper]
|
||||
except KeyError:
|
||||
raise ValueError(f"Unknown speculative algorithm name: {name}")
|
||||
pass
|
||||
spec = _get_registered_spec(upper)
|
||||
if spec is not None:
|
||||
return spec
|
||||
raise ValueError(f"Unknown speculative algorithm name: {name}")
|
||||
|
||||
@classmethod
|
||||
def register(
|
||||
cls,
|
||||
name: str,
|
||||
*,
|
||||
supports_overlap: bool = False,
|
||||
validate_server_args: Optional[ServerArgsValidator] = None,
|
||||
spec_class: Type[CustomSpecAlgo] = CustomSpecAlgo,
|
||||
) -> Callable[[WorkerFactory], WorkerFactory]:
|
||||
"""Decorator to register a plugin speculative algorithm. The factory
|
||||
takes ``server_args`` and returns the worker class. Pass a
|
||||
``CustomSpecAlgo`` subclass via ``spec_class`` to override any
|
||||
``is_*()`` / ``create_worker`` method.
|
||||
|
||||
Example:
|
||||
@SpeculativeAlgorithm.register("MY_SPEC", supports_overlap=False)
|
||||
def _factory(server_args):
|
||||
return MySpecWorker
|
||||
"""
|
||||
return _register_algorithm(
|
||||
name,
|
||||
supports_overlap=supports_overlap,
|
||||
validate_server_args=validate_server_args,
|
||||
spec_class=spec_class,
|
||||
)
|
||||
|
||||
def is_none(self) -> bool:
|
||||
return self == SpeculativeAlgorithm.NONE
|
||||
|
||||
@@ -0,0 +1,123 @@
|
||||
"""Internal storage backing ``SpeculativeAlgorithm.register``. Plugins
|
||||
should use that classmethod API; do not import from this module directly.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Callable, Dict, Optional, Type
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
WorkerFactory = Callable[["ServerArgs"], Type]
|
||||
ServerArgsValidator = Callable[["ServerArgs"], None]
|
||||
|
||||
|
||||
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_spec_v2`` follows ``supports_overlap``.
|
||||
"""
|
||||
|
||||
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_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_dflash(self) -> bool:
|
||||
return False
|
||||
|
||||
def is_standalone(self) -> bool:
|
||||
return False
|
||||
|
||||
def is_ngram(self) -> bool:
|
||||
return False
|
||||
|
||||
def supports_spec_v2(self) -> bool:
|
||||
return self.supports_overlap
|
||||
|
||||
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."
|
||||
)
|
||||
return self.factory(server_args)
|
||||
|
||||
|
||||
_REGISTRY: Dict[str, CustomSpecAlgo] = {}
|
||||
|
||||
# Builtin enum members + the NEXTN alias; plugins cannot shadow these.
|
||||
_RESERVED_NAMES = frozenset(
|
||||
{"DFLASH", "EAGLE", "EAGLE3", "NEXTN", "STANDALONE", "NGRAM", "NONE"}
|
||||
)
|
||||
|
||||
|
||||
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.")
|
||||
|
||||
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())
|
||||
@@ -0,0 +1,234 @@
|
||||
"""Unit tests for the speculative algorithm plugin registry."""
|
||||
|
||||
import unittest
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.speculative.spec_registry import (
|
||||
_REGISTRY,
|
||||
_RESERVED_NAMES,
|
||||
CustomSpecAlgo,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
|
||||
|
||||
|
||||
class _RegistryIsolated(CustomTestCase):
|
||||
"""Snapshot and restore the global registry so tests don't leak."""
|
||||
|
||||
def setUp(self):
|
||||
self._snapshot = _REGISTRY.copy()
|
||||
_REGISTRY.clear()
|
||||
|
||||
def tearDown(self):
|
||||
_REGISTRY.clear()
|
||||
_REGISTRY.update(self._snapshot)
|
||||
|
||||
|
||||
class TestFromString(_RegistryIsolated):
|
||||
def test_none_input_returns_none_member(self):
|
||||
self.assertIs(SpeculativeAlgorithm.from_string(None), SpeculativeAlgorithm.NONE)
|
||||
|
||||
def test_builtin_name_returns_enum(self):
|
||||
self.assertIs(
|
||||
SpeculativeAlgorithm.from_string("EAGLE"), SpeculativeAlgorithm.EAGLE
|
||||
)
|
||||
self.assertIs(
|
||||
SpeculativeAlgorithm.from_string("NGRAM"), SpeculativeAlgorithm.NGRAM
|
||||
)
|
||||
|
||||
def test_builtin_name_is_case_insensitive(self):
|
||||
self.assertIs(
|
||||
SpeculativeAlgorithm.from_string("eagle"), SpeculativeAlgorithm.EAGLE
|
||||
)
|
||||
|
||||
def test_unknown_name_raises(self):
|
||||
with self.assertRaisesRegex(ValueError, "Unknown speculative algorithm"):
|
||||
SpeculativeAlgorithm.from_string("NOT_REGISTERED")
|
||||
|
||||
def test_registered_plugin_returns_custom_spec(self):
|
||||
@SpeculativeAlgorithm.register("MY_FOO")
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
algo = SpeculativeAlgorithm.from_string("MY_FOO")
|
||||
self.assertIsInstance(algo, CustomSpecAlgo)
|
||||
self.assertEqual(algo.name, "MY_FOO")
|
||||
|
||||
def test_registered_plugin_lookup_is_case_insensitive(self):
|
||||
@SpeculativeAlgorithm.register("MY_FOO")
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
self.assertIs(
|
||||
SpeculativeAlgorithm.from_string("my_foo"),
|
||||
SpeculativeAlgorithm.from_string("MY_FOO"),
|
||||
)
|
||||
|
||||
|
||||
class TestRegister(_RegistryIsolated):
|
||||
def test_register_returns_factory_unchanged(self):
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
decorated = SpeculativeAlgorithm.register("MY_FOO")(_factory)
|
||||
self.assertIs(decorated, _factory)
|
||||
|
||||
def test_two_distinct_registrations_are_independent(self):
|
||||
@SpeculativeAlgorithm.register("FOO")
|
||||
def _foo_factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
@SpeculativeAlgorithm.register("BAR")
|
||||
def _bar_factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
foo = SpeculativeAlgorithm.from_string("FOO")
|
||||
bar = SpeculativeAlgorithm.from_string("BAR")
|
||||
self.assertIsNot(foo, bar)
|
||||
self.assertNotEqual(foo, bar)
|
||||
self.assertEqual(foo.name, "FOO")
|
||||
self.assertEqual(bar.name, "BAR")
|
||||
|
||||
def test_duplicate_name_raises(self):
|
||||
@SpeculativeAlgorithm.register("MY_FOO")
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "already registered"):
|
||||
|
||||
@SpeculativeAlgorithm.register("MY_FOO")
|
||||
def _factory2(server_args):
|
||||
return MagicMock
|
||||
|
||||
def test_reserved_name_raises(self):
|
||||
for reserved in _RESERVED_NAMES:
|
||||
with self.assertRaisesRegex(ValueError, "reserved"):
|
||||
SpeculativeAlgorithm.register(reserved)
|
||||
|
||||
def test_register_is_case_insensitive_on_collision(self):
|
||||
@SpeculativeAlgorithm.register("MY_FOO")
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
with self.assertRaisesRegex(ValueError, "already registered"):
|
||||
|
||||
@SpeculativeAlgorithm.register("my_foo")
|
||||
def _factory2(server_args):
|
||||
return MagicMock
|
||||
|
||||
|
||||
class TestCustomSpecAlgoInterface(_RegistryIsolated):
|
||||
"""CustomSpecAlgo must duck-type SpeculativeAlgorithm enum values."""
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
|
||||
@SpeculativeAlgorithm.register("MY_FOO", supports_overlap=False)
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
self.algo = SpeculativeAlgorithm.from_string("MY_FOO")
|
||||
|
||||
def test_is_predicates_all_false_except_speculative(self):
|
||||
self.assertFalse(self.algo.is_none())
|
||||
self.assertFalse(self.algo.is_eagle())
|
||||
self.assertFalse(self.algo.is_eagle3())
|
||||
self.assertFalse(self.algo.is_dflash())
|
||||
self.assertFalse(self.algo.is_standalone())
|
||||
self.assertFalse(self.algo.is_ngram())
|
||||
self.assertTrue(self.algo.is_speculative())
|
||||
|
||||
def test_supports_spec_v2_follows_supports_overlap(self):
|
||||
# Plugin registered with supports_overlap=False -> not spec_v2.
|
||||
self.assertFalse(self.algo.supports_spec_v2())
|
||||
|
||||
@SpeculativeAlgorithm.register("MY_V2", supports_overlap=True)
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
v2 = SpeculativeAlgorithm.from_string("MY_V2")
|
||||
self.assertTrue(v2.supports_spec_v2())
|
||||
|
||||
def test_create_worker_calls_factory(self):
|
||||
server_args = MagicMock()
|
||||
server_args.disable_overlap_schedule = True
|
||||
worker_cls = self.algo.create_worker(server_args)
|
||||
self.assertIs(worker_cls, MagicMock)
|
||||
|
||||
def test_create_worker_raises_on_overlap_mismatch(self):
|
||||
server_args = MagicMock()
|
||||
server_args.disable_overlap_schedule = False
|
||||
with self.assertRaisesRegex(ValueError, "does not support overlap"):
|
||||
self.algo.create_worker(server_args)
|
||||
|
||||
|
||||
class TestValidatorHook(_RegistryIsolated):
|
||||
def test_validator_invocation_is_caller_driven(self):
|
||||
validator = MagicMock()
|
||||
|
||||
@SpeculativeAlgorithm.register("MY_FOO", validate_server_args=validator)
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
algo = SpeculativeAlgorithm.from_string("MY_FOO")
|
||||
self.assertIs(algo.validate_server_args, validator)
|
||||
# Callers (e.g. ServerArgs.__post_init__) must invoke the hook themselves;
|
||||
# CustomSpecAlgo does not call it from create_worker.
|
||||
validator.assert_not_called()
|
||||
|
||||
|
||||
class TestSubclassOverride(_RegistryIsolated):
|
||||
"""Plugins can subclass CustomSpecAlgo to override is_*() / create_worker."""
|
||||
|
||||
def test_subclass_overrides_is_eagle(self):
|
||||
class EagleLike(CustomSpecAlgo):
|
||||
def is_eagle(self) -> bool:
|
||||
return True
|
||||
|
||||
@SpeculativeAlgorithm.register(
|
||||
"MY_LIKE_EAGLE", supports_overlap=True, spec_class=EagleLike
|
||||
)
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
algo = SpeculativeAlgorithm.from_string("MY_LIKE_EAGLE")
|
||||
self.assertIsInstance(algo, EagleLike)
|
||||
self.assertIsInstance(algo, CustomSpecAlgo)
|
||||
self.assertTrue(algo.is_eagle())
|
||||
# Other predicates default to False
|
||||
self.assertFalse(algo.is_ngram())
|
||||
self.assertFalse(algo.is_dflash())
|
||||
|
||||
def test_subclass_overrides_create_worker(self):
|
||||
class CustomDispatch(CustomSpecAlgo):
|
||||
def create_worker(self, server_args):
|
||||
return "custom-dispatched"
|
||||
|
||||
@SpeculativeAlgorithm.register("MY_CUSTOM", spec_class=CustomDispatch)
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
algo = SpeculativeAlgorithm.from_string("MY_CUSTOM")
|
||||
# Custom dispatch bypasses default overlap check
|
||||
self.assertEqual(algo.create_worker(MagicMock()), "custom-dispatched")
|
||||
|
||||
|
||||
class TestCrossTypeIdentity(_RegistryIsolated):
|
||||
"""A plugin algo and a builtin enum value must never compare equal."""
|
||||
|
||||
def test_plugin_not_equal_to_builtin(self):
|
||||
@SpeculativeAlgorithm.register("MY_FOO")
|
||||
def _factory(server_args):
|
||||
return MagicMock
|
||||
|
||||
algo = SpeculativeAlgorithm.from_string("MY_FOO")
|
||||
self.assertNotEqual(algo, SpeculativeAlgorithm.EAGLE)
|
||||
self.assertNotEqual(algo, SpeculativeAlgorithm.NONE)
|
||||
self.assertIsNot(algo, SpeculativeAlgorithm.EAGLE)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main(verbosity=3)
|
||||
Reference in New Issue
Block a user