Repository navigation
[Perf] Fork-safe import: no CUDA context at import time, lighter argument parsing #40201
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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", | ||
| ] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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", | ||
| ] |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -95,7 +95,6 @@ | |
| 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) | ||
|
Collaborator
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fixed. The helper returns |
||
| 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 | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. changed to a short note |
||
|
|
||
| encoded_image = torch.frombuffer(image_bytes, dtype=torch.uint8) | ||
| image_tensor = decode_jpeg(encoded_image, device="cuda") | ||
| return image_tensor | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The choices now prefer the live registry whenever it is already in
sys.modulesand fall back to the name list otherwise.I kept the name lists because they are worth about 535 ms of imports in every process. Let me know if you prefer to drop it.