Multi platform Plugin

python/sglang/cli/serve.py

@@ -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:

python/sglang/launch_server.py

@@ -56,6 +56,10 @@ def run_server(server_args):
         stacklevel=1,
     )
 
+    from sglang.srt.plugins import load_plugins
+
+    load_plugins()
+
     server_args = prepare_server_args(sys.argv[1:])
 
     try:

python/sglang/srt/compilation/backend.py

@@ -22,6 +22,7 @@
 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,

python/sglang/srt/entrypoints/engine.py

@@ -84,6 +84,7 @@
 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 @@ def __init__(self, **kwargs):
         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 @@ def _launch_subprocesses(
         # 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)
 

python/sglang/srt/environ.py

@@ -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

python/sglang/srt/layers/utils/multi_platform.py

@@ -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 @@
 
 
 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 @@ def forward_cpu(self, *args, **kwargs):
         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:

python/sglang/srt/managers/scheduler.py

@@ -204,6 +204,7 @@
 )
 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
@@ -3742,6 +3743,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,

python/sglang/srt/mem_cache/memory_pool.py

@@ -54,6 +54,7 @@
     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 @@ def __init__(
         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 @@ def __init__(
 
             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 @@ def __init__(
 
             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,
                 )

python/sglang/srt/model_executor/model_runner.py

@@ -150,6 +150,7 @@
 )
 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 @@
     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 @@ def initialize(self, pre_model_load_memory: float):
         # 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 @@ def initialize(self, pre_model_load_memory: float):
         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 @@ def model_load_weights(model, iter):
         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,23 +2550,29 @@ def init_device_graphs(self):
         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",
             },
         )
         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

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -282,7 +282,63 @@ def _init_pools(self: ModelRunner):
 
         # 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 @@ def _init_pools(self: ModelRunner):
         # 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
             ):