[yeighta]

Overview

Qwen4Exp PLE cpu_offload currently requires UVA. When UVA is unavailable, the NVIDIA path now keeps the table in host memory and stages each step's rows through the staging path that #59958 added for disk offload, instead of failing at init. This lets Qwen3.8-Flash-Next run on one H100 under NVIDIA Confidential Computing, where #59981 turns UVA off.

Closes #62815. Related: #57224 (root cause: under CC, pinned host memory reports is_pinned() == False and UVA views of it go stale), #59981 (CC detection that disables UVA), #59958 (PLE disk offload, whose staging this PR shares), #54371 (UVA PLE offload).

Not a duplicate. #59958, #58815 and #58439 read PLE rows from the checkpoint files (disk / page cache). This PR keeps the table in anonymous host memory and needs neither UVA nor file-backed reads; in a confidential VM, storage I/O patterns are visible to the host, so disk-backed tables are not an option there. Since #59958 has landed, this PR reuses its per-step staging and only adds the host-memory row source.

Claims

• With cpu_offload on and is_uva_available() false, Qwen4Exp (NVIDIA) selects a new Qwen4ExpPLEHostMemoryEmbedding instead of raising Engram CPU offload requires UVA support. BF16, FP8 and NVFP4 tables are supported.
• The per-step staging of disk offload moves into a shared Qwen4ExpPLEHostStagedEmbedding base, unchanged in behaviour. Disk offload and the new host-memory table differ only in where _read_rows reads the rows from.
• The forward only reads persistent device buffers, so the default FULL_AND_PIECEWISE CUDA graphs are kept.
• No change when UVA is available: the pinned UVA path and its prefetch are untouched.

Validation

All on one H100 NVL 94 GB in CC mode (Azure Standard_NCC40ads_H100_v5, driver 580.95.05, CUDA 13.0).

