Repository navigation
Conversation
The Run:ai streamer reads every tensor into one reusable pageable host buffer, and the weight loaders then copy out of it with param.copy_(), i.e. a driver-staged pageable memcpy per tensor. On an 8xB300 loading a 433 GiB NVFP4 checkpoint (232k tensors) as four TP=2 replicas that copy path ran at about 2.4 GB/s per worker and was the longest span of a tuned cold start. Add a `pinned_staging` key to the loader's extra config. When set, the streamer's staging buffer is page-locked right after stream_files() (pre-faulted with MADV_POPULATE_WRITE from several threads, then cudaHostRegister), the streamer copies each tensor to the current device out of that pinned buffer, and the loaders only ever see device tensors so every TP narrow is a device-to-device copy. The buffer is bounded to 16 GB unless `memory_limit` is set, is unregistered when the iterator is exhausted, and any failure to pin falls back to the stock path with a warning. Nothing changes when the key is unset. Measured on vLLM 0.20.2 (GLM-5.2-NVFP4, runai-model-streamer 0.16.1, 4x TP=2 on 8xB300): "Model loading took" 39-42 s vs 75-79 s (-48 %), four replicas serving at 133 s instead of 170 s from script start, every GPU parameter bit-identical to the stock loader, throughput unchanged. At TP=8 the gain is small (every rank reads the whole checkpoint through the page cache, so the ceiling is a memcpy). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Tom <tom@getimpala.ai>
…r level-2 sleep (NVFP4 MoE + MLA) After sleep(level=2) discards the weights pool, reload_weights on a ModelOpt NVFP4 MoE model (GLM-5.2-NVFP4 on the FlashInfer TRT-LLM backend, with MLA) brought the server back healthy and serving garbage. Two independent causes, both fixed here behind env flags so the default path is unchanged: 1. The fused-MoE kernel is rebuilt on every process_weights_after_loading. Its quant config holds the activation gscales (a1/a2) that are derived at load time and are not layer parameters, so the layerwise reload's copy into kernel tensors never covers them: the CUDA graphs captured at startup keep reading the old tensors, whose memory the sleep discarded. Eager steps use the new tensors, graph replays the stale ones. With VLLM_RELOAD_KEEP_MOE_KERNEL=1 a reload keeps the kernel the graphs were captured with and copies the freshly derived quant-config tensors into the ones it references. (The MLA W_UV/W_UK_T half of this is already upstream via vllm-project#48251, and vllm-project#49489 made the shared scales writable.) 2. The layerwise reload processes a layer once the elements its loaders reported reach the layer's declared size. NVFP4 MoE scale loaders are called more than once per parameter, so the count reaches the declared size before the last scale tensors have arrived (the Run:ai streamer also yields out of file order); the layer is quantized early and the trailing loads are dropped as "excessive loading". The per-call cap in get_numel_loaded cannot see repeated calls. With VLLM_LAYERWISE_RELOAD_MEASURED_TOTALS=1 the loaders' element counts are recorded during the first load (record_first_load_numel, wrapped around load_weights in BaseModelLoader.load_model) and used as the exact trigger total on reload; a ready layer is also only processed once a load for another layer arrives. A CPU unit test shows the stock trigger leaving a derived tensor stale and the measured trigger not. Measured on vLLM 0.20.2, 4x TP=2 on 8xB300: after wake_up + reload every parameter, buffer and derived tensor hashes identical to a fresh load, 10 rounds out of 10; reload through the pinned Run:ai path 113 -> 55 s per pair. This is what makes a GPU-state snapshot of a sleeping server usable (17 GB CRIU images instead of 491 GB per TP=2 pair). Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Tom <tom@getimpala.ai>
This was referenced Oct 13, 2026
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: 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. 🚀 |
Contributor
|
Documentation preview: https://vllm--61304.org.readthedocs.build/en/61304/ |
mergify
Bot
added
documentation
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Overview
Stacked on #61303 (its commit is the first of the two here; review the second commit).
After
sleep(level=2)discards the weights pool,reload_weightson a ModelOpt NVFP4 MoE model (GLM-5.2-NVFP4, FlashInfer TRT-LLM backend, MLA) brings the server back healthy and serving garbage: well-formed completions, wrong content. Two independent causes, fixed behind env flags so the default path is unchanged:VLLM_RELOAD_KEEP_MOE_KERNEL=1: on reload, keep the ModelOpt NVFP4 MoE kernel the CUDA graphs were captured with and copy the freshly derived quant-config tensors into the ones it references, instead of building a new kernel.VLLM_LAYERWISE_RELOAD_MEASURED_TOTALS=1: trigger a layer's layerwise processing from the element count its loaders reported during the first load, not from the declared parameter sizes, and process a ready layer only once a load for another layer arrives.Claims
wake_up+reload_weightsevery parameter, buffer and derived tensor hashes identical to a fresh load. Measured on vLLM 0.20.2 (this PR is the port tomain): 10 rounds out of 10 on 4x TP=2 GLM-5.2-NVFP4 on an 8xB300, content-verified completions; without the fixes, 0 of 10.Validation
test_layerwise_reload_measured_first_load_totalsintests/model_executor/model_loader/test_reload.py: a layer whose scale loader is called twice per expert (so its loaders report 8 elements where the layer declares 6), streamed with the over-counted parameter last. With the stock trigger the layer is processed two loads early and a tensor derived inprocess_weights_after_loadingstays stale ([222, 0]); with the measured total it is correct ([222, 242]). Both variants run and pass; the rest oftest_reload.pythat runs on CPU still passes (38 passed, the remaining cases need the GPU fixtures).pre-commitruff-check,ruff-format,typospass on the changed files.main. The 10/10 result is from the same two fixes applied to 0.20.2. Parts of the 0.20.2 patch are already upstream and were dropped from this port: [Bugfix][Attention] Preserve post-load tensors across weight reloads #48251 (MLAW_UV/W_UK_Trefreshed in place), [Bugfix] Make shared NVFP4 MoE scales writable #49489 (shared NVFP4 MoE scales writable, which also removes the need for an overlap-tolerant copy), [Bugfix][Rocm]Aiter MoE re-uses existing tensor addresses after weight update. #40390 (replace_parameter(prefer_copy=True)), [Feature] Offload KV-init runtime state on sleep #59158 (KV-init runtime state offloaded on sleep). I will re-run the wake+reload checksum drill onmainas soon as I have the node again and post here; draft until then.Details
Cause 1, stale derived tensors.
ModelOptNvFp4FusedMoE._build_moe_kernelruns on everyprocess_weights_after_loading, so a reload builds a newFusedMoEQuantConfigand kernel. The activation gscales (a1/a2) in that config are derived at load time and are not layer parameters, so the layerwise reload's copy into kernel tensors does not cover them: the captured graphs keep pointers to the old tensors, whose memory the level-2 sleep discarded. Eager steps use the new tensors, graph replays the stale ones. The fix keeps the kernel on reload and refreshes its quant-config tensors in place (_refresh_moe_quant_config_inplace, per-descriptor, same shape/dtype/device required, otherwise rebind with a warning). The new kernel'sfused_experts.process_weights_after_loading(layer)still runs so the layer's parameters are converted; the layerwise machinery then copies them back into their original storage as before.Cause 2, early trigger. The layerwise reload processes a layer once
load_numel >= get_layer_size(layer). NVFP4 MoE scale loaders are called more than once per parameter (one call per expert per projection), and the Run:ai streamer yields tensors out of file order, so the count reaches the declared size before the last scale tensors have arrived; the layer is quantized early and the trailing loads are dropped as "Excessive loading" (theLimitations(4)case inlayerwise.py). The per-call cap inget_numel_loadedcannot see repeated calls.record_first_load_numel(model)wraps every weight loader during the firstload_weights(only when the flag is set), records the per-layer total inLayerReloadingInfo.first_load_numel(kept acrossreset()), andinitialize_online_processinguses it asload_numel_totalwhen reloading. As a second line of defence a ready layer is processed only when a load for another layer arrives (PENDING_LAYERS), flushed infinalize_layerwise_processing.Both are opt-in because they have been validated on one model family and one backend; if maintainers would rather have the measured-total trigger on by default, I am happy to drop the flag.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI. (CPU unit test added; hardware re-run on
mainpending.)Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
🤖 Generated with Claude Code