[Perf] Fork-safe import: no CUDA context at import time, lighter argument parsing by xenshinu · Pull Request #40201 · sgl-project/sglang

xenshinu · 2026-09-18T16:26:59Z

Following #39683, decompose it into a stack of PRs for better management.

WHAT IS IN THIS PR:

Motivation: importing sglang created a CUDA context in every process through device probes, and server_args pulled the parser registries and sglang.kernels into every process that parses arguments. A process with a CUDA context cannot be a fork() parent, which rules out forkserver-based worker startup (PR 2), and the import cost lands on the launcher's critical path (PR 3).

Modifications: get_device_sm() answers from NVML while torch.cuda is uninitialized; get_device_module() picks the module from the platform checks; GraphCaptureContext.stream becomes a string annotation; DeepGEMM reads SGLANG_JIT_DEEPGEMM_PRECOMPILE at call time. CLI choices for --reasoning-parser / --tool-call-parser come from dependency-free name lists asserted in sync with the registries; the kv-canary enum, torchvision and HfApi are imported where used; flashinfer's JIT logger level is set by name. Always on; behaviour unchanged. Note for reviewers: sgl_kernel still initializes CUDA on import on the default path, so the launcher keeps a context until that is fixed upstream; PR 2 answers those probes from NVML inside the forkserver.

Accuracy: no model code touched. Speed: import sglang.srt.server_args no longer imports the parser registries (about 3 s of imports on our hosts); end-to-end gains are measured in PR 3. Tests: function_call/test_parser_names.py, parser/test_reasoning_parser_names.py.

OLD DESCRIPTION:

Motivation

A cold sglang serve spends 15-25 s in process init before any model work starts: the launcher imports the server stack, then every spawned worker re-imports the same sglang/torch stack, and on the DP path the DP controller re-imports once more before it spawns the schedulers. Every process also created a CUDA context at import time, which wastes a context in the launcher, detokenizer and DP controller and rules out fork()-based worker startup.

This PR makes the import fork-safe and, behind three opt-in environment variables, lets workers fork from a preloaded forkserver and start before the launcher imports its HTTP stack. With the variables unset every code path runs exactly as before.

Time from sglang serve to /health, Qwen3-30B-A3B (bf16) on H200, mean of three warm passes per cell (the first, cold pass that fills the weight page cache and the kernel JIT caches is excluded; one pass ran with the launch order reversed; all agree within 1 s):

config	knobs off	forkserver + pre-spawn	+ daemon
Single GPU	47.3 s	39.8 s	39.5 s
DP2	56.0 s	40.9 s	40.5 s
TP2	47.6 s	40.5 s	39.9 s
EP2 (DeepEP low-latency, DP attention)	46.6 s	32.1 s	31.4 s

Where the time goes, from the rank-0 log stages:

stage	knobs off	forkserver + pre-spawn
process init (launcher imports, spawn, worker import, scheduler setup)	15-23 s	8-9 s	this PR
weight loading (61 GB bf16)	6-10 s	6-10 s	unchanged
CUDA graph capture, prefill + decode	11-18 s	11-18 s	unchanged
distributed init, KV pool, tokenizer manager + HTTP server	4-6 s	4-6 s	unchanged

Modifications

Four commits, each usable on its own:

No CUDA at import (always on). get_device_sm() answers from NVML while torch.cuda is uninitialized, get_device_module() picks the module from the platform checks, GraphCaptureContext.stream becomes a string annotation, DeepGEMM reads SGLANG_JIT_DEEPGEMM_PRECOMPILE at call time. A process without a CUDA context can be a fork() parent.
Early forkserver (SGLANG_EARLY_FORKSERVER=1, new entrypoints/early_forkserver.py). The multiprocessing forkserver starts at CLI entry with the scheduler, DP controller and detokenizer modules preloaded; workers fork from it instead of re-importing. The child re-applies the launcher's environment and stdio, re-derives the HF offline flags and reports the launcher as its parent. python -m sglang.srt.entrypoints.early_forkserver runs a long-lived daemon that publishes its address in SGLANG_FORKSERVER_FILE.
Pre-spawn (SGLANG_PRESPAWN_WORKERS=1, new entrypoints/prespawn.py, entrypoints/worker_launch.py, managers/process_entry.py). run_server() resolves the config and spawns the schedulers (or the DP controller) right after argument parsing, before importing the HTTP and tokenizer stack; Engine._launch_subprocesses adopts them. Weight-cache daemons, the engine-info bootstrap server and Ray stay on the normal path.
Lighter launcher path (sglang/__init__.py, server_args.py, cli/utils.py, utils/common.py). Under the pre-spawn flag the package applies the transformers patches when transformers is imported and resolves the frontend API lazily (import sglang: 4.6 s -> 0.24 s); by default it imports eagerly as before. CLI choices for the parsers come from dependency-free name lists asserted in sync with the registries; a few imports move to their call sites.

