[Perf] Fork-safe import: no CUDA context at import time, lighter argument parsing (#40201)

This commit is contained in:
Xueshen Liu
2026-09-21 10:48:35 +08:00
committed by GitHub
parent 176dbcb85d
commit ab03a8e7eb
11 changed files with 381 additions and 12 deletions
+12 -2
View File
@@ -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:
@@ -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
@@ -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",
]
@@ -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",
]
+51 -6
View File
@@ -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. "
+62 -2
View File
@@ -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
@@ -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
@@ -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()
@@ -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()
@@ -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()
+71
View File
@@ -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()