[Draft][MoE][Offload] Row-level expert weight cache provider (moe_expert_cache_provider=row)

[01554]

Upstream extraction, unit 3 (notes/upstream-expert-provider-contract-draft-2026-09-09.md). Draft; two reviews required before ready (astra_mid: review 1, astra: review 2 incl. GPU).

Base: lab/upstream-37190-base d38903c = upstream/main a97dacb + PR vllm-project#37190 c66868c (notes/upstream-37190-base-merge-2026-09-09.md). Head: 60ba5b9 (last GPU-verified production state: b5ac54e; 60ba5b9 adds docstrings and moves one log line). Dependencies: vllm-project#37190 (provider surface). The native consumer path (not in this PR) additionally depends on #44.

Commits on this head

• Knob: OffloadConfig.moe_expert_cache_provider (cached default, row) wired EngineArgs -> CLI -> config, hash-covered (tests).

• RowCacheWeightProvider: CachedWeightProvider surface plus NVFP4 global-scale slot buffers, owner-stream contract, forward-scoped map, LRU with in-forward victim protection, staged copies on a provider-owned copy stream (release event on the owner stream -> copy stream waits -> six-tensor copies -> ready event -> owner waits), invalidate ordered against in-flight copies. Expert maps are uploaded as one pinned non-blocking copy each; per-element device writes stalled the host until the owner stream drained (d26e289).

• (b1) CachedWeightProvider: optional keyword-only w13_scale_2/w2_scale_2 with buf_*_scale_2 slot buffers; ExpertWeightResult and the path without them unchanged.

• (c) RoutedExperts: provider selection by the knob; _maybe_init_expert_lru_cache(scale_suffix, *, scale_2_suffix) repoints per-expert global-scale parameters at slot buffers before the quant config is built; run_with_expert_cache typed against a Protocol.

• Load path: with the cache enabled, ModelOptNvFp4FusedMoE.create_weights allocates the per-expert weights and block scales in CPU pinned memory; prepare_nvfp4_moe_layer_for_marlin converts host-resident sources expert_chunk (64) experts at a time on the accelerator into pinned host outputs (device sources: unchanged single pass); the providers alias the pinned results.

• (b2) ModelOptNvFp4FusedMoE: supports_expert_lru_cache True only for the Marlin backend (expert-major final layout, globals read through the quant config); cache installed after the backend conversion, before quant config/kernel; apply through run_with_expert_cache; monolithic rejects a cache.

• Loader interaction: device_loading_context restores CPU-resident parameters by name after processing; the cache marks its scale slot buffers device-resident (DEVICE_RESIDENT_ATTR) so they are not pulled back. Ordinary replacements keep the previous contract; the context no longer retains moved Parameter objects.

• Breakable CUDA graphs: for models that are not torch-compiled, the MoE custom ops are eager break points for layers holding a provider (uncached MoE layers stay captured), and both guards accept VLLM_USE_BREAKABLE_CUDAGRAPH=1 as the alternative to compile splitting. The legacy layer-name placeholder is resolved once per op.

• Byte accounting at cache setup and periodic row-provider stats.

• Stream ownership in the row provider is tracked, not enforced: the release event is recorded on the stream of the previous accepted prepare() (where that forward read the slots) and the current stream becomes the owner after waiting on the ready event; resident-map writes are serialized with the slot copies (copy stream in prepare(), owner stream in invalidate()). Contract: sequential host calls, all consumer reads enqueued on the current stream before the next prepare().

Tests

• CPU/mock: tests/config/test_moe_expert_cache_provider.py, tests/kernels/moe_provider/test_row_provider_cpu.py, test_row_provider_staging_mock.py (event call order).
• CUDA: tests/kernels/moe_provider/test_row_provider_cuda.py (eviction reuse with a pre-captured graph reader kept unfinished while the next prepare is enqueued; previous reader output oracle-checked; six-tensor bytes on the reused slot; every route planned once; other-stream rejection; invalidate->prepare ordering), tests/kernels/moe/test_expert_lru_cache.py (scale_2 lifecycle/validation added), tests/kernels/moe/test_nvfp4_expert_cache_consumer.py (Marlin capability-gated; both providers x device/host sources; distinct per-expert globals; overflow every forward; per-pass g1/g2_alphas[slot] equal to the uncached layer's final globals of the mapped expert; live-buffer identity; output match; chunked host repack byte-equal to the device path for N=128 and N=96), tests/kernels/moe/test_nvfp4_marlin_repack_args.py (CPU: non-positive expert_chunk rejected before side effects).

Executed: Mac: ruff, format, mypy, pre-commit; CPU+mock provider tests (9 passed, run outside the repo conftest). GPU (astra, overlay on this base, head d26e289): tests/kernels/moe_provider + test_expert_lru_cache.py 51 passed. GPU on the exact a97 build (astra, head b2b5ec1): tests/kernels/moe_provider, test_expert_lru_cache.py, test_nvfp4_expert_cache_consumer.py, test_nvfp4_marlin_repack_args.py = 63 passed, 0 skipped (results/upstream-provider-exact-b2b5ec199799-20260909). Real checkpoint (astra, mid cross-checked): integration head lab/upstream-provider-ple-mmap (this PR merged onto lab/upstream-ple-mmap-base a7a3d7a = upstream/main a97dacb + proposed sync PLE vllm-project#54129; 4c79d36 for the row run, fcf878a for the cached run) on RTX PRO 6000 Blackwell Max-Q 96 GiB limited to 48 GiB by balloon; Marlin backend; --moe-expert-cache-size 208 --moe-expert-cache-split token; chunk 512; max_num_seqs 1; breakable CUDA graphs auto-on (compilation mode NONE), cudagraph_mode PIECEWISE, capture sizes [1, 2]; sync PLE mmap; no speculation. Both providers: load 35.54 GiB, 48 layers cached at 208/512, KV 9.86 GiB, captures at 1 and 2 tokens recorded twice each (profiling and actual capture; 98 graphs / 97 eager breaks per capture), two probe requests answered correctly (a 64-token explanation and a 5-token arithmetic answer 3288). Host cgroup peak reached the 100 GiB budget without OOM; capacity samples (25 at load, 4 during generation) recorded no error. Results: cached results/upstream-provider-checkpoint-fcf878ad5781-cached-20260909 (head 0b6fe94), row results/upstream-provider-checkpoint-4c79d366113b-row-20260909 (head b5ac54e), each with an independent review file. Exact suite at 4c79d36: 71 passed, 0 skipped. Limits: short smokes only; the servers were stopped externally after the requests and the logs keep the shutdown EngineDeadError/semaphore warnings (not a clean shutdown claim); no speed comparison between providers, no long-output quality, no per-step replay counters, no concurrent-consumer guarantee (sequential-consumer stream contract, see the class docstring).

Equal-capacity pair on the same head (astra, mid cross-checked; integration head e5c920b; same launch except --moe-expert-cache-provider; two fixed tokenized prompts A 1070 / B 753 tokens, max_tokens 2048, one request each, no repeats; results/upstream-provider-pair-e5c920b0fdda-20260909):

provider	A: output / TTFT / client decode	B: output / TTFT / client decode
cached	926 tok / 10.16 s / 18.24 tok/s	1616 tok / 6.36 s / 18.84 tok/s
row	878 tok / 9.39 s / 19.76 tok/s	1306 tok / 6.14 s / 19.66 tok/s

Single observations with differing outputs (greedy paths diverge between providers), so this is a check, not a speed claim: the row provider's staged copies show a few percent higher client decode rate at N=208. Both are far below the lab fork's post-fix figure (66 tok/s at a 26 GiB tier with the in-graph planner); that is a different configuration and not a single-variable comparison. Capacity samples recorded no guard error (cached 26 startup + 75 generation, row 27 + 62); host cgroup peak reached the 100 GiB budget in both runs (memory reclaim events 4012 / 4138) without OOM; both containers exited 0.

Port source for the internals: lab/flashnext-cpu-base 5bf2151 (global_pool, device_lru, staging, promote, device_record).

🤖 Generated with Claude Code

https://claude.ai/code/session_016QWXP5rMj1rGh9xasXNLyT