IMPORTANT note on the copy of the engine's launch code. worker_launch.py duplicates the launch preparation and the scheduler spawn loop of engine.py verbatim, against the "avoid duplication" style rule. This is deliberate for a first PR: pre-spawn cannot import engine.py (that import is what it overlaps), and we did not want to restructure the engine here, so engine.py keeps its code (19 added lines ignoring whitespace; the original preparation block moves into an else: branch, see git diff -w) and test_worker_launch.py fails when the copy drifts from the engine or when the light modules import the server stack.
I am happy to remove the duplication based on the reviewers' opinion: the launch path moves into worker_launch.py and engine.py re-exports it. That variant exists and can go into this PR or a follow-up.

New environment variables (documented in environment_variables.mdx): SGLANG_EARLY_FORKSERVER (default off), SGLANG_FORKSERVER_FILE (default /tmp/sglang_forkserver.json), SGLANG_PRESPAWN_WORKERS (default off), SGLANG_MP_START_METHOD (set by start_early(), not meant to be set by hand).

Accuracy Tests

No change to model forward or sampling code. Greedy output for "The capital of France is" is identical across all 48 launches of the matrix above and, in an earlier run, identical to upstream main in the same environment.

Speed Tests and Profiling

See the tables and figure in Motivation (4 configurations x 3 settings x 3 warm passes on one H200 host; experimental/startup/startup_matrix.sh and startup_matrix_report.py in the foundry-org repo). Also on the branch: import sglang 4.6 s -> 0.24 s under the pre-spawn flag, import sglang.srt.server_args 5.3 s -> 3.0 s. With the knobs on, pre-spawn fires about 6.5 s after interpreter start: 1.7 s of imports and 4.5-5 s of ServerArgs.resolve_once(), which this PR keeps ahead of the spawn so the workers receive the same record as on the normal path.

Unit tests (CPU, base-a-test-cpu): entrypoints/test_early_forkserver.py, entrypoints/test_prespawn.py, entrypoints/test_worker_launch.py (copy matches the engine source-for-source; light modules never import the server stack), cli/test_sglang_init.py (eager default; lazy mode stays free of torch/transformers), function_call/test_parser_names.py, parser/test_reasoning_parser_names.py.

How to reproduce

Unit tests (about 15 s):

cd test/registered/unit
python3 -m unittest entrypoints.test_early_forkserver entrypoints.test_prespawn entrypoints.test_worker_launch cli.test_sglang_init function_call.test_parser_names parser.test_reasoning_parser_names

Launcher import cost with and without the lazy package init:

time python3 -c "import sglang"                             # eager (default)
SGLANG_PRESPAWN_WORKERS=1 time python3 -c "import sglang"   # lazy

Time to /health for one configuration under the three settings (set MODEL, ARGS such as --dp 2 or --tp 2, and a free PORT):

serve_and_time() {   # usage: serve_and_time "<ENV assignments>"
  t0=$(date +%s.%N)
  env $1 sglang serve --model-path $MODEL --port $PORT --host 127.0.0.1 $ARGS > server.log 2>&1 &
  until curl -fsS -m 2 http://127.0.0.1:$PORT/health >/dev/null 2>&1; do sleep 0.5; done
  python3 -c "import time; print(f'health after {time.time()-$t0:.1f} s')"
  curl -s http://127.0.0.1:$PORT/generate -H "Content-Type: application/json" -d '{"text":"The capital of France is","sampling_params":{"max_new_tokens":8,"temperature":0}}'
  pkill -f "port $PORT"; sleep 3
}
MODEL=Qwen/Qwen3-30B-A3B PORT=31000 ARGS=""
serve_and_time ""                                                     # knobs off
serve_and_time "SGLANG_EARLY_FORKSERVER=1 SGLANG_PRESPAWN_WORKERS=1"  # forkserver + pre-spawn
SGLANG_FORKSERVER_FILE=/tmp/sglang_forkserver.json python3 -m sglang.srt.entrypoints.early_forkserver &   # daemon, once per host
sleep 8
serve_and_time "SGLANG_EARLY_FORKSERVER=1 SGLANG_PRESPAWN_WORKERS=1 SGLANG_FORKSERVER_FILE=/tmp/sglang_forkserver.json"

