[Draft][MoE][Offload] Expert pool: shared GPU bank of NVFP4 Marlin expert rows with a device-side planner (moe_expert_pool_rows)

[01554]

Draft for team review (author claude, review 1 mid, GPU/review 2 astra). Fork-only; nothing is proposed upstream from here without the user's go.

What

Opt-in expert pool for MoE models whose experts do not fit the GPU: --moe-expert-pool-rows N (OffloadConfig.moe_expert_pool_rows, default 0 = unchanged behaviour).

• Expert tensors (NVFP4 weights and block scales) are allocated in pinned host memory; the loader still moves one layer at a time to the device for the Marlin conversion and restores the converted tensors to pinned memory, which becomes the pool's host source.
• After every layer's process_weights_after_loading, install_expert_pool allocates one VRAM bank shared by all MoE layers (N rows per layer resident at start, plus top_k staging rows), keyed by (layer, expert).
• Per forward a device-side step program (LRU over the bank, promotion policy in device controls) plans hits, misses and staging; row copies are fixed-grid kernels from the host source's accelerator view; the Marlin consumer runs on the bank with logical alignment (moe_align_block_size(..., ignore_invalid_experts=True)) and a physical-row remap. The device-planned decode path runs no host code between the router and the Marlin call, so the MoE op stays inside CUDA graphs (FULL_DECODE_ONLY works); the Python dispatch around it and the prefill/partition path are ordinary eager code.
• Batches wider than the planner's lanes take a bank + host-view partition path.
• The placement is frozen (gate closed) through profiling and graph capture and opened at the end of compile_or_warm_up_model after a table check and a sampled bank/host comparison.

Supported: ModelOpt NVFP4 with the Marlin MoE backend, no EP/DP/SP; rejected at layer construction otherwise.

Base and dependencies

