chore: import upstream snapshot with attribution

tests/model_executor/model_loader/fastsafetensors_loader/test_fastsafetensors_loader.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.platforms import current_platform
+
+test_model = "openai-community/gpt2"
+
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+# Create a sampling params object.
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="fastsafetensors requires NVIDIA/AMD GPUs",
+)
+def test_model_loader_download_files(vllm_runner):
+    with vllm_runner(test_model, load_format="fastsafetensors") as llm:
+        deserialized_outputs = llm.generate(prompts, sampling_params)
+        assert deserialized_outputs

tests/model_executor/model_loader/fastsafetensors_loader/test_weight_utils.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import glob
+import tempfile
+
+import huggingface_hub.constants
+import pytest
+import torch
+
+from vllm.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf,
+    fastsafetensors_weights_iterator,
+    safetensors_weights_iterator,
+)
+from vllm.platforms import current_platform
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda_alike(),
+    reason="fastsafetensors requires NVIDIA/AMD GPUs",
+)
+@pytest.mark.parametrize("queue_size", [0, 1])
+def test_fastsafetensors_model_loader(monkeypatch, queue_size):
+    monkeypatch.setenv("VLLM_FASTSAFETENSORS_QUEUE_SIZE", str(queue_size))
+    with tempfile.TemporaryDirectory() as tmpdir:
+        huggingface_hub.constants.HF_HUB_OFFLINE = False
+        download_weights_from_hf(
+            "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
+        )
+        safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
+        assert len(safetensors) > 0
+
+        fastsafetensors_tensors = {}
+        hf_safetensors_tensors = {}
+
+        for name, tensor in fastsafetensors_weights_iterator(safetensors, True):
+            fastsafetensors_tensors[name] = tensor
+
+        for name, tensor in safetensors_weights_iterator(safetensors, True):
+            hf_safetensors_tensors[name] = tensor
+
+        assert len(fastsafetensors_tensors) == len(hf_safetensors_tensors)
+
+        for name, fastsafetensors_tensor in fastsafetensors_tensors.items():
+            fastsafetensors_tensor = fastsafetensors_tensor.to("cpu")
+            assert fastsafetensors_tensor.dtype == hf_safetensors_tensors[name].dtype
+            assert fastsafetensors_tensor.shape == hf_safetensors_tensors[name].shape
+            assert torch.all(fastsafetensors_tensor.eq(hf_safetensors_tensors[name]))

tests/model_executor/model_loader/instanttensor_loader/test_instanttensor_loader.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.platforms import current_platform
+
+test_model = "openai-community/gpt2"
+
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+# Create a sampling params object.
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="InstantTensor requires NVIDIA GPUs",
+)
+def test_model_loader_download_files(vllm_runner):
+    with vllm_runner(test_model, load_format="instanttensor") as llm:
+        deserialized_outputs = llm.generate(prompts, sampling_params)
+        assert deserialized_outputs

tests/model_executor/model_loader/instanttensor_loader/test_weight_utils.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import glob
+import tempfile
+
+import huggingface_hub.constants
+import pytest
+import torch
+
+from vllm.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf,
+    instanttensor_weights_iterator,
+    safetensors_weights_iterator,
+)
+from vllm.platforms import current_platform
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(),
+    reason="InstantTensor requires NVIDIA GPUs",
+)
+def test_instanttensor_model_loader():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        huggingface_hub.constants.HF_HUB_OFFLINE = False
+        download_weights_from_hf(
+            "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
+        )
+        safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
+        assert len(safetensors) > 0
+
+        instanttensor_tensors = {}
+        hf_safetensors_tensors = {}
+
+        for name, tensor in instanttensor_weights_iterator(safetensors, True):
+            # Copy the tensor immediately as it is a reference to the internal
+            # buffer of instanttensor.
+            instanttensor_tensors[name] = tensor.to("cpu")
+
+        for name, tensor in safetensors_weights_iterator(safetensors, True):
+            hf_safetensors_tensors[name] = tensor
+
+        assert len(instanttensor_tensors) == len(hf_safetensors_tensors)
+
+        for name, instanttensor_tensor in instanttensor_tensors.items():
+            assert instanttensor_tensor.dtype == hf_safetensors_tensors[name].dtype
+            assert instanttensor_tensor.shape == hf_safetensors_tensors[name].shape
+            assert torch.all(instanttensor_tensor.eq(hf_safetensors_tensors[name]))
+
+
+if __name__ == "__main__":
+    test_instanttensor_model_loader()

tests/model_executor/model_loader/runai_streamer_loader/conftest.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
+from vllm.v1.executor import UniProcExecutor
+from vllm.v1.worker.worker_base import WorkerWrapperBase
+
+
+# This is a dummy executor for patching in test_runai_model_streamer_s3.py.
+# We cannot use vllm_runner fixture here, because it spawns worker process.
+# The worker process reimports the patched entities, and the patch is not applied.
+class RunaiDummyExecutor(UniProcExecutor):
+    def _init_executor(self) -> None:
+        distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())
+
+        local_rank = 0
+        rank = 0
+        is_driver_worker = True
+
+        device_info = self.vllm_config.device_config.device.__str__().split(":")
+        if len(device_info) > 1:
+            local_rank = int(device_info[1])
+
+        worker_rpc_kwargs = dict(
+            vllm_config=self.vllm_config,
+            local_rank=local_rank,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            is_driver_worker=is_driver_worker,
+        )
+
+        self.driver_worker = WorkerWrapperBase()
+
+        self.collective_rpc("init_worker", args=([worker_rpc_kwargs],))
+        self.collective_rpc("init_device")

tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_loader.py

@@ -0,0 +1,124 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+import types
+from unittest.mock import patch
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.config.load import LoadConfig
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader import runai_streamer_loader as rsl
+
+load_format = "runai_streamer"
+test_model = "openai-community/gpt2"
+# TODO(amacaskill): Replace with a GKE owned GCS bucket.
+test_gcs_model = "gs://vertex-model-garden-public-us/codegemma/codegemma-2b/"
+
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+# Create a sampling params object.
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
+
+
+def get_runai_model_loader():
+    load_config = LoadConfig(load_format=load_format)
+    return get_model_loader(load_config)
+
+
+def test_get_model_loader_with_runai_flag():
+    model_loader = get_runai_model_loader()
+    assert model_loader.__class__.__name__ == "RunaiModelStreamerLoader"
+
+
+def test_runai_model_loader_download_files(vllm_runner):
+    with vllm_runner(test_model, load_format=load_format) as llm:
+        deserialized_outputs = llm.generate(prompts, sampling_params)
+        assert deserialized_outputs
+
+
+@pytest.mark.skip(
+    reason="Temporarily disabled due to GCS access issues. "
+    "TODO: Re-enable this test once the underlying issue is resolved."
+)
+def test_runai_model_loader_download_files_gcs(
+    vllm_runner, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "fake-project")
+    monkeypatch.setenv("RUNAI_STREAMER_GCS_USE_ANONYMOUS_CREDENTIALS", "true")
+    monkeypatch.setenv(
+        "CLOUD_STORAGE_EMULATOR_ENDPOINT", "https://storage.googleapis.com"
+    )
+    with vllm_runner(test_gcs_model, load_format=load_format) as llm:
+        deserialized_outputs = llm.generate(prompts, sampling_params)
+        assert deserialized_outputs
+
+
+def test_runai_passes_revision_by_name():
+    # revision must reach download_safetensors_index_file_from_hf as the
+    # ``revision`` keyword, not the positional ``subfolder`` slot.
+    fake_self = types.SimpleNamespace(
+        load_config=types.SimpleNamespace(download_dir="/cache", ignore_patterns=[])
+    )
+    with (
+        patch.object(rsl, "is_runai_obj_uri", return_value=False),
+        patch.object(rsl, "download_weights_from_hf", return_value="/folder"),
+        patch.object(
+            rsl, "list_safetensors", return_value=["/folder/model.safetensors"]
+        ),
+        patch.object(rsl, "download_safetensors_index_file_from_hf") as mock_idx,
+    ):
+        rsl.RunaiModelStreamerLoader._prepare_weights(fake_self, "org/model", "myrev")
+
+    mock_idx.assert_called_once()
+    assert mock_idx.call_args.kwargs.get("revision") == "myrev"
+    assert "myrev" not in mock_idx.call_args.args
+
+
+def _runai_loader(extra):
+    return rsl.RunaiModelStreamerLoader(
+        LoadConfig(load_format="runai_streamer", model_loader_extra_config=extra)
+    )
+
+
+@pytest.mark.parametrize(
+    "extra, match",
+    [
+        ({"typo_key": 1}, "Unexpected extra config"),
+        ({"distributed": "yes"}, "distributed must be a bool"),
+        ({"concurrency": "16"}, "concurrency must be a positive integer"),
+        ({"concurrency": -1}, "concurrency must be a positive integer"),
+    ],
+)
+def test_runai_rejects_invalid_extra_config(extra, match):
+    # The loader used to silently drop unknown keys / wrong types / negatives.
+    with pytest.raises(ValueError, match=match):
+        _runai_loader(extra)
+
+
+def test_runai_accepts_valid_extra_config():
+    with patch.dict(os.environ, {}, clear=False):
+        os.environ.pop("RUNAI_STREAMER_CONCURRENCY", None)
+        os.environ.pop("RUNAI_STREAMER_MEMORY_LIMIT", None)
+        loader = _runai_loader(
+            {"distributed": True, "concurrency": 16, "memory_limit": 1024}
+        )
+        assert loader._is_distributed is True
+        assert os.environ["RUNAI_STREAMER_CONCURRENCY"] == "16"
+        assert os.environ["RUNAI_STREAMER_MEMORY_LIMIT"] == "1024"
+
+
+def test_runai_invalid_extra_config_leaves_environ_untouched():
+    # A later invalid key must not leave an earlier valid key applied to
+    # os.environ (all values are validated before any global mutation).
+    with patch.dict(os.environ, {}, clear=False):
+        os.environ.pop("RUNAI_STREAMER_CONCURRENCY", None)
+        with pytest.raises(ValueError, match="memory_limit must be an integer >= -1"):
+            _runai_loader({"concurrency": 16, "memory_limit": -5})
+        assert "RUNAI_STREAMER_CONCURRENCY" not in os.environ

