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

.buildkite/test_areas/models_basic.yaml

@@ -71,8 +71,9 @@ steps:
   - tests/models/test_deepseek_v41_replay_batch.py
   - tests/models/transformers/test_layer_registry_e2e.py
   - tests/models/test_deepseek_v41_replay_start.py
+  - tests/models/test_deepseek_v4_vl_stream_weights.py
   commands:
-    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py
+    - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v4_vl_stream_weights.py
   mirror:
     amd:
       label: ":amd: (MI355) Basic Models (Other)"

.pre-commit-config.yaml

@@ -49,7 +49,7 @@ repos:
         -o, requirements/test/cuda.txt,
         --index-strategy, unsafe-best-match,
         --torch-backend, cu130,
-        --python-platform, x86_64-manylinux_2_28,
+        --python-platform, x86_64-manylinux_2_31,
         --python-version, "3.12",
       ]
       files: ^requirements/(common|cuda|test/cuda)\.(in|txt)$
@@ -61,7 +61,7 @@ repos:
         -c, requirements/rocm.txt,
         -o, requirements/test/rocm.txt,
         --index-strategy, unsafe-best-match,
-        --python-platform, x86_64-manylinux_2_28,
+        --python-platform, x86_64-manylinux_2_31,
         --python-version, "3.12",
         # Exclude torch and CUDA/NVIDIA packages
         --no-emit-package, torch,
@@ -142,7 +142,7 @@ repos:
         -o, requirements/test/cpu.txt,
         --index-strategy, unsafe-best-match,
         --torch-backend, cpu,
-        --python-platform, x86_64-manylinux_2_28,
+        --python-platform, x86_64-manylinux_2_31,
         --python-version, "3.12",
       ]
       files: ^requirements/(common|cpu|test/(cuda|cpu))\.(in|txt)$

docker/Dockerfile

@@ -966,7 +966,7 @@ RUN set -eux; \
 # These are stable packages that don't depend on vLLM itself
 # From versions.json: .openai_server_extras.timm, .openai_server_extras.runai_model_streamer
 ARG TIMM_VERSION=">=1.0.17"
-ARG RUNAI_MODEL_STREAMER_VERSION=">=0.15.7"
+ARG RUNAI_MODEL_STREAMER_VERSION=">=0.17.0"
 RUN --mount=type=cache,target=/opt/uv/cache \
     uv pip install --system accelerate 'modelscope<1.38' \
         "timm${TIMM_VERSION}" "runai-model-streamer[s3,gcs,azure]${RUNAI_MODEL_STREAMER_VERSION}"

docker/versions.json

@@ -83,7 +83,7 @@
       "default": ">=1.0.17"
     },
     "RUNAI_MODEL_STREAMER_VERSION": {
-      "default": ">=0.15.7"
+      "default": ">=0.17.0"
     }
   }
 }

docs/models/extensions/runai_model_streamer.md

