From ab03a8e7eb82d35907bfbeb645bf0af8b4ce290e Mon Sep 17 00:00:00 2001 From: Xueshen Liu Date: Sun, 20 Sep 2026 22:48:35 -0400 Subject: [PATCH] [Perf] Fork-safe import: no CUDA context at import time, lighter argument parsing (#40201) --- python/sglang/cli/utils.py | 14 +++- .../sglang/srt/distributed/parallel_state.py | 4 +- .../sglang/srt/function_call/parser_names.py | 51 +++++++++++++ .../srt/parser/reasoning_parser_names.py | 39 ++++++++++ python/sglang/srt/server_args.py | 57 +++++++++++++-- python/sglang/srt/utils/common.py | 64 ++++++++++++++++- .../srt/utils/hf_transformers_patches.py | 5 +- .../test_function_call_parser.py | 11 +++ .../unit/parser/test_reasoning_parser.py | 10 +++ .../unit/server_args/test_server_args.py | 67 +++++++++++++++++ test/registered/unit/utils/test_common.py | 71 +++++++++++++++++++ 11 files changed, 381 insertions(+), 12 deletions(-) create mode 100644 python/sglang/srt/function_call/parser_names.py create mode 100644 python/sglang/srt/parser/reasoning_parser_names.py diff --git a/python/sglang/cli/utils.py b/python/sglang/cli/utils.py index f27ad6cf0..5df252a3e 100644 --- a/python/sglang/cli/utils.py +++ b/python/sglang/cli/utils.py @@ -4,8 +4,6 @@ import os import subprocess from functools import lru_cache -from huggingface_hub import HfApi - from sglang.srt.environ import envs from sglang.utils import ( has_diffusion_overlay_registry_match, @@ -24,7 +22,17 @@ def _is_overlay_diffusion_model(model_path: str) -> bool: return has_diffusion_overlay_registry_match(model_path, _load_overlay_registry()) +def _diffusion_deps_available() -> bool: + # Locating diffusers is cheap; importing the registry costs ~2 s and then + # fails anyway without it. A false positive is caught by the caller. + import importlib.util + + return importlib.util.find_spec("diffusers") is not None + + def _is_diffusion_model_from_registry(model_path: str) -> bool: + if not _diffusion_deps_available(): + return False try: from sglang.multimodal_gen.registry import is_registered_diffusion_model_path except ImportError: @@ -49,6 +57,8 @@ def _is_diffusers_model_dir(model_dir: str) -> bool: def _is_gated_diffusion_repo(repo_id: str) -> bool: """Query HF model card metadata to check if a gated repo is a diffusers model.""" try: + from huggingface_hub import HfApi # lazy: ~0.3 s at CLI entry otherwise + info = HfApi().model_info(repo_id) return getattr(info, "library_name", None) == "diffusers" except Exception: diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 1d03cdc64..3dfb63849 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -107,7 +107,9 @@ def get_torch_distributed_pg_options(group_name=None): @dataclass class GraphCaptureContext: - stream: torch.get_device_module().Stream + # Evaluating torch.get_device_module() at import marks the process unsafe + # to fork, and a child then fails in cuInit; torch.Stream is its base. + stream: torch.Stream @dataclass diff --git a/python/sglang/srt/function_call/parser_names.py b/python/sglang/srt/function_call/parser_names.py new file mode 100644 index 000000000..54ca98352 --- /dev/null +++ b/python/sglang/srt/function_call/parser_names.py @@ -0,0 +1,51 @@ +"""Tool-call parser names, kept in a dependency-free module. + +`server_args` needs these for the `--tool-call-parser` CLI choices, and +importing `function_call_parser` for them costs seconds (OpenAI protocol +models -> xgrammar -> transformers -> torch.distributed). A unit test checks +this list against `FunctionCallParser.ToolCallParserEnum`; add a name here +when you register a parser. +""" + +TOOL_CALL_PARSER_NAMES = [ + "apertus2509", + "cohere_command4", + "deepseekv3", + "deepseekv31", + "deepseekv32", + "deepseekv4", + "deepseekv41", + "dots", + "glm", + "glm45", + "glm47", + "gpt-oss", + "k2_horizon", + "kimi_k2", + "kimi_k3", + "lfm2", + "ling3", + "llama3", + "mimo", + "minicpm5", + "mistral", + "muse", + "poolside_v1", + "pythonic", + "qwen", + "qwen25", + "qwen3_coder", + "spark25", + "step3", + "step3p5", + "minimax-m2", + "minimax-m3", + "nanbeige", + "trinity", + "interns1", + "hermes", + "hunyuan", + "gigachat3", + "gemma4", + "inkling", +] diff --git a/python/sglang/srt/parser/reasoning_parser_names.py b/python/sglang/srt/parser/reasoning_parser_names.py new file mode 100644 index 000000000..bf0005cdd --- /dev/null +++ b/python/sglang/srt/parser/reasoning_parser_names.py @@ -0,0 +1,39 @@ +"""Reasoning parser names, kept in a dependency-free module (see +`function_call/parser_names.py`). A unit test checks that this list matches +`ReasoningParser.DetectorMap`; add a name here when you register a parser. +""" + +REASONING_PARSER_NAMES = [ + "apertus2509", + "deepseek-r1", + "deepseek-v3", + "deepseek-v4", + "deepseek-v41", + "dots", + "glm45", + "ling3", + "hunyuan", + "gpt-oss", + "k2_horizon", + "kimi", + "kimi_k2", + "kimi_k3", + "mimo", + "muse", + "poolside_v1", + "qwen3", + "qwen3-thinking", + "minimax", + "minimax-append-think", + "minimax-m3", + "nanbeige", + "step3", + "step3p5", + "mistral", + "nemotron_3", + "granite_thinking_parser", + "interns1", + "gemma4", + "inkling", + "cohere_command4", +] diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 1fafc9035..10d9ff06d 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -37,14 +37,15 @@ import argparse import copy import dataclasses import functools +import importlib import logging +import sys import tempfile import uuid from typing import Any, NoReturn import msgspec -from sglang.kernels.ops.kv_canary.consts import RealKvHashMode from sglang.srt.arg_groups.arg_utils import ( add_cli_args_from_dataclass, is_record, @@ -60,14 +61,41 @@ from sglang.srt.arg_groups.overrides import ( resolving_view, ) from sglang.srt.environ import envs -from sglang.srt.function_call.function_call_parser import FunctionCallParser -from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.runtime_context import get_platform, publish from sglang.srt.speculative.decoupled_spec_io import DecoupledSpecIpcConfig from sglang.srt.utils.network import NetworkAddress, get_free_port, wait_port_available logger = logging.getLogger(__name__) + +def _reasoning_parser_choices(): + # Importing the registry here costs seconds in every process that parses + # arguments; a plugin that registered a parser has already imported it. + module = sys.modules.get("sglang.srt.parser.reasoning_parser") + if module is not None: + return list(module.ReasoningParser.DetectorMap) + from sglang.srt.parser.reasoning_parser_names import REASONING_PARSER_NAMES + + return list(REASONING_PARSER_NAMES) + + +def _tool_call_parser_choices(): + module = sys.modules.get("sglang.srt.function_call.function_call_parser") + if module is not None: + return list(module.FunctionCallParser.ToolCallParserEnum) + from sglang.srt.function_call.parser_names import TOOL_CALL_PARSER_NAMES + + return list(TOOL_CALL_PARSER_NAMES) + + +def _real_kv_hash_modes(): + # Lazy: this pulls the whole sglang.kernels package (~2 s) into every + # process that imports server_args, most of which never use it. + from sglang.kernels.ops.kv_canary.consts import RealKvHashMode + + return list(RealKvHashMode) + + # Re-exported. These were importable from this module while the field # declarations that used them lived here; the declarations moved to # `arg_groups/fields/` but out-of-tree code -- and `tokenizer_control_mixin` @@ -172,6 +200,23 @@ from sglang.srt.utils.common import ( # noqa: F401 nullable_str, ) +# Re-exported like the imports above, but resolved on first use: importing them +# eagerly is what the choices helpers avoid, and most processes never read them. +_LAZY_REEXPORTS = { + "FunctionCallParser": "sglang.srt.function_call.function_call_parser", + "ReasoningParser": "sglang.srt.parser.reasoning_parser", + "RealKvHashMode": "sglang.kernels.ops.kv_canary.consts", +} + + +def __getattr__(name: str) -> Any: + module_name = _LAZY_REEXPORTS.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(importlib.import_module(module_name), name) + globals()[name] = value + return value + def _plain(value: Any) -> Any: """``asdict``'s conversion, applied to one value: a record -- Struct or @@ -362,7 +407,7 @@ class ServerArgs: help="Choose the kernels for sampling layers.", ) - reasoning_parser_choices = list(ReasoningParser.DetectorMap.keys()) + reasoning_parser_choices = _reasoning_parser_choices() parser.add_argument( "--reasoning-parser", type=str, @@ -372,7 +417,7 @@ class ServerArgs: f"Use 'auto' to detect from chat template. " f"Options include: {reasoning_parser_choices}.", ) - tool_call_parser_choices = list(FunctionCallParser.ToolCallParserEnum.keys()) + tool_call_parser_choices = _tool_call_parser_choices() parser.add_argument( "--tool-call-parser", type=str, @@ -386,7 +431,7 @@ class ServerArgs: "--kv-canary-real-data", type=str, default=_declared_default("kv_canary_real_data"), - choices=[m.name.lower() for m in RealKvHashMode], + choices=[m.name.lower() for m in _real_kv_hash_modes()], help=( "Check the real KV-cache in the canary. " "'none' (default) disables the feature. " diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py index 1db8f8e6d..b706544ee 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py @@ -95,7 +95,6 @@ from starlette.routing import Mount from torch import nn from torch.library import Library from torch.utils._contextlib import _DecoratorContextManager -from torchvision.io import decode_jpeg from typing_extensions import Literal from sglang.srt.environ import envs @@ -266,8 +265,10 @@ def _check_cuda_device_version( ): if not is_cuda(): return False + # get_device_sm() answers from NVML while torch.cuda is uninitialized, so + # the platform probes evaluated at import time do not create a CUDA context. return ( - torch.cuda.get_device_capability()[0] in device_capability_majors + get_device_sm() // 10 in device_capability_majors and tuple(map(int, torch.version.cuda.split(".")[:2])) >= cuda_version ) @@ -582,6 +583,16 @@ def get_dispatch_device_backend(): @lru_cache(maxsize=1) def get_device_module(): + # Resolve from the platform checks: torch.get_device_module() with no + # argument initializes the CUDA runtime, which poisons fork() startup. + if is_cuda() or is_hip(): + return torch.cuda + if is_npu(): + return torch.npu + if is_xpu(): + return torch.xpu + if is_musa(): + return torch.musa return torch.get_device_module() @@ -630,8 +641,55 @@ def get_amdgpu_memory_capacity(): ) +def _get_device_sm_via_nvml() -> Optional[int]: + # Compute capability of torch device 0, read while torch.cuda stays + # uninitialized; None when NVML cannot answer and the caller falls back. + try: + import pynvml + except ImportError: + logger.debug("get_device_sm: pynvml is not installed, using torch.cuda") + return None + # Private torch API, read defensively: it maps the torch ordinal to the NVML + # index under CUDA_VISIBLE_DEVICES / MIG; absent or failing -> fall back. + getter = getattr(torch.cuda, "_get_nvml_device_index", None) + if getter is None: + logger.debug( + "get_device_sm: torch.cuda._get_nvml_device_index is missing, " + "using torch.cuda" + ) + return None + try: + idx = getter(0) + except Exception: + logger.debug( + "get_device_sm: torch.cuda._get_nvml_device_index(0) failed, " + "using torch.cuda", + exc_info=True, + ) + return None + try: + pynvml.nvmlInit() + try: + handle = pynvml.nvmlDeviceGetHandleByIndex(idx) + major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle) + finally: + pynvml.nvmlShutdown() + return major * 10 + minor + except Exception: + logger.debug( + "get_device_sm: NVML query failed, using torch.cuda", exc_info=True + ) + return None + + def get_device_sm(): if torch.cuda.is_available() or is_musa(): + # Called at import time (e.g. by the DeepGEMM configurer): initializing + # torch.cuda here would create a context and poison fork() startup. + if not is_musa() and not torch.cuda.is_initialized(): + sm = _get_device_sm_via_nvml() + if sm is not None: + return sm major, minor = torch.cuda.get_device_capability() return major * 10 + minor return 0 @@ -1859,6 +1917,8 @@ def _load_image( ) return decode_jpeg_with_fancy_upsampling(image_bytes) + from torchvision.io import decode_jpeg # lazy: ~1 s of torch._dynamo + encoded_image = torch.frombuffer(image_bytes, dtype=torch.uint8) image_tensor = decode_jpeg(encoded_image, device="cuda") return image_tensor diff --git a/python/sglang/srt/utils/hf_transformers_patches.py b/python/sglang/srt/utils/hf_transformers_patches.py index d06403e2f..df6809cd6 100644 --- a/python/sglang/srt/utils/hf_transformers_patches.py +++ b/python/sglang/srt/utils/hf_transformers_patches.py @@ -23,8 +23,11 @@ all patches. It is safe to import multiple times -- patches are idempotent. """ import inspect +import logging -from sglang.srt.utils import logger +# Plain logger: importing sglang.srt.utils here pulls torch/transformers/triton +# into every `import sglang` (this module runs from sglang/__init__.py). +logger = logging.getLogger(__name__) _applied = False diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py index 5b1101c9e..359ae726d 100644 --- a/test/registered/unit/function_call/test_function_call_parser.py +++ b/test/registered/unit/function_call/test_function_call_parser.py @@ -34,6 +34,7 @@ from sglang.srt.function_call.lfm2_detector import Lfm2Detector from sglang.srt.function_call.ling3_detector import Ling3Detector from sglang.srt.function_call.llama32_detector import Llama32Detector from sglang.srt.function_call.mistral_detector import MistralDetector +from sglang.srt.function_call.parser_names import TOOL_CALL_PARSER_NAMES from sglang.srt.function_call.pythonic_detector import PythonicDetector from sglang.srt.function_call.qwen3_coder_detector import Qwen3CoderDetector from sglang.srt.function_call.utils import get_schema_properties @@ -6013,5 +6014,15 @@ class TestTopLevelCompositeToolSchema(unittest.TestCase): self.assertEqual(json.loads(arguments), self.expected) +class TestToolCallParserNames(unittest.TestCase): + def test_matches_registry(self): + # `server_args` builds the --tool-call-parser choices from this list to + # keep the registry, and its dependencies, out of argument parsing. + self.assertEqual( + sorted(TOOL_CALL_PARSER_NAMES), + sorted(FunctionCallParser.ToolCallParserEnum), + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/parser/test_reasoning_parser.py b/test/registered/unit/parser/test_reasoning_parser.py index 5ef4a98dd..acf2152d7 100644 --- a/test/registered/unit/parser/test_reasoning_parser.py +++ b/test/registered/unit/parser/test_reasoning_parser.py @@ -20,6 +20,7 @@ from sglang.srt.parser.reasoning_parser import ( Qwen3Detector, ReasoningParser, ) +from sglang.srt.parser.reasoning_parser_names import REASONING_PARSER_NAMES from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -1726,5 +1727,14 @@ class TestGraniteThinkingDetector(CustomTestCase): self.assertEqual(normal, "truncated") +class TestReasoningParserNames(CustomTestCase): + def test_matches_registry(self): + # `server_args` builds the --reasoning-parser choices from this list to + # keep the registry, and its dependencies, out of argument parsing. + self.assertEqual( + sorted(REASONING_PARSER_NAMES), sorted(ReasoningParser.DetectorMap) + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index 08a67e2e4..fed3493e5 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -4002,5 +4002,72 @@ class TestDcpCommBackendDefault(CustomTestCase): ) +class TestParserChoices(CustomTestCase): + """The choices come from dependency-free name lists, but `cli/serve.py` + loads plugins before parsing, so a plugin's parser must still be accepted.""" + + def test_a_plugin_registered_parser_is_accepted(self): + from sglang.srt.function_call.function_call_parser import FunctionCallParser + from sglang.srt.parser.reasoning_parser import ReasoningParser + + ReasoningParser.DetectorMap["plugin-reasoning"] = object + FunctionCallParser.ToolCallParserEnum["plugin-toolcall"] = object + try: + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args( + [ + "--model-path", + "dummy-model", + "--reasoning-parser", + "plugin-reasoning", + "--tool-call-parser", + "plugin-toolcall", + ] + ) + finally: + del ReasoningParser.DetectorMap["plugin-reasoning"] + del FunctionCallParser.ToolCallParserEnum["plugin-toolcall"] + self.assertEqual(args.reasoning_parser, "plugin-reasoning") + self.assertEqual(args.tool_call_parser, "plugin-toolcall") + + def test_name_lists_are_used_when_the_registries_are_not_imported(self): + from sglang.srt.function_call.parser_names import TOOL_CALL_PARSER_NAMES + from sglang.srt.parser.reasoning_parser_names import REASONING_PARSER_NAMES + + with patch.dict(server_args_module.sys.modules): + server_args_module.sys.modules.pop( + "sglang.srt.parser.reasoning_parser", None + ) + server_args_module.sys.modules.pop( + "sglang.srt.function_call.function_call_parser", None + ) + self.assertEqual( + server_args_module._reasoning_parser_choices(), + list(REASONING_PARSER_NAMES), + ) + self.assertEqual( + server_args_module._tool_call_parser_choices(), + list(TOOL_CALL_PARSER_NAMES), + ) + + +class TestLazyReexports(CustomTestCase): + def test_the_names_that_lost_their_eager_import_are_still_attributes(self): + # Out-of-tree code reaches these through `sglang.srt.server_args`; they + # now resolve through the module __getattr__ instead of a top import. + from sglang.kernels.ops.kv_canary.consts import RealKvHashMode + from sglang.srt.function_call.function_call_parser import FunctionCallParser + from sglang.srt.parser.reasoning_parser import ReasoningParser + + self.assertIs(server_args_module.FunctionCallParser, FunctionCallParser) + self.assertIs(server_args_module.ReasoningParser, ReasoningParser) + self.assertIs(server_args_module.RealKvHashMode, RealKvHashMode) + + def test_an_unknown_attribute_still_raises(self): + with self.assertRaises(AttributeError): + server_args_module.NotAThing + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/utils/test_common.py b/test/registered/unit/utils/test_common.py index f93a65ecd..2b28b089d 100644 --- a/test/registered/unit/utils/test_common.py +++ b/test/registered/unit/utils/test_common.py @@ -1,9 +1,12 @@ +import sys import unittest from array import array +from unittest import mock import torch from sglang.srt.utils.common import ( + _get_device_sm_via_nvml, flatten_arrays_to_int64_tensor, get_device_sm_nvidia_smi, get_nvidia_driver_version_str, @@ -144,5 +147,73 @@ class TestGetDeviceSmNvidiaSmi(CustomTestCase): subprocess.run = original +class _FakePynvml: + """Records the NVML index it was asked for, so a test can tell which + physical GPU the helper would have reported.""" + + def __init__(self, capability=(9, 0)): + self.capability = capability + self.requested_index = None + self.initialized = False + + def nvmlInit(self): + self.initialized = True + + def nvmlShutdown(self): + pass + + def nvmlDeviceGetHandleByIndex(self, index): + self.requested_index = index + return f"handle-{index}" + + def nvmlDeviceGetCudaComputeCapability(self, handle): + return self.capability + + +class TestGetDeviceSmViaNvml(CustomTestCase): + """The torch ordinal and the NVML index differ under CUDA_VISIBLE_DEVICES + and MIG; without that mapping the helper must return None, not GPU 0.""" + + def test_torch_exposes_the_mapping_api(self): + # The cases below install the private attribute themselves, so they stay + # green on a torch that dropped it while the helper silently falls back. + self.assertTrue(hasattr(torch.cuda, "_get_nvml_device_index")) + + def test_maps_the_torch_ordinal_to_the_nvml_index(self): + fake = _FakePynvml(capability=(9, 0)) + with ( + mock.patch.dict(sys.modules, {"pynvml": fake}), + mock.patch.object( + torch.cuda, "_get_nvml_device_index", lambda index: 3, create=True + ), + ): + self.assertEqual(_get_device_sm_via_nvml(), 90) + self.assertEqual(fake.requested_index, 3) + + def test_returns_none_when_the_mapping_api_is_absent(self): + fake = _FakePynvml() + saved = torch.cuda.__dict__.pop("_get_nvml_device_index", None) + try: + with mock.patch.dict(sys.modules, {"pynvml": fake}): + self.assertIsNone(_get_device_sm_via_nvml()) + finally: + if saved is not None: + torch.cuda._get_nvml_device_index = saved + self.assertFalse(fake.initialized, "must not query NVML without the mapping") + + def test_returns_none_when_the_mapping_api_raises(self): + fake = _FakePynvml() + + def boom(index): + raise RuntimeError("no such device") + + with ( + mock.patch.dict(sys.modules, {"pynvml": fake}), + mock.patch.object(torch.cuda, "_get_nvml_device_index", boom, create=True), + ): + self.assertIsNone(_get_device_sm_via_nvml()) + self.assertFalse(fake.initialized, "must not query NVML without the mapping") + + if __name__ == "__main__": unittest.main()