[MoE][Offload] Shared GPU expert pool with a device-side planner for NVFP4 Marlin experts (moe_expert_pool_rows)

[01554]

Purpose

Serve MoE models whose expert weights do not fit in GPU memory with a GPU-resident, device-managed expert pool, so that decode stays inside CUDA graphs. Discussed in #56175.

--moe-expert-pool-rows N (OffloadConfig.moe_expert_pool_rows, default 0 = unchanged behaviour):

• NVFP4 expert tensors are allocated in pinned host memory at load time; the loader's existing per-layer device round trip does the Marlin conversion and restores the converted tensors to pinned memory.
• After loading, one GPU bank shared by all MoE layers is allocated (initially N rows per layer) and a Marlin consumer is bound to it; the placement is frozen through profiling and graph capture and opened at the end of warm-up.
• The captured decode path runs a device-side LRU step (hits/misses, promotions across layers), fixed-grid row copies from the host source, and one Marlin call with logical alignment and a physical-row remap. No host-side routing readback or cache planning between routing and the GEMM, so the MoE layers stay inside the CUDA graph (FULL_DECODE_ONLY).
• Wider batches take a bank + host-view partition path.

Scope: ModelOpt NVFP4 with the Marlin MoE backend, no EP/DP/SP; backend and parallelism are checked at layer construction, cross-layer shape/top-k consistency at pool installation.

Files: vllm/model_executor/layers/fused_moe/expert_pool/ (tables, copy, pool, layer, install), hooks in routed_experts.py, modelopt.py, model_loader/utils.py, gpu_worker.py, config/CLI knob, docs/features/moe_expert_pool.md, benchmarks/expert_pool/ (reproducible client and frozen prompts).

Test Plan

pytest -v -s tests/kernels/expert_pool tests/config/test_moe_expert_pool_rows.py
python -m unittest discover -s benchmarks/expert_pool -p 'test_*.py'

tests/kernels/expert_pool/test_pool_marlin_cuda.py runs two layers through one bank with the real Marlin kernel (decode with cross-layer eviction, a wide batch through the partition path) and compares against the uncached layers. The other tests are CPU-only (planner tables, copy helpers, installation guards, config/CLI). CI: added to the MoE Kernels shard.

End to end: Qwen3.8-Flash-Next NVFP4 on a 48 GiB GPU budget (RTX PRO 6000 limited by a separate memory-holding process, host 100 GiB), see benchmarks/expert_pool/README.md for the exact launch.

Test Result

• Unit/kernel tests on this branch (production code at feba2ca, identical to the head except one docs commit): pytest -v -s tests/kernels/expert_pool tests/config/test_moe_expert_pool_rows.py -> 40 passed, 0 skipped (12.65 s); unittest discover -s benchmarks/expert_pool -> 2 passed. Environment boundary: run as the Python sources of this branch over a wheel built from base a97dacb (no C++ changes between a97dacb and the rebase base 285cbce in csrc/, CMakeLists.txt, _custom_ops.py; requirements/cuda.txt moved flashinfer 0.6.18 -> 0.6.18.post1, which the a97 wheel environment does not reflect). Not done for this head: a full wheel build at 285cbce, the Buildkite CI run (it will run on this PR), and a re-measurement of the speed numbers below. Previous head 5fbc240 on base a97dacb: 49 passed + 5 subtests on the GPU host, pre-commit incl. mypy passed.
• Decode, 4096 context, one request, same frozen prompts (client-observed, prefill excluded):
  • existing --offload-backend uva --cpu-offload-gb 40 --cpu-offload-params experts, eager: 7.4 tok/s (1 run)
  • this PR, --moe-expert-pool-rows 258 (about a 32 GiB bank), FULL_DECODE_ONLY, with deferred PLE rows (Add opt-in single-token deferred mmap PLE rows for FULL graphs 01554/vllm#46 on [Model] Support disk-backed (mmap) PLE table for Qwen3.8-Flash-Next (VLLM_PLE_MMAP) #54129): 63.2 tok/s, median of 3 fresh server launches (63.2 / 62.7 / 63.7)
• Quality, 32768 context: two SWELancer tasks pass the official grader; six long-output prompts show no repetition, degeneration or restarts before stopping or reaching the output cap.
• Long inputs, --max-model-len 262144 with --moe-expert-pool-rows 208: 4K / 64K / 128K / 256K-token documents with five key/value facts, 5/5 retrieved at every length; prefill 81 / 107 / 105 / 77 tok/s (input tokens / TTFT, one run each).

Numbers above were measured on the integrated head 7dedc6d8d = this PR at 5fbc240 + PLE mmap support (#54129, needed to load this model's PLE table within the budget) + deferred PLE rows (01554#46 at 4f859de, part of the measured configuration, not required for loading). The pool itself depends on neither.

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[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--56177.org.readthedocs.build/en/56177/

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

🚀

[mergify]

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

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