[Perf][Model Loader] DeepSeek V4.1 Flash mixed RunAI loading

[SherifWaly]

Overview

Consolidate the DeepSeek V4.1 Flash RunAI loading stack: stream the VL wrapper, keep ordinary weights distributed, read shared CPU Engram tables once, and preserve/select the DSpark S3 weight source. This draft includes ported prior work and the experimental integration measured on eight H200s.

Claims

• Keep TP1/DP8/EP8 ordinary loading distributed while oversized Engram shards load into CPU staging. Shared embeddings/scales have one reader; the small non-shared tensors in those shards still reach every GPU.
• Pass the reader process group to the streamer before metadata; only the single shared Engram CPU reader passes None. Independent readers retain shared filesystem queue sizing.
• Preserve streamed tensor lifetimes without materializing the full language checkpoint. Borrowed GPU tensors keep vLLM's normal clone; shared CPU embedding sources are consumed synchronously.
• Load the same-checkpoint DSpark draft from S3 and select its three required shards from the verified local index.
• Require runai-model-streamer>=0.17.0 in the Python extra and Docker default, with matching 0.17.0 CI requirements and cloud adapters. Released Linux wheels require glibc 2.30; CUDA/CPU/ROCm lock generation targets 2.31, the lowest compatible target supported by the pinned uv. XPU already targets 2.39.

Validation

Real DeepSeek V4.1 Flash revision 5254448, eight H200s, TP1/DP8/EP8, Engram CPU offload and shared DP storage, DSpark5, RunAI 0.17.0, spans, plus the separate streamer rank-discovery fix model-streamer #252. Ordinary staging remained 4 GiB/rank (32 GiB node total under 0.17.0); oversized CPU staging remained 100 GiB.

Loading scope	Earlier S3 integration	Combined patched integration
Ordinary 286.1 GiB	585 s	37.7 s
Engram 189.1 GiB	140 s	88.4 s
Target payload, ordinary plus Engram	725 s	126 s
Draft-containing 7.4 GiB	—	approximately 6 s

This is the combined stack/configuration result, not the isolated effect of this PR: released 0.17.0 already supplies native read coalescing/ring improvements, and different nodes were used. One full model run per integration configuration; no repeated same-node A/B startup campaign is claimed.

The patched integration reached Ready with zero restarts. Full vLLM model loading, including approximately 150 s of host-table initialization, took 291 s. Cold container-to-Ready took approximately 15m28s, including regression preflight, setup and kernel/graph warmup. The two-minute payload time is not a two-minute server startup claim. Cached NFS/page-cache default-loader controls took 110–113 s for target weights and 265–279 s for full model loading; this does not claim a speedup over cached filesystem loading.

• Focused wrapper/loader/draft regressions: 12 passed, 27 deselected, 9.62 s.
• Real eight-H200 shared Engram toy loading: all eight ranks passed, including FP8/scales, follower synchronization and independent GPU weights in the CPU shard.
• Updated process-group integration: 4 passed, one unrelated GPT-2 download test deselected, 5.38 s. Covers the three distinct mixed-loader paths and real native borrowed-buffer cloning.
• Updated streamer dependency: 45 passed + 7 subtests, no skips, 3.50 s. Real eight-process Gloo coverage verifies metadata/payload reader sizing, idle peers, a non-contiguous subgroup, sequential handles and caller-group lifetime; cleanup-failure coverage verifies environment restoration.
• Changed-file pre-commit checks and Linux-target mypy completed. Native macOS mypy reports two pre-existing Linux-only os.posix_fadvise/os.POSIX_FADV_DONTNEED errors; equivalent Linux-target checks passed.
• Dependency bump: CUDA/CPU/ROCm/XPU locks regenerated with uv 0.11.1; prek run --files passed all applicable hooks, including the four dependency resolvers and Docker version consistency. Existing unrelated pins were retained; regeneration also restored the openai-harmony==0.0.8 transitive entry required by the already-pinned gpt-oss package.

The updated process-group tests ran on an isolated CPU-only Linux pod with the released native core 0.17.0 and the modified Python SDK from model-streamer #252, now at 0e0ebbb. This revision requires that matching Python SDK; the released 0.17.0 Python API does not yet accept process_group. The final dependency version must be set after the companion API is released. These tests do not establish a new model-loading or NFS performance result, and no revised serving deployment was validated.

Latest integration reproduction with that matching SDK installed:

.venv/bin/python -m pytest -q -s \
  tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py \
  -k 'not test_runai_model_loader'

Focused reproduction in a CUDA vLLM environment:

VLLM_RUNAI_BORROWED_WEIGHTS=0 .venv/bin/python -m pytest \
  tests/models/test_deepseek_v4_vl_stream_weights.py \
  tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py \
  tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py \
  tests/v1/worker/test_dspark_utils.py \
  -q -k 'streams_language_model_group or clones_reused_buffers or explicit_cpu_shards or index_selection or inherits_only_shared_checkpoint_uri'

The hardware results came from the original integration on fe4df3df0 (PyTorch 2.13.0/CUDA 13). This branch is rebased on main; a complete model rerun on the rebased head, HTTP text/image/tool checks, draft acceptance and quality evaluation remain pending. No output-quality equivalence is claimed. The new wrapper test is ported with its Buildkite entry; existing model-executor/worker jobs cover the extended suites. Full-model CI additionally needs the checkpoint, eight GPUs and sufficient host memory.

Details

The reproduction example explicitly selects the Python API frontend with VLLM_USE_RUST_FRONTEND=0. It supports an explicit local tokenizer metadata path for S3 deployment. This is independent of RunAI tensor loading. A separate eight-B200 canary using preserved image digest e5c9e6bb82f645e8fca04e4d395743ee5ed473a62b98942d65f83c35d0c69301 and the cached filesystem loader reached Ready with zero restarts and passed 16 serving checks: health/model discovery, text, SSE streaming, image input, forced/automatic tool calls, concurrent requests directed to DP ranks 0–7, and per-engine metric deltas of exactly one request per rank. It validates the existing Python frontend and DP header routing; it does not validate this PR's loading changes or its rebased head. The H200 Python/S3 canary remains unvalidated because BAR's requested SUNK QoS has reached its 80-GPU quota.

Included prior work and attribution:

• Ported VL streaming and its regression/CI wiring from [Model] DeepSeek-V4/V4.1 VL: stream language-model weights instead of sorting the checkpoint #58857, by @gitbisector; original author credit is retained in the commit trailer.
• Adapted the model-owned draft shard-selection approach from [Feature][Model Loader] Apply model-owned MTP shard selection to RunAI ModelStreamer #60759, by @zhy76. This integration uses a verified local index and a DeepSeek V4.1 DSpark unused-weight predicate; it is not the full generic metadata-selection implementation from that PR.
• Related draft-source fix Fix/spec draft inherit model weights #48023, by @kunling-anyscale, and issue [Bug]: runai_streamer + MTP drafter fails to load weights from model_streamer local cache #42060. The recipe uses method=dspark, not method=mtp. This branch fixes the observed DSpark config-only-cache failure in its loader using a shallow copy of the resolved config; it does not cherry-pick the older generic speculative-config patch.

These pieces are deliberately included in one reviewable draft alongside the CPU Engram integration. Refreshed duplicate checks found the existing PRs above; their overlap is explicit. #61303 optimizes pinned GPU staging and #57195/#61553 change shared-storage inference/topology support; they do not implement this mixed RunAI single-reader path.

The current loader's sorted VL input retains the stream; the GPU streamer cannot safely stage the oversized CPU tables on an H200. Splitting GPU/CPU phases addresses that destination mismatch. Shared followers receive meta placeholders so the existing single-writer weight loaders retain equal call counts/barriers; non-shared tensors are broadcast as raw bytes. The streamer dependency scopes reader discovery before metadata: ordinary readers supply their process group, while the lone CPU reader passes None to avoid discovery collectives.

Before merge:

• Replace checkpoint-specific file-list/index opt-ins and private storage-owner discovery with a reviewed configuration/metadata contract; validate missing/stale index and shard failures.
• Harden collective/error ordering, tensor lifetime and configuration validation across supported topologies; test TP/PP, reload and multi-node compatibility explicitly.
• Re-run the model on the final head, complete request/acceptance/eval validation and human line review.
• Resolve/release the separate streamer dependency. It is not vendored into this vLLM diff.

All experiment switches default off. The documented initial scope is one eight-H200 DP/EP server. See docs/models/extensions/runai_model_streamer.md for the exact environment and staging configuration. This draft adds loader-time synchronization, not steady-state inference synchronization, and remains experimental until the items above are resolved.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Initial integration is measured; generic metadata/configuration and compatibility review remain pending.

• Testing and Validation: Regressions, eight-rank toy loading and real-model startup recorded; final-head request/acceptance/eval validation remains pending.

• Code Quality and Style: Applicable checks completed, docs and ported CI wiring included; platform limitation documented.

• Pull Request Contents: Scope, port attribution, dependencies, measurements, root causes and draft limitations are explicit.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing

[mergify]

Documentation preview: https://vllm--61715.org.readthedocs.build/en/61715/

[mergify]

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

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

[gitbisector]

process_group is now passed to stream_files() on every RunAI load, but the released 0.17.0 that this PR pins doesn't accept it, so every --load-format runai_streamer load fails:

TypeError: SafetensorsStreamer.stream_files() got an unexpected keyword argument 'process_group'

Tested on GB10 (torch 2.14) against base 7840bede98, with Qwen3-0.6B over runai_streamer:

• base + 0.17.0: loads in 16.8 s, correct greedy output.

• PR + 0.17.0: engine init fails with the TypeError above. Five tests fail on the PR only. These three pass once the Not able to used qlora models with vllm  #252 files below are in place:

  • test_weight_utils.py::test_runai_model_loader
  • test_weight_utils.py::test_runai_safetensors_weights_iterator_clones_reused_buffers
  • test_runai_model_streamer_loader.py::test_runai_model_loader_download_files

  The other two, in test_runai_model_streamer_s3.py, fail on the SDK's mock (see below).

• PR + the two changed files from model-streamer Not able to used qlora models with vllm  #252 (0e0ebbb) on top of 0.17.0: loads in 14.6 s with the same output, and the focused tests pass 13/13.

Until the SDK API is released, could process_group be passed only on the mixed CPU/GPU path, or only when stream_files() accepts it? Then ordinary RunAI loads keep working on 0.17.0.

Separately, #252 also needs stream_files() in streamer_mock.py (StreamerPatcher.MockSafetensorsStreamer) to accept process_group. With the #252 files in place, these two still fail on unexpected keyword argument 'process_group':

• test_runai_model_streamer_s3.py::test_runai_model_loader_download_files_s3_mocked_with_patch
• test_runai_model_streamer_s3.py::test_runai_reload_weights_path_from_mocked_s3