tests/model_executor/model_loader/runai_streamer_loader/test_runai_model_streamer_s3.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+from runai_model_streamer.safetensors_streamer.streamer_mock import StreamerPatcher
+
+from vllm.engine.arg_utils import EngineArgs
+
+from .conftest import RunaiDummyExecutor
+
+load_format = "runai_streamer"
+test_model = "openai-community/gpt2"
+
+
+def test_runai_model_loader_download_files_s3_mocked_with_patch(
+    vllm_runner,
+    tmp_path: Path,
+    monkeypatch,
+):
+    patcher = StreamerPatcher(str(tmp_path))
+
+    test_mock_s3_model = "s3://my-mock-bucket/gpt2/"
+
+    # Download model from HF
+    mock_model_dir = f"{tmp_path}/gpt2"
+    snapshot_download(repo_id=test_model, local_dir=mock_model_dir)
+
+    monkeypatch.setattr(
+        "vllm.transformers_utils.runai_utils.runai_list_safetensors",
+        patcher.shim_list_safetensors,
+    )
+    monkeypatch.setattr(
+        "vllm.transformers_utils.runai_utils.runai_pull_files",
+        patcher.shim_pull_files,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.model_loader.weight_utils.SafetensorsStreamer",
+        patcher.create_mock_streamer,
+    )
+
+    engine_args = EngineArgs(
+        model=test_mock_s3_model,
+        load_format=load_format,
+        tensor_parallel_size=1,
+    )
+
+    vllm_config = engine_args.create_engine_config()
+
+    executor = RunaiDummyExecutor(vllm_config)
+    executor.driver_worker.load_model()

tests/model_executor/model_loader/runai_streamer_loader/test_runai_utils.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import glob
+import hashlib
+import os
+import tempfile
+
+import huggingface_hub.constants
+
+from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
+from vllm.transformers_utils.runai_utils import (
+    ObjectStorageModel,
+    is_runai_obj_uri,
+    list_safetensors,
+)
+
+
+def test_is_runai_obj_uri():
+    assert is_runai_obj_uri("gs://some-gcs-bucket/path")
+    assert is_runai_obj_uri("s3://some-s3-bucket/path")
+    assert is_runai_obj_uri("az://some-azure-container/path")
+    assert not is_runai_obj_uri("nfs://some-nfs-path")
+
+
+def test_runai_list_safetensors_local():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        huggingface_hub.constants.HF_HUB_OFFLINE = False
+        download_weights_from_hf(
+            "openai-community/gpt2",
+            allow_patterns=["*.safetensors", "*.json"],
+            cache_dir=tmpdir,
+        )
+        safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
+        assert len(safetensors) > 0
+        parentdir = [os.path.dirname(safetensor) for safetensor in safetensors][0]
+        files = list_safetensors(parentdir)
+        assert len(safetensors) == len(files)
+
+
+def test_runai_pull_files_gcs(monkeypatch):
+    monkeypatch.setenv("RUNAI_STREAMER_GCS_USE_ANONYMOUS_CREDENTIALS", "true")
+    # Bypass default project lookup by setting GOOGLE_CLOUD_PROJECT
+    monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "fake-project")
+    filename = "LT08_L1GT_074061_20130309_20170505_01_T2_MTL.txt"
+    gcs_bucket = "gs://gcp-public-data-landsat/LT08/01/074/061/LT08_L1GT_074061_20130309_20170505_01_T2/"
+    gcs_url = f"{gcs_bucket}/{filename}"
+    model = ObjectStorageModel(gcs_url)
+    model.pull_files(gcs_bucket, allow_pattern=[f"*{filename}"])
+    # To re-generate / change URLs:
+    #   gsutil ls -L gs://<gcs-url> | grep "Hash (md5)" | tr -d ' ' \
+    #     | cut -d":" -f2 | base64 -d | xxd -p
+    expected_checksum = "f60dea775da1392434275b311b31a431"
+    hasher = hashlib.new("md5")
+    with open(os.path.join(model.dir, filename), "rb") as f:
+        # Read the file in chunks to handle large files efficiently
+        for chunk in iter(lambda: f.read(4096), b""):
+            hasher.update(chunk)
+    actual_checksum = hasher.hexdigest()
+    assert actual_checksum == expected_checksum

tests/model_executor/model_loader/runai_streamer_loader/test_weight_utils.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import glob
+import tempfile
+
+import huggingface_hub.constants
+import torch
+from safetensors.torch import save_file
+
+from vllm.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf,
+    runai_safetensors_weights_iterator,
+    safetensors_weights_iterator,
+)
+
+
+def test_runai_safetensors_weights_iterator_clones_reused_buffers(
+    tmp_path, monkeypatch
+):
+    monkeypatch.setenv("RUNAI_STREAMER_MEMORY_LIMIT", "0")
+    weights_file = tmp_path / "model.safetensors"
+    expected_tensors = {
+        "first": torch.tensor([1.0, 2.0]),
+        "second": torch.tensor([3.0, 4.0]),
+    }
+    save_file(expected_tensors, weights_file)
+
+    actual_tensors = dict(
+        runai_safetensors_weights_iterator([str(weights_file)], False)
+    )
+
+    assert actual_tensors.keys() == expected_tensors.keys()
+    assert actual_tensors["first"].data_ptr() != actual_tensors["second"].data_ptr()
+    for name, expected_tensor in expected_tensors.items():
+        assert torch.equal(actual_tensors[name], expected_tensor)
+
+
+def test_runai_model_loader():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        huggingface_hub.constants.HF_HUB_OFFLINE = False
+        download_weights_from_hf(
+            "openai-community/gpt2", allow_patterns=["*.safetensors"], cache_dir=tmpdir
+        )
+        safetensors = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
+        assert len(safetensors) > 0
+
+        runai_model_streamer_tensors = {}
+        hf_safetensors_tensors = {}
+
+        for name, tensor in runai_safetensors_weights_iterator(safetensors, True):
+            runai_model_streamer_tensors[name] = tensor
+
+        for name, tensor in safetensors_weights_iterator(safetensors, True):
+            hf_safetensors_tensors[name] = tensor
+
+        assert len(runai_model_streamer_tensors) == len(hf_safetensors_tensors)
+
+        for name, runai_tensor in runai_model_streamer_tensors.items():
+            assert runai_tensor.dtype == hf_safetensors_tensors[name].dtype
+            assert runai_tensor.shape == hf_safetensors_tensors[name].shape
+            assert torch.all(runai_tensor.eq(hf_safetensors_tensors[name]))
+
+
+if __name__ == "__main__":
+    test_runai_model_loader()

tests/model_executor/model_loader/tensorizer_loader/conftest.py

@@ -0,0 +1,92 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Callable
+
+import pytest
+
+from vllm import LLM, EngineArgs
+from vllm.distributed import cleanup_dist_env_and_memory
+from vllm.model_executor.model_loader import tensorizer as tensorizer_mod
+from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
+from vllm.utils.network_utils import get_distributed_init_method, get_ip, get_open_port
+from vllm.v1.executor import UniProcExecutor
+from vllm.v1.worker.worker_base import WorkerWrapperBase
+
+MODEL_REF = "facebook/opt-125m"
+
+
+@pytest.fixture()
+def model_ref():
+    return MODEL_REF
+
+
+@pytest.fixture(autouse=True)
+def allow_insecure_serialization(monkeypatch):
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+
+
+@pytest.fixture(autouse=True)
+def cleanup():
+    cleanup_dist_env_and_memory(shutdown_ray=True)
+
+
+@pytest.fixture()
+def just_serialize_model_tensors(model_ref, monkeypatch, tmp_path):
+    def noop(*args, **kwargs):
+        return None
+
+    args = EngineArgs(model=model_ref)
+    tc = TensorizerConfig(tensorizer_uri=f"{tmp_path}/model.tensors")
+
+    monkeypatch.setattr(tensorizer_mod, "serialize_extra_artifacts", noop)
+
+    tensorizer_mod.tensorize_vllm_model(args, tc)
+    yield tmp_path
+
+
+@pytest.fixture(autouse=True)
+def tensorizer_config():
+    config = TensorizerConfig(tensorizer_uri="vllm")
+    return config
+
+
+@pytest.fixture()
+def model_path(model_ref, tmp_path):
+    yield tmp_path / model_ref / "model.tensors"
+
+
+def assert_from_collective_rpc(engine: LLM, closure: Callable, closure_kwargs: dict):
+    res = engine.collective_rpc(method=closure, kwargs=closure_kwargs)
+    return all(res)
+
+
+# This is an object pulled from tests/v1/engine/test_engine_core.py
+# Modified to strip the `load_model` method from its `_init_executor`
+# method. It's purely used as a dummy utility to run methods that test
+# Tensorizer functionality
+class DummyExecutor(UniProcExecutor):
+    def _init_executor(self) -> None:
+        """Initialize the worker and load the model."""
+        self.driver_worker = WorkerWrapperBase(rpc_rank=0)
+        distributed_init_method = get_distributed_init_method(get_ip(), get_open_port())
+        local_rank = 0
+        # set local rank as the device index if specified
+        device_info = self.vllm_config.device_config.device.__str__().split(":")
+        if len(device_info) > 1:
+            local_rank = int(device_info[1])
+        rank = 0
+        is_driver_worker = True
+        kwargs = dict(
+            vllm_config=self.vllm_config,
+            local_rank=local_rank,
+            rank=rank,
+            distributed_init_method=distributed_init_method,
+            is_driver_worker=is_driver_worker,
+        )
+        self.mm_receiver_cache = None
+        self.collective_rpc("init_worker", args=([kwargs],))
+        self.collective_rpc("init_device")
+
+    def shutdown(self):
+        if hasattr(self, "thread_pool"):
+            self.thread_pool.shutdown(wait=False)

tests/model_executor/model_loader/tensorizer_loader/test_tensorizer.py