Base lab/upstream-base-20260909 = upstream main a97dacb. No unmerged dependency (the earlier pool head on PR vllm-project#37190 is replaced by this standalone form; the host cache providers, breakable-graph break points, loader marker and chunked host repack are dropped, see the lab ledger notes/upstream-pool-standalone-ledger-2026-09-09.md).

Tests

• Exact suite on the integrated head 7dedc6d (this branch at 5fbc240 + fork pull/46 at 4f859de) on the GPU host: 49 passed + 5 subtests, 0 skipped (results/goal66-minimal-7dedc6d8-exact-20260909). This includes tests/kernels/expert_pool/test_pool_marlin_cuda.py (two layers sharing one bank through the real Marlin consumer, with the host-source fixture run under the loader's device_loading_context), the guard, shared-experts and config tests.
• CPU locally: test_pool_tables.py + test_pool_layer_helpers.py, 18 passed. pre-commit (ruff, format, mypy 3.10, config docstring check) passed.
• Environment as actually used on the host: the a97 base wheel installed with --no-deps, this head's Python sources overlaid, and the _ple_memops extension built in addition; the candidate head is not rebuilt as a full wheel per run.

Measurements

Purpose: run Qwen3.8-Flash-Next NVFP4 (48 layers x 512 experts) on an RTX 6000 Ada (48 GB). Development and measurement used an RTX PRO 6000 Blackwell Max-Q (96 GiB) with a separate process holding GPU memory so that vLLM had 48 GiB available; the container's host memory limit was 100 GiB.

Code measured

configuration	code	commits
baseline: existing UVA CPU offload	upstream main + the PLE mmap PR vllm-project#54129 (needed to load this model's PLE table from disk)	main 1970f3ed4, vllm-project#54129 branch 50a061f7
proposed: expert pool	upstream main + this PR + fork pull/46 (deferred PLE rows, on vllm-project#54129)	main a97dacb7, this PR 5fbc240ba, pull/46 4f859de9d; combined verification commit 7dedc6d8d

The combined commit identifies the code that was measured; it is not the head of either PR. The server ran the candidate's Python sources over a wheel built from a97dacb7, plus the _ple_memops extension built separately.

Speed (max-model-len 4096, one request at a time)

Two frozen prompts are sent in order to a freshly started server: a warmup request (1070 input tokens) and the measured request (753 input tokens); both ask for a fix proposal and a verification plan for a bug report. temperature=0, top_p=1, seed=0, max_tokens=2048, thinking disabled. Decode speed = (completion_tokens - 1) / (time between the first and last text event), measured by the client. Prompts, the client, and the exact launch commands are in benchmarks/expert_pool/ on this branch.

configuration	expert placement	execution	decode speed of the measured request	samples
baseline UVA offload	--offload-backend uva --cpu-offload-gb 40 --cpu-offload-params experts	--enforce-eager, PLE mmap	7.42 tok/s	1 (after 1 warmup)
proposed	--moe-expert-pool-rows 258 (about 32 GiB bank)	cudagraph_mode=FULL_DECODE_ONLY, PLE mmap + deferred rows (VLLM_PLE_MMAP_DEFERRED=1)	63.19 tok/s	median of 3 fresh server launches: 63.19 / 62.71 / 63.65

Both rows differ in source version and execution mode, as listed. Common launch arguments:

--model /model --served-model-name flashnext --tensor-parallel-size 1
--quantization modelopt --dtype bfloat16 --moe-backend marlin
--language-model-only --max-model-len 4096 --max-num-seqs 1
--max-num-batched-tokens 512 --no-enable-flashinfer-autotune
--gpu-memory-utilization 0.4548806288994517 --safetensors-load-strategy lazy
--reasoning-parser qwen3 --generation-config vllm

Common environment: VLLM_USE_V2_MODEL_RUNNER=1 VLLM_PLE_MMAP=1 VLLM_PLE_MMAP_PREWARM=0 VLLM_PLE_MMAP_PINNED=0. The proposed row adds --compilation-config '{"cudagraph_mode":"FULL_DECODE_ONLY"}' --default-chat-template-kwargs '{"enable_thinking":false}'.

The acceptance line for this work was 62.985 tok/s (our internal reference build's 66.30 minus 5%); the median above is a check against that line, not a tuned number.

Quality (max-model-len 32768, separate launch, proposed configuration)

• SWELancer, two tasks with the official grader: both passed (500/500 and 1000/1000). One task stopped on its own; the other reached the 300 s agent budget and was graded on its committed patch.
• Six long-output prompts (design document, JSON array of 120 objects, module + tests, refactor, tool-call sequence, multilingual), temperature 0: no repetition, degeneration or restart observed. Three hit the output token cap (the JSON array ends inside object 64, the generated tests end mid-method, the design document lost only an appendix). The refactor completed on its own but calls an undefined repository method. These checks establish that the served model does not degrade under this path; they do not certify the functional correctness of the generated content.

Long inputs (separate capacity configuration: --moe-expert-pool-rows 208, --max-model-len 262144)

At N=258 the 262144 context does not fit (KV needs 6.55 GiB, 1.24 GiB free), so the long-input runs use N=208 (about 6.2 GiB moved from the bank to KV); everything else as above. Each length is one request: a document with five key/value facts at 5/25/50/75/95 %, exact input length built from token ids, max_tokens=512, answer requested as JSON. Client-observed; prefill = input tokens / TTFT.

input tokens	TTFT	prefill	decode after prefill	5 values correct
4096	50.35 s	81 tok/s	42.2 tok/s	5/5
65536	614.87 s	107 tok/s	64.9 tok/s	5/5
131072	1250.71 s	105 tok/s	70.7 tok/s	5/5
261632	3395.10 s	77 tok/s	42.5 tok/s	5/5

The 4K/64K/128K requests were sent in order to one server (prefix caching may apply); 261632 was a fresh single request. Results: results/goal66-long-context-262k-n208-20260909 and -resumed-20260909 in the lab repo; client scripts/long_context_pool_check.py there.

For reference, the same frozen inputs on unmodified FreeToken af71ba4 (its own offload backend, resident expert bytes matched to N=208, KV for 262656 tokens, prefill chunk 512, single runs): TTFT 54.85 / 720.58 / 1441.27 / 2884.37 s, i.e. prefill 75 / 91 / 91 / 91 tok/s, decode 39.4 / 67.3 / 71.6 / 38.3 tok/s, 5/5 at every length. Up to 128K the pool's TTFT is shorter; at 261632 it is 17.7 % longer. Per-chunk cost is the candidate reading (both engines move expert bytes over PCIe per 512-token chunk); no transfer volume was measured and the cause is not isolated. Not tested on either engine: long inputs combined with long outputs.

Capacity

29 startup + 243 in-flight samples with no errors under the 48 GiB limit, minimum free 4449 / 4295 MiB, final exit 0, container OOMKilled=false. Retained observations: one allocator OOM retry during startup, and an EngineDeadError with a held semaphore after the external shutdown.

🤖 Generated with Claude Code

https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT

[01554]

Review 1 — static extraction review of 69fd13b (base a97dacb), with integrated PLE head 9a67b70.

No required changes found in the reviewed extraction diff. Relative to the accepted bc8 configuration, the expert_pool package changes only the installation attribute names and documentation; planner, copy and consumer code are unchanged. The integrated PLE model and runner paths are unchanged.

Checked the default-off pool knob, nonnegative validation, hash inclusion, CLI and engine-config pass-through, early Marlin/EP/DP/SP restrictions, pending installation guard, shared-expert overlap behavior, and post-warm-up gate boundary. The standard loader runs per-layer finalization inside device_loading_context, restores CPU parameters, then installs the pool after all layers. Removing the chunked host repack path is consistent with that standard loader path; this is a source-reading conclusion, not a claim about every loader. Pool-owned pinned source normalization and lifetime references remain. Removing the loader marker/helper restores upstream behavior; the final load peak and temporary Parameter lifetime still need verification on this exact assembly.

Optional follow-ups: directly exercise the new early DP/SP and rows < top_k guards in CPU tests. The CLI test currently constructs OffloadConfig manually, so it should not be described as executing create_engine_config (the production pass-through was checked statically).

I have not executed full-environment tests for this head. Exact regression tests and final 48 GiB capacity, fresh-three-run speed, and 32K quality acceptance remain pending. Earlier bc8 acceptance does not establish acceptance of this minimal assembly. Keep Draft pending both reviewers and final validation.

[01554]

Review 1 follow-up: checked e5cc8d6 against 69fd13b. Only two test files changed; production code is unchanged. The added tests directly exercise the supported early-guard case and six rejection cases, and the CLI test now correctly states that create_engine_config is not executed. No required changes found; the static review carries forward to this head. Execution results and final integrated acceptance remain pending.

[01554]

Review 1 correction after exact-environment execution: independently checked the 9a67 test log (40 passed, 2 failed). The migrated CUDA fixture still calls process_weights_after_loading directly with CPU expert tensors, relying on the host-repack path removed by this extraction. I missed this fixture dependency in the initial static review. It must use the standard device_loading_context around finalization, verify device processing and CPU/pinned restoration, then reach install and the pool oracle. The CLI test also needs a local model path to avoid unrelated offline snapshot lookup. Both issues remain in the e5cc8 test version. These are required test fixes before acceptance; the pool oracle has not been reached in this run, and these failures do not establish a production runtime defect. Author and GPU reviewer have been notified.