[Perf] Lazy-load built-in model definitions and nixl_ep at startup

python/sglang/srt/layers/moe/token_dispatcher/nixl.py

@@ -2,6 +2,7 @@
 
 import logging
 from enum import Enum, auto
+from functools import cache
 
 import torch
 import torch.distributed as dist
@@ -28,24 +29,29 @@
     get_resources,
 )
 
-try:
-    from nixl_ep import Buffer
+logger = logging.getLogger(__name__)
 
-    use_nixl = True
-except ImportError:
-    use_nixl = False
+NixlEPDispatchOutput = DeepEPLLDispatchOutput
+NixlEPCombineInput = DeepEPLLCombineInput
 
-try:
-    from nixl_ep import topk_idx_t as NIXL_EP_TOPK_INDICES_DTYPE
-except ImportError:
-    NIXL_EP_TOPK_INDICES_DTYPE = torch.int64
 
-assert isinstance(NIXL_EP_TOPK_INDICES_DTYPE, torch.dtype)
+@cache
+def _load_nixl_ep() -> tuple[type, torch.dtype]:
+    try:
+        from nixl_ep import Buffer
+    except ImportError as exc:
+        raise ImportError(
+            "NixlEP is not installed. Please install NixlEP package from "
+            "https://github.com/ai-dynamo/nixl."
+        ) from exc
 
-logger = logging.getLogger(__name__)
+    try:
+        from nixl_ep import topk_idx_t
+    except ImportError:
+        topk_idx_t = torch.int64
 
-NixlEPDispatchOutput = DeepEPLLDispatchOutput
-NixlEPCombineInput = DeepEPLLCombineInput
+    assert isinstance(topk_idx_t, torch.dtype)
+    return Buffer, topk_idx_t
 
 
 class NixlEPBuffer:
@@ -125,6 +131,7 @@ def get_nixl_buffer(
                 cls._update_connections(state, state.scale_to)
             return state.buffer
 
+        Buffer, _ = _load_nixl_ep()
         state.hidden_size = hidden_size
         state.num_max_dispatch_tokens_per_rank = num_max_dispatch_tokens_per_rank
         state.num_experts = num_experts
@@ -207,11 +214,7 @@ def __init__(
         params_dtype: torch.dtype,
         deepep_mode: DeepEPMode,
     ):
-        if not use_nixl:
-            raise ImportError(
-                "NixlEP is not installed. Please install NixlEP package from "
-                "https://github.com/ai-dynamo/nixl."
-            )
+        _, self.topk_indices_dtype = _load_nixl_ep()
 
         self.group = group
         self.router_topk = router_topk
@@ -296,7 +299,7 @@ def dispatch_a(
     ):
         buffer = self._get_buffer()
         topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
-        topk_ids = topk_ids.to(NIXL_EP_TOPK_INDICES_DTYPE)
+        topk_ids = topk_ids.to(self.topk_indices_dtype)
         state = NixlEPBuffer._state()
         dispatch_ep_size = state.dispatch_ep_size
         num_local_experts = state.num_local_experts

python/sglang/srt/model_loader/utils.py

@@ -221,8 +221,9 @@ def get_model_architecture(model_config: ModelConfig) -> Tuple[Type[nn.Module],
     ):
         architectures = ["QuantMixtralForCausalLM"]
 
-    supported_archs = ModelRegistry.get_supported_archs()
-    is_native_supported = any(arch in supported_archs for arch in architectures)
+    is_native_supported = any(
+        ModelRegistry.is_arch_supported(arch) for arch in architectures
+    )
 
     if model_config.model_impl == ModelImpl.MINDSPORE:
         architectures = ["MindSporeForCausalLM"]

python/sglang/srt/models/registry.py

@@ -20,10 +20,19 @@
 class _ModelRegistry:
     # Keyed by model_arch
     models: Dict[str, Union[Type[nn.Module], str]] = field(default_factory=dict)
+    _default_models_loaded: bool = True
+
+    def _load_default_models(self):
+        if self._default_models_loaded:
+            return
+
+        self.models = import_model_classes("sglang.srt.models") | self.models
+        self._default_models_loaded = True
 
     def register(
         self, package_name: str, overwrite: bool = False, strict: bool = False
     ):
+        self._load_default_models()
         new_models = import_model_classes(package_name, strict=strict)
         if overwrite:
             self.models.update(new_models)
@@ -36,8 +45,12 @@ def register(
                 self.models[arch] = cls
 
     def get_supported_archs(self) -> AbstractSet[str]:
+        self._load_default_models()
         return self.models.keys()
 
+    def is_arch_supported(self, architecture: str) -> bool:
+        return architecture in self.models or architecture in self.get_supported_archs()
+
     def _raise_for_unsupported(self, architectures: List[str]):
         all_supported_archs = self.get_supported_archs()
 
@@ -81,6 +94,16 @@ def resolve_model_cls(
         self,
         architectures: Union[str, List[str]],
     ) -> Tuple[Type[nn.Module], str]:
+        # Fast startup path: resolve registered models without importing defaults.
+        if architectures:
+            first_arch = (
+                architectures if isinstance(architectures, str) else architectures[0]
+            )
+            model_cls = self.models.get(first_arch)
+            if model_cls is not None:
+                return model_cls, first_arch
+
+        self._load_default_models()
         architectures = self._normalize_archs(architectures)
 
         for arch in architectures:
@@ -128,7 +151,9 @@ def import_model_classes(package_name: str, strict: bool = False):
 
 
 ModelRegistry = _ModelRegistry()
-ModelRegistry.register("sglang.srt.models")
-
 if external_pkg := envs.SGLANG_EXTERNAL_MODEL_PACKAGE.get():
-    ModelRegistry.register(external_pkg, overwrite=True)
+    # Defer default model imports to speed up startup when using external models.
+    ModelRegistry._default_models_loaded = False
+    ModelRegistry.models.update(import_model_classes(external_pkg))
+else:
+    ModelRegistry.register("sglang.srt.models")

test/registered/unit/models/test_model_registry.py

@@ -0,0 +1,91 @@
+"""Unit tests for lazy default-model loading in ``_ModelRegistry``.
+
+With ``SGLANG_EXTERNAL_MODEL_PACKAGE`` set, the registry holds only the external
+package's models at import time and imports the built-in ``sglang.srt.models``
+package on the first lookup that needs it; an out-of-tree architecture must
+resolve without paying for that import.
+"""
+
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=20, suite="base-a-test-cpu")
+
+import unittest
+from unittest.mock import patch
+
+import torch.nn as nn
+
+from sglang.srt.models import registry as registry_module
+from sglang.srt.models.registry import _ModelRegistry
+from sglang.test.test_utils import CustomTestCase
+
+
+class _DefaultModel(nn.Module):
+    pass
+
+
+class _ExternalModel(nn.Module):
+    pass
+
+
+class _ExternalOverride(nn.Module):
+    pass
+
+
+DEFAULT_PACKAGE = "sglang.srt.models"
+FAKE_PACKAGES = {
+    DEFAULT_PACKAGE: {"DefaultArch": _DefaultModel, "SharedArch": _DefaultModel},
+    "another.models": {"DefaultArch": _ExternalModel},
+}
+EXTERNAL_MODELS = {"ExternalArch": _ExternalModel, "SharedArch": _ExternalOverride}
+
+
+class TestLazyDefaultModels(CustomTestCase):
+    def setUp(self):
+        self.import_calls = []
+        patcher = patch.object(
+            registry_module, "import_model_classes", side_effect=self._fake_import
+        )
+        self.addCleanup(patcher.stop)
+        patcher.start()
+        # The state ``registry.py`` builds when SGLANG_EXTERNAL_MODEL_PACKAGE is set.
+        self.registry = _ModelRegistry(
+            models=dict(EXTERNAL_MODELS), _default_models_loaded=False
+        )
+
+    def _fake_import(self, package_name, strict=False):
+        self.import_calls.append(package_name)
+        return dict(FAKE_PACKAGES[package_name])
+
+    def test_external_arch_resolves_without_importing_defaults(self):
+        self.assertTrue(self.registry.is_arch_supported("ExternalArch"))
+        model_cls, arch = self.registry.resolve_model_cls(["ExternalArch"])
+        self.assertIs(model_cls, _ExternalModel)
+        self.assertEqual(arch, "ExternalArch")
+        self.assertEqual(self.import_calls, [])
+
+    def test_default_arch_imports_defaults_once(self):
+        model_cls, arch = self.registry.resolve_model_cls(["DefaultArch"])
+        self.assertIs(model_cls, _DefaultModel)
+        self.assertEqual(arch, "DefaultArch")
+        self.assertFalse(self.registry.is_arch_supported("MissingArch"))
+        self.registry.resolve_model_cls("DefaultArch")
+        self.assertEqual(self.import_calls, [DEFAULT_PACKAGE])
+
+    def test_external_models_override_defaults(self):
+        self.assertEqual(
+            set(self.registry.get_supported_archs()),
+            {"ExternalArch", "SharedArch", "DefaultArch"},
+        )
+        self.assertIs(
+            self.registry.resolve_model_cls("SharedArch")[0], _ExternalOverride
+        )
+
+    def test_register_checks_duplicates_against_defaults(self):
+        with self.assertRaises(ValueError):
+            self.registry.register("another.models")
+        self.assertEqual(self.import_calls, [DEFAULT_PACKAGE, "another.models"])
+
+
+if __name__ == "__main__":
+    unittest.main()