[wyj-IIRtyj]

Overview

Three independent follow-ups to vllm-project#56177 (on top of its head 3e13f466), from serving GLM-5.3-Flash NVFP4 with the expert pool at 1M context on 2x RTX PRO 6000 (TP=2). Each commit stands alone and can be cherry-picked separately.

Claims

• Speculative decoding (MTP) works with the pool (796df0733). MoE layers that do not use the NVFP4 method, such as the BF16 MTP head of an NVFP4 checkpoint, now keep their experts on the GPU (one warning) instead of failing the layer guard (... got UnquantizedFusedMoEMethod). The decode lanes are sized for 1 + num_speculative_tokens tokens per request; previously, for example, --max-num-seqs 1 with 3 speculative tokens sent every verify step (4 tokens) down the partition path.
• Banks larger than Marlin's 32-bit weight index are rejected at installation (d2a05d42b). Marlin MoE adds the int64 expert offset of its weight operand into int b_gl_rd, which is counted in 16-byte units (csrc/moe/marlin_moe_wna16/marlin_template.h L698/L706, unchanged on main). Any bank weight tensor above 32 GiB is therefore read out of bounds. --moe-expert-pool-rows 240 on GLM-5.3-Flash TP=2 (over 10,000 bank rows of 4 MiB w13, about 39.5 GiB) died with an illegal memory access in profile_run. A direct kernel check, shown below, puts the bound exactly at 8192 rows of 4 MiB. Such a configuration now fails at installation with use --moe-expert-pool-rows <= 193.
• Prefill staging (25bbec966). TTFT at 81.6K input tokens drops from 67.6 s to 13.0 s and at 12.5K from 9.6 s to 1.8 s, with decode unchanged. On by default; --no-moe-expert-pool-prefill-staging turns it off.

Validation

Setup: nvidia/GLM-5.3-Flash-NVFP4, 2x RTX PRO 6000 Blackwell (SM120, PCIe, GPU-GPU P2P), TP=2, --moe-backend marlin --moe-expert-pool-rows 193 --kv-cache-dtype fp8 --max-model-len 1048576 --max-num-seqs 4, FULL_DECODE_ONLY, MTP with 3 speculative tokens ("moe_backend": "triton" for the BF16 draft). Measurements use single requests, synthetic prompts and a streaming client.

The end-to-end numbers come from vLLM 0.32.0 with the pool applied. The expert_pool/ package and its tests are byte-identical to 3e13f466, and the hook files were ported. That build ran the same staging and MTP logic, but enabled staging through an environment variable rather than the new flag, and had no bank-size check. They were not measured on a wheel built from this branch.

Prefill path	--max-num-batched-tokens	TTFT 12.5K	TTFT 81.6K	GPU power in prefill
host view (3e13f466)	4096	9.5-9.8 s	67.6 s	~180 W
staging (this PR)	8192 [1]	1.81-1.83 s	13.0 s	~390 W
staging (this PR) + custom all-reduce	9216	1.80-1.84 s	10.3 s	~470 W

[1] On this hybrid model, align-mode prefix caching clips non-final prefill chunks to multiples of the 4608-token mamba block, so 8192 actually runs 4608-token steps.

• Root cause, from nvidia-smi dmon: each 4096-token step pulled ~141 GB over PCIe for a ~28 GB non-resident set. Marlin reads its B operand once per M-block, so every non-resident expert crossed PCIe about 5 times.
• Staging with synchronous copies (no prefetch) measured 3.60 s / 18.9 s at 4096 and 2.39-2.47 s / 14.7 s at 8192 [1]. At the same step size, prefetching the next layer gives the 1.81-1.83 s / 13.0 s above. Most of the gain comes from staging itself; prefetch cuts TTFT by a further 12-25%.
• Decode (tok/s, 3 runs x 512 tokens each; short / Chinese prose / 10K context): host view 104.7 / 107.8 / 102.0, staging 107.7 / 111.9 / 109.0. These are within run-to-run variation (about ±5%).
• MTP 3 since the last restart: 61% of draft tokens accepted (83% / 60% / 41% by position), 2.84 tokens per verify step, 105-114 tok/s single stream. No MTP-off baseline was measured on this setup.
• Smoke checks on every build: arithmetic and a tool call (get_weather) through the OpenAI API.

Tests, on one RTX PRO 6000 (SM120). Same environment boundary as in vllm-project#56177: the Python sources of this branch over the vLLM 0.32.0 wheel. The expert_pool/ package is byte-identical to this branch, and the hook files are ported.

pytest -v tests/config/test_moe_expert_pool_rows.py tests/kernels/expert_pool
43 passed in 24.65s   (39 existing + 4 new)

