[Offloader] Prefetch weight offloading for large MoE models, plus a schedule planner

[lucifer1004]

Purpose

Extends the prefetch weight offloader (#29941) so DeepSeek-V4 class MoE models
run on hardware their expert weights do not fit on, and adds a way to choose an
offload schedule without launching one server per candidate.

Applying #29941 to DSv4 exposed three gaps: expert weights are heterogeneous
after quantization, H2D copies contend with tensor-parallel collectives on the
same links, and nothing reported which units were selected or how many bytes
each moved.

Scope. The mechanism is model-agnostic: units are discovered through
named_modules/named_parameters and selected by parameter-name segment, with
no model-class dispatch and no architecture-specific branches. The selectors
cover fused (experts.w13_weight/w2_weight) and unfused
(experts.gate_up_proj/down_proj) MoE layouts as well as feed_forward.w*
and attention. DeepSeek-V4 is what it was developed and measured against, not
what it is limited to — but it is the only architecture I have evidence for,
so treat other models as untested rather than unsupported.

Commit 1 is @Tom-Zheng's, submitted on his behalf with his authorization and
unmodified. Commits 2–3 are mine.

Not a duplicate

Checked weight offloading, prefetch offload, offload selectors,
offload NUMA, prefetch slab pinned, offloader diagnostics manifest,
offload planner, weight offload configuration.

• [Perf] Batch Weight Prefetching via cuMemcpyBatchAsync to Reduce Latency #41474 (batch prefetch via cuMemcpyBatchAsync) is closest: it changes
  the copy submission API, this changes what is copied and when. Intents
  compose and I have not duplicated its swap_blocks_batch change, but both
  touch prefetch.py and will conflict textually — happy to rebase either way.
• [Config] Migrate weight-offloading env vars to OffloadConfig #46337 (offload env vars → OffloadConfig) is a June draft; that
  migration landed independently, so commit 2 follows the existing convention.
• [Bugfix] Align MLA indexer block table with MTP speculative decode #43970 and [Bugfix][Kernel] Fix persistent top-k histogram reuse after short rows #49139 touch neighbouring DSv4 code, not the offloader.
  Hybrid KV offload: planner, MultiConnector, and mamba alignment for hybrid models #38261 is a KV offload planner — a different subsystem.

What is in it

1. Prefetch offloading for DSv4-Flash (@Tom-Zheng) — semantic parameter
selectors (routed_experts, shared_experts, dense_mlp, attention),
pooled slab and storage-group runtime buffers with a direct fallback for
layouts that cannot be pooled, a planner resolving which ordered module
positions become offload units, NUMA-aware process binding, and a read-only
post-init manifest describing the resulting schedule.

2. Bounded slabs and collective-aware pacing — the slab allocator could
request a pinned host block per layout with no upper bound; now bounded. H2D
copies are chunked and gated while a collective is in flight, behind
--offload-comm-aware, so an unpaced copy cannot delay a collective every rank
is blocked on. The offloader carries the flag and exposes gates_collectives,
so low-level distributed code reads a boolean instead of vllm.envs and the
common no-op case costs no context manager. Pacing engages only at TP > 1.

Also moves the offloader tests from a new top-level tests/weight_offload/
into tests/model_executor/offloader/. Not cosmetic: the old location was in
no CI area's file list and outside the mypy hook's scope, so those tests were
neither run by CI nor type-checked. The Model Executor area now also triggers
on vllm/config/offload.py.

3. A schedule planner — python -m vllm.benchmarks.weight_offload.planner

Choosing offload_group_size / offload_num_in_group /
offload_prefetch_step is guesswork, and a schedule that does not fit says so
only by failing to start. Two exact quantities decide viability — bytes that
stay resident, and bytes re-copied H2D each forward — but both need
per-position sizes that only a running engine knows, which is what commit 1's
manifest reports. Serve once with VLLM_PREFETCH_LOG_SCHEDULE=1, then score
every other schedule offline.

Results are deltas against that run, so schedule-independent terms (KV,
non-offloadable weights, allocator slack, graph reserve) cancel and the
operator enters nothing. It models memory, not latency, and says so: no
bandwidth setting, because bytes over bandwidth gives the PCIe time a schedule
occupies, not whether that time is hidden. This is also why the manifest lands
with it — it is the only artifact that makes the identities checkable.

Evidence

pytest tests/model_executor/offloader   # 123 passed, 1 skipped
pytest tests/benchmarks/weight_offload  # 10 passed
pre-commit run --from-ref origin/main --to-ref HEAD

All clean on main @ 3a79957b62; the three commits cherry-pick onto it
without conflicts.

Serving, TP8

DeepSeek-V4-Pro, TP8, RTX PRO 6000, offload_group_size=3 offload_num_in_group=1 offload_prefetch_step=3 offload_params=experts,
gpu_memory_utilization=0.90, 8000-token prompts with one output token at
concurrency 1/2/4/8, cold prefix cache. Each arm is a separate server with the
torch.compile cache disabled, so no arm inherits another's compiled graphs.
Every flag in the baseline case already parses on 3a79957b62.

arm	tree	TTFT @1 (ms)	steady req/s
A	main @ 3a79957b62	does not reach readiness	—
B	+ commit 1	1956.17	0.5165
C	+ commits 1–3, pacing off	1955.87	0.5165
C	+ commits 1–3, pacing on	1766.73 / 1764.43	0.5725 / 0.5726

• A cannot start. Weights load, then CUDA graph capture aborts under
  determine_available_memory with cudaErrorStreamCaptureUnjoined at
  breakable_cudagraph.py capture_end, on six of eight ranks — precisely
  what commit 1's join_after_forward hook prevents. Commit 1 is a
  prerequisite for this configuration, not a speedup on a working one.
• Pacing off is free: B and C differ by 0.02% at every concurrency, two
  independent instances on different trees.
• Pacing on is worth ~10%: −9.7% TTFT and +10.9% throughput against the
  mean of the pacing-off arms, uniform across the sweep, at ~75× the observed
  run-to-run spread (≤0.13%).

KV holds 10,171 tokens here, barely more than one 8000-token request, so
concurrency above 1 measures queueing — TTFT @1 and steady throughput are the
load-bearing numbers.

Memory and transfer identities

Same model and schedule. Byte counts from the post-init manifest, transfer
figures from an Nsight Systems trace over explicit GPU-time windows. Measured
2026-08-04 on this code before the rebase onto current main; the rebase touched
no selector or buffer arithmetic.

quantity	value
ordered candidate positions	61
runtime units selected (predicted / observed)	20 / 20
offloaded bytes per rank	33,690,746,880 (31.38 GiB)
runtime buffer bytes per rank	5,053,612,032 (4.71 GiB)
net GPU bytes freed per rank	28,637,134,848 (26.67 GiB)
steady-request H2D per rank	33.6911 GB, 0.001% from the analytical identity
per-rank H2D bandwidth	24.68 – 27.28 GB/s
copy union overlapped by compute	≥ 99.813%
steady copy-only exposure	2.253 – 2.305 ms/rank

Manifest against the analytical model: −46,880 B offloaded and −12,032 B
runtime buffer, both inside the log rounding interval. Reading that same
manifest, the planner reproduces all of it — 20 units, identical byte counts,
zero delta scoring the recorded schedule against itself — and reduces 10,446
enumerated schedules to 331 distinct outcomes.

Caveats: per-position compute was not measured (aggregate divided uniformly
gives 40.36 ms/position, a proxy). Synchronous TP is bounded by the slowest
participant, 24.68 GB/s, not the 25.98 GB/s mean. These figures are
instrumented and the latency table is not — separate measurements.

One interaction worth knowing

At gpu_memory_utilization=0.952 both B and C reach readiness and then die
identically on the first request: FlashInfer's CUTLASS fused-MoE workspace
requests 1.32 GiB from FusedMoeRunner::getWorkspaceInfo after memory
profiling has already passed. Offloading converts weight memory into headroom
that gets respent on KV, so a workspace allocated outside vLLM's profiling
becomes marginal exactly when offloading is doing its job. Not a defect in
these commits and the fix is headroom rather than code, but it is the first
failure an operator will hit.

AI assistance

Produced with AI assistance. Every changed line has been reviewed by the
submitter, and the tests and measurements above were executed on the hardware
described, with results reported as measured.

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

[mergify]

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

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

[Big2Wheel]

I found #51710 while doing the final duplicate/overlap check for #54975. The focused PR fixes the same non-divisible static-buffer slot reuse problem that your new PrefetchRuntimeController._build_prefetch_after() also handles.

The focused fix was independently derived from vllm-project/vllm-ascend#10932 and validated with real Qwen3-0.6B weights on Ascend 910B3 in eager and ACLGraph modes. Since #51710 is a much broader draft, #54975 isolates only this correctness fix plus six small ring-shape regression cases so it can potentially land first.

I called out the overlap explicitly in the PR description. If #54975 lands first, the corresponding scheduling portion here can be rebased onto it or dropped while retaining the broader runtime-controller design. Please let me know if you would prefer a different coordination path.

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