From a3a77ad6d7160fd74104b690708a899fdb9c551b Mon Sep 17 00:00:00 2001 From: metamergebot Date: Thu, 24 Sep 2026 00:36:20 -0700 Subject: [PATCH] Lazy-load built-in model definitions and nixl_ep at startup When SGLANG_EXTERNAL_MODEL_PACKAGE is set, register only the external package at import time and import sglang.srt.models on the first lookup that needs it, keeping external models as overrides. Resolve nixl_ep inside a cached helper at first use instead of at module import. Co-authored-by: Yongji Wu <30348494+libertyeagle@users.noreply.github.com> --- .../srt/layers/moe/token_dispatcher/nixl.py | 41 +++++---- python/sglang/srt/model_loader/utils.py | 5 +- python/sglang/srt/models/registry.py | 31 ++++++- .../unit/models/test_model_registry.py | 91 +++++++++++++++++++ 4 files changed, 144 insertions(+), 24 deletions(-) create mode 100644 test/registered/unit/models/test_model_registry.py diff --git a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py b/python/sglang/srt/layers/moe/token_dispatcher/nixl.py index 6f71d4a38736..27e29b27262e 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/nixl.py +++ b/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 diff --git a/python/sglang/srt/model_loader/utils.py b/python/sglang/srt/model_loader/utils.py index c82e5d60fb31..55496273bda3 100644 --- a/python/sglang/srt/model_loader/utils.py +++ b/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"] diff --git a/python/sglang/srt/models/registry.py b/python/sglang/srt/models/registry.py index b83e5bb0792a..506be2362a32 100644 --- a/python/sglang/srt/models/registry.py +++ b/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") diff --git a/test/registered/unit/models/test_model_registry.py b/test/registered/unit/models/test_model_registry.py new file mode 100644 index 000000000000..75d55ef13e54 --- /dev/null +++ b/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()