Repository navigation
[MoE] Stage UVA-offloaded expert weights in VRAM for large steps - #62060
mushroomcarbon wants to merge 1 commit into
Conversation
With UVA weight offloading (--cpu-offload-gb, --offload-backend uva) the fused-MoE kernel reads offloaded expert weights from host memory in place, and it reads an expert once per tile of the tokens routed to it. In a large step every offloaded expert therefore crosses the interconnect several times; small steps leave many experts unrouted, and there reading in place is the cheaper option. Add UVAOffloadConfig.cpu_offload_stage_min_tokens (--cpu-offload-stage-min-tokens), off by default. Steps with at least that many tokens copy each layer's offloaded w1/w2 into a reusable device buffer (one per slot and shape, shared by all layers, ordered by the stream) right before the expert kernel, which then reads the copy. "auto" stages every step that runs without CUDA graphs (more than max_cudagraph_capture_size tokens, and at least 512); it is resolved in VllmConfig once the capture sizes are known. Steps inside a CUDA graph capture never stage, so decode keeps the zero-copy path. The buffers are allocated by the memory profiling run, so the KV cache is sized with them counted. Experts that do not read the w1/w2 arguments (FlashInfer MoE-EP, b12x) opt out through FusedMoEExperts.reads_weight_args. CUDA only for now: ROCm kernels read per-tensor dispatch attributes that a staging buffer does not carry. The threshold is left out of the compilation hash: it acts inside the opaque moe_forward op. Assisted-by: Claude Opus 5.5 Signed-off-by: Andrew Goh <andrew.goh@zillion.network>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
👋 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. 🚀 |
Overview
With UVA weight offloading (
--offload-backend uva --cpu-offload-gb N --cpu-offload-params experts), the fused-MoE kernel reads offloaded experts from host memory in place, once per M-tile (64 tokens with the default config here) of the tokens routed to each expert, so in prefill every offloaded expert crosses PCIe several times per step. This adds an opt-in--cpu-offload-stage-min-tokens(default 0 = off;auto= every step that runs without CUDA graphs): such steps copy each layer's offloadedw1/w2into a reusable device buffer right before the expert kernel, which reads the copy. Decode and every graph-captured step keep today's zero-copy reads.Claims
vLLM 0.31.0 with this patch against stock 0.31.0; Qwen3-30B-A3B-Instruct-2507-FP8, one RTX 4090 (PCIe 4.0 x16), 16 GiB of experts offloaded.
autoin one build gives byte-identical greedy output on 18/18 requests; GSM8K 0.896 vs 0.900 (±0.019) against stock.Validation
Run on the nightly wheel of this PR's base commit with the patched files applied (CUDA 13.0, torch 2.14).
auto/0/N resolution (inVllmConfig, excluded fromcompute_hash); pass-through below the threshold, for non-offloaded weights and when off; an identical copy that keeps strides, one buffer reused across layers, separate buffers per slot; no staging or allocation inside a CUDA graph capture; and throughFusedMoEKernel.apply(Triton experts): at the threshold the experts receive the staging buffers (checked bydata_ptr) and match the same weights in VRAM, one token below they receive the UVA views.TitanML/tiny-mixtral(eager, every step staged; CUDA graphs, prefill staged) andAdamji/tiny-qwen3-moe-fp8-per-block(CUDA graphs), threshold off vs on: identical greedy output.VLLM_GPU_SYNC_CHECK=error: no sync error. Pre-commit hooks pass.auto)--numa-bind, chunk 2048, 1 runServer command and kernel microbenchmark
One step through all 48 MoE layers (checkpoint FP8 block shapes, random weights, routing from recorded chat traffic): the MoE part is 2.4× faster at 2048 tokens, 6.9× at 8192, 11.3× at 16384. UVA reads 2.32× / 7.75× / 15.05× the offloaded bytes per step; the copy reads them once at 26.3 GB/s. Copying loses below 256–512 tokens (0.91× at 256) and wins from 512 (1.09×), which is what
auto's floor of 512 is based on.Clients: the agent rows replay 16 long ShareGPT conversations as sessions (turn 0 sends the first ≥16k tokens of the conversation, each later turn appends the dataset's next exchange, so prompts are 16–19k tokens and each extends the previous one; 4 turns, 64 output tokens per turn, 2 sessions in flight) with
agent_sessions.py+agent_load.py(exact commands in their headers); the decode rows arevllm bench serve --dataset-name sharegpt --sharegpt-output-len 256 --ignore-eos --max-concurrency 1|4; GSM8K islm_eval --model local-completions --tasks gsm8k --limit 250(greedy).Details
Root cause.
fused_moe_kernelgives each program one M-tile of one expert and streams that expert's weights, so every tile re-reads the expert. In VRAM that is cheap; over UVA each re-read crosses PCIe (0.97–0.98 host reads per tile). A tunedBLOCK_SIZE_M128 config helps UVA 1.4–1.9× but copying in is still 1.1–6× faster, and it slows the all-in-VRAM case 8–12% at ≥ 2048 tokens.Design.
BaseOffloader.maybe_stage(weight, num_tokens, slot)returns the weight unchanged except inUVAOffloader, whoseUVAExpertStagercopies a weight carrying the UVA marker when the step meets the threshold and the stream is not capturing.FusedMoEKernelModularImpl.applycalls it forw1/w2after_prepare, unless the experts setreads_weight_args = False(FlashInfer MoE-EP and b12x run on weights prepared at load time). It is a copy, not a cache: no routing inspection (no device-to-host sync), no eviction, any layout the kernel reads, nothing to invalidate after EPLB or weight updates. One buffer per (slot, shape, stride, dtype, device) is shared by all layers; copy and kernel run on the same stream, so layer l+1's copy waits for layer l's kernel (also under DBO, whose micro-batches share the stream).autoresolves inVllmConfig.__post_init__tomax(max_cudagraph_capture_size + 1, 512); an explicit N at or below the capture size warns, since captured steps never stage. The threshold is left out ofcompute_hash(it acts inside the opaquemoe_forwardop). Like the other offload options, the flag is documented through the CLI help; there is no offload docs page to extend.Why opt-in. It shrinks the KV cache of every
--cpu-offload-gbuser by the buffers (several GiB for large BF16 MoE layers), and it is unmeasured on NVLink-C2C systems and APUs, where reading in place may already be fast enough.Limitations. CUDA only: ROCm AITER kernels read a per-tensor attribute (
is_shuffled) that the buffer does not carry, and a backend keyed on the weight'sdata_ptrwould see one address for every layer (none in the modular path is); such backends must opt out viareads_weight_args. Measured with Triton experts and FP8 block weights only; FlashInfer, DeepGEMM and CUTLASS experts untested. Only the modular fused-MoE path; monolithic kernels read in place as before. Under DBO each micro-batch copies the layer. With all2all EP the threshold counts tokens before dispatch. The buffers are not in the sleep-mode pool, and come on top of--kv-cache-memory-bytes. Copy and compute are sequential: double-buffering would save 4% at 2048-token steps and 12–19% at 8k–16k (microbenchmark), for a second buffer; left as a follow-up.Not a duplicate. #37190 (open) is an opt-in GPU expert cache: it inspects routing each forward (
topk_ids.unique().tolist(), a sync per layer), copies missing experts one by one, evicts with LFRU, and runs the MoE layer outside CUDA graphs (it capscudagraph_modeat PIECEWISE). RFC #38256 asks to load each distinct expert once in batched prefill; this is the narrow form of that for the existing UVA path. Residency work (#54610, #56177, #57794, #57943) decides which experts live in VRAM; this composes with it, since only weights carrying the UVA marker are staged. #51710 extends the prefetch backend, which copies every step including decode. #54972 edits a differentUVAOffloadermethod. Searched 2026-10-09 and rechecked 2026-10-16: none stages UVA experts for large steps.AI assistance. Claude Opus 5.5 assisted with the code, tests, benchmarks and this description; the tests and benchmarks above were run on my hardware. I reviewed every changed line and the results, and can answer for it; the commit carries an
Assisted-by:trailer.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.
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.