Run each setting once to warm the page cache and kernel caches, then measure. The stage breakdown comes from the rank-0 log timestamps in server.log relative to t0: Init torch distributed begin (end of process init), Load weight begin/end, KV Cache is allocated, first Capture ... CUDA graph begin to last ... end, first GET /health ... 200.

Follow-up work

Graph capture is now the bottleneck. After this PR, weight loading (6-10 s) and CUDA graph capture (11-18 s) dominate every bar above. The foundry integration (CUDA graph save/restore across processes) restores the captured graphs from an archive in a few seconds instead of re-capturing them on every launch, and it will be another PR once all tests are done.
Make the daemon count. A warm daemon saves only 0.5-1 s here because the fork waits for ServerArgs.resolve_once() (4.5-5 s), which this PR keeps ahead of the spawn so workers get the same record as the normal path. Two options keep a single resolution and still move the fork to about 2 s: make resolution cheaper (most of it is imports of the model stack, recoverable with lazy imports), or spawn the workers first and have the child block on the launcher's resolved record over the pipe before it does anything. With either, the daemon's instant fork becomes the critical-path win it was designed to be.
Smaller items measured on the research branch: memoize get_tokenizer() in the scheduler process (the tokenizer is loaded twice, about 0.6 s), lazy LoRA drainer/overlap-loader imports in the scheduler (about 0.5 s).
De-duplicate the launch code by moving it into worker_launch.py, if the reviewers prefer that over the guarded copy.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

CI States

Latest PR Test (Base): ✅ Run #35521688622
Latest PR Test (Extra): ✅ Run #35521688477
Latest PR Test (AMD ROCm 10): ❌ Run #35521688646

ispobock · 2026-09-19T14:19:57Z

@xenshinu the CI tests failed, please check https://github.com/sgl-project/sglang/actions/runs/35368544693

ispobock · 2026-09-19T14:25:32Z

+        return None
+    try:
+        idx = 0
+        getter = getattr(torch.cuda, "_get_nvml_device_index", None)


_get_nvml_device_index is a private torch API that may be absent. Add a note so it is not changed to a direct attribute access.

ispobock · 2026-09-19T14:25:32Z

 @dataclass
 class GraphCaptureContext:
-    stream: torch.get_device_module().Stream
+    # String annotation on purpose: evaluating torch.get_device_module() here


trim to one or two lines, also for other comments. ref https://github.com/sgl-project/sglang/tree/main/.claude/rules

ispobock · 2026-09-19T14:25:32Z

                return decode_jpeg_with_fancy_upsampling(image_bytes)
+            # Lazy: torchvision (and torch._dynamo behind it) costs ~1.3 s at
+            # import and is only needed for GPU JPEG decode.
+            from torchvision.io import decode_jpeg


add a short note on why this import is lazy, like the HfApi import.

ispobock · 2026-09-19T14:44:34Z

+                idx = 0
+        pynvml.nvmlInit()
+        try:
+            handle = pynvml.nvmlDeviceGetHandleByIndex(idx)


_get_nvml_device_index missing or raising leaves idx=0, which queries physical GPU 0 and returns the wrong SM under CUDA_VISIBLE_DEVICES/MIG. Return None and fall back to torch, with a test.

Fixed. The helper returns None when _get_nvml_device_index is missing or raises, and get_device_sm() falls back to torch.cuda.get_device_capability().

ispobock · 2026-09-19T14:44:34Z

+        from sglang.srt.function_call.parser_names import TOOL_CALL_PARSER_NAMES
+        from sglang.srt.parser.reasoning_parser_names import REASONING_PARSER_NAMES
+
+        reasoning_parser_choices = list(REASONING_PARSER_NAMES)


Static choices drop parsers a plugin registers before parsing (engine.py:1090). Is that path still supported? If so keep the dynamic choices and add a test.

Thanks for pointing this out!
The choices now prefer the live registry whenever it is already in sys.modules and fall back to the name list otherwise.
I kept the name lists because they are worth about 535 ms of imports in every process, but I am good with dropping it if it is not worthy chasing.

xenshinu · 2026-09-20T03:18:46Z

                )

                return decode_jpeg_with_fancy_upsampling(image_bytes)
+            from torchvision.io import decode_jpeg  # lazy: ~1 s of torch._dynamo


changed to a short note

xenshinu · 2026-09-20T03:23:05Z


+
+def _reasoning_parser_choices():
+    # Importing the registry here costs seconds in every process that parses


The choices now prefer the live registry whenever it is already in sys.modules and 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.

