Skip to content

feat: retain GPU weights across vLLM and SGLang engine reloads - #5640

Merged
qinxuye merged 16 commits into
xorbitsai:mainfrom
qinxuye:feat/model-weight-cache-reload
Oct 6, 2026
Merged

qinxuye merged 16 commits into
xorbitsai:mainfrom
qinxuye:feat/model-weight-cache-reload

Conversation

@qinxuye

@qinxuye qinxuye commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

Changing engine limits for a large model currently requires terminating the model and loading its weights again. This adds an opt-in enable_weight_cache=true launch option and a reload operation that keeps GPU weights in an upstream IPC cache daemon while rebuilding the inference engine, KV cache and CUDA graphs.

POST /v1/models/{uid}/reload accepts a partial model_config patch and returns HTTP 202 with an operation ID. Status and editable configuration are available through GET endpoints. Reload drains existing requests, including streams, rejects new requests with 503 during reconfiguration, preserves the model UID and GPU reservation, and attempts to restore the previous configuration if the replacement engine fails. Committed limits are persisted for worker recovery and the active autostart launch entry. Existing requests remain cancellable during drain; model termination interrupts reload and forcibly removes a rebuilding model pool while retaining and cleaning up native engine child process identities. Full termination releases the daemon; a parent watcher also cleans it up if the model process dies.

Sync/async clients, xinference reload, a launch switch and a running-model parameter dialog expose the operation. Documentation and all maintained translations are included.

The initial scope is Linux CUDA/ROCm, vLLM >= 0.31.0 or SGLang >= 0.5.21, one worker, one replica and fixed tensor parallelism. Editable fields cover concurrency, token/context limits, memory fractions and graph controls. Parallelism, dtype, quantization, model path and weight layout cannot change. This integration rejects non-LLM models, replica scale-up, GGUF, an explicit xoscar executor, PD, LoRA, speculative decoding and weight offloading. Multi-GPU vLLM caching uses its native multiprocessing executor. Unsupported quantization methods remain subject to upstream IPC checks.

The daemon uses a private socket namespace and remains independent of engine lifetime. vLLM IPC loading disables disk fallback. SGLang uses its standalone daemon and client loader; the original sockets and rank liveness are checked before and after loading to reject a lost cache. Upstream references: vLLM preload, SGLang weight cache.

Validation:

  • 118 model/reload/lifecycle/routing tests passed with one platform skip. Tests include the real vLLM/SGLang model classes with mocked upstream engine construction, health-task and loading-thread re-entry, graceful and forced daemon cleanup, parent death, termination during reload, stale actor ownership, abort during drain and scale-up rejection.
  • 46 API/client/autostart tests and 220 broader worker/model/RPC/history tests passed in separate suites. The final control/worker/API/client/autostart suite passed 208 tests, including a real native-child orphan regression and abort via the actual actor RPC metadata envelope. The autostart test uses the real SQLite launch store and verifies committed-only persistence. Cancelled reload job ownership remains safe across immediate same-UID relaunch.
  • Four React dialog integration tests passed, including hiding an earlier finished operation's error when opening a new dialog. Backend metadata exposes worker count for the existing sharded-model gate.
  • Full pre-commit run --all-files, ESLint, TypeScript, focused Prettier checks and the production static frontend build passed; 47 static-export/backend-serving tests passed.
  • All nine changed PO catalogs previously passed msgfmt --check --check-format; python doc/build_i18n.py --all passed, and the changed page rendered in each maintained locale. The review fixes do not change documentation messages.
  • Earlier CI failures were repaired, including SGLang process fixtures, POSIX-only tests, decimal activated-size parsing, shared FlagEmbedding constraints, rerank tokenizer detection, OPT FPQuant prerequisites and speculative fixtures. Current main supplies the shared repairs; the branch avoids duplicate constraints/fixtures. Linux Python 3.10 passed on the preceding merge head before the review-fix push. After merging current main (f37c131ba), 478 targeted tests passed with one platform skip, and full pre-commit passed. The previous macOS CI process-group cleanup failure is fixed by checking for live members after a permission error; zombie-only groups complete cleanup, while live-group permission errors still propagate. On that head, Linux 3.10–3.14, macOS 3.10/3.14, Metal, GPU, lint and docs passed. Two independent Windows timing assertions were repaired in a tests-only follow-up (07ff4ad4d): diffusion fatal-error cleanup waits for its completion event, and executor shared-deadline checks use a deterministic clock while preserving timeout exhaustion assertions. The two test files passed 69 checks; an injected delay after fatal-error publication also passed. The latest matrix is queued at CI run 37482108712; Windows success on this head remains pending.

Real GPU validation on two RTX 3090 Ti 24 GiB GPUs (CUDA 13, PyTorch 2.13.0):

  • vLLM 0.31.0 and SGLang 0.5.21 passed real Qwen2.5-0.5B-Instruct reload checks with one and two GPU ranks. Greedy outputs were unchanged, concurrency changes took effect, and the original daemon processes and socket identities survived engine reconstruction. Full unload returned GPU usage to the existing idle service contexts.
  • The real REST pipeline passed three reloads covering concurrency, context length, memory utilization and CUDA graph controls. Each operation returned 202, served status transitions, rejected inference with 503 while reloading, and resumed with the same model UID/output. Quantization/parallelism changes and an insufficient memory budget were rejected without interrupting serving.
  • These runs uncovered and fixed SGLang daemon YAML flag names and deprecated graph-flag parsing, vLLM memory accounting for the independently retained weights, and long Linux Unix socket paths. Numeric launch arguments retain CLI string compatibility. The public memory fraction remains unchanged; the internal engine budget reserves the daemon's per-device allocation first. GPU smoke tests use bounded KV pools.
  • Existing request-limit and document-analysis fixture regressions introduced by the reload guard are fixed; their 47 focused checks passed.

