Multi platform Plugin (#21388)
Co-authored-by: root <root@tjzj-inf-sci-k8s-bzz2-0183.tjzj.baidu.com> Co-authored-by: Alex Nails <alex.nails@radixark.ai> Co-authored-by: Alex Nails <alexj.nails@gmail.com> Co-authored-by: root <root@tjzj-inf-sci-k8s-bzz2-0000.tjzj.baidu.com> Co-authored-by: Mick <mickjagger19@icloud.com>
This commit is contained in:
Baidu-AIAK
co-authored by
root
Alex Nails
Alex Nails
root
Mick
root
Alex Nails
Alex Nails
root
Mick
parent
ebcc2b3eec
commit
7ca3566130
@@ -86,6 +86,10 @@ def serve(args, extra_argv):
|
||||
)
|
||||
return
|
||||
|
||||
from sglang.srt.plugins import load_plugins
|
||||
|
||||
load_plugins()
|
||||
|
||||
model_type, dispatch_argv = _extract_model_type_override(extra_argv)
|
||||
model_path = get_model_path(dispatch_argv)
|
||||
try:
|
||||
|
||||
@@ -56,6 +56,10 @@ if __name__ == "__main__":
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
from sglang.srt.plugins import load_plugins
|
||||
|
||||
load_plugins()
|
||||
|
||||
server_args = prepare_server_args(sys.argv[1:])
|
||||
|
||||
try:
|
||||
|
||||
@@ -22,6 +22,7 @@ from sglang.srt.compilation.cuda_piecewise_backend import CUDAPiecewiseBackend
|
||||
from sglang.srt.compilation.npu_piecewise_backend import NPUPiecewiseBackend
|
||||
from sglang.srt.compilation.pass_manager import PostGradPassManager
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.platforms import current_platform
|
||||
from sglang.srt.utils.common import is_npu
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@@ -48,7 +49,12 @@ def make_backend(
|
||||
sglang_backend,
|
||||
):
|
||||
|
||||
backend_cls = CUDAPiecewiseBackend if not is_npu() else NPUPiecewiseBackend
|
||||
if current_platform.is_out_of_tree():
|
||||
backend_cls = current_platform.get_piecewise_backend_cls()
|
||||
elif is_npu():
|
||||
backend_cls = NPUPiecewiseBackend
|
||||
else:
|
||||
backend_cls = CUDAPiecewiseBackend
|
||||
return backend_cls(
|
||||
graph,
|
||||
compile_config,
|
||||
|
||||
@@ -84,6 +84,7 @@ from sglang.srt.managers.scheduler import run_scheduler_process
|
||||
from sglang.srt.managers.template_manager import TemplateManager
|
||||
from sglang.srt.managers.tokenizer_manager import TokenizerManager
|
||||
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
|
||||
from sglang.srt.plugins import load_plugins
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs
|
||||
from sglang.srt.utils import (
|
||||
MultiprocessingSerializer,
|
||||
@@ -167,6 +168,10 @@ class Engine(EngineScoreMixin, EngineBase):
|
||||
Please refer to `ServerArgs` for the documentation.
|
||||
"""
|
||||
|
||||
# Ensure plugins are loaded before ServerArgs construction,
|
||||
# so hooks on ServerArgs.__post_init__ fire correctly.
|
||||
load_plugins()
|
||||
|
||||
# Parse server_args
|
||||
if "server_args" in kwargs:
|
||||
# Directly load server_args
|
||||
@@ -647,6 +652,11 @@ class Engine(EngineScoreMixin, EngineBase):
|
||||
# Configure global environment
|
||||
configure_logger(server_args)
|
||||
_set_envs_and_config(server_args)
|
||||
|
||||
# Defensive: ensure plugins loaded (may already be loaded by
|
||||
# Engine.__init__ or CLI entry).
|
||||
load_plugins()
|
||||
|
||||
server_args.check_server_args()
|
||||
_set_gc(server_args)
|
||||
|
||||
|
||||
@@ -554,6 +554,10 @@ class Envs:
|
||||
# Sglang Cache Dir
|
||||
SGLANG_CACHE_DIR = EnvStr(os.path.expanduser("~/.cache/sglang"))
|
||||
|
||||
# Plugin system
|
||||
SGLANG_PLATFORM = EnvStr("")
|
||||
SGLANG_PLUGINS = EnvStr("")
|
||||
|
||||
|
||||
envs = Envs()
|
||||
EnvField._allow_set_name = False
|
||||
|
||||
@@ -1,8 +1,9 @@
|
||||
from typing import Callable
|
||||
from typing import Callable, ClassVar
|
||||
|
||||
from torch import nn
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.platforms import current_platform
|
||||
from sglang.srt.utils import (
|
||||
cpu_has_amx_support,
|
||||
is_cpu,
|
||||
@@ -23,6 +24,15 @@ _is_musa = is_musa()
|
||||
|
||||
|
||||
class MultiPlatformOp(nn.Module):
|
||||
|
||||
# OOT forward registry: maps dispatch_key -> {op_cls -> forward_fn}
|
||||
_oot_forward_registry: ClassVar[dict[str, dict[type, Callable]]] = {}
|
||||
|
||||
@classmethod
|
||||
def register_oot_forward(cls, op_cls: type, fn: Callable, platform_key: str):
|
||||
"""Register an OOT forward implementation for a specific op class and platform."""
|
||||
cls._oot_forward_registry.setdefault(platform_key, {})[op_cls] = fn
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._forward_method: Callable = self.dispatch_forward()
|
||||
@@ -100,6 +110,17 @@ class MultiPlatformOp(nn.Module):
|
||||
return self.forward_native(*args, **kwargs)
|
||||
|
||||
def dispatch_forward(self):
|
||||
# OOT platform dispatch: check registry then method lookup
|
||||
if current_platform.is_out_of_tree():
|
||||
key = current_platform.get_dispatch_key_name()
|
||||
oot = self._oot_forward_registry.get(key, {})
|
||||
if type(self) in oot:
|
||||
return oot[type(self)].__get__(self)
|
||||
method = getattr(self, f"forward_{key}", None)
|
||||
if method is not None:
|
||||
return method
|
||||
return self.forward_native
|
||||
|
||||
if _is_cuda:
|
||||
return self.forward_cuda
|
||||
elif _is_hip:
|
||||
|
||||
@@ -204,6 +204,7 @@ from sglang.srt.observability.scheduler_metrics_mixin import (
|
||||
)
|
||||
from sglang.srt.observability.trace import process_tracing_init, trace_set_thread_info
|
||||
from sglang.srt.parser.reasoning_parser import ReasoningParser
|
||||
from sglang.srt.plugins import load_plugins
|
||||
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||
from sglang.srt.server_args import PortArgs, ServerArgs, get_global_server_args
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
@@ -3743,6 +3744,8 @@ def run_scheduler_process(
|
||||
dp_rank: Optional[int],
|
||||
pipe_writer,
|
||||
):
|
||||
# Load plugins so hooks can override Scheduler and its dependencies.
|
||||
load_plugins()
|
||||
dp_rank = configure_scheduler_process(
|
||||
server_args,
|
||||
gpu_id,
|
||||
|
||||
@@ -54,6 +54,7 @@ from sglang.srt.mem_cache.utils import (
|
||||
set_mla_kv_buffer_triton_fp8_quant,
|
||||
set_mla_kv_scale_buffer_triton,
|
||||
)
|
||||
from sglang.srt.platforms import current_platform
|
||||
from sglang.srt.utils import (
|
||||
cpu_has_amx_support,
|
||||
is_cpu,
|
||||
@@ -780,8 +781,12 @@ class MHATokenToKVPool(KVCache):
|
||||
self._create_buffers()
|
||||
|
||||
self.device_module = torch.get_device_module(self.device)
|
||||
|
||||
_use_alt_stream = _is_cuda or current_platform.is_cuda_alike()
|
||||
self.alt_stream = (
|
||||
self.device_module.Stream() if _is_cuda and enable_alt_stream else None
|
||||
self.device_module.Stream()
|
||||
if _use_alt_stream and enable_alt_stream
|
||||
else None
|
||||
)
|
||||
|
||||
if enable_kv_cache_copy:
|
||||
@@ -1262,7 +1267,9 @@ class HybridLinearKVPool(KVCache):
|
||||
|
||||
TokenToKVPoolClass = MHATokenToKVPool
|
||||
|
||||
if _is_npu:
|
||||
if current_platform.is_out_of_tree():
|
||||
TokenToKVPoolClass = current_platform.get_mha_kv_pool_cls()
|
||||
elif _is_npu:
|
||||
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
|
||||
NPUMHATokenToKVPool,
|
||||
)
|
||||
@@ -1283,7 +1290,9 @@ class HybridLinearKVPool(KVCache):
|
||||
|
||||
TokenToKVPoolClass = MLATokenToKVPool
|
||||
|
||||
if _is_npu:
|
||||
if current_platform.is_out_of_tree():
|
||||
TokenToKVPoolClass = current_platform.get_mla_kv_pool_cls()
|
||||
elif _is_npu:
|
||||
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
|
||||
NPUMLATokenToKVPool,
|
||||
)
|
||||
|
||||
@@ -150,6 +150,7 @@ from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
|
||||
)
|
||||
from sglang.srt.model_loader.utils import set_default_torch_dtype
|
||||
from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.platforms import current_platform
|
||||
from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
|
||||
from sglang.srt.server_args import (
|
||||
ServerArgs,
|
||||
@@ -207,6 +208,8 @@ if _is_npu:
|
||||
from sglang.srt.hardware_backend.npu.utils import init_npu_backend
|
||||
|
||||
init_npu_backend()
|
||||
elif current_platform.is_out_of_tree():
|
||||
current_platform.init_backend()
|
||||
|
||||
MLA_ATTENTION_BACKENDS = [
|
||||
"aiter",
|
||||
@@ -702,6 +705,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
# Init routed experts capturer
|
||||
self.init_routed_experts_capturer()
|
||||
|
||||
# TODO: Refactor device-specific init branches into platform interface (separate PR).
|
||||
# Must be called BEFORE init_device_graphs() so CUDA graph capture
|
||||
# runs with aux hidden state capture enabled.
|
||||
self.init_aux_hidden_state_capture()
|
||||
@@ -714,6 +718,13 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
elif self.device in ["npu", "cpu"]:
|
||||
self.init_attention_backend()
|
||||
self.init_device_graphs()
|
||||
elif current_platform.is_out_of_tree():
|
||||
self.init_attention_backend()
|
||||
if current_platform.support_cuda_graph():
|
||||
self.init_device_graphs()
|
||||
else:
|
||||
self.graph_runner = None
|
||||
self.graph_mem_usage = 0
|
||||
else:
|
||||
self.graph_runner = None
|
||||
self.graph_mem_usage = 0
|
||||
@@ -1483,7 +1494,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
self.server_args.load_format = load_format
|
||||
self.load_config = load_config
|
||||
|
||||
if recapture_cuda_graph and (self.device == "cuda" or self.device == "musa"):
|
||||
if recapture_cuda_graph and (
|
||||
self.device == "cuda"
|
||||
or self.device == "musa"
|
||||
or (
|
||||
current_platform.is_out_of_tree()
|
||||
and current_platform.support_cuda_graph()
|
||||
)
|
||||
):
|
||||
self.init_device_graphs()
|
||||
|
||||
logger.info("Update weights end.")
|
||||
@@ -2532,8 +2550,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
tic = time.perf_counter()
|
||||
before_mem = get_available_gpu_memory(self.device, self.gpu_id)
|
||||
graph_backend = defaultdict(
|
||||
lambda: "cuda graph",
|
||||
lambda: f"{current_platform.device_name} graph",
|
||||
{
|
||||
"cuda": "cuda graph",
|
||||
"musa": "cuda graph",
|
||||
"cpu": "cpu graph",
|
||||
"npu": "npu graph",
|
||||
},
|
||||
@@ -2541,14 +2561,18 @@ class ModelRunner(ModelRunnerKVCacheMixin):
|
||||
logger.info(
|
||||
f"Capture {graph_backend[self.device]} begin. This can take up to several minutes. avail mem={before_mem:.2f} GB"
|
||||
)
|
||||
graph_runners = defaultdict(
|
||||
lambda: CudaGraphRunner,
|
||||
{
|
||||
"cpu": CPUGraphRunner,
|
||||
"npu": NPUGraphRunner,
|
||||
},
|
||||
)
|
||||
self.graph_runner = graph_runners[self.device](self)
|
||||
if current_platform.is_out_of_tree():
|
||||
GraphRunnerCls = current_platform.get_graph_runner_cls()
|
||||
self.graph_runner = GraphRunnerCls(self)
|
||||
else:
|
||||
graph_runners = defaultdict(
|
||||
lambda: CudaGraphRunner,
|
||||
{
|
||||
"cpu": CPUGraphRunner,
|
||||
"npu": NPUGraphRunner,
|
||||
},
|
||||
)
|
||||
self.graph_runner = graph_runners[self.device](self)
|
||||
|
||||
after_mem = get_available_gpu_memory(self.device, self.gpu_id)
|
||||
self.graph_mem_usage = before_mem - after_mem
|
||||
|
||||
@@ -282,7 +282,63 @@ class ModelRunnerKVCacheMixin:
|
||||
|
||||
# Initialize token_to_kv_pool
|
||||
is_nsa_model = is_deepseek_nsa(self.model_config.hf_config)
|
||||
if self.server_args.attention_backend == "ascend" and not self.mambaish_config:
|
||||
|
||||
# Check out-of-tree platform (plugin system) first
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if current_platform.is_out_of_tree() and not self.mambaish_config:
|
||||
if self.use_mla_backend and is_nsa_model:
|
||||
PoolCls = current_platform.get_nsa_kv_pool_cls()
|
||||
self.token_to_kv_pool = PoolCls(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
kv_lora_rank=self.model_config.kv_lora_rank,
|
||||
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
|
||||
layer_num=self.num_effective_layers,
|
||||
device=self.device,
|
||||
kv_cache_dim=self.calculate_mla_kv_cache_dim(),
|
||||
enable_memory_saver=self.server_args.enable_memory_saver,
|
||||
start_layer=self.start_layer,
|
||||
end_layer=self.end_layer,
|
||||
index_head_dim=get_nsa_index_head_dim(self.model_config.hf_config),
|
||||
)
|
||||
elif self.use_mla_backend:
|
||||
PoolCls = current_platform.get_mla_kv_pool_cls()
|
||||
self.token_to_kv_pool = PoolCls(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
kv_lora_rank=self.model_config.kv_lora_rank,
|
||||
qk_rope_head_dim=self.model_config.qk_rope_head_dim,
|
||||
index_head_dim=(
|
||||
self.model_config.index_head_dim if is_nsa_model else None
|
||||
),
|
||||
layer_num=self.num_effective_layers,
|
||||
device=self.device,
|
||||
enable_memory_saver=self.server_args.enable_memory_saver,
|
||||
start_layer=self.start_layer,
|
||||
end_layer=self.end_layer,
|
||||
)
|
||||
else:
|
||||
PoolCls = current_platform.get_mha_kv_pool_cls()
|
||||
self.token_to_kv_pool = PoolCls(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
head_num=self.model_config.get_num_kv_heads(
|
||||
get_attention_tp_size()
|
||||
),
|
||||
head_dim=self.model_config.head_dim,
|
||||
layer_num=self.num_effective_layers,
|
||||
device=self.device,
|
||||
enable_memory_saver=self.server_args.enable_memory_saver,
|
||||
start_layer=self.start_layer,
|
||||
end_layer=self.end_layer,
|
||||
)
|
||||
elif (
|
||||
self.server_args.attention_backend == "ascend" and not self.mambaish_config
|
||||
):
|
||||
if self.is_hybrid_swa:
|
||||
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
|
||||
NPUMHATokenToKVPool,
|
||||
@@ -513,7 +569,17 @@ class ModelRunnerKVCacheMixin:
|
||||
# Initialize token_to_kv_pool_allocator
|
||||
need_sort = self.server_args.disaggregation_mode in ("decode", "prefill")
|
||||
if self.token_to_kv_pool_allocator is None:
|
||||
if _is_npu and (
|
||||
if current_platform.is_out_of_tree():
|
||||
AllocatorCls = current_platform.get_paged_allocator_cls()
|
||||
self.token_to_kv_pool_allocator = AllocatorCls(
|
||||
self.max_total_num_tokens,
|
||||
page_size=self.page_size,
|
||||
dtype=self.kv_cache_dtype,
|
||||
device=self.device,
|
||||
kvcache=self.token_to_kv_pool,
|
||||
need_sort=need_sort,
|
||||
)
|
||||
elif _is_npu and (
|
||||
self.server_args.attention_backend == "ascend"
|
||||
or self.hybrid_gdn_config is not None
|
||||
):
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
"""
|
||||
SGLang Platform Discovery and Lazy Initialization.
|
||||
|
||||
Provides `current_platform` as a module-level lazy singleton. On first access,
|
||||
it discovers platform plugins via entry_points and instantiates the appropriate
|
||||
SRTPlatform subclass.
|
||||
|
||||
Usage:
|
||||
from sglang.srt.platforms import current_platform
|
||||
print(current_platform.device_name)
|
||||
"""
|
||||
|
||||
import logging
|
||||
import pkgutil
|
||||
from importlib.metadata import entry_points
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.platforms.interface import SRTPlatform
|
||||
from sglang.srt.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_current_platform: SRTPlatform | None = None
|
||||
|
||||
|
||||
def _resolve_platform() -> SRTPlatform:
|
||||
"""
|
||||
Discover and instantiate the active platform.
|
||||
|
||||
Discovery flow:
|
||||
1. Branch on SGLANG_PLATFORM:
|
||||
|
||||
SGLANG_PLATFORM set (front-loading filter):
|
||||
- Enumerate entry_points without importing any plugin modules
|
||||
- Only ep.load() + activate() the named plugin
|
||||
- Other plugins are never imported (avoids pulling their dependencies)
|
||||
- Plugin name not found → RuntimeError
|
||||
- activate() returns None → RuntimeError (hardware unavailable)
|
||||
|
||||
SGLANG_PLATFORM unset (auto-discover):
|
||||
- Import and activate all discovered plugins
|
||||
- 0 activated → fallback base SRTPlatform
|
||||
- 1 activated → use it
|
||||
- N activated → RuntimeError (must set SGLANG_PLATFORM)
|
||||
|
||||
SGLANG_PLATFORM matches against entry_point names.
|
||||
"""
|
||||
selected = envs.SGLANG_PLATFORM.get()
|
||||
|
||||
if selected:
|
||||
# Front-loading filter: only import and activate the specified plugin.
|
||||
# Other plugins' modules are never loaded — avoids pulling their deps.
|
||||
discovered = entry_points(group=PLATFORM_PLUGINS_GROUP)
|
||||
ep_map = {ep.name: ep for ep in discovered}
|
||||
|
||||
if selected not in ep_map:
|
||||
available = ", ".join(f"'{n}'" for n in ep_map) if ep_map else "none"
|
||||
raise RuntimeError(
|
||||
f"SGLANG_PLATFORM={selected!r} not found in discovered platform plugins "
|
||||
f"(available: {available}). Install the plugin with 'pip install -e' "
|
||||
f"to register its entry_points."
|
||||
)
|
||||
|
||||
try:
|
||||
plugin_fn = ep_map[selected].load()
|
||||
result = plugin_fn()
|
||||
except Exception:
|
||||
logger.exception("Failed to activate platform plugin: %s", selected)
|
||||
raise
|
||||
|
||||
if result is None:
|
||||
raise RuntimeError(
|
||||
f"Platform plugin {selected!r} is installed but activate() returned None "
|
||||
f"(hardware not available on this machine?)."
|
||||
)
|
||||
logger.info("OOT platform plugin activated: %s -> %s", selected, result)
|
||||
return _load_platform_class(result)()
|
||||
|
||||
# Auto-discover: import and activate all plugins, expect exactly one
|
||||
all_plugins = load_plugins_by_group(PLATFORM_PLUGINS_GROUP)
|
||||
|
||||
activated: dict[str, str] = {}
|
||||
for name, (plugin_fn, _dist) in all_plugins.items():
|
||||
try:
|
||||
result = plugin_fn()
|
||||
if result is not None:
|
||||
activated[name] = result
|
||||
logger.info("OOT platform plugin activated: %s -> %s", name, result)
|
||||
except Exception:
|
||||
logger.exception("Failed to activate platform plugin: %s", name)
|
||||
|
||||
if len(activated) == 0:
|
||||
logger.warning("No platform detected. Using base SRTPlatform with defaults.")
|
||||
return SRTPlatform()
|
||||
|
||||
if len(activated) == 1:
|
||||
name, qualname = next(iter(activated.items()))
|
||||
return _load_platform_class(qualname)()
|
||||
|
||||
# Multiple activated without SGLANG_PLATFORM
|
||||
names_str = ", ".join(f"'{n}'" for n in activated)
|
||||
raise RuntimeError(
|
||||
f"Multiple platform plugins activated: {names_str}. "
|
||||
f"Set SGLANG_PLATFORM to select one."
|
||||
)
|
||||
|
||||
|
||||
def _load_platform_class(qualname: str) -> type:
|
||||
"""Load an SRTPlatform subclass from its fully-qualified class name."""
|
||||
cls = pkgutil.resolve_name(qualname)
|
||||
if not isinstance(cls, type) or not issubclass(cls, SRTPlatform):
|
||||
raise TypeError(
|
||||
f"Expected an SRTPlatform subclass, got {type(cls)}: {qualname}"
|
||||
)
|
||||
return cls
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
"""Lazy initialization of current_platform on first access."""
|
||||
if name == "current_platform":
|
||||
global _current_platform
|
||||
if _current_platform is None:
|
||||
_current_platform = _resolve_platform()
|
||||
return _current_platform
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -0,0 +1,244 @@
|
||||
"""
|
||||
Shared device abstraction for SGLang platforms.
|
||||
|
||||
DeviceMixin provides the common device identity queries and operations
|
||||
shared between the SRT (LLM inference) and Multimodal (diffusion)
|
||||
platform hierarchies. Concrete per-device mixins (e.g. MyDeviceMixin)
|
||||
implement the abstract operations; subsystem-specific platforms
|
||||
(SRTPlatform, MMPlatform) inherit DeviceMixin and add their own methods.
|
||||
|
||||
Hierarchy example (OOT plugin)::
|
||||
|
||||
DeviceMixin
|
||||
├── MyDeviceMixin(DeviceMixin) # vendor-specific device operations
|
||||
├── SRTPlatform(DeviceMixin) # + graph runner, KV pool, …
|
||||
│ └── MySRTPlatform(SRTPlatform, MyDeviceMixin)
|
||||
└── MMPlatform(DeviceMixin) # + attention backend, VAE, …
|
||||
└── MyMMPlatform(MMPlatform, MyDeviceMixin)
|
||||
|
||||
Method status annotations:
|
||||
|
||||
- ``[Active]`` — SGLang core calls this method through ``current_platform``.
|
||||
OOT implementations take effect immediately.
|
||||
- ``[Planned]`` — Reserved interface. SGLang core still uses hardcoded calls
|
||||
(e.g. ``torch.cuda.empty_cache()``). OOT implementations will NOT take
|
||||
effect until the core is migrated in a future PR.
|
||||
"""
|
||||
|
||||
import enum
|
||||
from typing import TYPE_CHECKING, NamedTuple, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
|
||||
class PlatformEnum(enum.Enum):
|
||||
"""Enumeration of known platform types.
|
||||
|
||||
Superset of both SRT and MM enums so that a single PlatformEnum can
|
||||
be shared across subsystems.
|
||||
"""
|
||||
|
||||
CUDA = enum.auto()
|
||||
ROCM = enum.auto()
|
||||
CPU = enum.auto()
|
||||
XPU = enum.auto()
|
||||
MUSA = enum.auto()
|
||||
NPU = enum.auto()
|
||||
TPU = enum.auto()
|
||||
MPS = enum.auto()
|
||||
OOT = enum.auto() # Out-of-tree (external plugin)
|
||||
UNSPECIFIED = enum.auto()
|
||||
|
||||
|
||||
class CpuArchEnum(enum.Enum):
|
||||
"""CPU architecture enumeration."""
|
||||
|
||||
X86 = enum.auto()
|
||||
ARM = enum.auto()
|
||||
UNSPECIFIED = enum.auto()
|
||||
|
||||
|
||||
class DeviceCapability(NamedTuple):
|
||||
"""Device compute capability (major, minor).
|
||||
|
||||
Uses NamedTuple for built-in comparison support:
|
||||
``DeviceCapability(9, 0) >= DeviceCapability(8, 9)`` works naturally.
|
||||
"""
|
||||
|
||||
major: int
|
||||
minor: int
|
||||
|
||||
def as_version_str(self) -> str:
|
||||
return f"{self.major}.{self.minor}"
|
||||
|
||||
def to_int(self) -> int:
|
||||
"""Express capability as ``<major><minor>`` (minor is single digit)."""
|
||||
assert 0 <= self.minor < 10
|
||||
return self.major * 10 + self.minor
|
||||
|
||||
|
||||
class DeviceMixin:
|
||||
"""Mixin providing device identity queries and basic device operations.
|
||||
|
||||
Class-level attributes (override in subclasses):
|
||||
_enum: PlatformEnum identifying this platform.
|
||||
device_name: Human-readable short name (e.g. "cuda", "npu").
|
||||
device_type: ``torch.device`` type string (e.g. "cuda", "npu").
|
||||
"""
|
||||
|
||||
_enum: PlatformEnum = PlatformEnum.UNSPECIFIED
|
||||
device_name: str = "unknown"
|
||||
device_type: str = "cpu"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Platform identity queries
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def is_cuda(self) -> bool:
|
||||
return self._enum == PlatformEnum.CUDA
|
||||
|
||||
def is_rocm(self) -> bool:
|
||||
return self._enum == PlatformEnum.ROCM
|
||||
|
||||
def is_cpu(self) -> bool:
|
||||
return self._enum == PlatformEnum.CPU
|
||||
|
||||
def is_xpu(self) -> bool:
|
||||
return self._enum == PlatformEnum.XPU
|
||||
|
||||
def is_musa(self) -> bool:
|
||||
return self._enum == PlatformEnum.MUSA
|
||||
|
||||
def is_npu(self) -> bool:
|
||||
return self._enum == PlatformEnum.NPU
|
||||
|
||||
def is_tpu(self) -> bool:
|
||||
return self._enum == PlatformEnum.TPU
|
||||
|
||||
def is_mps(self) -> bool:
|
||||
return self._enum == PlatformEnum.MPS
|
||||
|
||||
def is_cuda_alike(self) -> bool:
|
||||
"""True for CUDA, ROCm, or MUSA (all expose CUDA-like APIs)."""
|
||||
return self._enum in (
|
||||
PlatformEnum.CUDA,
|
||||
PlatformEnum.ROCM,
|
||||
PlatformEnum.MUSA,
|
||||
)
|
||||
|
||||
def is_out_of_tree(self) -> bool:
|
||||
"""True for externally-registered OOT platforms."""
|
||||
return self._enum == PlatformEnum.OOT
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Active methods — core calls these through current_platform.
|
||||
# OOT implementations take effect immediately.
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_device_total_memory(self, device_id: int = 0) -> int:
|
||||
"""[Active] Get total device memory in bytes."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_current_memory_usage(
|
||||
self, device: Optional["torch.device"] = None
|
||||
) -> float:
|
||||
"""[Active] Get current peak memory usage in bytes."""
|
||||
raise NotImplementedError
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Planned methods — reserved interface. Core still uses hardcoded
|
||||
# calls (e.g. torch.cuda.*). OOT implementations will NOT take
|
||||
# effect until the core is migrated in a future PR.
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
# ---- Device management ----
|
||||
|
||||
def get_device(self, local_rank: int) -> "torch.device":
|
||||
"""[Planned] Return ``torch.device`` for the given local rank."""
|
||||
raise NotImplementedError
|
||||
|
||||
def set_device(self, device: "torch.device") -> None:
|
||||
"""[Planned] Set the current device."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_device_name(self, device_id: int = 0) -> str:
|
||||
"""[Planned] Get human-readable device name."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_device_uuid(self, device_id: int = 0) -> str:
|
||||
"""[Planned] Get unique device identifier string."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_device_capability(self, device_id: int = 0) -> Optional["DeviceCapability"]:
|
||||
"""[Planned] Get device compute capability. None if N/A."""
|
||||
raise NotImplementedError
|
||||
|
||||
def empty_cache(self) -> None:
|
||||
"""[Planned] Release cached device memory. No-op for CPU-like platforms."""
|
||||
pass
|
||||
|
||||
def synchronize(self) -> None:
|
||||
"""[Planned] Synchronize device operations. No-op for CPU-like platforms."""
|
||||
pass
|
||||
|
||||
# ---- Memory ----
|
||||
|
||||
def get_available_memory(self, device_id: int = 0) -> tuple[int, int]:
|
||||
"""[Planned] Return ``(free_bytes, total_bytes)``."""
|
||||
raise NotImplementedError
|
||||
|
||||
# ---- Distributed ----
|
||||
|
||||
def get_torch_distributed_backend_str(self) -> str:
|
||||
"""[Planned] Return the torch.distributed backend string (e.g. "nccl", "hccl")."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_communicator_class(self) -> type | None:
|
||||
"""[Planned] Return platform-specific communicator class, or None for default."""
|
||||
return None
|
||||
|
||||
# ---- Misc ----
|
||||
|
||||
@classmethod
|
||||
def inference_mode(cls):
|
||||
"""[Planned] Return inference mode context manager."""
|
||||
import torch
|
||||
|
||||
return torch.inference_mode(mode=True)
|
||||
|
||||
@classmethod
|
||||
def seed_everything(cls, seed: int | None = None) -> None:
|
||||
"""[Planned] Set random seeds for reproducibility across all libraries."""
|
||||
if seed is not None:
|
||||
import random
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
|
||||
def verify_quantization(self, quant: str) -> None:
|
||||
"""[Planned] Validate that a quantization method is supported. No-op by default."""
|
||||
pass
|
||||
|
||||
@classmethod
|
||||
def get_cpu_architecture(cls) -> "CpuArchEnum":
|
||||
"""[Planned] Detect CPU architecture."""
|
||||
import platform as _platform
|
||||
|
||||
machine = _platform.machine().lower()
|
||||
if machine in ("x86_64", "amd64", "i386", "i686"):
|
||||
return CpuArchEnum.X86
|
||||
elif machine in ("arm64", "aarch64"):
|
||||
return CpuArchEnum.ARM
|
||||
return CpuArchEnum.UNSPECIFIED
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Dunder helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}(device={self.device_name})"
|
||||
@@ -0,0 +1,133 @@
|
||||
"""
|
||||
SGLang SRT Hardware Platform Abstraction.
|
||||
|
||||
Defines SRTPlatform — the base class for SRT (LLM inference) platform
|
||||
backends. SRTPlatform inherits DeviceMixin for shared device operations
|
||||
and adds SRT-specific subsystem factory methods, capability flags, and
|
||||
configuration lifecycle hooks.
|
||||
|
||||
Out-of-tree platforms register via setuptools entry_points under the
|
||||
"sglang.platform_plugins" group and should subclass SRTPlatform.
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum
|
||||
|
||||
if TYPE_CHECKING:
|
||||
pass
|
||||
|
||||
# Re-export for convenience
|
||||
__all__ = ["SRTPlatform", "PlatformEnum"]
|
||||
|
||||
|
||||
class SRTPlatform(DeviceMixin):
|
||||
"""
|
||||
Base class for SRT hardware platform backends.
|
||||
|
||||
Inherits device identity queries and operations from DeviceMixin.
|
||||
Adds SRT-specific factory methods, capability flags, and lifecycle hooks.
|
||||
|
||||
OOT platforms should subclass SRTPlatform and override the methods
|
||||
relevant to their hardware.
|
||||
"""
|
||||
|
||||
# SRT-specific class-level attribute
|
||||
supported_quantization: list[str] = []
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Configuration lifecycle
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def apply_server_args_defaults(self, server_args) -> None:
|
||||
"""Apply platform-specific default values to server arguments.
|
||||
|
||||
Called after ServerArgs is parsed.
|
||||
"""
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Subsystem factory methods
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_default_attention_backend(self) -> str:
|
||||
"""Return the default attention backend name for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_graph_runner_cls(self) -> type:
|
||||
"""Return the graph runner class for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_mha_kv_pool_cls(self) -> type:
|
||||
"""Return the MHA KV pool class for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_mla_kv_pool_cls(self) -> type:
|
||||
"""Return the MLA KV pool class for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_nsa_kv_pool_cls(self) -> type:
|
||||
"""Return the NSA KV pool class for this platform (DeepSeek V3.2)."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_paged_allocator_cls(self) -> type:
|
||||
"""Return the paged allocator class for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_compile_backend(self, mode: str | None = None) -> str:
|
||||
"""Return the compilation backend identifier.
|
||||
|
||||
``mode`` is an optional hint for the platform (e.g. "npugraph_ex").
|
||||
"""
|
||||
return "inductor"
|
||||
|
||||
def get_piecewise_backend_cls(self) -> type:
|
||||
"""Return the piecewise compilation backend class for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Capability flags (safe conservative defaults)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def supports_fp8(self) -> bool:
|
||||
"""Whether this platform supports FP8 quantization."""
|
||||
return False
|
||||
|
||||
def is_pin_memory_available(self) -> bool:
|
||||
"""Whether pinned memory is available on this platform."""
|
||||
return True
|
||||
|
||||
def support_cuda_graph(self) -> bool:
|
||||
"""Whether this platform supports device graph capture and replay.
|
||||
Controls CUDA graph (CudaGraphRunner) for the decode path.
|
||||
OOT platforms that support graph-style capture should return True.
|
||||
"""
|
||||
return False
|
||||
|
||||
def support_piecewise_cuda_graph(self) -> bool:
|
||||
"""Whether this platform supports piecewise CUDA graph.
|
||||
|
||||
Controls PiecewiseCudaGraphRunner for the prefill/extend path
|
||||
(torch.compile backend).
|
||||
"""
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Initialization
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def init_backend(self) -> None:
|
||||
"""One-time backend initialization. Called in each worker."""
|
||||
pass
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# MultiPlatformOp integration
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def get_dispatch_key_name(self) -> str:
|
||||
"""Return the dispatch key name for MultiPlatformOp.
|
||||
|
||||
Determines which ``forward_<key>()`` method is selected.
|
||||
E.g. "cuda", "npu", "hip", "xpu", "cpu".
|
||||
"""
|
||||
return "native"
|
||||
@@ -0,0 +1,141 @@
|
||||
"""
|
||||
SGLang Unified Plugin Framework.
|
||||
|
||||
Supports two types of plugins via setuptools entry_points:
|
||||
1. Hardware Platform Plugins (sglang.srt.platforms) - register custom hardware platforms
|
||||
2. General Plugins (sglang.srt.plugins) - inject hooks into functions/methods, replace classes, etc.
|
||||
|
||||
Plugins are discovered automatically when installed via pip.
|
||||
- Platform plugins: use ``SGLANG_PLATFORM`` to select when multiple are installed.
|
||||
- General plugins: use ``SGLANG_PLUGINS`` (comma-separated) to restrict which are loaded.
|
||||
"""
|
||||
|
||||
import logging
|
||||
from collections.abc import Callable
|
||||
from importlib.metadata import entry_points
|
||||
from typing import Any
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.plugins.hook_registry import (
|
||||
HookRegistry,
|
||||
HookSource,
|
||||
_current_plugin_source,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Entry point group names
|
||||
PLATFORM_PLUGINS_GROUP = "sglang.srt.platforms"
|
||||
GENERAL_PLUGINS_GROUP = "sglang.srt.plugins"
|
||||
|
||||
# Guard against multiple loads in the same process
|
||||
_plugins_loaded = False
|
||||
|
||||
|
||||
def load_plugins_by_group(
|
||||
group: str,
|
||||
excluded_dists: set[str] | None = None,
|
||||
) -> dict[str, tuple[Callable[[], Any], str | None]]:
|
||||
"""
|
||||
Discover and load plugins registered under the given entry point group.
|
||||
|
||||
Args:
|
||||
group: The setuptools entry_point group name.
|
||||
excluded_dists: Distribution names to skip. Plugins from these
|
||||
distributions are never ``ep.load()``-ed (avoids importing
|
||||
their modules and pulling hardware-specific dependencies).
|
||||
|
||||
Returns:
|
||||
Dictionary mapping plugin name to ``(callable, dist_name)``.
|
||||
"""
|
||||
# SGLANG_PLUGINS whitelist (comma-separated plugin names)
|
||||
allowed_set: set[str] | None = None
|
||||
allowed_str = envs.SGLANG_PLUGINS.get()
|
||||
if allowed_str:
|
||||
allowed_set = {x.strip() for x in allowed_str.split(",") if x.strip()}
|
||||
|
||||
discovered = entry_points(group=group)
|
||||
if len(discovered) == 0:
|
||||
logger.debug("No plugins found for group %s.", group)
|
||||
return {}
|
||||
|
||||
logger.info("Available plugins for group %s:", group)
|
||||
for ep in discovered:
|
||||
logger.info(" - %s -> %s", ep.name, ep.value)
|
||||
|
||||
plugins: dict[str, tuple[Callable[[], Any], str | None]] = {}
|
||||
for ep in discovered:
|
||||
if allowed_set is not None and ep.name not in allowed_set:
|
||||
logger.info("Skipping plugin %s (not in SGLANG_PLUGINS)", ep.name)
|
||||
continue
|
||||
dist_name = ep.dist.name if ep.dist else None
|
||||
if excluded_dists and dist_name in excluded_dists:
|
||||
logger.info(
|
||||
"Skipping plugin %s (dist %s excluded by SGLANG_PLATFORM)",
|
||||
ep.name,
|
||||
dist_name,
|
||||
)
|
||||
continue
|
||||
try:
|
||||
func = ep.load()
|
||||
plugins[ep.name] = (func, dist_name)
|
||||
logger.info("Loaded plugin %s from group %s", ep.name, group)
|
||||
except Exception:
|
||||
logger.exception("Failed to load plugin %s from group %s", ep.name, group)
|
||||
|
||||
return plugins
|
||||
|
||||
|
||||
def _get_excluded_dists() -> set[str]:
|
||||
"""Compute dist names to skip when ``SGLANG_PLATFORM`` is set.
|
||||
|
||||
Returns dist names that provide a platform plugin but are NOT the one
|
||||
selected by ``SGLANG_PLATFORM``. This prevents unselected platform
|
||||
packages from registering hooks that pull their hardware dependencies.
|
||||
"""
|
||||
selected = envs.SGLANG_PLATFORM.get()
|
||||
if not selected:
|
||||
return set()
|
||||
platform_eps = entry_points(group=PLATFORM_PLUGINS_GROUP)
|
||||
return {ep.dist.name for ep in platform_eps if ep.dist and ep.name != selected}
|
||||
|
||||
|
||||
def load_plugins():
|
||||
"""
|
||||
Load and execute all general plugins, then apply registered hooks.
|
||||
|
||||
Idempotent - safe to call multiple times. General plugins are functions
|
||||
whose side effects (registering hooks, replacing classes, etc.) are the
|
||||
desired behavior. Return values are ignored.
|
||||
|
||||
When ``SGLANG_PLATFORM`` is set, general plugins from unselected platform
|
||||
packages are automatically skipped (avoids pulling their dependencies).
|
||||
|
||||
After all plugins execute, ``HookRegistry.apply_hooks()`` is called
|
||||
automatically so callers only need this single function call.
|
||||
|
||||
This should be called early in every process (main, engine core, workers).
|
||||
"""
|
||||
global _plugins_loaded
|
||||
if _plugins_loaded:
|
||||
return
|
||||
_plugins_loaded = True
|
||||
|
||||
plugins = load_plugins_by_group(
|
||||
GENERAL_PLUGINS_GROUP,
|
||||
excluded_dists=_get_excluded_dists(),
|
||||
)
|
||||
|
||||
for name, (func, dist_name) in plugins.items():
|
||||
source = HookSource(plugin_name=name, dist_name=dist_name)
|
||||
token = _current_plugin_source.set(source)
|
||||
try:
|
||||
func()
|
||||
logger.info("Executed general plugin: %s", name)
|
||||
except Exception:
|
||||
logger.exception("Failed to execute general plugin: %s", name)
|
||||
finally:
|
||||
_current_plugin_source.reset(token)
|
||||
|
||||
# Apply all registered hooks (idempotent — already-patched targets are skipped).
|
||||
HookRegistry.apply_hooks()
|
||||
@@ -0,0 +1,430 @@
|
||||
"""
|
||||
Hook registry for SGLang plugins.
|
||||
|
||||
Provides before/after/around/replace hooks that can be applied to any
|
||||
function, method, or class in the sglang codebase. Hooks are registered
|
||||
during plugin loading and applied before the engine starts.
|
||||
|
||||
Usage:
|
||||
from sglang.srt.plugins.hook_registry import HookRegistry, HookType
|
||||
|
||||
def my_timer(original_fn, *args, **kwargs):
|
||||
start = time.perf_counter()
|
||||
result = original_fn(*args, **kwargs)
|
||||
print(f"Elapsed: {time.perf_counter() - start:.3f}s")
|
||||
return result
|
||||
|
||||
HookRegistry.register(
|
||||
"sglang.srt.managers.scheduler.Scheduler.schedule",
|
||||
my_timer,
|
||||
HookType.AROUND,
|
||||
)
|
||||
"""
|
||||
|
||||
import contextvars
|
||||
import functools
|
||||
import logging
|
||||
import pkgutil
|
||||
import sys
|
||||
import types
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable
|
||||
from enum import Enum
|
||||
from typing import NamedTuple
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HookSource(NamedTuple):
|
||||
"""Identifies which plugin registered a hook."""
|
||||
|
||||
plugin_name: str # entry_point name, e.g. "xpu_hooks"
|
||||
dist_name: str | None # distribution name, e.g. "sglang_xpu_platform"
|
||||
|
||||
|
||||
# Set by load_plugins() around each plugin's func() call, read by register().
|
||||
_current_plugin_source: contextvars.ContextVar[HookSource | None] = (
|
||||
contextvars.ContextVar("_current_plugin_source", default=None)
|
||||
)
|
||||
|
||||
|
||||
def _format_source(source: HookSource | None) -> str:
|
||||
"""Format source info for log messages."""
|
||||
if source is None:
|
||||
return "unknown"
|
||||
if source.dist_name:
|
||||
return f"plugin={source.plugin_name}, dist={source.dist_name}"
|
||||
return f"plugin={source.plugin_name}"
|
||||
|
||||
|
||||
class HookType(Enum):
|
||||
"""Types of hooks that can be applied to functions or classes."""
|
||||
|
||||
BEFORE = "before" # Execute before original; can modify args
|
||||
AFTER = "after" # Execute after original; can modify return value
|
||||
AROUND = "around" # Wrap original; full control over execution
|
||||
REPLACE = "replace" # Replace the original function or class entirely
|
||||
|
||||
|
||||
class HookRegistry:
|
||||
"""
|
||||
Global registry for function/method/class hooks.
|
||||
|
||||
Thread safety: All registration should happen during load_plugins()
|
||||
phase (single-threaded). apply_hooks() should be called once before the
|
||||
engine starts serving requests.
|
||||
"""
|
||||
|
||||
_hooks: dict[str, list[tuple[HookType, Callable, HookSource | None]]] = defaultdict(
|
||||
list
|
||||
)
|
||||
_patched: set[str] = set()
|
||||
|
||||
@classmethod
|
||||
def register(
|
||||
cls,
|
||||
target: str,
|
||||
hook: Callable,
|
||||
hook_type: HookType = HookType.AFTER,
|
||||
*,
|
||||
source: HookSource | None = None,
|
||||
):
|
||||
"""
|
||||
Register a hook on a target function, method, or class.
|
||||
|
||||
Args:
|
||||
target: Fully-qualified dotted path to the target.
|
||||
e.g. "sglang.srt.managers.scheduler.Scheduler.schedule"
|
||||
or "sglang.srt.managers.scheduler.Scheduler" (class)
|
||||
hook: The hook callable (function or class). Signature depends on hook_type:
|
||||
- BEFORE: fn(*args, **kwargs) -> (args, kwargs) or None
|
||||
- AFTER: fn(result, *args, **kwargs) -> new_result or None
|
||||
- AROUND: fn(original_fn, *args, **kwargs) -> result
|
||||
- REPLACE: fn(*args, **kwargs) -> result (function replacement)
|
||||
MyClass (class replacement)
|
||||
hook_type: Type of hook (default: AFTER).
|
||||
source: Optional source info. If None, auto-read from context var
|
||||
set by ``load_plugins()``.
|
||||
|
||||
Raises:
|
||||
TypeError: If a class is passed with a hook_type other than REPLACE.
|
||||
"""
|
||||
if isinstance(hook, type) and hook_type != HookType.REPLACE:
|
||||
raise TypeError(
|
||||
f"Class {hook.__name__} can only be used with HookType.REPLACE, "
|
||||
f"got HookType.{hook_type.name}. "
|
||||
f"Use a function for BEFORE/AFTER/AROUND hooks."
|
||||
)
|
||||
resolved_source = source or _current_plugin_source.get()
|
||||
# Warn on duplicate REPLACE for the same target
|
||||
if hook_type == HookType.REPLACE:
|
||||
existing_replace = [
|
||||
(h, src) for ht, h, src in cls._hooks[target] if ht == HookType.REPLACE
|
||||
]
|
||||
if existing_replace:
|
||||
prev, prev_src = existing_replace[-1]
|
||||
prev_name = getattr(prev, "__qualname__", None) or repr(prev)
|
||||
new_name = getattr(hook, "__qualname__", None) or repr(hook)
|
||||
logger.warning(
|
||||
"Multiple REPLACE hooks on '%s': previous (%s [%s]) will be "
|
||||
"overridden by (%s [%s]). The last registered REPLACE takes effect.",
|
||||
target,
|
||||
prev_name,
|
||||
_format_source(prev_src),
|
||||
new_name,
|
||||
_format_source(resolved_source),
|
||||
)
|
||||
cls._hooks[target].append((hook_type, hook, resolved_source))
|
||||
logger.debug(
|
||||
"Registered %s hook on %s [%s]",
|
||||
hook_type.value,
|
||||
target,
|
||||
_format_source(resolved_source),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def apply_hooks(cls):
|
||||
"""
|
||||
Apply all registered hooks to their target functions/classes.
|
||||
|
||||
This performs the actual monkey-patching. Should be called once after
|
||||
all plugins have been loaded and before the engine starts.
|
||||
|
||||
Targets with class REPLACE hooks are applied first, so that
|
||||
subsequent method-level hooks (AROUND, BEFORE, AFTER) on child
|
||||
attributes resolve against the *replaced* class rather than the
|
||||
original.
|
||||
"""
|
||||
sorted_items = sorted(cls._hooks.items(), key=cls._target_sort_key)
|
||||
for target, hooks in sorted_items:
|
||||
if target in cls._patched:
|
||||
continue
|
||||
try:
|
||||
cls._apply_target(target, hooks)
|
||||
cls._patched.add(target)
|
||||
except Exception:
|
||||
logger.exception("Failed to apply hooks to %s", target)
|
||||
|
||||
@staticmethod
|
||||
def _target_sort_key(item):
|
||||
"""Sort key: class REPLACE targets (tier 0) before all others (tier 1).
|
||||
|
||||
This ensures that when a class is replaced, subsequent method-level
|
||||
hooks on ``ClassName.method`` resolve against the replacement class.
|
||||
"""
|
||||
_target, hooks = item
|
||||
has_class_replace = any(
|
||||
isinstance(h, type) and ht == HookType.REPLACE for ht, h, _ in hooks
|
||||
)
|
||||
return (0 if has_class_replace else 1, _target)
|
||||
|
||||
@classmethod
|
||||
def _apply_target(cls, target: str, hooks: list):
|
||||
"""Resolve target, build wrapper chain, and replace the original."""
|
||||
parts = target.rsplit(".", 1)
|
||||
if len(parts) != 2:
|
||||
raise ValueError(
|
||||
f"Invalid target path (need at least module.attr): {target}"
|
||||
)
|
||||
|
||||
obj_path, attr_name = parts
|
||||
obj = pkgutil.resolve_name(obj_path)
|
||||
|
||||
# Check if the original is a classmethod or staticmethod by
|
||||
# inspecting __dict__ before getattr() triggers the descriptor
|
||||
# protocol (which would lose the wrapper type for classmethod).
|
||||
original = getattr(obj, attr_name)
|
||||
is_classmethod = False
|
||||
is_staticmethod = False
|
||||
if isinstance(obj, type):
|
||||
raw_attr = obj.__dict__.get(attr_name)
|
||||
if isinstance(raw_attr, classmethod):
|
||||
is_classmethod = True
|
||||
original = raw_attr.__func__
|
||||
elif isinstance(raw_attr, staticmethod):
|
||||
is_staticmethod = True
|
||||
original = raw_attr.__func__
|
||||
|
||||
# Cross-target conflict detection: if the parent object is a class
|
||||
# that was already class-REPLACE'd, and the replacement class defines
|
||||
# its own version of this method, a method REPLACE here will silently
|
||||
# override the replacement class's implementation.
|
||||
if isinstance(obj, type) and obj_path in cls._patched:
|
||||
has_method_replace = any(ht == HookType.REPLACE for ht, _, _ in hooks)
|
||||
if has_method_replace and attr_name in obj.__dict__:
|
||||
replace_sources = [
|
||||
_format_source(src)
|
||||
for ht, _, src in hooks
|
||||
if ht == HookType.REPLACE
|
||||
]
|
||||
logger.warning(
|
||||
"Method REPLACE on '%s' will override the class REPLACE's "
|
||||
"own implementation of '%s'. If this is unintended, remove "
|
||||
"the method REPLACE and modify the replacement class "
|
||||
"directly, or use AROUND to wrap it. (from: %s)",
|
||||
target,
|
||||
attr_name,
|
||||
", ".join(replace_sources),
|
||||
)
|
||||
|
||||
# Guard: if the target is a class, only REPLACE is safe. Wrapping a
|
||||
# class in a function would break isinstance/issubclass/inheritance.
|
||||
if isinstance(original, type):
|
||||
bad = [ht for ht, _, _ in hooks if ht != HookType.REPLACE]
|
||||
if bad:
|
||||
raise TypeError(
|
||||
f"Target '{target}' is a class. Only HookType.REPLACE is "
|
||||
f"allowed for class targets (got {bad[0].value}). "
|
||||
f"To hook a method, use '{target}.<method_name>' instead."
|
||||
)
|
||||
|
||||
# Warn about risky hook combinations
|
||||
hook_types = [ht for ht, _, _ in hooks]
|
||||
around_count = hook_types.count(HookType.AROUND)
|
||||
has_replace = HookType.REPLACE in hook_types
|
||||
has_others = any(ht != HookType.REPLACE for ht in hook_types)
|
||||
|
||||
if around_count > 1:
|
||||
around_sources = [
|
||||
_format_source(src) for ht, _, src in hooks if ht == HookType.AROUND
|
||||
]
|
||||
logger.warning(
|
||||
"Multiple AROUND hooks on '%s' (%d hooks, from: %s). If any AROUND hook "
|
||||
"skips calling original_fn, inner hooks will be bypassed.",
|
||||
target,
|
||||
around_count,
|
||||
", ".join(around_sources),
|
||||
)
|
||||
if has_replace and has_others:
|
||||
logger.info(
|
||||
"Target '%s' has both REPLACE and %s hooks. "
|
||||
"REPLACE will be applied first, then wrapped by other hooks.",
|
||||
target,
|
||||
", ".join(
|
||||
sorted({ht.value for ht in hook_types if ht != HookType.REPLACE})
|
||||
),
|
||||
)
|
||||
|
||||
# Build the wrapper chain.
|
||||
# Sort: REPLACE hooks first (stable sort preserves registration order
|
||||
# within the same type). This ensures AROUND/BEFORE/AFTER always wrap
|
||||
# the replaced function, regardless of registration order.
|
||||
sorted_hooks = sorted(
|
||||
hooks, key=lambda h: (0 if h[0] == HookType.REPLACE else 1)
|
||||
)
|
||||
wrapped = original
|
||||
for hook_type, hook, _src in sorted_hooks:
|
||||
if isinstance(hook, type) and hook_type == HookType.REPLACE:
|
||||
# Class replacement: direct substitution to preserve type identity.
|
||||
# This keeps isinstance(), issubclass(), and inheritance working.
|
||||
wrapped = hook
|
||||
else:
|
||||
wrapped = _wrap_fn(wrapped, hook, hook_type)
|
||||
|
||||
# Restore classmethod/staticmethod decorator if the original had one.
|
||||
if is_classmethod:
|
||||
wrapped = classmethod(wrapped)
|
||||
logger.debug("Preserved @classmethod decorator for %s", target)
|
||||
elif is_staticmethod:
|
||||
wrapped = staticmethod(wrapped)
|
||||
logger.debug("Preserved @staticmethod decorator for %s", target)
|
||||
|
||||
setattr(obj, attr_name, wrapped)
|
||||
|
||||
# Propagate the patch to all other modules that imported the original
|
||||
# via ``from source_module import name``. Python's ``from X import Y``
|
||||
# copies the reference at import time; patching X alone leaves
|
||||
# importers with a stale binding.
|
||||
if wrapped is not original:
|
||||
extra = _propagate_patch(original, wrapped, obj)
|
||||
if extra:
|
||||
logger.debug(
|
||||
"Propagated patch for %s to %d additional module(s)",
|
||||
target,
|
||||
extra,
|
||||
)
|
||||
|
||||
sources = sorted({_format_source(src) for _, _, src in hooks})
|
||||
logger.info(
|
||||
"Applied %d hook(s) to %s (from: %s)",
|
||||
len(hooks),
|
||||
target,
|
||||
", ".join(sources),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def reset(cls):
|
||||
"""Reset all hooks and patches. Primarily for testing."""
|
||||
cls._hooks.clear()
|
||||
cls._patched.clear()
|
||||
|
||||
|
||||
def _propagate_patch(original: object, wrapped: object, source_module: object) -> int:
|
||||
"""Propagate a monkey-patch to all modules holding a stale ``from X import Y`` binding.
|
||||
|
||||
After ``setattr(source_module, name, wrapped)`` updates the defining module,
|
||||
other modules that did ``from source_module import name`` still hold a direct
|
||||
reference to the old *original* object. This walks ``sys.modules`` and
|
||||
replaces every such stale binding with *wrapped*.
|
||||
|
||||
Returns the number of additional module attributes that were patched.
|
||||
"""
|
||||
patched_count = 0
|
||||
for mod in list(sys.modules.values()):
|
||||
if mod is source_module or mod is None:
|
||||
continue
|
||||
if not isinstance(mod, types.ModuleType):
|
||||
continue
|
||||
try:
|
||||
mod_vars = vars(mod)
|
||||
except TypeError:
|
||||
continue
|
||||
for attr_name, attr_value in list(mod_vars.items()):
|
||||
if attr_value is original:
|
||||
try:
|
||||
setattr(mod, attr_name, wrapped)
|
||||
patched_count += 1
|
||||
except (AttributeError, TypeError):
|
||||
pass
|
||||
return patched_count
|
||||
|
||||
|
||||
def _wrap_fn(original_fn: Callable, hook: Callable, hook_type: HookType) -> Callable:
|
||||
"""Create a wrapper function based on the hook type."""
|
||||
if hook_type == HookType.REPLACE:
|
||||
|
||||
@functools.wraps(original_fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
return hook(*args, **kwargs)
|
||||
|
||||
wrapper.__wrapped__ = original_fn
|
||||
return wrapper
|
||||
|
||||
elif hook_type == HookType.BEFORE:
|
||||
|
||||
@functools.wraps(original_fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
result = hook(*args, **kwargs)
|
||||
if result is not None:
|
||||
args, kwargs = result
|
||||
return original_fn(*args, **kwargs)
|
||||
|
||||
wrapper.__wrapped__ = original_fn
|
||||
return wrapper
|
||||
|
||||
elif hook_type == HookType.AFTER:
|
||||
|
||||
@functools.wraps(original_fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
result = original_fn(*args, **kwargs)
|
||||
modified = hook(result, *args, **kwargs)
|
||||
return modified if modified is not None else result
|
||||
|
||||
wrapper.__wrapped__ = original_fn
|
||||
return wrapper
|
||||
|
||||
elif hook_type == HookType.AROUND:
|
||||
|
||||
@functools.wraps(original_fn)
|
||||
def wrapper(*args, **kwargs):
|
||||
return hook(original_fn, *args, **kwargs)
|
||||
|
||||
wrapper.__wrapped__ = original_fn
|
||||
return wrapper
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown hook type: {hook_type}")
|
||||
|
||||
|
||||
def plugin_hook(
|
||||
target: str,
|
||||
type: HookType = HookType.AFTER,
|
||||
) -> Callable:
|
||||
"""Decorator that registers a function or class as a hook on *target*.
|
||||
|
||||
Usage::
|
||||
|
||||
# Function hook (AROUND)
|
||||
@plugin_hook("sglang.srt.managers.scheduler.Scheduler.schedule",
|
||||
type=HookType.AROUND)
|
||||
def my_timer(original_fn, *args, **kwargs):
|
||||
start = time.perf_counter()
|
||||
result = original_fn(*args, **kwargs)
|
||||
print(f"Elapsed: {time.perf_counter() - start:.3f}s")
|
||||
return result
|
||||
|
||||
# Class replacement (REPLACE)
|
||||
@plugin_hook("sglang.srt.managers.scheduler.Scheduler",
|
||||
type=HookType.REPLACE)
|
||||
class MyScheduler(Scheduler):
|
||||
...
|
||||
|
||||
The decorated function/class is returned unchanged so it can still be
|
||||
used directly if needed.
|
||||
"""
|
||||
|
||||
def decorator(hook: Callable) -> Callable:
|
||||
HookRegistry.register(target, hook, type)
|
||||
return hook
|
||||
|
||||
return decorator
|
||||
@@ -790,6 +790,11 @@ class ServerArgs:
|
||||
self._handle_mps_backends()
|
||||
self._handle_xpu_backends()
|
||||
|
||||
# Allow OOT platform plugins to apply server args defaults.
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
current_platform.apply_server_args_defaults(self)
|
||||
|
||||
# Handle piecewise CUDA graph.
|
||||
self._handle_piecewise_cuda_graph()
|
||||
|
||||
@@ -1157,6 +1162,12 @@ class ServerArgs:
|
||||
# 5. Non-CUDA hardware (AMD, NPU, CPU, MPS, XPU, etc.)
|
||||
if is_hip() or is_npu() or is_cpu() or is_mps() or is_xpu():
|
||||
self.disable_piecewise_cuda_graph = True
|
||||
# 5b. OOT platforms that don't support piecewise cuda graph
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if current_platform.is_out_of_tree():
|
||||
if not current_platform.support_piecewise_cuda_graph():
|
||||
self.disable_piecewise_cuda_graph = True
|
||||
# 6. MoE A2A backend
|
||||
if self.moe_a2a_backend != "none":
|
||||
self.disable_piecewise_cuda_graph = True
|
||||
@@ -2326,6 +2337,12 @@ class ServerArgs:
|
||||
2.2 We will use Flashinfer backend on blackwell.
|
||||
2.3 Otherwise, we will use triton backend.
|
||||
"""
|
||||
# OOT platforms provide their own default attention backend.
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if current_platform.is_out_of_tree():
|
||||
return current_platform.get_default_attention_backend()
|
||||
|
||||
# Whisper requires flashinfer for cross-attention CUDA graph support
|
||||
if "WhisperForConditionalGeneration" in (
|
||||
model_config.hf_config.architectures or []
|
||||
|
||||
@@ -590,6 +590,18 @@ def get_available_gpu_memory(
|
||||
free_gpu_memory, total_gpu_memory = torch.musa.mem_get_info()
|
||||
elif device == "mps":
|
||||
free_gpu_memory = psutil.virtual_memory().available
|
||||
else:
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if not current_platform.is_out_of_tree():
|
||||
raise ValueError(
|
||||
f"Unsupported device type: {device!r}. "
|
||||
"If this is an OOT platform, ensure it is properly registered "
|
||||
"via the 'sglang.platform_plugins' entry point."
|
||||
)
|
||||
total_mem = current_platform.get_device_total_memory(gpu_id)
|
||||
used_mem = current_platform.get_current_memory_usage()
|
||||
free_gpu_memory = total_mem - used_mem
|
||||
|
||||
if distributed:
|
||||
tensor = torch.tensor(free_gpu_memory, dtype=torch.float32)
|
||||
@@ -1671,6 +1683,14 @@ def get_mtgpu_memory_capacity():
|
||||
|
||||
|
||||
def get_device_memory_capacity(device: str = None):
|
||||
# OOT platforms provide their own memory query via the platform class.
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if current_platform.is_out_of_tree():
|
||||
mem_bytes = current_platform.get_device_total_memory()
|
||||
if mem_bytes:
|
||||
return mem_bytes / (1 << 20) # bytes -> MiB
|
||||
return None
|
||||
if is_cuda():
|
||||
gpu_mem = get_nvgpu_memory_capacity()
|
||||
elif is_hip():
|
||||
@@ -1913,6 +1933,12 @@ def get_device_capability(device_id: int = 0) -> Tuple[int, int]:
|
||||
|
||||
|
||||
def get_compiler_backend(mode=None) -> str:
|
||||
# OOT platforms provide their own compile backend.
|
||||
from sglang.srt.platforms import current_platform
|
||||
|
||||
if current_platform.is_out_of_tree():
|
||||
return current_platform.get_compile_backend(mode)
|
||||
|
||||
if hasattr(torch, "hpu") and torch.hpu.is_available():
|
||||
return "hpu_backend"
|
||||
|
||||
|
||||
Reference in New Issue
Block a user