@@ -0,0 +1,560 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import asyncio
+import gc
+import json
+import os
+import pathlib
+import subprocess
+import sys
+from typing import Any
+
+import pytest
+import torch
+
+import vllm.model_executor.model_loader.tensorizer
+from tests.utils import VLLM_PATH, RemoteOpenAIServer
+from vllm import LLM, SamplingParams
+from vllm.engine.arg_utils import EngineArgs
+from vllm.model_executor.model_loader.tensorizer import (
+    TensorizerConfig,
+    TensorSerializer,
+    is_vllm_tensorized,
+    open_stream,
+    tensorize_vllm_model,
+)
+from vllm.model_executor.model_loader.tensorizer_loader import (
+    BLACKLISTED_TENSORIZER_ARGS,
+)
+from vllm.utils.import_utils import PlaceholderModule
+
+from .conftest import DummyExecutor, assert_from_collective_rpc
+
+try:
+    import tensorizer
+    from tensorizer import EncryptionParams
+except ImportError:
+    tensorizer = PlaceholderModule("tensorizer")  # type: ignore[assignment]
+    EncryptionParams = tensorizer.placeholder_attr("EncryptionParams")
+
+
+class TensorizerCaughtError(Exception):
+    pass
+
+
+EXAMPLES_PATH = VLLM_PATH / "examples"
+
+pytest_plugins = ("pytest_asyncio",)
+
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+# Create a sampling params object.
+sampling_params = SamplingParams(temperature=0.8, top_p=0.95, seed=0)
+
+
+def patch_init_and_catch_error(self, obj, method_name, expected_error: type[Exception]):
+    original = getattr(obj, method_name, None)
+    if original is None:
+        raise ValueError("Method '{}' not found.".format(method_name))
+
+    def wrapper(*args, **kwargs):
+        try:
+            return original(*args, **kwargs)
+        except expected_error as err:
+            raise TensorizerCaughtError from err
+
+    setattr(obj, method_name, wrapper)
+
+    self.load_model()
+
+
+def assert_specific_tensorizer_error_is_raised(
+    executor,
+    obj: Any,
+    method_name: str,
+    expected_error: type[Exception],
+):
+    with pytest.raises(TensorizerCaughtError):
+        executor.collective_rpc(
+            patch_init_and_catch_error,
+            args=(
+                obj,
+                method_name,
+                expected_error,
+            ),
+        )
+
+
+def is_curl_installed():
+    try:
+        subprocess.check_call(["curl", "--version"])
+        return True
+    except (subprocess.CalledProcessError, FileNotFoundError):
+        return False
+
+
+def write_keyfile(keyfile_path: str):
+    encryption_params = EncryptionParams.random()
+    pathlib.Path(keyfile_path).parent.mkdir(parents=True, exist_ok=True)
+    with open(keyfile_path, "wb") as f:
+        f.write(encryption_params.key)
+
+
+@pytest.mark.skipif(not is_curl_installed(), reason="cURL is not installed")
+def test_deserialized_encrypted_vllm_model_has_same_outputs(
+    model_ref, vllm_runner, tmp_path, model_path
+):
+    args = EngineArgs(model=model_ref)
+    with vllm_runner(model_ref) as vllm_model:
+        key_path = tmp_path / model_ref / "model.key"
+        write_keyfile(key_path)
+
+        outputs = vllm_model.generate(prompts, sampling_params)
+
+    config_for_serializing = TensorizerConfig(
+        tensorizer_uri=str(model_path), encryption_keyfile=str(key_path)
+    )
+
+    tensorize_vllm_model(args, config_for_serializing)
+
+    config_for_deserializing = TensorizerConfig(
+        tensorizer_uri=str(model_path), encryption_keyfile=str(key_path)
+    )
+
+    with vllm_runner(
+        model_ref,
+        load_format="tensorizer",
+        model_loader_extra_config=config_for_deserializing,
+    ) as loaded_vllm_model:  # noqa: E501
+        deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
+        # noqa: E501
+
+        assert outputs == deserialized_outputs
+
+
+def test_deserialized_hf_model_has_same_outputs(
+    hf_runner, vllm_runner, tmp_path, model_ref, model_path
+):
+    with hf_runner(model_ref) as hf_model:
+        max_tokens = 50
+        outputs = hf_model.generate_greedy(prompts, max_tokens=max_tokens)
+        with open_stream(model_path, "wb+") as stream:
+            serializer = TensorSerializer(stream)
+            serializer.write_module(hf_model.model)
+
+    with vllm_runner(
+        model_ref,
+        load_format="tensorizer",
+        model_loader_extra_config=TensorizerConfig(
+            tensorizer_uri=str(model_path),
+            num_readers=1,
+        ),
+    ) as loaded_hf_model:
+        deserialized_outputs = loaded_hf_model.generate_greedy(
+            prompts, max_tokens=max_tokens
+        )
+
+        assert outputs == deserialized_outputs
+
+
+def test_load_without_tensorizer_load_format(vllm_runner, capfd, model_ref):
+    model = None
+    try:
+        model = vllm_runner(
+            model_ref, model_loader_extra_config=TensorizerConfig(tensorizer_uri="test")
+        )
+        pytest.fail("Expected RuntimeError for extra config keys")
+    except RuntimeError:
+        out, err = capfd.readouterr()
+        combined_output = out + err
+        assert (
+            "ValueError: Unexpected extra config keys for load format auto"
+        ) in combined_output
+    finally:
+        del model
+        gc.collect()
+        torch.accelerator.empty_cache()
+
+
+def test_raise_value_error_on_invalid_load_format(vllm_runner, capfd, model_ref):
+    model = None
+    try:
+        model = vllm_runner(
+            model_ref,
+            load_format="safetensors",
+            model_loader_extra_config=TensorizerConfig(tensorizer_uri="test"),
+        )
+        pytest.fail("Expected RuntimeError for extra config keys")
+    except RuntimeError:
+        out, err = capfd.readouterr()
+
+        combined_output = out + err
+        assert (
+            "ValueError: Unexpected extra config keys for load format safetensors"
+        ) in combined_output
+    finally:
+        del model
+        gc.collect()
+        torch.accelerator.empty_cache()
+
+
+@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Requires 2 GPUs")
+def test_tensorizer_with_tp_path_without_template(vllm_runner, capfd):
+    try:
+        model_ref = "EleutherAI/pythia-1.4b"
+        tensorized_path = f"s3://tensorized/{model_ref}/fp16/model.tensors"
+
+        vllm_runner(
+            model_ref,
+            load_format="tensorizer",
+            model_loader_extra_config=TensorizerConfig(
+                tensorizer_uri=tensorized_path,
+                num_readers=1,
+                s3_endpoint="object.ord1.coreweave.com",
+            ),
+            tensor_parallel_size=2,
+            disable_custom_all_reduce=True,
+        )
+    except RuntimeError:
+        out, err = capfd.readouterr()
+        combined_output = out + err
+        assert (
+            "ValueError: For a sharded model, tensorizer_uri "
+            "should include a string format template like '%04d' "
+            "to be formatted with the rank "
+            "of the shard"
+        ) in combined_output
+
+
+@pytest.mark.skipif(torch.accelerator.device_count() < 2, reason="Requires 2 GPUs")
+def test_deserialized_encrypted_vllm_model_with_tp_has_same_outputs(
+    vllm_runner, tmp_path
+):
+    model_ref = "EleutherAI/pythia-1.4b"
+    # record outputs from un-sharded un-tensorized model
+    with vllm_runner(
+        model_ref,
+        disable_custom_all_reduce=True,
+        enforce_eager=True,
+    ) as base_model:
+        outputs = base_model.generate(prompts, sampling_params)
+
+    # load model with two shards and serialize with encryption
+    model_path = str(tmp_path / model_ref / "model-%02d.tensors")
+    key_path = tmp_path / (model_ref + ".key")
+
+    tensorizer_config = TensorizerConfig(
+        tensorizer_uri=model_path,
+        encryption_keyfile=str(key_path),
+    )
+
+    tensorize_vllm_model(
+        engine_args=EngineArgs(
+            model=model_ref,
+            tensor_parallel_size=2,
+            disable_custom_all_reduce=True,
+            enforce_eager=True,
+        ),
+        tensorizer_config=tensorizer_config,
+    )
+    assert os.path.isfile(model_path % 0), "Serialization subprocess failed"
+    assert os.path.isfile(model_path % 1), "Serialization subprocess failed"
+
+    with vllm_runner(
+        model_ref,
+        tensor_parallel_size=2,
+        load_format="tensorizer",
+        disable_custom_all_reduce=True,
+        enforce_eager=True,
+        model_loader_extra_config=tensorizer_config,
+    ) as loaded_vllm_model:
+        deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
+
+    assert outputs == deserialized_outputs
+
+
+@pytest.mark.flaky(reruns=3)
+def test_vllm_tensorized_model_has_same_outputs(
+    model_ref, vllm_runner, tmp_path, model_path
+):
+    gc.collect()
+    torch.accelerator.empty_cache()
+    config = TensorizerConfig(tensorizer_uri=str(model_path))
+    args = EngineArgs(model=model_ref)
+
+    with vllm_runner(model_ref) as vllm_model:
+        outputs = vllm_model.generate(prompts, sampling_params)
+
+    tensorize_vllm_model(args, config)
+    assert is_vllm_tensorized(config)
+
+    with vllm_runner(
+        model_ref, load_format="tensorizer", model_loader_extra_config=config
+    ) as loaded_vllm_model:
+        deserialized_outputs = loaded_vllm_model.generate(prompts, sampling_params)
+        # noqa: E501
+
+        assert outputs == deserialized_outputs
+
+
+def test_load_with_just_model_tensors(just_serialize_model_tensors, model_ref):
+    # For backwards compatibility, ensure Tensorizer can be still be loaded
+    # for inference by passing the model reference name, not a local/S3 dir,
+    # and the location of the model tensors
+
+    model_dir = just_serialize_model_tensors
+
+    extra_config = {"tensorizer_uri": f"{model_dir}/model.tensors"}
+
+    ## Start OpenAI API server
+    args = [
+        "--load-format",
+        "tensorizer",
+        "--model-loader-extra-config",
+        json.dumps(extra_config),
+    ]
+
+    with RemoteOpenAIServer(model_ref, args):
+        # This test only concerns itself with being able to load the model
+        # and successfully initialize the server
+        pass
+
+
+def test_assert_serialization_kwargs_passed_to_tensor_serializer(tmp_path):
+    serialization_params = {
+        "limit_cpu_concurrency": 2,
+    }
+    model_ref = "facebook/opt-125m"
+    model_path = tmp_path / (model_ref + ".tensors")
+    config = TensorizerConfig(
+        tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
+    )
+    llm = LLM(
+        model=model_ref,
+    )
+
+    def serialization_test(self, *args, **kwargs):
+        # This is performed in the ephemeral worker process, so monkey-patching
+        # will actually work, and cleanup is guaranteed so don't
+        # need to reset things
+
+        original_dict = serialization_params
+        to_compare = {}
+
+        original = tensorizer.serialization.TensorSerializer.__init__
+
+        def tensorizer_serializer_wrapper(self, *args, **kwargs):
+            nonlocal to_compare
+            to_compare = kwargs.copy()
+            return original(self, *args, **kwargs)
+
+        tensorizer.serialization.TensorSerializer.__init__ = (
+            tensorizer_serializer_wrapper
+        )
+
+        tensorizer_config = TensorizerConfig(**kwargs["tensorizer_config"])
+        self.save_tensorized_model(
+            tensorizer_config=tensorizer_config,
+        )
+        return to_compare | original_dict == to_compare
+
+    kwargs = {"tensorizer_config": config.to_serializable()}
+
+    assert assert_from_collective_rpc(llm, serialization_test, kwargs)
+
+
+def test_assert_deserialization_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
+    deserialization_kwargs = {
+        "num_readers": "bar",  # illegal value
+    }
+
+    serialization_params = {
+        "limit_cpu_concurrency": 2,
+    }
+
+    model_ref = "facebook/opt-125m"
+    model_path = tmp_path / (model_ref + ".tensors")
+    config = TensorizerConfig(
+        tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
+    )
+
+    args = EngineArgs(model=model_ref)
+    tensorize_vllm_model(args, config)
+
+    loader_tc = TensorizerConfig(
+        tensorizer_uri=str(model_path),
+        deserialization_kwargs=deserialization_kwargs,
+    )
+
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        load_format="tensorizer",
+        model_loader_extra_config=loader_tc.to_serializable(),
+    )
+
+    vllm_config = engine_args.create_engine_config()
+    executor = DummyExecutor(vllm_config)
+
+    assert_specific_tensorizer_error_is_raised(
+        executor,
+        tensorizer.serialization.TensorDeserializer,
+        "__init__",
+        TypeError,
+    )
+
+
+def test_assert_stream_kwargs_passed_to_tensor_deserializer(tmp_path, capfd):
+    deserialization_kwargs = {
+        "num_readers": 1,
+    }
+
+    serialization_params = {
+        "limit_cpu_concurrency": 2,
+    }
+
+    model_ref = "facebook/opt-125m"
+    model_path = tmp_path / (model_ref + ".tensors")
+    config = TensorizerConfig(
+        tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
+    )
+
+    args = EngineArgs(model=model_ref)
+    tensorize_vllm_model(args, config)
+
+    stream_kwargs = {"mode": "foo"}
+
+    loader_tc = TensorizerConfig(
+        tensorizer_uri=str(model_path),
+        deserialization_kwargs=deserialization_kwargs,
+        stream_kwargs=stream_kwargs,
+    )
+
+    engine_args = EngineArgs(
+        model="facebook/opt-125m",
+        load_format="tensorizer",
+        model_loader_extra_config=loader_tc.to_serializable(),
+    )
+
+    vllm_config = engine_args.create_engine_config()
+    executor = DummyExecutor(vllm_config)
+
+    assert_specific_tensorizer_error_is_raised(
+        executor,
+        vllm.model_executor.model_loader.tensorizer,
+        "open_stream",
+        ValueError,
+    )
+
+
+@pytest.mark.asyncio
+async def test_serialize_and_serve_entrypoints(tmp_path):
+    model_ref = "facebook/opt-125m"
+
+    suffix = "test"
+    try:
+        result = subprocess.run(
+            [
+                sys.executable,
+                f"{VLLM_PATH}/examples/features/tensorize_vllm_model.py",
+                "--model",
+                model_ref,
+                "serialize",
+                "--serialized-directory",
+                str(tmp_path),
+                "--suffix",
+                suffix,
+                "--serialization-kwargs",
+                '{"limit_cpu_concurrency": 4}',
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+    except subprocess.CalledProcessError as e:
+        print("Tensorizing failed.")
+        print("STDOUT:\n", e.stdout)
+        print("STDERR:\n", e.stderr)
+        raise
+
+    assert "Successfully serialized" in result.stdout
+
+    # Next, try to serve with vllm serve
+    model_uri = tmp_path / "vllm" / model_ref / suffix / "model.tensors"
+
+    model_loader_extra_config = {
+        "tensorizer_uri": str(model_uri),
+        "stream_kwargs": {
+            "force_http": False,
+        },
+        "deserialization_kwargs": {
+            "verify_hash": True,
+            "num_readers": 8,
+        },
+    }
+
+    cmd = [
+        "-m",
+        "vllm.entrypoints.cli.main",
+        "serve",
+        "--host",
+        "localhost",
+        "--load-format",
+        "tensorizer",
+        model_ref,
+        "--model-loader-extra-config",
+        json.dumps(model_loader_extra_config, indent=2),
+    ]
+
+    proc = await asyncio.create_subprocess_exec(
+        sys.executable,
+        *cmd,
+        stdout=asyncio.subprocess.PIPE,
+        stderr=asyncio.subprocess.STDOUT,
+    )
+
+    assert proc.stdout is not None
+    fut = proc.stdout.readuntil(b"Application startup complete.")
+
+    try:
+        await asyncio.wait_for(fut, 180)
+    except asyncio.TimeoutError:
+        pytest.fail("Server did not start successfully")
+    finally:
+        proc.terminate()
+    await proc.communicate()
+
+
+@pytest.mark.parametrize("illegal_value", BLACKLISTED_TENSORIZER_ARGS)
+def test_blacklisted_parameter_for_loading(tmp_path, vllm_runner, capfd, illegal_value):
+    serialization_params = {
+        "limit_cpu_concurrency": 2,
+    }
+
+    model_ref = "facebook/opt-125m"
+    model_path = tmp_path / (model_ref + ".tensors")
+    config = TensorizerConfig(
+        tensorizer_uri=str(model_path), serialization_kwargs=serialization_params
+    )
+
+    args = EngineArgs(model=model_ref)
+    tensorize_vllm_model(args, config)
+
+    loader_tc = {"tensorizer_uri": str(model_path), illegal_value: "foo"}
+
+    try:
+        vllm_runner(
+            model_ref,
+            load_format="tensorizer",
+            model_loader_extra_config=loader_tc,
+        )
+    except RuntimeError:
+        out, err = capfd.readouterr()
+        combined_output = out + err
+        assert (
+            f"ValueError: {illegal_value} is not an allowed Tensorizer argument."
+        ) in combined_output

tests/model_executor/model_loader/test_ep_weight_filter.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for EP weight filtering during model loading."""
+
+import glob
+import tempfile
+
+import huggingface_hub.constants
+import pytest
+import torch
+
+from vllm.model_executor.model_loader.ep_weight_filter import (
+    compute_local_expert_ids,
+    parse_expert_id,
+    should_skip_weight,
+)
+from vllm.model_executor.model_loader.weight_utils import (
+    safetensors_weights_iterator,
+)
+
+# ---------------------------------------------------------------------------
+# Unit tests for parse_expert_id
+# ---------------------------------------------------------------------------
+
+
+class TestParseExpertId:
+    def test_routed_expert(self):
+        name = "model.layers.0.mlp.experts.42.gate_proj.weight"
+        assert parse_expert_id(name) == 42
+
+    def test_large_expert_id(self):
+        name = "model.layers.60.mlp.experts.383.down_proj.weight"
+        assert parse_expert_id(name) == 383
+
+    def test_shared_expert(self):
+        # Shared experts use a different naming convention in most models
+        name = "model.layers.0.mlp.shared_experts.gate_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_attention_weight(self):
+        name = "model.layers.0.self_attn.q_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_embedding(self):
+        name = "model.embed_tokens.weight"
+        assert parse_expert_id(name) is None
+
+    def test_layernorm(self):
+        name = "model.layers.0.input_layernorm.weight"
+        assert parse_expert_id(name) is None
+
+    def test_fused_3d_expert(self):
+        # 3D fused-expert tensors (e.g. gpt-oss) have no numeric expert id.
+        # They must NOT be filtered — slicing happens later in weight_loader.
+        name = "model.layers.0.mlp.experts.gate_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_fused_3d_expert_down_proj(self):
+        name = "model.layers.10.mlp.experts.down_proj.weight"
+        assert parse_expert_id(name) is None
+
+    def test_expert_scale(self):
+        # NVFP4 quantized models have scale tensors for experts
+        name = "model.layers.5.mlp.experts.100.gate_proj.weight_scale"
+        assert parse_expert_id(name) == 100
+
+    def test_expert_zero_id(self):
+        name = "model.layers.0.mlp.experts.0.up_proj.weight"
+        assert parse_expert_id(name) == 0
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for compute_local_expert_ids
+# ---------------------------------------------------------------------------
+
+
+class TestComputeLocalExpertIds:
+    def test_ep_disabled(self):
+        assert compute_local_expert_ids(64, ep_size=1, ep_rank=0) is None
+
+    def test_even_split(self):
+        # 64 experts, EP=8 → 8 per rank
+        ids = compute_local_expert_ids(64, ep_size=8, ep_rank=0)
+        assert ids == set(range(0, 8))
+
+        ids = compute_local_expert_ids(64, ep_size=8, ep_rank=7)
+        assert ids == set(range(56, 64))
+
+    def test_uneven_split(self):
+        # 10 experts, EP=3 → ranks get 4, 3, 3
+        ids_0 = compute_local_expert_ids(10, ep_size=3, ep_rank=0)
+        ids_1 = compute_local_expert_ids(10, ep_size=3, ep_rank=1)
+        ids_2 = compute_local_expert_ids(10, ep_size=3, ep_rank=2)
+
+        assert len(ids_0) == 4
+        assert len(ids_1) == 3
+        assert len(ids_2) == 3
+        # All experts covered, no overlap
+        assert ids_0 | ids_1 | ids_2 == set(range(10))
+        assert ids_0.isdisjoint(ids_1)
+        assert ids_1.isdisjoint(ids_2)
+
+    def test_384_experts_ep8(self):
+        # Kimi-K2.5 config: 384 experts, EP=8
+        for rank in range(8):
+            ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
+            assert len(ids) == 48
+
+        # All experts covered
+        all_ids = set()
+        for rank in range(8):
+            ids = compute_local_expert_ids(384, ep_size=8, ep_rank=rank)
+            all_ids |= ids
+        assert all_ids == set(range(384))
+
+    def test_384_experts_ep16(self):
+        for rank in range(16):
+            ids = compute_local_expert_ids(384, ep_size=16, ep_rank=rank)
+            assert len(ids) == 24
+
+    def test_384_experts_ep24(self):
+        # 384 / 24 = 16 exactly
+        for rank in range(24):
+            ids = compute_local_expert_ids(384, ep_size=24, ep_rank=rank)
+            assert len(ids) == 16
+
+    # round_robin placement tests
+
+    def test_round_robin_basic(self):
+        # 8 experts, EP=2: rank 0 → {0,2,4,6}, rank 1 → {1,3,5,7}
+        rr = "round_robin"
+        ids_0 = compute_local_expert_ids(8, 2, 0, placement=rr)
+        ids_1 = compute_local_expert_ids(8, 2, 1, placement=rr)
+        assert ids_0 == {0, 2, 4, 6}
+        assert ids_1 == {1, 3, 5, 7}
+
+    def test_round_robin_full_coverage(self):
+        # 384 experts, EP=8: all experts covered, no overlap
+        rr = "round_robin"
+        all_ids: set[int] = set()
+        for rank in range(8):
+            ids = compute_local_expert_ids(384, 8, rank, placement=rr)
+            assert ids is not None and len(ids) == 48
+            assert all_ids.isdisjoint(ids)
+            all_ids |= ids
+        assert all_ids == set(range(384))
+
+    def test_round_robin_uneven(self):
+        # 10 experts, EP=3: rank 0→{0,3,6,9}, rank 1→{1,4,7}, rank 2→{2,5,8}
+        rr = "round_robin"
+        ids_0 = compute_local_expert_ids(10, 3, 0, placement=rr)
+        ids_1 = compute_local_expert_ids(10, 3, 1, placement=rr)
+        ids_2 = compute_local_expert_ids(10, 3, 2, placement=rr)
+        assert ids_0 == {0, 3, 6, 9}
+        assert ids_1 == {1, 4, 7}
+        assert ids_2 == {2, 5, 8}
+        assert ids_0 | ids_1 | ids_2 == set(range(10))
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for should_skip_weight
+# ---------------------------------------------------------------------------
+
+
+class TestShouldSkipWeight:
+    def setup_method(self):
+        # Simulate EP=8, rank=0 → experts 0-47
+        self.local_ids = compute_local_expert_ids(384, ep_size=8, ep_rank=0)
+
+    def test_no_filter(self):
+        assert not should_skip_weight("anything", None)
+
+    def test_dense_not_skipped(self):
+        assert not should_skip_weight(
+            "model.layers.0.self_attn.q_proj.weight", self.local_ids
+        )
+
+    def test_local_expert_not_skipped(self):
+        assert not should_skip_weight(
+            "model.layers.0.mlp.experts.10.gate_proj.weight", self.local_ids
+        )
+
+    def test_remote_expert_skipped(self):
+        assert should_skip_weight(
+            "model.layers.0.mlp.experts.200.gate_proj.weight", self.local_ids
+        )
+
+    def test_boundary_expert(self):
+        # Expert 47 is local (last one), 48 is not
+        assert not should_skip_weight(
+            "model.layers.0.mlp.experts.47.gate_proj.weight", self.local_ids
+        )
+        assert should_skip_weight(
+            "model.layers.0.mlp.experts.48.gate_proj.weight", self.local_ids
+        )
+
+    def test_shared_expert_not_skipped(self):
+        assert not should_skip_weight(
+            "model.layers.0.mlp.shared_experts.gate_proj.weight", self.local_ids
+        )
+
+    def test_embedding_not_skipped(self):
+        assert not should_skip_weight("model.embed_tokens.weight", self.local_ids)
+
+    def test_fused_3d_expert_not_skipped(self):
+        # 3D fused-expert tensors (gpt-oss style) have no numeric id.
+        # Must not be skipped — weight_loader handles slicing later.
+        assert not should_skip_weight(
+            "model.layers.0.mlp.experts.gate_proj.weight", self.local_ids
+        )
+
+
+# ---------------------------------------------------------------------------
+# Integration test: safetensors_weights_iterator with EP filtering
+# ---------------------------------------------------------------------------
+
+
+class TestSafetensorsWeightsIteratorWithEpFilter:
+    """Verify that EP filtering produces a strict subset of unfiltered loading
+    and that all expected dense + local expert weights are present."""
+
+    @pytest.fixture(scope="class")
+    def gpt2_files(self):
+        """Download GPT-2 safetensors to a temp dir (shared across class)."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            huggingface_hub.constants.HF_HUB_OFFLINE = False
+            from vllm.model_executor.model_loader.weight_utils import (
+                download_weights_from_hf,
+            )
+
+            download_weights_from_hf(
+                "openai-community/gpt2",
+                allow_patterns=["*.safetensors"],
+                cache_dir=tmpdir,
+            )
+            files = glob.glob(f"{tmpdir}/**/*.safetensors", recursive=True)
+            assert len(files) > 0
+            yield files
+
+    def test_no_filter_returns_all(self, gpt2_files):
+        """With local_expert_ids=None, all weights are returned (no MoE)."""
+        all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
+        filtered_weights = dict(
+            safetensors_weights_iterator(gpt2_files, False, local_expert_ids=None)
+        )
+        assert set(all_weights.keys()) == set(filtered_weights.keys())
+
+    def test_empty_filter_skips_experts_only(self, gpt2_files):
+        """GPT-2 has no expert weights, so even an empty local_expert_ids
+        set should return all weights (all are dense)."""
+        all_weights = dict(safetensors_weights_iterator(gpt2_files, False))
+        filtered_weights = dict(
+            safetensors_weights_iterator(gpt2_files, False, local_expert_ids=set())
+        )
+        # GPT-2 has no experts, so nothing should be filtered
+        assert set(all_weights.keys()) == set(filtered_weights.keys())
+
+
+class TestEpFilterOnSyntheticMoeWeights:
+    """Create synthetic safetensors files with expert-like naming and verify
+    that the filter correctly skips non-local experts."""
+
+    @pytest.fixture
+    def synthetic_moe_files(self, tmp_path):
+        """Create synthetic safetensors with expert-patterned tensor names."""
+        from safetensors.torch import save_file
+
+        tensors = {}
+        # Dense weights
+        tensors["model.embed_tokens.weight"] = torch.randn(100, 64)
+        tensors["model.layers.0.self_attn.q_proj.weight"] = torch.randn(64, 64)
+        tensors["model.layers.0.input_layernorm.weight"] = torch.randn(64)
+        # Expert weights: 8 experts
+        for expert_id in range(8):
+            tensors[f"model.layers.0.mlp.experts.{expert_id}.gate_proj.weight"] = (
+                torch.randn(128, 64)
+            )
+            tensors[f"model.layers.0.mlp.experts.{expert_id}.up_proj.weight"] = (
+                torch.randn(128, 64)
+            )
+            tensors[f"model.layers.0.mlp.experts.{expert_id}.down_proj.weight"] = (
+                torch.randn(64, 128)
+            )
+        # Shared expert (should never be filtered)
+        tensors["model.layers.0.mlp.shared_experts.gate_proj.weight"] = torch.randn(
+            128, 64
+        )
+
+        filepath = str(tmp_path / "model-00001-of-00001.safetensors")
+        save_file(tensors, filepath)
+        return [filepath], tensors
+
+    def test_no_filter_returns_all(self, synthetic_moe_files):
+        files, expected = synthetic_moe_files
+        loaded = dict(safetensors_weights_iterator(files, False))
+        assert set(loaded.keys()) == set(expected.keys())
+
+    def test_ep2_rank0_gets_half_experts(self, synthetic_moe_files):
+        files, expected = synthetic_moe_files
+        # EP=2, rank=0 → experts 0-3
+        local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
+        loaded = dict(
+            safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+        )
+
+        # Should have all dense + shared + experts 0-3 only
+        for name in loaded:
+            eid = parse_expert_id(name)
+            if eid is not None:
+                assert eid in local_ids, f"Non-local expert {eid} was loaded"
+
+        # Check expert count: 4 experts × 3 weights = 12
+        expert_names = [n for n in loaded if parse_expert_id(n) is not None]
+        assert len(expert_names) == 4 * 3
+
+        # Check all dense weights present
+        assert "model.embed_tokens.weight" in loaded
+        assert "model.layers.0.self_attn.q_proj.weight" in loaded
+        assert "model.layers.0.input_layernorm.weight" in loaded
+        assert "model.layers.0.mlp.shared_experts.gate_proj.weight" in loaded
+
+    def test_ep2_rank1_gets_other_half(self, synthetic_moe_files):
+        files, expected = synthetic_moe_files
+        local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=1)
+        loaded = dict(
+            safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+        )
+
+        expert_names = [n for n in loaded if parse_expert_id(n) is not None]
+        assert len(expert_names) == 4 * 3
+        for name in expert_names:
+            assert parse_expert_id(name) in local_ids
+
+    def test_ep8_each_rank_gets_one_expert(self, synthetic_moe_files):
+        files, _ = synthetic_moe_files
+        all_expert_names = set()
+        for rank in range(8):
+            local_ids = compute_local_expert_ids(8, ep_size=8, ep_rank=rank)
+            loaded = dict(
+                safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+            )
+            expert_names = {n for n in loaded if parse_expert_id(n) is not None}
+            # 1 expert × 3 weights
+            assert len(expert_names) == 3
+            all_expert_names |= expert_names
+
+        # All 8 experts × 3 weights covered across ranks
+        assert len(all_expert_names) == 24
+
+    def test_tensor_values_match(self, synthetic_moe_files):
+        """Filtered tensors have identical values to unfiltered ones."""
+        files, _ = synthetic_moe_files
+        all_weights = dict(safetensors_weights_iterator(files, False))
+
+        local_ids = compute_local_expert_ids(8, ep_size=2, ep_rank=0)
+        filtered = dict(
+            safetensors_weights_iterator(files, False, local_expert_ids=local_ids)
+        )
+
+        for name, tensor in filtered.items():
+            assert torch.equal(tensor, all_weights[name]), f"Tensor mismatch for {name}"

tests/model_executor/model_loader/test_filter_duplicate_safetensors.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+import os
+import tempfile
+
+import pytest
+
+from vllm.model_executor.model_loader.weight_utils import (
+    filter_duplicate_safetensors_files,
+)
+
+
+def test_filter_duplicate_safetensors_files_missing_weight():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        existing_file = os.path.join(tmpdir, "model-00001-of-00002.safetensors")
+        with open(existing_file, "wb") as f:
+            f.write(b"")
+
+        existing_file2 = os.path.join(tmpdir, "model-00002-of-00002.safetensors")
+        with open(existing_file2, "wb") as f:
+            f.write(b"")
+
+        index_file = os.path.join(tmpdir, "model.safetensors.index.json")
+        index_content = {
+            "weight_map": {
+                "layer.0.weight": "model-00001-of-00002.safetensors",
+                "layer.1.weight": "model-00002-of-00002.safetensors",
+                "layer.2.weight": "model-00003-of-00002.safetensors",
+            }
+        }
+        with open(index_file, "w") as f:
+            json.dump(index_content, f)
+
+        hf_weights_files = [
+            os.path.join(tmpdir, "model-00001-of-00002.safetensors"),
+            os.path.join(tmpdir, "model-00002-of-00002.safetensors"),
+        ]
+
+        with pytest.raises(FileNotFoundError) as exc_info:
+            filter_duplicate_safetensors_files(
+                hf_weights_files=hf_weights_files,
+                hf_folder=tmpdir,
+                index_file="model.safetensors.index.json",
+            )
+
+        assert "model-00003-of-00002.safetensors" in str(exc_info.value)
+
+
+def test_filter_duplicate_safetensors_files_all_exist():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        existing_files = []
+        for i in range(1, 3):
+            file_path = os.path.join(tmpdir, f"model-0000{i}-of-00002.safetensors")
+            with open(file_path, "wb") as f:
+                f.write(b"")
+            existing_files.append(file_path)
+
+        index_file = os.path.join(tmpdir, "model.safetensors.index.json")
+        index_content = {
+            "weight_map": {
+                "layer.0.weight": "model-00001-of-00002.safetensors",
+                "layer.1.weight": "model-00002-of-00002.safetensors",
+            }
+        }
+        with open(index_file, "w") as f:
+            json.dump(index_content, f)
+
+        filter_duplicate_safetensors_files(
+            hf_weights_files=existing_files,
+            hf_folder=tmpdir,
+            index_file="model.safetensors.index.json",
+        )
+
+
+if __name__ == "__main__":
+    test_filter_duplicate_safetensors_files_missing_weight()
+    test_filter_duplicate_safetensors_files_all_exist()

tests/model_executor/model_loader/test_modelexpress_loader.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import sys
+from types import ModuleType, SimpleNamespace
+
+import pytest
+from torch import nn
+
+from vllm.config import VllmConfig
+from vllm.config.load import LoadConfig
+from vllm.model_executor.model_loader import get_model_loader
+from vllm.model_executor.model_loader.modelexpress_loader import (
+    ModelExpressModelLoader,
+)
+
+
+class FakeModelexpressLoader:
+    calls: list[tuple[str, tuple, dict]] = []
+    loaded_model: nn.Module
+
+    def __init__(self, load_config: LoadConfig):
+        self.load_config = load_config
+
+    def download_model(self, *args, **kwargs):
+        self.calls.append(("download_model", args, kwargs))
+
+    def load_weights(self, *args, **kwargs):
+        self.calls.append(("load_weights", args, kwargs))
+
+    def load_model(self, *args, **kwargs):
+        self.calls.append(("load_model", args, kwargs))
+        return self.loaded_model
+
+
+def _install_fake_modelexpress(monkeypatch):
+    FakeModelexpressLoader.calls = []
+    FakeModelexpressLoader.loaded_model = nn.Module()
+
+    for name in [
+        "modelexpress",
+        "modelexpress.engines",
+        "modelexpress.engines.vllm",
+    ]:
+        monkeypatch.setitem(sys.modules, name, ModuleType(name))
+
+    module = ModuleType("modelexpress.engines.vllm.loader")
+    module.__dict__["MxModelLoader"] = FakeModelexpressLoader
+    monkeypatch.setitem(sys.modules, module.__name__, module)
+
+
+def test_modelexpress_load_format_resolves_to_modelexpress_loader(monkeypatch):
+    _install_fake_modelexpress(monkeypatch)
+
+    loader = get_model_loader(LoadConfig(load_format="modelexpress"))
+
+    assert isinstance(loader, ModelExpressModelLoader)
+
+
+def test_modelexpress_loader_delegates_to_modelexpress(monkeypatch):
+    _install_fake_modelexpress(monkeypatch)
+    loader = ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
+    model = nn.Module()
+    model_config = SimpleNamespace()
+    vllm_config = SimpleNamespace()
+
+    loader.download_model(model_config)
+    loader.load_weights(model, model_config)
+    FakeModelexpressLoader.loaded_model.train()
+    result = loader.load_model(
+        vllm_config=vllm_config,
+        model_config=model_config,
+        prefix="model",
+    )
+
+    assert result is FakeModelexpressLoader.loaded_model
+    assert not result.training
+    assert FakeModelexpressLoader.calls == [
+        ("download_model", (model_config,), {}),
+        ("load_weights", (model, model_config), {}),
+        (
+            "load_model",
+            (),
+            {
+                "vllm_config": vllm_config,
+                "model_config": model_config,
+                "prefix": "model",
+            },
+        ),
+    ]
+
+
+def test_modelexpress_loader_missing_modelexpress_error(monkeypatch):
+    import importlib
+
+    def missing_modelexpress(name):
+        raise ModuleNotFoundError(name=name)
+
+    monkeypatch.setattr(importlib, "import_module", missing_modelexpress)
+
+    with pytest.raises(ImportError, match="requires the ModelExpress Python package"):
+        ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
+
+
+def test_modelexpress_loader_preserves_internal_import_errors(monkeypatch):
+    import importlib
+
+    def missing_dependency(name):
+        raise ModuleNotFoundError(name="not_modelexpress_dependency")
+
+    monkeypatch.setattr(importlib, "import_module", missing_dependency)
+
+    with pytest.raises(ModuleNotFoundError) as exc_info:
+        ModelExpressModelLoader(LoadConfig(load_format="modelexpress"))
+    assert exc_info.value.name == "not_modelexpress_dependency"
+
+
+def test_modelexpress_load_format_allows_object_storage_model_weights():
+    model_config = SimpleNamespace(
+        architecture="UnknownForTest",
+        config_updated=False,
+        convert_type=None,
+        is_hybrid=False,
+        model="test-model",
+        model_weights="s3://bucket/model",
+    )
+    vllm_config = object.__new__(VllmConfig)
+    vllm_config.model_config = model_config
+    vllm_config.load_config = LoadConfig(load_format="modelexpress")
+
+    vllm_config.try_verify_and_update_config()
+
+    assert vllm_config.load_config.load_format == "modelexpress"

tests/model_executor/model_loader/test_registry.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from torch import nn
+
+from vllm.config import ModelConfig
+from vllm.config.load import LoadConfig
+from vllm.model_executor.model_loader import get_model_loader, register_model_loader
+from vllm.model_executor.model_loader.base_loader import BaseModelLoader
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
+
+
+@register_model_loader("custom_load_format")
+class CustomModelLoader(BaseModelLoader):
+    def __init__(self, load_config: LoadConfig) -> None:
+        super().__init__(load_config)
+
+    def download_model(self, model_config: ModelConfig) -> None:
+        pass
+
+    def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None:
+        pass
+
+
+def test_register_model_loader():
+    load_config = LoadConfig(load_format="custom_load_format")
+    assert isinstance(get_model_loader(load_config), CustomModelLoader)
+
+
+def test_invalid_model_loader():
+    with pytest.raises(ValueError):
+
+        @register_model_loader("invalid_load_format")
+        class InValidModelLoader:
+            pass
+
+
+def test_default_loader_rejects_zero_num_threads():
+    # num_threads=0 used to fail late in ThreadPoolExecutor ("max_workers must be > 0").
+    with pytest.raises(ValueError, match="num_threads"):
+        DefaultModelLoader(
+            LoadConfig(
+                model_loader_extra_config={
+                    "enable_multithread_load": True,
+                    "num_threads": 0,
+                }
+            )
+        )
+
+
+def test_default_loader_rejects_multithread_with_non_lazy_strategy():
+    # The multi-thread loader ignores safetensors_load_strategy; reject the
+    # combination instead of silently dropping the requested strategy.
+    with pytest.raises(ValueError, match="does not support"):
+        DefaultModelLoader(
+            LoadConfig(
+                safetensors_load_strategy="torchao",
+                model_loader_extra_config={"enable_multithread_load": True},
+            )
+        )
+
+
+def test_default_loader_explicit_safetensors_does_not_misread_pt(tmp_path):
+    # Explicit safetensors must not fall back to a .pt and open it as safetensors.
+    (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00")
+    loader = DefaultModelLoader(LoadConfig(load_format="safetensors"))
+    with pytest.raises(RuntimeError, match="Cannot find any model weights"):
+        loader._prepare_weights(
+            str(tmp_path),
+            None,
+            None,
+            fall_back_to_pt=True,
+            allow_patterns_overrides=None,
+        )
+
+
+def test_default_loader_hf_still_falls_back_to_pt(tmp_path):
+    # Control: load_format="hf" still picks up .pt weights via fallback.
+    (tmp_path / "model.pt").write_bytes(b"\x00\x00\x00\x00")
+    loader = DefaultModelLoader(LoadConfig(load_format="hf"))
+    _, files, use_safetensors = loader._prepare_weights(
+        str(tmp_path),
+        None,
+        None,
+        fall_back_to_pt=True,
+        allow_patterns_overrides=None,
+    )
+    assert use_safetensors is False
+    assert any(f.endswith("model.pt") for f in files)

tests/model_executor/model_loader/test_reload.py

@@ -0,0 +1,498 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import gc
+import inspect
+from weakref import WeakKeyDictionary, ref
+
+import pytest
+import torch
+from torch.nn.parameter import UninitializedParameter
+
+import vllm.model_executor.model_loader.reload.meta as reload_meta
+from vllm.model_executor.layers.linear import QKVParallelLinear
+from vllm.model_executor.model_loader.reload.layerwise import (
+    finalize_layerwise_reload,
+    initialize_layerwise_reload,
+    record_metadata_for_reloading,
+)
+from vllm.model_executor.model_loader.reload.meta import (
+    capture_layer_to_meta,
+    get_numel_loaded,
+    materialize_layer,
+    materialize_meta_tensor,
+    restore_layer_on_meta,
+    to_meta_tensor,
+)
+from vllm.model_executor.model_loader.reload.types import LayerReloadingInfo
+from vllm.model_executor.model_loader.reload.utils import get_layer_tensors
+from vllm.model_executor.model_loader.weight_utils import (
+    composed_weight_loader,
+    default_weight_loader,
+)
+from vllm.platforms import current_platform
+
+
+def _fp8_reload_unsupported() -> bool:
+    """Whether the FP8 reload/online-quantize tests should be skipped.
+
+    ``supports_fp8()`` returns True on MI250 (gfx90a) because the general
+    quantization paths upcast FP8 weights, but gfx90a has no native FP8 and
+    cannot run these reload models, so treat it as unsupported here.
+    """
+    if not current_platform.supports_fp8():
+        return True
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx90a
+
+        return on_gfx90a()
+    return False
+
+
+class _AliasedBufferLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        weight = torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        self.weight = torch.nn.Parameter(weight)
+        self.register_buffer(
+            "weight_view", self.weight.detach().view(-1), persistent=False
+        )
+
+
+class _ParentAliasedChildBufferLayer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.ones(1))
+        self.conv1d = torch.nn.Linear(3, 2, bias=False)
+        self.conv1d.weight.data.copy_(
+            torch.arange(6, dtype=torch.float32).reshape(2, 3)
+        )
+        self.register_buffer(
+            "conv_weights", self.conv1d.weight.detach().view(-1), persistent=False
+        )
+
+
+class _AliasedBufferWithUninitializedChildLayer(_AliasedBufferLayer):
+    def __init__(self):
+        super().__init__()
+        self.child = torch.nn.Module()
+        self.child.register_parameter(
+            "lazy_weight", UninitializedParameter(requires_grad=False)
+        )
+
+
+def test_move_metatensors():
+    tensor = torch.empty((1, 2, 3))
+    meta_tensor = to_meta_tensor(tensor)
+    materialized_tensor = materialize_meta_tensor(meta_tensor)
+
+    assert meta_tensor.device.type == "meta"
+    assert tensor.device == materialized_tensor.device
+
+    assert tensor.dtype == meta_tensor.dtype == materialized_tensor.dtype
+    assert tensor.shape == meta_tensor.shape == materialized_tensor.shape
+    assert tensor.__class__ == meta_tensor.__class__ == materialized_tensor.__class__
+    assert tensor.__dict__ == meta_tensor.__dict__ == materialized_tensor.__dict__
+
+
+def test_reload_lifecycle():
+    layer = torch.nn.Linear(2, 3)
+    info = LayerReloadingInfo(
+        restore_metadata=capture_layer_to_meta(layer),
+        restore_device=torch.device("cpu"),
+    )
+
+    restore_layer_on_meta(layer, info)
+    for name, tensor in get_layer_tensors(layer).items():
+        meta_tensor = getattr(layer, name)
+        assert tensor.dtype == meta_tensor.dtype
+        assert tensor.shape == meta_tensor.shape
+        assert tensor.__class__ == meta_tensor.__class__
+        assert tensor.__dict__ == meta_tensor.__dict__
+
+    materialize_layer(layer, info)
+    for name, tensor in get_layer_tensors(layer).items():
+        materialized_tensor = getattr(layer, name)
+        assert tensor.dtype == materialized_tensor.dtype
+        assert tensor.shape == materialized_tensor.shape
+        assert tensor.__class__ == materialized_tensor.__class__
+        assert tensor.__dict__ == materialized_tensor.__dict__
+
+
+def test_materialize_layer_preserves_non_meta_tensors():
+    """Ensure that materialize_layer does not overwrite non meta tensors."""
+    layer = torch.nn.Linear(2, 3, bias=True)
+
+    # Create a non meta bias tensor and meta weight, which can happen with FP8
+    bias_values = torch.ones(3)
+    layer.bias.data.copy_(bias_values)
+    layer.weight = torch.nn.Parameter(layer.weight.data.to("meta"))
+
+    assert layer.weight.is_meta
+    assert not layer.bias.is_meta
+
+    # materialize the layer weights after the bias is initialized
+    info = LayerReloadingInfo(
+        restore_metadata=({}, {}),
+        restore_device=torch.device("cpu"),
+    )
+    materialize_layer(layer, info)
+
+    # Ensure the weight materialized off meta
+    assert not layer.weight.is_meta
+    assert layer.weight.device.type == "cpu"
+
+    # Ensure that the bias is (still) not meta and values are unchanged
+    assert not layer.bias.is_meta
+    assert torch.equal(layer.bias.data, bias_values)
+
+
+def test_model_cleanup(dist_init, default_vllm_config):
+    layer = QKVParallelLinear(2, 3, 4)
+    assert layer.weight.weight_loader.__self__ is layer
+    info = LayerReloadingInfo(
+        restore_metadata=capture_layer_to_meta(layer),
+        restore_device=torch.device("cpu"),
+    )
+
+    mock_info_dict: WeakKeyDictionary[torch.nn.Module, LayerReloadingInfo] = (
+        WeakKeyDictionary()
+    )
+    mock_info_dict[layer] = info
+    layer_ref = ref(layer)
+
+    del layer
+    gc.collect()
+
+    assert layer_ref() is None
+    assert len(mock_info_dict) == 0
+
+
+def test_get_numel_loaded():
+    param = torch.empty(10, device="meta")
+    loaded_weight = torch.empty(10)
+
+    def complex_weight_loader(param, loaded_weight):
+        param[:3] = loaded_weight[:3]
+        param[5:8] = loaded_weight[5:8]
+        return "value"
+
+    args = inspect.signature(complex_weight_loader).bind(param, loaded_weight)
+    num_loaded, ret = get_numel_loaded(complex_weight_loader, args)
+    assert num_loaded == 6
+    assert ret == "value"
+
+
+def test_get_numel_loaded_caps_at_param_size():
+    # composed_weight_loader copies into the param twice (the load and the
+    # in-place post-load transform), but only param.numel() distinct elements
+    # are loaded. get_numel_loaded must not double-count, otherwise a layer's
+    # loaded-element total can be reached early and trailing params get dropped.
+    param = torch.empty(10)
+    loaded_weight = torch.ones(10)
+    loader = composed_weight_loader(default_weight_loader, lambda x: x + 1)
+
+    args = inspect.signature(loader).bind(param, loaded_weight)
+    num_loaded, _ = get_numel_loaded(loader, args)
+    assert num_loaded == 10
+
+
+class _ComposedLoaderLayer(torch.nn.Module):
+    """Mimics a Mamba2 mixer's equal-numel direct params (A, D, dt_bias).
+
+    ``A`` uses ``composed_weight_loader`` (an extra in-place transform copy),
+    matching ``MambaMixer2`` where ``A`` is loaded as ``-exp(A_log)``.
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.A = torch.nn.Parameter(torch.empty(4, dtype=torch.float32))
+        self.D = torch.nn.Parameter(torch.ones(4))
+        self.dt_bias = torch.nn.Parameter(torch.ones(4))
+        self.A.weight_loader = composed_weight_loader(
+            default_weight_loader, lambda x: -torch.exp(x.float())
+        )
+        self.D.weight_loader = default_weight_loader
+        self.dt_bias.weight_loader = default_weight_loader
+
+
+def test_layerwise_reload_composed_loader_does_not_drop_params(monkeypatch):
+    # Regression test: a composed_weight_loader param (A) used to double-count
+    # its elements, finalizing the layer before the trailing param (D) was
+    # loaded and leaving it as uninitialized materialized memory.
+    layer = _ComposedLoaderLayer()
+    model = torch.nn.Sequential(layer)
+
+    def materialize_with_sentinel(meta_tensor):
+        tensor = torch.empty_strided(
+            size=tuple(meta_tensor.size()),
+            stride=tuple(meta_tensor.stride()),
+            dtype=meta_tensor.dtype,
+            requires_grad=False,
+        )
+        tensor.fill_(float("nan"))
+        tensor.__class__ = meta_tensor.__class__
+        tensor.__dict__ = meta_tensor.__dict__.copy()
+        return tensor
+
+    monkeypatch.setattr(
+        reload_meta, "materialize_meta_tensor", materialize_with_sentinel
+    )
+
+    loaded = {
+        "A": torch.full((4,), 0.5),
+        "dt_bias": torch.full((4,), 3.0),
+        "D": torch.full((4,), 7.0),
+    }
+
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    # Mimic real load_weights: resolve params once, then load in checkpoint
+    # order with D last (the param that was dropped).
+    params = dict(layer.named_parameters())
+    for name in ("A", "dt_bias", "D"):
+        param = params[name]
+        param.weight_loader(param, loaded[name])
+    finalize_layerwise_reload(model, model_config=None)
+
+    assert torch.equal(layer.A, -torch.exp(loaded["A"]))
+    assert torch.equal(layer.dt_bias, loaded["dt_bias"])
+    assert torch.equal(layer.D, loaded["D"])
+
+
+def test_layerwise_reload_skips_non_persistent_parameter_alias_buffers(monkeypatch):
+    layer = _AliasedBufferLayer()
+    model = torch.nn.Sequential(layer)
+    loaded_weight = torch.full_like(layer.weight, 7.0)
+
+    def materialize_with_sentinel(meta_tensor):
+        tensor = torch.empty_strided(
+            size=tuple(meta_tensor.size()),
+            stride=tuple(meta_tensor.stride()),
+            dtype=meta_tensor.dtype,
+            requires_grad=False,
+        )
+        tensor.fill_(-123.0)
+        tensor.__class__ = meta_tensor.__class__
+        tensor.__dict__ = meta_tensor.__dict__.copy()
+        return tensor
+
+    monkeypatch.setattr(
+        reload_meta, "materialize_meta_tensor", materialize_with_sentinel
+    )
+
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    layer.weight.weight_loader(layer.weight, loaded_weight)
+    finalize_layerwise_reload(model, model_config=None)
+
+    assert torch.equal(layer.weight, loaded_weight)
+    assert layer.weight_view.untyped_storage().data_ptr() == (
+        layer.weight.untyped_storage().data_ptr()
+    )
+
+
+def test_capture_layer_to_meta_skips_uninitialized_parameter_storage_ptrs():
+    layer = _AliasedBufferWithUninitializedChildLayer()
+
+    _, buffers = capture_layer_to_meta(layer)
+
+    assert "weight_view" not in buffers
+
+
+def test_layerwise_reload_skips_child_parameter_alias_buffers(monkeypatch):
+    layer = _ParentAliasedChildBufferLayer()
+    model = torch.nn.Sequential(layer)
+    loaded_conv = torch.full_like(layer.conv1d.weight, 7.0)
+    loaded_scale = torch.full_like(layer.scale, 3.0)
+
+    def materialize_with_sentinel(meta_tensor):
+        tensor = torch.empty_strided(
+            size=tuple(meta_tensor.size()),
+            stride=tuple(meta_tensor.stride()),
+            dtype=meta_tensor.dtype,
+            requires_grad=False,
+        )
+        tensor.fill_(-123.0)
+        tensor.__class__ = meta_tensor.__class__
+        tensor.__dict__ = meta_tensor.__dict__.copy()
+        return tensor
+
+    monkeypatch.setattr(
+        reload_meta, "materialize_meta_tensor", materialize_with_sentinel
+    )
+
+    record_metadata_for_reloading(model)
+    initialize_layerwise_reload(model)
+    layer.conv1d.weight.weight_loader(layer.conv1d.weight, loaded_conv)
+    layer.scale.weight_loader(layer.scale, loaded_scale)
+    finalize_layerwise_reload(model, model_config=None)
+
+    assert torch.equal(layer.conv1d.weight, loaded_conv)
+    assert torch.equal(layer.conv_weights, loaded_conv.view(-1))
+    assert layer.conv_weights.untyped_storage().data_ptr() == (
+        layer.conv1d.weight.untyped_storage().data_ptr()
+    )
+
+
+@pytest.mark.parametrize(
+    "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
+)
+@pytest.mark.parametrize(
+    "base_model,mul_model,add_model",
+    [
+        pytest.param(
+            "Qwen/Qwen3-0.6B",
+            "inference-optimization/Qwen3-0.6B-debug-multiply",
+            "inference-optimization/Qwen3-0.6B-debug-add",
+            marks=[pytest.mark.slow_test],
+        ),
+        pytest.param(
+            "inference-optimization/Qwen3-0.6B-FP8_BLOCK",
+            "inference-optimization/Qwen3-0.6B-debug-multiply-FP8_BLOCK",
+            "inference-optimization/Qwen3-0.6B-debug-add-FP8_BLOCK",
+            marks=[pytest.mark.slow_test],
+        ),
+        pytest.param(
+            "inference-optimization/Qwen3-0.6B-W4A16-G128",
+            "inference-optimization/Qwen3-0.6B-debug-multiply-W4A16-G128",
+            "inference-optimization/Qwen3-0.6B-debug-add-W4A16-G128",
+            marks=[pytest.mark.slow_test],
+        ),
+        pytest.param(
+            "inference-optimization/DeepSeek-V3-debug-empty",
+            "inference-optimization/DeepSeek-V3-debug-multiply",
+            "inference-optimization/DeepSeek-V3-debug-add",
+            marks=[pytest.mark.slow_test],
+        ),
+        pytest.param(
+            "inference-optimization/DeepSeek-V3-debug-empty-FP8_DYNAMIC",
+            "inference-optimization/DeepSeek-V3-debug-multiply-FP8_DYNAMIC",
+            "inference-optimization/DeepSeek-V3-debug-add-FP8_DYNAMIC",
+        ),
+        pytest.param(
+            "inference-optimization/DeepSeek-V3-debug-empty-NVFP4A16",
+            "inference-optimization/DeepSeek-V3-debug-multiply-NVFP4A16",
+            "inference-optimization/DeepSeek-V3-debug-add-NVFP4A16",
+            marks=[pytest.mark.slow_test],
+        ),
+    ],
+)
+def test_reload_weights(base_model, mul_model, add_model, tp_size, vllm_runner):
+    if current_platform.device_count() < tp_size:
+        pytest.skip(reason="Not enough CUDA devices")
+
+    if "FP8" in base_model and _fp8_reload_unsupported():
+        pytest.skip(reason="Requires FP8 support")
+
+    with vllm_runner(
+        model_name=base_model,
+        tensor_parallel_size=tp_size,
+        enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
+        enable_prefix_caching=False,
+        max_model_len=16,
+        max_num_seqs=1,
+    ) as llm:
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
+        mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
+        add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
+        assert mul_perp < add_perp
+
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})
+        mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
+        add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
+        assert add_perp < mul_perp
+
+
+def test_kv_scale_reload(vllm_runner):
+    """Test reloading a checkpoint that contains k_scale/v_scale weights."""
+    if _fp8_reload_unsupported():
+        pytest.skip(reason="Requires FP8 support")
+
+    model = "nm-testing/Llama-3.2-1B-Instruct-FP8-KV"
+
+    # Load dummy weights, then reload real checkpoint
+    with vllm_runner(
+        model_name=model,
+        load_format="dummy",
+        enable_prefix_caching=False,
+        max_model_len=16,
+        max_num_seqs=1,
+    ) as llm:
+        llm.collective_rpc(
+            "update_config",
+            kwargs={"overrides": {"load_config": {"load_format": "auto"}}},
+        )
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": model})
+        reloaded_perp = llm.generate_prompt_perplexity(
+            ["The capital of France is the city of Paris"],
+            mask=["The capital of France is"],
+        )[0]
+
+    assert reloaded_perp < 10
+
+
+@pytest.mark.parametrize(
+    "tp_size", [pytest.param(1), pytest.param(2, marks=[pytest.mark.slow_test])]
+)
+@pytest.mark.parametrize(
+    "base_model,mul_model,add_model,quantization",
+    [
+        pytest.param(
+            "Qwen/Qwen3-0.6B",
+            "inference-optimization/Qwen3-0.6B-debug-multiply",
+            "inference-optimization/Qwen3-0.6B-debug-add",
+            "fp8",
+        ),
+        pytest.param(
+            "inference-optimization/DeepSeek-V3-debug-empty",
+            "inference-optimization/DeepSeek-V3-debug-multiply",
+            "inference-optimization/DeepSeek-V3-debug-add",
+            "fp8",
+            marks=[pytest.mark.slow_test],
+        ),
+        pytest.param(
+            "Qwen/Qwen3-0.6B",
+            "inference-optimization/Qwen3-0.6B-debug-multiply",
+            "inference-optimization/Qwen3-0.6B-debug-add",
+            "mxfp8",
+            marks=[pytest.mark.slow_test],
+        ),
+        pytest.param(
+            "inference-optimization/DeepSeek-V3-debug-empty",
+            "inference-optimization/DeepSeek-V3-debug-multiply",
+            "inference-optimization/DeepSeek-V3-debug-add",
+            "mxfp8",
+            marks=[
+                pytest.mark.slow_test,
+                pytest.mark.xfail(reason="mxfp4 & mla is not supported yet"),
+            ],
+        ),
+    ],
+)
+def test_online_quantize_reload(
+    base_model, mul_model, add_model, quantization, tp_size, vllm_runner
+):
+    if current_platform.device_count() < tp_size:
+        pytest.skip(reason="Not enough GPU devices")
+
+    if quantization == "fp8" and _fp8_reload_unsupported():
+        pytest.skip(reason="Requires FP8 support")
+
+    with vllm_runner(
+        model_name=base_model,
+        quantization=quantization,
+        tensor_parallel_size=tp_size,
+        enable_expert_parallel=(tp_size > 1 and "DeepSeek" in base_model),
+        enable_prefix_caching=False,
+        max_model_len=16,
+        max_num_seqs=1,
+    ) as llm:
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": mul_model})
+        mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
+        add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
+        assert mul_perp < add_perp
+
+        llm.collective_rpc("reload_weights", kwargs={"weights_path": add_model})
+        mul_perp = llm.generate_prompt_perplexity(["3 4 = 12"], mask=["3 4 ="])[0]
+        add_perp = llm.generate_prompt_perplexity(["3 4 = 7"], mask=["3 4 ="])[0]
+        assert add_perp < mul_perp

tests/model_executor/model_loader/test_sharded_state_loader.py

@@ -0,0 +1,165 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import fnmatch
+import multiprocessing as mp
+import os
+import shutil
+from tempfile import TemporaryDirectory
+
+import pytest
+import torch
+from huggingface_hub import snapshot_download
+
+from vllm import LLM, SamplingParams
+from vllm.model_executor.model_loader import ShardedStateLoader
+from vllm.platforms import current_platform
+
+prompts = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+# Create a sampling params object.
+sampling_params = SamplingParams(
+    temperature=0,
+    max_tokens=256,
+    ignore_eos=True,
+)
+
+
+def test_filter_subtensors():
+    state_dict = {
+        "a": torch.empty(2),
+        "b": torch.empty((2, 4)),
+        "c": torch.empty((2, 4, 8)),
+    }
+    state_dict.update(
+        {
+            "x": state_dict["b"],
+            "y": state_dict["c"][1, 2, :],
+            "z": state_dict["c"][1, :, 4],
+        }
+    )
+    filtered_state_dict = ShardedStateLoader._filter_subtensors(state_dict)
+    assert tuple(filtered_state_dict.keys()) == ("a", "b", "c")
+    for key, tensor in filtered_state_dict.items():
+        # NOTE: don't use `equal` here, as the tensor might contain NaNs
+        assert tensor is state_dict[key]
+
+
+@pytest.fixture(scope="module")
+def llama_3p2_1b_files():
+    input_dir = snapshot_download(
+        "meta-llama/Llama-3.2-1B-Instruct", ignore_patterns=["*.bin*", "original/*"]
+    )
+
+    yield input_dir
+
+
+def _run_writer(input_dir, output_dir, weights_patterns, **kwargs):
+    llm_sharded_writer = LLM(model=input_dir, **kwargs)
+
+    # Dump worker states to output directory
+    llm_sharded_writer.llm_engine.engine_core.save_sharded_state(path=output_dir)
+
+    # Copy metadata files to output directory
+    for file in os.listdir(input_dir):
+        if os.path.isdir(os.path.join(input_dir, file)):
+            shutil.copytree(
+                os.path.join(input_dir, file), os.path.join(output_dir, file)
+            )
+        elif not any(fnmatch.fnmatch(file, ext) for ext in weights_patterns):
+            shutil.copy(os.path.join(input_dir, file), output_dir)
+
+
+def _run_generate(input_dir, queue: mp.Queue, **kwargs):
+    llm = LLM(model=input_dir, **kwargs)
+    gen = llm.generate(prompts, sampling_params)
+    queue.put([g.outputs[0].__dict__ for g in gen])
+    queue.close()
+    queue.join_thread()
+
+
+@pytest.mark.parametrize("enable_lora", [False, True])
+@pytest.mark.parametrize("tp_size", [1, 2])
+def test_sharded_state_loader(
+    enable_lora, tp_size, num_gpus_available, llama_3p2_1b_files
+):
+    if num_gpus_available < tp_size:
+        pytest.skip(f"Not enough GPUs for tensor parallelism {tp_size}")
+
+    weights_patterns = ("*.safetensors",)
+    gpu_memory_utilization = 0.8
+    input_dir = llama_3p2_1b_files
+    ctx = mp.get_context("spawn")
+
+    platform_args = {}
+    if current_platform.is_rocm() or current_platform.is_xpu():
+        platform_args["max_num_seqs"] = 1
+
+    # Run in separate processes for memory & CUDA isolation
+    with TemporaryDirectory() as output_dir:
+        p = ctx.Process(
+            target=_run_writer,
+            args=(input_dir, output_dir, weights_patterns),
+            kwargs=dict(
+                tensor_parallel_size=tp_size,
+                gpu_memory_utilization=gpu_memory_utilization,
+                enforce_eager=True,
+                **platform_args,
+            ),
+        )
+        p.start()
+        p.join()
+
+        queue = ctx.Queue()
+
+        p = ctx.Process(
+            target=_run_generate,
+            args=(input_dir, queue),
+            kwargs=dict(
+                enable_lora=enable_lora,
+                gpu_memory_utilization=gpu_memory_utilization,
+                tensor_parallel_size=tp_size,
+                **platform_args,
+            ),
+        )
+        p.start()
+        # Call queue.get() before p.join() to prevent deadlock:
+        # If p.join() is called before queue.get() and the queue is full,
+        # the child process may block while writing to the queue and never
+        # terminate, causing the parent to wait indefinitely on p.join().
+        # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814
+        out_before = queue.get()
+        p.join()
+        queue.close()
+        queue.join_thread()
+
+        queue = ctx.Queue()
+
+        p = ctx.Process(
+            target=_run_generate,
+            args=(output_dir, queue),
+            kwargs=dict(
+                enable_lora=enable_lora,
+                gpu_memory_utilization=gpu_memory_utilization,
+                tensor_parallel_size=tp_size,
+                load_format="sharded_state",
+                **platform_args,
+            ),
+        )
+        p.start()
+        # Call queue.get() before p.join() to prevent deadlock:
+        # If p.join() is called before queue.get() and the queue is full,
+        # the child process may block while writing to the queue and never
+        # terminate, causing the parent to wait indefinitely on p.join().
+        # See: https://github.com/vllm-project/vllm/pull/22371#discussion_r2257773814
+        out_after = queue.get()
+        p.join()
+        queue.close()
+        queue.join_thread()
+
+        assert out_before == out_after