ispobock · 2026-09-20T07:58:31Z

/rerun-failed-ci

Importing sglang created a CUDA context in every process (launcher, detokenizer, DP controller): get_device_sm() and the platform probes that call it went through torch.cuda.get_device_capability(), and get_device_module() asked torch for the current accelerator. Besides the wasted init, a process with a CUDA context cannot serve as a fork() parent. - get_device_sm(): answer from NVML while torch.cuda is uninitialized, and fall back to torch when torch's private ordinal-to-NVML index mapping is unavailable, rather than reporting physical GPU 0 under CUDA_VISIBLE_DEVICES or MIG. - get_device_module(): pick the module from the platform checks. - GraphCaptureContext.stream: annotate with torch.Stream. Evaluating torch.get_device_module() in the class body marks the process unsafe to fork -- a child then dies in cuInit -- while torch.cuda.is_initialized() stays False, so nothing reports it until a worker fails. This removes sglang's own import-time initializer. sgl_kernel still creates a context when it is imported, so the launcher keeps one today; what changes is that the modules a forkserver preloads no longer do.

Importing sglang.srt.server_args pulled the parser registries (OpenAI protocol models, xgrammar, torch.distributed) and sglang.kernels (via the kv-canary enum) into every process that parses arguments, and a few utilities imported heavy optional packages at module level. - CLI choices for --reasoning-parser / --tool-call-parser come from dependency-free name lists (function_call/parser_names.py, parser/reasoning_parser_names.py), checked against the registries by unit tests. cli/serve.py loads plugins before it parses arguments, and a plugin that registers a parser has imported the registry to do so, so the live registry is preferred whenever it is already in sys.modules. - The kv-canary enum is imported when the parser is built. FunctionCallParser, ReasoningParser and RealKvHashMode stay importable from server_args through a module __getattr__, like the other re-exports above it. - cli/utils.py looks for diffusers before importing sglang.multimodal_gen, and imports HfApi where it is used. - utils/common.py imports torchvision where the GPU JPEG decode uses it: at module level it pulls torch._dynamo into every process that imports this module, and no other path here needs it. - hf_transformers_patches.py uses a plain logger instead of the one from sglang.srt.utils, which pulls torch into every `import sglang`. Behaviour is unchanged; only import order and cost move.

xenshinu requested review from BBuf, Edwardf0t1, Fridge003, HaiShaw, JustinTong0323, Ying1123, ch-wan, ispobock, merrymercy and yizhang2077 as code owners September 18, 2026 16:27

ispobock added run-ci CI: run the baseline test suite on this PR bypass-fastfail run-ci-extra CI: also run the extra suite (requires run-ci) labels Sep 19, 2026

ispobock self-assigned this Sep 19, 2026

ispobock reviewed Sep 19, 2026

View reviewed changes

xenshinu force-pushed the startup/1-fork-safe-import branch from 73d0ef3 to ecc68ef Compare September 20, 2026 03:16

xenshinu commented Sep 20, 2026

View reviewed changes

ispobock approved these changes Sep 20, 2026

View reviewed changes

xenshinu added 2 commits September 20, 2026 12:04

xenshinu force-pushed the startup/1-fork-safe-import branch from ecc68ef to 8560eb9 Compare September 20, 2026 16:06

ispobock merged commit ab03a8e into sgl-project:main Sep 21, 2026
227 of 247 checks passed

McZyWu mentioned this pull request Sep 21, 2026

[NPU] Add opt-in MLA prefix FIA and Gemma norm fallback #40131

Open

6 tasks done

mmangkad mentioned this pull request Sep 21, 2026

[Fix] Add gigachat35 to the tool-call and reasoning parser name lists #40554

Merged

xenshinu mentioned this pull request Sep 21, 2026

[Perf] Fork workers from a preloaded forkserver (SGLANG_ENABLE_EARLY_FORKSERVER) #40596

Open

5 tasks done

This was referenced Oct 9, 2026

[CLI] Keep the diffusion pipeline runtime out of backend auto-detection #43317

Open

[CLI] Avoid loading diffusion runtime during backend auto-detection #39407

Open



		def _reasoning_parser_choices():
		# Importing the registry here costs seconds in every process that parses

Conversation

xenshinu commented Sep 18, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

WHAT IS IN THIS PR:

OLD DESCRIPTION:

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

How to reproduce

Follow-up work

Checklist

CI States

Uh oh!

ispobock commented Sep 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ispobock commented Sep 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

xenshinu commented Sep 18, 2026 •

edited by github-actions Bot

Loading