Unit tests (current head): nightly 5689c279 + this PR, python -m pytest tests/models/qwen4_exp/test_ple.py tests/models/qwen4_exp/test_config.py on that GPU: 140 passed, 8 failed, 2 skipped. The new and renamed cases pass: the graph-replay comparison against the resident table for both row sources (disk and host memory) in BF16, FP8 and NVFP4, the ETP owner masking, and the fallback selection. The 8 failures are the same set without this PR: 7 pinned/UVA tests fail because UVA views of pinned memory go stale under CC (this run has no #59981), and 1 config test needs PP=2 on a one-GPU node.

The end-to-end results below were measured on the first version of this PR (nightly 8e2beaa3 + #59981 f37260ad, V2 model runner, VnimanieAI/Qwen3.8-Flash-Next-W4A16), which gathered the rows with the same index_select but had its own staging code. I will re-run them on the final head of #59981.

Correctness: per-token prompt logprobs of 30 fixed texts (5,756 tokens). The reference is a naive lookup inside the forward (rows fetched synchronously with F.embedding(ids.cpu(), weight), --enforce-eager). The serving engine is not bit-deterministic here (two reference runs differ), so the comparison is against that run-to-run noise:

Run	Mean NLL	Mean |Δlogprob| vs reference run 1
Reference, run 1 / run 2	0.682 / 0.683	— / 0.053
This PR (CUDA graphs on), run 1 / run 2	0.684 / 0.682	0.051 / 0.051
Control: staged rows left at zero	0.926	0.394

Performance: vllm bench serve, random dataset, --ignore-eos, one run per point, total output tok/s / mean TPOT ms / mean TTFT ms:

Concurrency	input 1024 / output 1024	input 4096 / output 512
1	122.6 / 8.0 / 148	114.8 / 8.1 / 330
8	638.3 / 12.0 / 584	486.4 / 13.0 / 1,756
32	1,293.2 / 23.5 / 1,317	883.3 / 30.8 / 2,761

There is no baseline on this hardware: without this PR the model does not start (the table does not fit on the GPU, and the UVA offload refuses to run).

Cost of the staging itself (gather + copies only, 16 heads × 160 dims BF16, median of 20 on the same VM): 0.14 ms for 32 tokens, 12–19 ms for an 8,192-token prefill chunk. The H2D runs at about 4.1 GB/s, the bounce-buffer limit under CC.

Commands

VLLM_USE_V2_MODEL_RUNNER=1 vllm serve VnimanieAI/Qwen3.8-Flash-Next-W4A16 \
  --trust-remote-code --reasoning-parser qwen3 \
  --max-model-len 65536 --max-num-seqs 64 --gpu-memory-utilization 0.92 \
  --moe-backend marlin --language-model-only

vllm bench serve --backend openai --endpoint /v1/completions --dataset-name random \
  --random-input-len 1024 --random-output-len 1024 --ignore-eos \
  --num-prompts $((3 * C)) --max-concurrency $C   # min 8 prompts

# prompt logprobs: /v1/completions with max_tokens=1, prompt_logprobs=0 and a fresh cache_salt per request

Two local workarounds unrelated to this PR were needed for this checkpoint on main (noted in #62815): the fused q/k/v + indexer projection (#59533) does not load a checkpoint whose indexer is BF16, and the MoE HUMMING conversion OOMs on a 94 GB GPU (--moe-backend marlin).

Details

Design. PLE lookups depend only on token IDs, so the rows a step needs are known before its forward. #59958 stages them in Qwen4ExpModelState.prepare_inputs: n-gram IDs to the host, one stream sync, distinct rows read per storage plane, one H2D per plane into fixed device buffers, NVFP4 decoded on the GPU. This PR moves that into Qwen4ExpPLEHostStagedEmbedding with one abstract _read_rows(name, ids, rows). Qwen4ExpPLEDiskOffloadEmbedding keeps its preadv reads there. Qwen4ExpPLEHostMemoryEmbedding gathers from the ETP-sharded host table with index_select on raw bytes, zeroes rows owned by other ETP ranks, and reduces after the lookup like the pinned-host table. stage_checkpoint_rows becomes stage_host_rows.

Tradeoff: one host sync per step on this path. The IDs hash the tokens sampled by the previous step on the GPU, so the read-back waits for that step to finish, which removes some async-scheduling overlap. This is the same sync disk offload already makes. Only the fallback pays it; the UVA path is unchanged. I could not measure the overlap loss on CC, since there is no UVA path to compare against there. VLLM_GPU_SYNC_CHECK=error could not be exercised on CC either: it first stops in #59981's staged input copy, because pinned host memory reports is_pinned() == False under CC.

Limitations.

• NVIDIA path only. The AMD Qwen4Exp path still requires UVA.
• ETP spanning DP ranks (embedding_across_dp) is rejected with NotImplementedError. TP/ETP within one DP rank is supported (rows owned by other ranks are zeroed, then reduced). Unlike disk offload, each rank keeps only its own rows in host memory.
• DeepSeek V4.1 Engram is not covered.
• The NVIDIA Qwen4Exp model never selects this path on stock main unless UVA is unavailable (for example with [Perf] V2 runner decode optimizations for NVIDIA Confidential Computing #59981 under CC, or where pinned memory is unavailable).

CI. The new tests live in tests/models/qwen4_exp/test_ple.py and test_config.py, which the existing Qwen4 Exp lanes run. The ETP owner-masking test also runs on CPU. CC hardware is not available in CI; the end-to-end results above were collected outside CI.

---

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.

🤖 Generated with Claude Code

[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.

[mergify]

Documentation preview: https://vllm--62847.org.readthedocs.build/en/62847/

[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.

🚀

[yeighta]

@gau-nernst @DarkLight1337 @ywang96 @AndreasKaratzas (auto-requested reviewers) could one of you add the verified label so pre-commit and CI can run? This is my first contribution, so I can't trigger CI myself.

The change only affects Qwen4Exp when cpu_offload is on and UVA is unavailable; the UVA path is unchanged. Validation on an H100 in Confidential Computing mode is in the description. Thanks!

[ZJY0516]

please wait for #59981

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @yeighta.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork