[JimmyToluene]

Overview

Routes the Qwen4Exp PLE shard copy through copy_weight_, the follow-up named in #58868. Fixes #58726.

@Willian-Zhang, as you asked in #58868.

Not a duplicate: when this was opened, no other open PR mentioned #58726 or changed ple.py. #65713 now applies the same swap to other loaders and doesn't touch this file. The open PLE storage PRs (#58439, #58815, #62847) avoid or stage this copy for offloaded tables; this one speeds up the existing device-resident path.

Claims

• On a DGX Spark serving a device-resident packed-NVFP4 PLE table (cpu_offload: false), main Loading weights took falls from about 213 s to about 76 s from a cold page cache, and start to /health from about 320 s to about 184 s.
• Output is unchanged: logprobs and greedy text were exactly equal to the old line in every run.
• Off integrated GPUs copy_weight_ is a plain copy_. A dtype-converting copy skips the touch.
• The old line also created a temporary GPU tensor through .to(device=…) before copying. This drops it, with or without a dtype conversion: in a small GB10 test, extra peak GPU memory went from the target's size to 0. I did not measure discrete GPUs.

Validation

End to end, this branch. vllm serve on one GB10, Qwen3.8-Flash-Next-NVFP4-LIL-qad5500, local NVMe, TP=1, four boots in the order old, new, new, old. Only ple.py differed. The checkpoint was evicted with posix_fadvise before each boot, and mincore reported 0% resident.

Run	Main Loading weights took	Start to /health	Output tok/s	Mean ITL
old line, run 1	212.12 s	318 s	50.41	67.11 ms
copy_weight_, run 1	74.66 s	179 s	50.53	67.39 ms
copy_weight_, run 2	77.50 s	188 s	50.96	67.17 ms
old line, run 2	213.11 s	322 s	48.57	67.44 ms

Output. Against the code it replaces, per tests/AGENTS.md: 16 fixed prompts with temperature=0 and seed=0. All 566 prompt-token logprobs and all 16 48-token greedy completions were exactly equal across all four boots.

Model evaluation not run. tests/evals/qwen4_exp targets Qwen3.8-Flash-Next-FP8 with --tensor-parallel-size 4 on H200/B200, which I don't have. This change does not alter the loaded weights, as the exact comparison above and the table check below show.

Table copy alone, this branch. The 128 weight and 128 weight_scale shards (28.8 GB), each copied by the real _copy_embedding_shard_ into a GPU tensor of the shard's dtype, cold cache, two runs per arm. The old line took 176.8 s and 178.9 s; copy_weight_ took 11.4 s and 10.2 s. Both tables compared equal with torch.equal after every run.

Rebase. The boots and table copy above ran just before a final rebase onto current main, which changed neither ple.py nor model_executor/utils.py. The tests below were rerun after it.

Earlier result. The 160 s → 14 s figure for the table within a stock nightly boot is in this comment on #58868.

Tests, from .venv per AGENTS.md, on the GB10 with its GPU:

.venv/bin/python -m pytest tests/models/qwen4_exp/test_ple.py -q
.venv/bin/python -m pytest tests/model_executor/test_utils.py -q
pre-commit run --files vllm/models/qwen4_exp/common/ple.py

• test_ple.py: 127 passed, 2 skipped.
• test_utils.py: 32 passed.
• pre-commit: passed.

Serve command and environment

The host has no nvcc, so serving ran in vllm/vllm-openai:nightly-e4a5b5537344 (torch 2.14.0+cu130) with this tree on PYTHONPATH. The baseline bind-mounted the previous ple.py over it.

vllm serve <checkpoint> --engram-config '{"cpu_offload":false}' \
  --kv-cache-dtype fp8 --kv-cache-memory-bytes 4G --max-model-len 32768 \
  --max-num-seqs 4 --gpu-memory-utilization 0.88 \
  --limit-mm-per-prompt '{"image":0,"video":0}' --disable-flashinfer-autotune

vllm bench serve --dataset-name random --random-input-len 256 \
  --random-output-len 64 --num-prompts 16 --max-concurrency 4 --seed 0

Autotune was off because, with the model loaded, it took this machine down to 1 GiB available.

Details

Root cause: _copy_embedding_shard_ copied each shard with target.copy_(source.to(device=target.device, dtype=target.dtype)), straight from the safetensors mmap view. For a same-dtype source, the page faults land inside the driver copy, which #58726 measured as several times slower than faulting on the CPU.

Scope: load_weights in nvidia/ngram_embedding.py sends split table shards through weight_loader to this copy. Disk-offloaded tables take accept_checkpoint_shard instead. The buffer.copy_ in the same function only handles the small hash buffers, so I left it alone.

Other backends: amd/ple_layer.py reaches the same copy. copy_weight_ only touches pages when is_integrated_gpu is true, and otherwise runs a plain copy_. I only ran this on a GB10.

Tradeoffs: this is load-time only, so steady-state inference and CPU/GPU synchronization are untouched. The touch reads the pages it is about to copy, and the lowest available memory during boot was 4–5 GiB in both arms.

CI coverage: CI has no integrated GPU, so it can't exercise the touch path here. test_ple.py runs on H200 and MI355 through models_basic.yaml, and #58868 tests the helper with the gate forced on. I added no test because this is one line of wiring. Validation outside CI is above.

AI assistance (Claude) was used for the measurements and this write-up; I reviewed the change and the results.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (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.

[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. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

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.

🚀

[Willian-Zhang]

Thanks @JimmyToluene! Loading the PLE table on GB10 dropping from ~177 s to ~11 s is exactly what we hoped unified memory could do once the page faults moved out of the driver copy. Great to see it land end to end. 🚀

Nit: the old line wasn't a plain copy_ either. .to(device=…) made a temporary GPU tensor first, so this also drops that on discrete GPUs. Worth a sentence in the description.

LGTM once test_ple.py passes in CI.

[JimmyToluene]

/ci run

[github-actions]

❌ @JimmyToluene, A reviewer with write access must run /ci run, approve the PR, or add the ready label first.

[JimmyToluene]

Nit: the old line wasn't a plain copy_ either. .to(device=…) made a temporary GPU tensor first, so this also drops that on discrete GPUs. Worth a sentence in the description.

Thanks @Willian-Zhang! I added the temporary GPU tensor point to the description; for a same-dtype source the temporary is gone, and a dtype-converting copy now stages through the CPU. I also rebased onto main, with no code change.

/ci run needs an approval from someone with write access, so test_ple.py hasn't run yet. @gau-nernst, could you review when you have time?

[gau-nernst]

/ci run

[github-actions]

✅ Triggered Buildkite CI #100119 for commit 197bdf01a2d6.

ℹ️ This PR is 1 commit behind upstream main. Merged upstream main into this branch as 197bdf01a2d6 before starting CI. Pull the branch before pushing more changes.

[github-actions]

✅ @JimmyToluene, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Run commands test your branch with every commit currently on its upstream target branch. If it is behind, the bot merges the target branch into it first (pull before pushing again). Append --allow-stale to a run command to test an outdated branch as is, at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[JimmyToluene]

Correction to my comment above: I measured this on a GB10, and the temporary GPU tensor is gone with or without a dtype conversion. The CPU staging isn't new. I updated the description.

Thanks @Willian-Zhang and @gau-nernst for the reviews and for getting this merged!