chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Auto Sync] Update scheduler.py, spec_info.py, run_suite.py... (20251027) (#12235)

Browse Source 
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: gongwei-130 <56567052+gongwei-130@users.noreply.github.com>
This commit is contained in:
Yineng Zhang
2025-10-27 17:21:08 -07:00
committed by GitHub
parent 326c84c493
commit a6ea3add76
4 changed files with 478 additions and 84 deletions
Download Patch File Download Diff File Expand all files Collapse all files

75
python/sglang/srt/managers/scheduler.py
View File

@@ -327,8 +327,28 @@ class Scheduler(
# Launch a draft worker for speculative decoding
self.launch_draft_worker(
gpu_id, tp_rank, moe_ep_rank, server_args, port_args, dp_rank
draft_worker_kwargs = dict(
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
server_args=server_args,
nccl_port=port_args.nccl_port,
target_worker=self.tp_worker,
dp_rank=dp_rank,
)
if server_args.speculative_draft_load_format is not None:
server_args.load_format = server_args.speculative_draft_load_format
logger.info(
f"Using draft model load_format: '{server_args.speculative_draft_load_format}'"
)
# Draft workers are looked up via `SpeculativeAlgorithm` registry; new
# algorithms should register their factory instead of patching this code.
if self.spec_algorithm.name in {"EAGLE", "EAGLE3"}:
draft_worker_kwargs["enable_overlap"] = self.enable_overlap
self.draft_worker = self.spec_algorithm.create_draft_worker(
**draft_worker_kwargs
)
# Dispatch the model worker
@@ -557,57 +577,6 @@ class Scheduler(
]
)
def launch_draft_worker(
self, gpu_id, tp_rank, moe_ep_rank, server_args, port_args, dp_rank
):
if server_args.speculative_draft_load_format is not None:
server_args.load_format = server_args.speculative_draft_load_format
logger.info(
f"Using draft model load_format: '{server_args.speculative_draft_load_format}'"
)
if self.spec_algorithm.is_eagle():
from sglang.srt.speculative.eagle_worker import EAGLEWorker
from sglang.srt.speculative.eagle_worker_v2 import EAGLEWorkerV2
WorkerClass = EAGLEWorkerV2 if self.enable_overlap else EAGLEWorker
self.draft_worker = WorkerClass(
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
server_args=server_args,
nccl_port=port_args.nccl_port,
target_worker=self.tp_worker,
dp_rank=dp_rank,
)
elif self.spec_algorithm.is_standalone():
from sglang.srt.speculative.standalone_worker import StandaloneWorker
self.draft_worker = StandaloneWorker(
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
server_args=server_args,
nccl_port=port_args.nccl_port,
target_worker=self.tp_worker,
dp_rank=dp_rank,
)
elif self.spec_algorithm.is_ngram():
from sglang.srt.speculative.ngram_worker import NGRAMWorker
self.draft_worker = NGRAMWorker(
gpu_id=gpu_id,
tp_rank=tp_rank,
moe_ep_rank=moe_ep_rank,
server_args=server_args,
nccl_port=port_args.nccl_port,
target_worker=self.tp_worker,
dp_rank=dp_rank,
)
else:
self.draft_worker = None
def init_sockets(self, server_args: ServerArgs, port_args: PortArgs):
context = zmq.Context(2)
self.idle_sleeper = None

337
python/sglang/srt/speculative/spec_info.py
View File

@@ -1,46 +1,321 @@
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
@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):