[Bugfix][Memory] Load weights outside the CuMem pool and re-home them for sleep mode

[JasonVranek]

Purpose

FIX #48680 (--enable-sleep-mode OOMs loading an NVFP4 (modelopt) 30B on 16 GB SM120 cards where the identical model loads fine without it).

Root cause, pinned down by per-module instrumentation on the reproducing hardware: weight loading with sleep mode runs the whole load — including process_weights_after_loading — inside a live CuMem MemPool. A live pool can neither release emptied segments nor reuse freed blocks for larger requests (pytorch/pytorch#159674). The NVFP4 MoE conversion replaces each layer's weights with slightly larger (padded) tensors, so the pool strands ~one full weight copy: instrumented, the pool grows 9.18 → 14.44 GiB during conversion while torch-live memory grows only 9.08 → 9.61 GiB. That +5.2 GiB is the OOM — not per-op transient peaks, which is why the transient-reduction reworks in #48701 didn't resolve it (validated on the same hardware in that thread).

Two further allocator behaviors constrain the fix (established with small standalone probes, happy to share):

1. While a pool context is active, torch.cuda.empty_cache() is a device-wide no-op from any thread: pool routing registers in the caching allocator's captures_underway machinery, and while non-empty the allocator refuses to release cached blocks. Freed memory cannot reach the driver mid-context.
2. Creating a second MemPool for the same tag (e.g. re-entering use_memory_pool per chunk) is unsound: the orphaned pool's destructor fires mid-load and trips the captures_underway.empty() internal assert. Releasing "empty" pool segments behind torch's back mid-context causes illegal memory access (torch still caches those blocks).

The fix

vllm/v1/worker/gpu_worker.py only. When enable_cumem_allocator is set:

1. Run the entire load + weight conversion with the default caching allocator — identical to the proven non-sleep path, where freed/transient memory is reusable and reclaimable.
2. empty_cache() to hand the conversion churn back to the driver.
3. Enter use_memory_pool(tag="weights") once and re-home the finished model storage-by-storage into the pool (sleep can only offload/discard pool-owned memory). Storage-granular copying preserves tensors that share a storage (tied weights).
4. Every ~1 GiB, suspend pool routing on the same pool (_cuda_endAllocateToPool without _cuda_releasePool), empty_cache(), resume (_cuda_beginAllocateCurrentThreadToPool). Without this the pool's growth has no physical memory to map on 16 GiB cards; with it, free VRAM holds constant while the pool grows (probe: 1 GiB/round for 10 rounds vs OOM at round 5 without).

Peak overhead vs the non-sleep path is a single chunk's double-residency. The non-sleep path is unchanged.

Test Plan

Reproducing hardware from #48680: 2x RTX 5070 Ti 16 GB (SM120), driver 595.71.05, torch 2.11.0+cu130, flashinfer 0.6.13. Original failing configuration: nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4, TP=2, --max-model-len 32768, --max-num-seqs 64, --gpu-memory-utilization 0.85, --enable-sleep-mode, VLLM_USE_FLASHINFER_MOE_FP4=1.

• Load with sleep mode; compare KV budget vs non-sleep control on the same stack.
• Greedy completion before sleep; POST /sleep?level=1; VRAM check; POST /wake_up; identical greedy completion after wake; longer generation.
• Sanity: non-sleep path byte-unchanged (gated behind isinstance(pool_ctx, nullcontext)).

Test Result

• Load completes: Sleep mode: re-homed 9.82 GiB of weights into the CuMem pool after out-of-pool loading on both TP workers. Previously: deterministic torch.OutOfMemoryError during conversion (168 MiB request, ~130 MiB free, 4.85–5.20 GiB stranded in private pools across [Bugfix][Memory] Reduce CuMem weight-load memory #48701's revisions).
• KV cache: 1,004,885 tokens with sleep mode vs 980,309 tokens non-sleep control — sleep mode now costs ~0 KV budget on this setup.
• POST /sleep?level=1: 14.3 → 2.1 GiB per GPU. POST /wake_up: restored; greedy completion byte-identical to pre-sleep; longer generation coherent.
• Fix is independent of the [Bugfix][Memory] Reduce CuMem weight-load memory #48701 flashinfer rework: validated with the stock file (results above) and with that PR's version (same load success, KV 953,002).

Caveats for review: (a) _cuda_endAllocateToPool / _cuda_beginAllocateCurrentThreadToPool are private torch APIs — though torch.cuda.use_mem_pool itself is implemented on exactly these; (b) tensors stored as plain module attributes (not in _parameters/_buffers) stay outside the pool and won't be offloaded on sleep — none found that matter on this model class; (c) wake_up(level=2) + reload paths deserve CI exercise.

Credit: the investigation in #48701 (@hubunt) — especially its revisions isolating the TRTLLM branch and the max_split_size_mb policy — is what let us falsify the transient-peak hypotheses on hardware and converge on the pool-ratchet mechanism.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