SGLang Plugin System

def activate() -> str | None:
    """Return the platform class qualname when this hardware is available."""

entry_points("sglang.srt.platforms")  → Enumerate ALL plugins by name (metadata only)
  │
  ├─ SGLANG_PLATFORM set (front-loading filter):
  │   ├─ Name not found in discovered → RuntimeError
  │   ├─ activate() returns non-None  → load that platform
  │   └─ activate() returns None      → RuntimeError (hardware unavailable)
  │
  └─ SGLANG_PLATFORM unset (auto-discover, activate all):
      ├─ 0 activated + SGLANG_USE_CPU_ENGINE=1 → fallback CpuSRTPlatform
      ├─ 0 activated + CUDA available          → fallback CudaSRTPlatform
      ├─ 0 activated + ROCm available          → fallback RocmSRTPlatform
      ├─ 0 activated + XPU available           → fallback XpuSRTPlatform
      ├─ 0 activated + none of the above       → fallback base SRTPlatform
      ├─ 1 activated → use it
      └─ N activated → RuntimeError (must set SGLANG_PLATFORM)

load_plugins()
  ├── _get_excluded_dists()                       → compute dists to skip (via SGLANG_PLATFORM)
  ├── load_plugins_by_group("sglang.srt.plugins",     → discover entry_points, filter by SGLANG_PLUGINS
  │     excluded_dists=...)                          skip plugins from unselected platform packages
  ├── for each plugin:                            → set _current_plugin_source context var
  │     func()                                      side effects (register hooks with source tracking)
  └── HookRegistry.apply_hooks()                  → monkey-patch targets

spawn unpickles SchedulerProcessSpec (stdlib types + opaque ServerArgs bytes)
  -> initialize_current_platform()
  -> load_plugins()                 # discover callbacks and register hooks only
  -> apply_plugin_hooks()           # target imports are allowed from here onward
  -> materialize ServerArgs
  -> resolve and invoke the patched run_scheduler_process

from sglang.multimodal_gen import DiffGenerator  # a proxy: imports nothing yet

if __name__ == "__main__":
    generator = DiffGenerator.from_pretrained(model_path="...")

my_platform_plugin/
├── pyproject.toml
└── my_platform_plugin/
    ├── __init__.py    # activate() function
    ├── device.py      # MyDeviceMixin
    └── platform.py    # MySRTPlatform

[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"

[project]
name = "my-platform-plugin"
version = "0.1.0"

[project.entry-points."sglang.srt.platforms"]
my_device = "my_platform_plugin:activate"

def activate():
    """Return fully-qualified class name to activate, or None to skip."""
    if _my_device_is_available():
        return "my_platform_plugin.platform.MySRTPlatform"
    return None

from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum

class MyDeviceMixin(DeviceMixin):
    _enum = PlatformEnum.OOT
    device_name = "my_device"
    device_type = "my_device"   # torch device type

    def set_device(self, device) -> None: ...
    def get_device_name(self, device_id=0) -> str: ...
    def get_device_total_memory(self, device_id=0) -> int: ...
    def get_current_memory_usage(self, device=None) -> float: ...
    def get_device_capability(self, device_id=0): ...
    def get_torch_distributed_backend_str(self) -> str: ...

from sglang.srt.platforms.interface import SRTPlatform
from my_platform_plugin.device import MyDeviceMixin

class MySRTPlatform(SRTPlatform, MyDeviceMixin):
    def get_default_attention_backend(self) -> str: ...
    def support_cuda_graph(self) -> bool: ...
    # ... override other methods as needed

pip install -e my_platform_plugin/
python -c "from sglang.srt.platforms import current_platform; print(current_platform)"

[project.entry-points."sglang.multimodal_gen.platforms"]
my_device = "my_platform_plugin:activate_diffusion"

[project.entry-points."sglang.multimodal_gen.plugins"]
my_device = "my_platform_plugin:register_hooks"

def activate_diffusion() -> str | None:
    try:
        import my_device_runtime
    except ImportError:
        return None
    if not my_device_runtime.is_available():
        return None
    return "my_platform_plugin.diffusion_platform.MyDiffusionPlatform"

from sglang.multimodal_gen.runtime.platforms import Platform, PlatformEnum

class MyDiffusionPlatform(Platform):
    _enum = PlatformEnum.OOT
    device_name = "my_device"
    device_type = "my_device"
    dispatch_key = "PrivateUse1"

    def get_dispatch_key_name(self) -> str:
        return "my_device"

class MyDiffusionPlatform(Platform):
    ...

    def init_backend(self) -> None:
        from my_platform_plugin.ops import rms_norm_forward
        from sglang.multimodal_gen.runtime.layers.custom_op import CustomOp
        from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm

        CustomOp.register_oot_forward(
            RMSNorm, fn=rms_norm_forward, platform_key="my_device"
        )

pip install -e my_platform_plugin/
SGLANG_DIFFUSION_PLATFORM_OVERRIDE=my_device python -c \
  "from sglang.multimodal_gen.runtime.platforms import current_platform; print(current_platform)"

# SRT plugin
from sglang.srt.plugins.hook_registry import HookType, plugin_hook

# Diffusion plugin
from sglang.multimodal_gen.runtime.platforms.plugins import HookType, plugin_hook

my_general_plugin/
├── pyproject.toml
└── my_general_plugin/
    └── __init__.py    # register() function

[build-system]
requires = ["setuptools"]
build-backend = "setuptools.build_meta"

[project]
name = "my-general-plugin"
version = "0.1.0"

[project.entry-points."sglang.srt.plugins"]
my_plugin = "my_general_plugin:register"

from sglang.srt.plugins.hook_registry import HookRegistry, HookType

def register():
    """Entry point called by load_plugins()."""
    HookRegistry.register(
        "sglang.srt.managers.scheduler.Scheduler.__init__",
        my_hook,
        HookType.AROUND,
    )

def my_hook(original_fn, self, *args, **kwargs):
    result = original_fn(self, *args, **kwargs)
    print(f"Scheduler initialized! gpu_id={self.gpu_id}")
    return result

pip install -e my_general_plugin/
sglang serve --model-path <model> [options]
# Look for "Scheduler initialized!" in logs

# --- Imperative API ---
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.get_next_batch_to_run",
    my_timer,
    HookType.AROUND,
)

# --- Decorator API ---
from sglang.srt.plugins.hook_registry import plugin_hook, HookType

@plugin_hook(
    "sglang.srt.managers.scheduler.Scheduler.get_next_batch_to_run",
    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) ---
from sglang.srt.plugins.hook_registry import plugin_hook, HookType
from sglang.srt.managers.scheduler import Scheduler

@plugin_hook(
    "sglang.srt.managers.scheduler.Scheduler",
    type=HookType.REPLACE,
)
class MyScheduler(Scheduler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        print("Enhanced scheduler initialized!")