The new staging test runs the real Marlin kernel. It asserts that staging happened (two layers staged, eight rows copied, resident rows untouched) and that the output is byte-identical to the host-view read.

Marlin bound, checked directly: one NVFP4 expert (K=4096, N=1024, a 4 MiB w13 row) is placed at row 0 and at the last row of rows-row weight tensors, and the same 8 tokens are routed to each copy through the pool's align-and-remap path. Routing to an all-zero row instead changes the output, which confirms the physical row id is honored.

rows	w13 tensor	last row offset	result
8192	32.000 GiB	0.9999 x 2^31 x 16 B	identical to row 0 (max abs diff 0)
8193	32.004 GiB	2^31 x 16 B	CUDA error: an illegal memory access was encountered

pre-commit run --from-ref 3e13f466 --to-ref HEAD: all hooks pass (ruff, typos, markdownlint, mypy, SPDX, config docstrings).

Details

Prefill staging (stage.py).

• Applies only to uncaptured partition steps of at least 512 tokens. Decode and every captured step are unchanged.
• Two device buffers, each the size of a whole layer, alternate by layer. They are full-layer because, under the global LRU, a layer can have any number of non-resident rows.
• The tables only change in decode steps. So at the first layer of a prefill step, every layer's non-resident map is read back once and turned into contiguous runs.
• Rows are copied by DMA (copy_ from the pinned source) on a side stream: layer i+1 is copied while layer i computes, ordered with events and wait_stream. A Triton gather kernel reached the same ~52 GB/s, but it competes with the expert GEMM for SMs.
• Marlin itself is unchanged: the staging buffer simply replaces the host view as the second partition's tensors.

Tradeoffs:

• Memory. Two full layers of expert weights (2 x 1.9 GiB per rank here). They are allocated before memory profiling, so the KV cache shrinks accordingly; here that is about 0.45M of 1.8M KV tokens.
• One host sync per prefill step, for the placement readback. VLLM_GPU_SYNC_CHECK=error would flag it. Decode is unaffected.
• PCIe cost per step remains. Each prefill step still moves the non-resident set over PCIe once. Long inputs therefore benefit from larger chunks; see [1] for hybrid models.

MTP. The layer guard now rejects only NVFP4 layers on a non-Marlin backend. Any other MoE layer is left out of the pool, which also covers FP8 draft heads. A model where no layer is pooled therefore loads with all experts on the GPU, behind a single warning.

Bank-size guard. Only the weight operands are checked. The scale and global-scale pointers are advanced per expert in 64-bit arithmetic, so they are not affected. A kernel-side fix would advance B the same way (pointer per expert instead of an int offset). It is not included here because it needs a rebuilt _moe_C, and on main no single MoE layer comes near 32 GiB.

Related work. These commits target vllm-project#56177's branch rather than main because they depend on it.

• Marlin offset overflow: I searched open and closed PRs and issues (b_gl_rd, B_expert_off, int32/32-bit index) and found none.
• [MoE] Stage UVA-offloaded expert weights in VRAM for large steps vllm-project/vllm#62060 stages expert weights in VRAM for large steps on the --offload-backend uva path. It copies right before the expert kernel and, as its description notes, stages only weights carrying the UVA marker. It therefore does not cover the pool's host view.
• [MoE][Lab] PREFILL_STAGE_ROWS: stage routed cold experts on device for multi-token forwards #38 staged the routed cold experts for prefill on the lab native chain. DMA copies ([MoE][Lab] PREFILL_STAGE_COPY=memcpy: DMA copies from the pinned host bank into the scratch #39) and a double buffer with next-layer prefetch ([MoE][Lab] PREFILL_STAGE_BUFFERS=2: double-buffered cold staging with next-layer prefetch #40) were measured slower there and not adopted. Here both help. The likely differences: this stage copies a layer's whole non-resident set, so consecutive rows coalesce into runs rather than one copy per routed row, and a 4608-token Marlin step runs long enough to hide the next layer's copy.
• Speculative decoding was handled in the lab line ([MoE][Lab] Pool decode path for speculative verify rows (SPEC_ROWS, ngram admitted) #27, [MoE][Lab] Draft-model load scope: keep draft models outside the expert tier #31). On [MoE][Offload] Shared GPU expert pool with a device-side planner for NVFP4 Marlin experts (moe_expert_pool_rows) vllm-project/vllm#56177 the layer guard still rejects the BF16 MTP head; the first commit is the minimal form for this branch.

---

AI assistance (Claude Code) was used to prepare these changes. I reviewed every changed line and ran the tests above.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.