@@ -85,6 +85,51 @@ vllm serve /home/meta-llama/Llama-3.2-3B-Instruct \
 !!! note
     For further instructions about tunable parameters and additional parameters configurable through environment variables, read the [Environment Variables Documentation](https://github.com/run-ai/runai-model-streamer/blob/master/docs/src/env-vars.md).
 
+### Experimental DeepSeek V4.1 shared CPU loading
+
+The mixed loading path keeps ordinary weights on the distributed GPU streamer
+and reads oversized Engram shards on CPU. With CPU Engram storage shared across
+DP replicas, one rank reads the shared tables and scales. Other ranks receive
+metadata placeholders and participate in the existing weight-loader barriers.
+Non-shared tensors in those shards are broadcast as CPU bytes and loaded on each
+rank. The CPU staging limit is restored after this phase.
+
+The `vllm[runai]` extra and Docker image require Run:ai Model Streamer 0.17.0
+or newer; its Linux wheels require glibc 2.30 or newer. This experimental path
+additionally requires the
+[non-distributed rank-discovery fix](https://github.com/dsx-ai-factory/model-streamer/pull/252).
+It was validated through startup with TP1/DP8/EP8 on eight co-located H200 GPUs.
+Request correctness, broader topologies and a generic configuration API remain
+pending. The wrapper and DSpark changes in this patch retain streamed tensor
+lifetimes and use the target's object-storage source for the draft.
+
+For checkpoint revision `5254448`, verify that the checkpoint index places the
+Engram tables in the two filenames below before enabling these settings:
+
+```bash
+export VLLM_USE_RUST_FRONTEND=0
+export RUNAI_STREAMER_PARTITION_POLICY=spans
+export VLLM_RUNAI_CPU_WEIGHT_FILES=model-00047-of-00048.safetensors,model-00048-of-00048.safetensors
+export VLLM_RUNAI_CPU_MEMORY_LIMIT=107374182400
+export VLLM_RUNAI_SINGLE_ENGRAM_READER=1
+export VLLM_RUNAI_BORROWED_WEIGHTS=1
+export VLLM_RUNAI_USE_LOCAL_WEIGHT_INDEX=1
+vllm serve s3://bucket/deepseek-v4-1-flash/ \
+    --load-format runai_streamer \
+    --model-loader-extra-config '{"distributed":true,"memory_limit":34359738368}' \
+    --tensor-parallel-size 1 --data-parallel-size 8 --enable-expert-parallel \
+    --engram-config '{"cpu_offload":true,"embedding_across_dp":true}' \
+    --speculative-config '{"method":"dspark","num_speculative_tokens":5}'
+```
+
+In streamer 0.17.0, 32 GiB is the distributed node staging budget, giving eight
+readers 4 GiB each. The CPU phase uses a separate 100 GiB budget. These limits
+control checkpoint staging, independently of KV cache and KV offload settings.
+Borrowed GPU tensors still receive vLLM's normal clone; shared CPU tensors are
+consumed synchronously. The local index must match the actual checkpoint and its
+shard headers. These environment flags are experimental; they are disabled
+unless explicitly set.
+
 ## Sharded Model Loading
 
 vLLM also supports loading sharded models using Run:ai Model Streamer. This is particularly useful for large models that are split across multiple files. To use this feature, use the `--load-format runai_streamer_sharded` flag:

requirements/rocm.txt

@@ -16,7 +16,7 @@ packaging>=24.2
 setuptools>=78.1.1,<80.0.0
 setuptools-scm>=8
 setuptools-rust>=1.9.0
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.0
 conch-triton-kernels==1.2.1
 timm>=1.0.17
 # amd-quark: required for Quark quantization on ROCm

requirements/test/cpu.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test/cuda.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu --python-platform x86_64-manylinux_2_28 --python-version 3.12
+#    uv pip compile requirements/test/cuda.in -o requirements/test/cpu.txt --index-strategy unsafe-best-match --torch-backend cpu --python-platform x86_64-manylinux_2_31 --python-version 3.12
 abi3info==2025.11.29
     # via torch-abi-audit
 absl-py==2.1.0
@@ -615,6 +615,8 @@ openai==2.44.0
     # via
     #   -r requirements/test/../common.txt
     #   gpt-oss
+openai-harmony==0.0.8
+    # via gpt-oss
 opencensus==0.11.4
     # via ray
 opencensus-context==0.1.3
@@ -829,6 +831,7 @@ pydantic==2.12.0
     #   model-hosting-container-standards
     #   mteb
     #   openai
+    #   openai-harmony
     #   oss-harmony
     #   pydantic-extra-types
     #   pydantic-settings
@@ -1011,13 +1014,13 @@ rpds-py==0.20.1
     #   referencing
 rsa==4.9.1
     # via google-auth
-runai-model-streamer==0.15.7
+runai-model-streamer==0.17.0
     # via -r requirements/test/cuda.in
-runai-model-streamer-azure==0.15.7
+runai-model-streamer-azure==0.17.0
     # via runai-model-streamer
-runai-model-streamer-gcs==0.15.7
+runai-model-streamer-gcs==0.17.0
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.7
+runai-model-streamer-s3==0.17.0
     # via runai-model-streamer
 s3transfer==0.10.3
     # via boto3

requirements/test/cuda.in

@@ -55,7 +55,7 @@ grpcio-reflection==1.78.0
 
 arctic-inference == 0.1.1; platform_machine == "x86_64" # Required for suffix decoding test
 numba == 0.65.0 # Required for N-gram speculative decoding
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.0
 fastsafetensors>=0.3.3
 instanttensor>=0.2.2; platform_machine == "x86_64"
 decord==0.6.0; platform_machine == "x86_64"

requirements/test/cuda.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_28 --python-version 3.12
+#    uv pip compile requirements/test/cuda.in -c requirements/cuda.txt -o requirements/test/cuda.txt --index-strategy unsafe-best-match --torch-backend cu130 --python-platform x86_64-manylinux_2_31 --python-version 3.12
 abi3info==2025.11.29
     # via torch-abi-audit
 absl-py==2.1.0
@@ -692,6 +692,8 @@ openai==2.44.0
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   gpt-oss
+openai-harmony==0.0.8
+    # via gpt-oss
 opencensus==0.11.4
     # via ray
 opencensus-context==0.1.3
@@ -923,6 +925,7 @@ pydantic==2.12.0
     #   model-hosting-container-standards
     #   mteb
     #   openai
+    #   openai-harmony
     #   oss-harmony
     #   pydantic-extra-types
     #   pydantic-settings
@@ -1107,13 +1110,13 @@ rpds-py==0.20.1
     #   referencing
 rsa==4.9.1
     # via google-auth
-runai-model-streamer==0.15.7
+runai-model-streamer==0.17.0
     # via -r requirements/test/cuda.in
-runai-model-streamer-azure==0.15.7
+runai-model-streamer-azure==0.17.0
     # via runai-model-streamer
-runai-model-streamer-gcs==0.15.7
+runai-model-streamer-gcs==0.17.0
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.7
+runai-model-streamer-s3==0.17.0
     # via runai-model-streamer
 s3transfer==0.10.3
     # via boto3

requirements/test/nightly-torch.txt

@@ -40,7 +40,7 @@ tritonclient>=2.51.0
 
 numba == 0.65.0 # Required for N-gram speculative decoding
 numpy
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.0
 fastsafetensors>=0.3.3
 instanttensor>=0.2.2
 pydantic>=2.12 # 2.11 leads to error on python 3.13

requirements/test/rocm.in

@@ -52,7 +52,7 @@ grpcio-reflection==1.78.0
 
 arctic-inference==0.1.1 # Required for suffix decoding test
 numba==0.65.0 # Required for N-gram speculative decoding
-runai-model-streamer[s3,gcs,azure]==0.15.7
+runai-model-streamer[s3,gcs,azure]==0.17.0
 fastsafetensors>=0.3.3
 instanttensor>=0.1.9
 decord==0.6.0

requirements/test/rocm.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_28 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13
+#    uv pip compile requirements/test/rocm.in -c requirements/rocm.txt -o requirements/test/rocm.txt --index-strategy unsafe-best-match --python-platform x86_64-manylinux_2_31 --python-version 3.12 --no-emit-package torch --no-emit-package torchvision --no-emit-package torchaudio --no-emit-package triton --no-emit-package cuda-bindings --no-emit-package cuda-pathfinder --no-emit-package cuda-toolkit --no-emit-package cupy-cuda12x --no-emit-package nvidia-cublas --no-emit-package nvidia-cuda-cupti --no-emit-package nvidia-cuda-nvrtc --no-emit-package nvidia-cuda-runtime --no-emit-package nvidia-cudnn --no-emit-package nvidia-cufft --no-emit-package nvidia-cufile --no-emit-package nvidia-curand --no-emit-package nvidia-cusolver --no-emit-package nvidia-cusparse --no-emit-package nvidia-cusparselt --no-emit-package nvidia-nccl --no-emit-package nvidia-nvjitlink --no-emit-package nvidia-nvshmem --no-emit-package nvidia-nvtx --no-emit-package nvidia-cublas-cu12 --no-emit-package nvidia-cuda-cupti-cu12 --no-emit-package nvidia-cuda-nvrtc-cu12 --no-emit-package nvidia-cuda-runtime-cu12 --no-emit-package nvidia-cudnn-cu12 --no-emit-package nvidia-cufft-cu12 --no-emit-package nvidia-cufile-cu12 --no-emit-package nvidia-curand-cu12 --no-emit-package nvidia-cusolver-cu12 --no-emit-package nvidia-cusparse-cu12 --no-emit-package nvidia-cusparselt-cu12 --no-emit-package nvidia-nccl-cu12 --no-emit-package nvidia-nvjitlink-cu12 --no-emit-package nvidia-nvshmem-cu12 --no-emit-package nvidia-nvtx-cu12 --no-emit-package nvidia-cublas-cu13 --no-emit-package nvidia-cuda-cupti-cu13 --no-emit-package nvidia-cuda-nvrtc-cu13 --no-emit-package nvidia-cuda-runtime-cu13 --no-emit-package nvidia-cudnn-cu13 --no-emit-package nvidia-cufft-cu13 --no-emit-package nvidia-cufile-cu13 --no-emit-package nvidia-curand-cu13 --no-emit-package nvidia-cusolver-cu13 --no-emit-package nvidia-cusparse-cu13 --no-emit-package nvidia-cusparselt-cu13 --no-emit-package nvidia-nccl-cu13 --no-emit-package nvidia-nvjitlink-cu13 --no-emit-package nvidia-nvshmem-cu13 --no-emit-package nvidia-nvtx-cu13
 absl-py==2.4.0
     # via rouge-score
 accelerate==1.13.0
@@ -658,6 +658,8 @@ openai==2.31.0
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   gpt-oss
+openai-harmony==0.0.8
+    # via gpt-oss
 opencensus==0.11.4
     # via ray
 opencensus-context==0.1.3
@@ -896,6 +898,7 @@ pydantic==2.12.5
     #   model-hosting-container-standards
     #   mteb
     #   openai
+    #   openai-harmony
     #   oss-harmony
     #   pydantic-extra-types
     #   pydantic-settings
@@ -1078,15 +1081,15 @@ rpds-py==0.30.0
     # via
     #   jsonschema
     #   referencing
-runai-model-streamer==0.15.7
+runai-model-streamer==0.17.0
     # via
     #   -c requirements/rocm.txt
     #   -r requirements/test/rocm.in
-runai-model-streamer-azure==0.15.7
+runai-model-streamer-azure==0.17.0
     # via runai-model-streamer
-runai-model-streamer-gcs==0.15.7
+runai-model-streamer-gcs==0.17.0
     # via runai-model-streamer
-runai-model-streamer-s3==0.15.7
+runai-model-streamer-s3==0.17.0
     # via runai-model-streamer
 s3transfer==0.16.0
     # via boto3

requirements/test/xpu.in

@@ -18,7 +18,7 @@ arctic-inference
 lm_eval[api]>=0.4.12
 nltk>=3.10.3 # pin transitive rouge-score/lm-eval dep for known GHSAs
 modelscope<1.38
-runai-model-streamer[s3,gcs,azure]>=0.15.7
+runai-model-streamer[s3,gcs,azure]>=0.17.0
 tensorizer==2.10.1
 transformers==5.19.0
 

requirements/test/xpu.txt

@@ -605,6 +605,8 @@ openai==3.2.0
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   gpt-oss
+openai-harmony==0.0.8
+    # via gpt-oss
 opencv-python-headless==5.0.0.93
     # via
     #   -c requirements/common.txt
@@ -786,6 +788,7 @@ pydantic==2.13.4
     #   model-hosting-container-standards
     #   mteb
     #   openai
+    #   openai-harmony
     #   oss-harmony
     #   pydantic-extra-types
     #   pydantic-settings
@@ -945,13 +948,13 @@ rpds-py==2026.6.3
     # via
     #   jsonschema
     #   referencing
-runai-model-streamer==0.16.1
+runai-model-streamer==0.17.0
     # via -r requirements/test/xpu.in
-runai-model-streamer-azure==0.16.1
+runai-model-streamer-azure==0.17.0
     # via runai-model-streamer
-runai-model-streamer-gcs==0.16.1
+runai-model-streamer-gcs==0.17.0
     # via runai-model-streamer
-runai-model-streamer-s3==0.16.1
+runai-model-streamer-s3==0.17.0
     # via runai-model-streamer
 s3transfer==0.19.2
     # via boto3

setup.py

@@ -1693,7 +1693,7 @@ def add_vllm_package_data(filename: str) -> None:
         "tensorizer": ["tensorizer==2.10.1"],
         "fastsafetensors": ["fastsafetensors >= 0.3.3"],
         "instanttensor": ["instanttensor >= 0.2.2"],
-        "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.15.7"],
+        "runai": ["runai-model-streamer[s3,gcs,azure] >= 0.17.0"],
         "audio": [
             "av",
             "scipy",

tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import json
 import os
 import types
 from unittest.mock import patch
@@ -164,6 +165,36 @@ def test_runai_get_all_weights_resolves_source(model_config, expected_source):
     assert weights == [(f"{expected_source}@myrev", None)]
 
 
+def test_runai_index_selection_keeps_mixed_shards_and_order(tmp_path):
+    files = ["s3://b/z.safetensors", "s3://b/a.safetensors", "s3://b/m.safetensors"]
+    (tmp_path / "model.safetensors.index.json").write_text(
+        json.dumps(
+            {
+                "weight_map": {
+                    "mtp.0.w": "z.safetensors",
+                    "layers.0.w": "a.safetensors",
+                    "mtp.1.w": "m.safetensors",
+                    "layers.1.w": "m.safetensors",
+                }
+            }
+        )
+    )
+    selected = rsl._select_indexed_weight_files(
+        files, str(tmp_path), lambda name: not name.startswith("mtp.")
+    )
+    assert selected == [files[0], files[2]]
+
+
+@pytest.mark.parametrize("index", [None, {}, {"weight_map": {"mtp.w": "missing"}}])
+def test_runai_index_selection_falls_back_when_metadata_unusable(tmp_path, index):
+    files = ["s3://b/a.safetensors"]
+    if index is not None:
+        (tmp_path / "model.safetensors.index.json").write_text(json.dumps(index))
+    assert (
+        rsl._select_indexed_weight_files(files, str(tmp_path), lambda _: False) == files
+    )
+
+
 @pytest.mark.parametrize(
     "base_model,mul_model,add_model",
     [