Measured Qwen2.5-7B-Instruct reload performance (TP=2, warm OS page cache):

Engine Ordinary backend reconstruction Retained-weight reload Time reduction
vLLM 0.31.0 20.853 s 18.258 s 12.44%
SGLang 0.5.21 75.276 s 73.279 s 2.65%

Each cell is the median of four accepted trials. Timing includes stopping/reconstructing the engine and the first successful eight-token greedy response. The ordinary baseline keeps the Python actor process alive, so it excludes actor creation and REST terminate/launch overhead. Both sides use fixed KV pools (vLLM: 512 MiB per GPU; SGLang: 2048 total tokens), context length 512, CUDA graphs disabled, and concurrency changing between 4 and 8. Phases run ordinary/cached/cached/ordinary, each after an excluded initial load and request. Outputs, original daemon PIDs and both socket inodes were checked throughout. GPU ownership is sampled every 0.5 seconds, with only two known 262 MiB idle service contexts allowed; foreign-job overlaps and conservatively misclassified teardown samples are excluded and retried.

vLLM cached loading reports 0.0 GiB of new model allocation. SGLang maps 200 tensors per rank without new post-quantization allocations and reduces the weight-loading stage from about 2.6 s to 0.12 s. The modest end-to-end gains on a warm 7B checkpoint reflect the remaining engine startup/readiness cost. Full termination leaves no task-owned GPU processes. No extrapolation to hundred-billion or trillion-parameter models is claimed.

The review-fix runtime also passed real TP=2 vLLM reloads without an explicit KV-byte allocation. The real NVML process records were offset to simulate a driver/container PID namespace mismatch; the cache reservation remained correct, two concurrency reloads retained both socket identities and greedy outputs, and an undersized memory budget was rejected before teardown. This exercises the accounting fallback on real GPUs; it is not a container deployment test. Seven production runtime files on the GPU host were SHA-256 verified against the prior review-fix head b6ae72448. The later main merge preserves SGLang Xavier/PD initialization and GC cleanup alongside cached reload; default backend-class re-entry and shutdown-failure tests verify the combined lifecycle. GPU timings retain their original revision scope.

The review-fix SGLang runtime passed two TP=1 reloads on GPU 1 with identical output, retained cache PID/socket inode and no foreign samples. Two attempted TP=2 reruns were rejected in full due to unrelated GPU work; their timings are excluded. The earlier TP=2 results remain historical. A final-worker real CUDA child cleanup test also passed: the child survived the pool parent kill, then WorkerActor.terminate_model removed its NVML entry without touching existing contexts.

To reproduce the opt-in correctness checks, set XINFERENCE_TEST_WEIGHT_CACHE_MODEL_PATH and optionally XINFERENCE_TEST_WEIGHT_CACHE_TP=2, then run pytest -v xinference/model/llm/tests/test_weight_cache_gpu.py.

@XprobeBot XprobeBot added this to the v3.x milestone Oct 5, 2026
@qinxuye
qinxuye requested a review from rogercloud October 6, 2026 07:30

@rogercloud rogercloud left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Major

  • weight_cache.py:87: NVML PID attribution can fail or raise in containers / on None usage (inline).
  • test_weight_cache.py:285: real vLLM/SGLang reload path not covered in default CI (inline).

Minor

  • supervisor.py:4626: replica lock held across drain + rebuild blocks terminate (inline).
  • worker.py:5835: broad except Exception resets state to ready (inline).
  • supervisor.py:2976: scale-up and non-LLM model types not rejected for enable_weight_cache (inline).
  • weight_cache.py:311: no graceful window / SIGKILL escalation in daemon stop (inline).
  • vllm/core.py:767: daemon starts before gguf/xoscar handling (inline).
  • weight_cache.py:414: get_reload_config races with in-place config mutation (inline).
  • weight_cache.py:270: vLLM preflight lacks per-rank liveness (inline).
  • worker.py:5848: abort cannot reach a reloading model (inline).
  • worker.py:5841: autostart entries keep pre-reload limits (inline).
  • running-model-detail/index.tsx:256: n_worker gate no-op; stale error in reload-dialog.tsx (inline).

Simplification

L355: delete: _weight_cache re-annotation/assignment duplicates the class default at L335. Delete the line.
net: -1 lines possible

Blocking: no — recommended event: APPROVE

Comment thread xinference/model/llm/weight_cache.py
Comment thread xinference/model/llm/tests/test_weight_cache.py
Comment thread xinference/core/supervisor.py Outdated
Comment thread xinference/core/worker.py Outdated
Comment thread xinference/core/supervisor.py
Comment thread xinference/model/llm/weight_cache.py
Comment thread xinference/model/llm/weight_cache.py
Comment thread xinference/core/worker.py
Comment thread xinference/core/worker.py
Comment thread frontend/src/components/pages/running-model-detail/index.tsx
@qinxuye
qinxuye merged commit 10f91bd into xorbitsai:main Oct 6, 2026
14 of 15 checks passed
@qinxuye
qinxuye deleted the feat/model-weight-cache-reload branch October 6, 2026 16:17
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants