vllm-snapshot

Fast model suspend / restore for vLLM sleep mode — swap large models on a single GPU in seconds instead of minutes, by snapshotting the built kernel-format weights and restoring them with a bulk cudaMemcpy instead of re-running the checkpoint loader.

Built and validated on an NVIDIA DGX Spark (GB10, 128GB unified memory, SM121, arm64) serving a 26B NVFP4 MoE model, but the mechanism is general.

The problem

vLLM sleep mode is great for time-sharing one GPU across models:

• POST /sleep?level=2 frees the GPU memory pool and keeps the warm engine + captured CUDA graphs (via the CuMemAllocator's virtual-address reservation).
• POST /wake_up brings it back.

But level 2 discards the weights — waking them requires reload_weights, which re-reads the checkpoint and re-processes it into kernel format. For a quantized MoE that reprocessing is the wall, and it is CPU-bound, not I/O bound, so faster loaders don't fix it:

reload path	time (26B NVFP4 MoE, GB10)
--load-format auto (safetensors)	82 s
--load-format instanttensor	~30 min (pathological layerwise path)
--load-format fastsafetensors	same wall (I/O was never the bottleneck; no GDS on-box)
sleep level 1 (offload weights to host)	57 s (and frees nothing on unified memory)
cold restart	135 s

Raw disk read of the model is only ~12 s (1.5 GB/s × 18 GB) — the rest is per-tensor processing.

The idea

The built, kernel-format weights already sit in GPU memory, in the CuMemAllocator's "weights"-tagged regions. So:

1. Snapshot those regions byte-for-byte to disk once (cudaMemcpy D→H, streamed one region at a time).
2. On wake, wake_up?tags=weights re-maps the physical pages at the same virtual addresses (fast — the mapping, not the copy, was L1's real cost).
3. Restore with a bulk cudaMemcpy H→D straight into those pointers.
4. wake_up?tags=kv_cache.

Because it captures raw memory regions, it round-trips 100% of inference state — params, non-persistent buffers, and post-load kernel tensors — that a state_dict() walk silently omits. (A state_dict()-based restore was faster to write and ran in ~3 s, but produced garbage on this quantized MoE precisely because of those omissions. That path is kept in the code, clearly marked, as a cautionary reference.)

Result

	time
snapshot_weights_raw (one-time, offline)	~86 s (19 GiB)
restore_weights_raw	~1.6 s (19 GiB @ ~12 GiB/s)
full suspend→restore swap	~9 s (vs 82 s / 30 min / 135 s)

Correct and reproducible (verified with real generations after restore).

Install

Runtime deps vllm and torch come from your vLLM environment/image — they are not installed by pip. Until a PyPI release lands, install from git:

pip install "git+https://github.com/dgr237/vllm-snapshot"   # or, in a checkout: pip install -e .

Installing registers a vllm.general_plugins entry point, so vLLM loads it automatically in every process (including workers) at startup. Two flags are required on the vLLM server:

VLLM_PLUGINS=vllm_snapshot,vllm_snapshot_endpoints \
vllm serve <model> \
  --enable-sleep-mode \        # arm the CuMemAllocator (required)
  ...                          # your usual args

The /snapshot, /suspend, /restore, /snapshot_status endpoints call the engine client directly in-process, so VLLM_SERVER_DEV_MODE is not required (and best left off — it also exposes a broad /collective_rpc surface). Set it only if you want the low-level HTTP routes below.

Environment:

• VLLM_SNAPSHOT_DIR — where snapshots live (default /snapshots).
• VLLM_SNAPSHOT_MIN_FREE_GIB — refuse to snapshot below this much free RAM (default 16). On unified-memory boxes this guard matters — a naive all-at-once host snapshot can OOM the whole pool.
• VLLM_SNAPSHOT_AUTOSUSPEND — if 1, the server suspends itself at startup (see "Boot-and-swap" below).

(vLLM logs Unknown vLLM environment variable for these VLLM_SNAPSHOT_* names — harmless; the plugin reads them directly.)

Usage

High-level endpoints (recommended)

Enable the optional HTTP surface by naming both plugins in VLLM_PLUGINS:

VLLM_PLUGINS=vllm_snapshot,vllm_snapshot_endpoints

Then:

curl -X POST localhost:8000/snapshot        # capture weights to disk (once; ?force=1 to redo)
curl -X POST localhost:8000/suspend         # snapshot (if needed) + verify + sleep 2
curl -X POST localhost:8000/restore         # wake weights + restore + wake kv
curl      localhost:8000/snapshot_status    # {is_sleeping, cached, restorable, reason, free_gib}

/suspend verifies a restorable snapshot before sleeping — it will never discard the weights unless it can bring them back (returns 409 otherwise).

Low-level RPCs (optional — requires VLLM_SERVER_DEV_MODE=1)

The plugin also patches the worker with methods callable over POST /collective_rpc: snapshot_weights_raw, verify_snapshot, restore_weights_raw, snapshot_status. The manual cycle:

curl -X POST localhost:8000/collective_rpc -d '{"method":"snapshot_weights_raw"}'  # once
curl -X POST 'localhost:8000/sleep?level=2'
curl -X POST 'localhost:8000/wake_up?tags=weights'
curl -X POST localhost:8000/collective_rpc -d '{"method":"restore_weights_raw"}'
curl -X POST 'localhost:8000/wake_up?tags=kv_cache'

See scripts/suspend_restore.sh and examples/compose-snippet.yaml.

Boot-and-swap: N warm models, one resident

VLLM_SNAPSHOT_AUTOSUSPEND=1 makes a server suspend itself at startup: after the engine is initialized (CUDA graphs captured), the plugin verifies a restorable snapshot — creating one first if needed and there is RAM headroom — then calls sleep 2. The server comes up healthy but empty, using only its CUDA context + captured graphs (~1-3 GiB), ready to /restore in seconds. If no snapshot can be made or verified it leaves the model awake rather than stranding it.

This enables a multi-model box: start N servers (each with its model + a pre-made snapshot) with VLLM_SNAPSHOT_AUTOSUSPEND=1; all boot suspended; a router /restores whichever is needed and /suspends the previous one. One vLLM engine serves one architecture — same-architecture models (e.g. different fine-tunes) can share a server by restoring different snapshots; different architectures need separate servers.

Note: creating a snapshot needs RAM headroom (the guard), so the usual pattern is to snapshot once (offline, at a lower --gpu-memory-utilization) and then run production with autosuspend where startup only has to verify + sleep.

Sequenced boot (avoid the simultaneous-load OOM): starting N servers at once would load all their weights at the same time and blow the pool. Because autosuspend finishes before a server answers its healthcheck, "healthy" means "booted + suspended (~2 GiB)", so chaining depends_on: {condition: service_healthy} loads models strictly one at a time. See examples/compose-multimodel.yaml.

Cross-restart reuse

Snapshots are reusable across process restarts. Each region is keyed by an order-independent content key (the lowest-offset named tensor in it, or anon:N for kernel-only regions), and the manifest records the model name, vLLM version, region count and per-region sizes. On restore the cache is validated against the live model and refuses rather than corrupts on any mismatch. A snapshot taken at one --gpu-memory-utilization restores correctly at another (weight regions are independent of the KV budget).

Compatibility

This plugin depends on vLLM internals — CuMemAllocator.pointer_to_data, the Worker class, and reload_weights — not just its public API, and that surface can change between vLLM releases.

• Tested against vLLM 0.26.x (the eugr/spark-vllm build) on an NVIDIA DGX Spark (GB10, SM121, arm64). The mechanism is general — not GB10-specific — but that is the only configuration validated end-to-end so far.
• Snapshots record the vLLM version + model and are validated on restore; a mismatch refuses rather than corrupts, so an upgrade can't silently restore a stale snapshot. The code paths may still need updating for a new vLLM — pin the version you tested and re-run the TESTING.md smoke after upgrading.
• Research prototype: it works and is reproducible, but it is not battle-tested across models/hardware. Region counts + swap times from other setups welcome.

Attention backend: serve with Triton, not FlashInfer

Serve any model you intend to snapshot/restore with --attention-backend TRITON_ATTN.

FlashInfer keeps non-torch attention workspace / plan buffers — GPU memory allocated outside vLLM's CuMemAllocator — which sleep(2) does not offload and wake does not re-establish. This snapshot only covers the allocator's weights-tagged regions, so after a restore the weights are bit-perfect but the FlashInfer attention state is garbage, and the model returns subtly corrupted output (fluent but wrong — not an obvious crash). Triton attention keeps no such sleep-fragile off-allocator state, so restore is bit-clean and deterministic.

This bites hardest on hybrid / linear-attention models (e.g. Qwen3.6 GDN), whose config tends to auto-select FlashInfer — but the rule is general: if you snapshot it, serve it with Triton attention. Root-caused by a greedy (temp 0) A/B before vs after restore combined with per-region weight checksums: the weights hashed identical while the output diverged, isolating the fault to the attention backend's off-allocator workspace rather than the weights this plugin saves. (Verified on vLLM 0.26.x; note the CLI flag --attention-backend TRITON_ATTN — the older VLLM_ATTENTION_BACKEND env var is not honored there.)

Security

The /snapshot, /suspend, /restore, /snapshot_status routes are unauthenticated — anyone who can reach the port can suspend or restore the model (a denial-of-service lever). Expose them only on a trusted network and put auth/allowlisting in front otherwise. The endpoint plugin is opt-in (VLLM_PLUGINS must name it), and the low-level /sleep / /collective_rpc routes stay off unless you set VLLM_SERVER_DEV_MODE=1.

Status / TODO

Research prototype. Working: raw snapshot/restore, content-keyed cross-restart reuse with validation, /snapshot + /suspend + /restore + /snapshot_status endpoints, autosuspend-on-startup, and the sequenced-boot multi-model pattern.

Investigated and rejected:

• --load-format dummy startup bypass — the idea was to skip the checkpoint read at boot and restore from the snapshot instead. It doesn't help here: (1) dummy boot was ~123s vs ~135s real — weight I/O was never the boot bottleneck; compile + CUDA-graph capture dominate, and dummy skips neither; (2) dummy allocates a different region structure (one fewer weights region) so a real-load snapshot won't verify into a dummy process. The autosuspend + sequenced-boot path (real-load once per boot, then suspend) is the working alternative.

Open:

• Router glue: a LiteLLM/router integration that /restores the target model and /suspends the previous one on switch.
• Weight-resident KV-swap (co-resident models, shift KV budget to the active one for copy-free swaps) — needs runtime KV resize, which vLLM has no API for (num_gpu_blocks is fixed at engine init and CUDA graphs bind the KV layout).

License

Apache-2.0.