Sync from v0.13

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,51 @@
+# 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",
+)
+def test_fastsafetensors_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
+
+        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]))
+
+
+if __name__ == "__main__":
+    test_fastsafetensors_model_loader()

tests/model_executor/model_loader/runai_streamer_loader/conftest.py

@@ -0,0 +1,39 @@
+# 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,
+        )
+
+        wrapper_kwargs = {
+            "vllm_config": self.vllm_config,
+        }
+
+        self.driver_worker = WorkerWrapperBase(**wrapper_kwargs)
+
+        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,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm import SamplingParams
+from vllm.config.load import LoadConfig
+from vllm.model_executor.model_loader import get_model_loader
+
+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
+
+
+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

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,59 @@
+# 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 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,44 @@
+# 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 vllm.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf,
+    runai_safetensors_weights_iterator,
+    safetensors_weights_iterator,
+)
+
+
+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,96 @@
+# 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(vllm_config=self.vllm_config, 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")
+
+    @property
+    def max_concurrent_batches(self) -> int:
+        return 2
+
+    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.cuda.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.cuda.empty_cache()
+
+
+@pytest.mark.skipif(torch.cuda.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.cuda.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.cuda.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/others/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_registry.py

@@ -0,0 +1,35 @@
+# 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
+
+
+@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

tests/model_executor/model_loader/test_sharded_state_loader.py

@@ -0,0 +1,157 @@
+# 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
+
+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")
+
+    # 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,
+            ),
+        )
+        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,
+            ),
+        )
+        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",
+            ),
+        )
+        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

tests/model_executor/test_eagle_quantization.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from vllm.config import LoadConfig, ModelConfig, SpeculativeConfig, VllmConfig
+from vllm.model_executor.models.utils import get_draft_quant_config
+from vllm.platforms import current_platform
+
+DEVICES = (
+    [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
+    if current_platform.is_cuda_alike()
+    else ["cpu"]
+)
+
+
+def test_get_draft_quant_config_with_draft_model():
+    mock_draft_model_config = Mock(spec=ModelConfig)
+    mock_load_config = Mock(spec=LoadConfig)
+    mock_speculative_config = Mock(spec=SpeculativeConfig)
+    mock_speculative_config.draft_model_config = mock_draft_model_config
+
+    mock_vllm_config = Mock(spec=VllmConfig)
+    mock_vllm_config.speculative_config = mock_speculative_config
+    mock_vllm_config.load_config = mock_load_config
+
+    mock_quant_config = Mock()
+    with patch.object(
+        VllmConfig, "get_quantization_config", return_value=mock_quant_config
+    ):
+        result = get_draft_quant_config(mock_vllm_config)
+
+        # Verify the function calls get_quantization_config with draft model config
+        VllmConfig.get_quantization_config.assert_called_once_with(
+            mock_draft_model_config, mock_load_config
+        )
+        assert result == mock_quant_config
+
+
+def test_get_draft_quant_config_without_draft_model():
+    mock_speculative_config = Mock(spec=SpeculativeConfig)
+    mock_speculative_config.draft_model_config = None
+
+    mock_vllm_config = Mock(spec=VllmConfig)
+    mock_vllm_config.speculative_config = mock_speculative_config
+    mock_vllm_config.load_config = Mock(spec=LoadConfig)
+
+    result = get_draft_quant_config(mock_vllm_config)
+
+    assert result is None
+
+
+@torch.inference_mode()
+@pytest.mark.parametrize("device", DEVICES)
+def test_fc_layer_quant_config_usage(dist_init, device) -> None:
+    import torch
+
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    if current_platform.is_cuda_alike():
+        torch.cuda.set_device(device)
+
+    torch.set_default_device(device)
+
+    input_size = 256
+    output_size = 128
+
+    fc_no_quant = ReplicatedLinear(
+        input_size=input_size,
+        output_size=output_size,
+        bias=False,
+        params_dtype=torch.float16,
+        quant_config=None,
+        prefix="fc",
+    )
+
+    assert fc_no_quant.quant_config is None
+    assert fc_no_quant.input_size == input_size
+    assert fc_no_quant.output_size == output_size
+
+    mock_quant_config = Mock()
+    fc_with_quant = ReplicatedLinear(
+        input_size=input_size,
+        output_size=output_size,
+        bias=False,
+        params_dtype=torch.float16,
+        quant_config=mock_quant_config,
+        prefix="fc",
+    )
+
+    assert fc_with_quant.quant_config == mock_quant_config
+
+    # Check forward pass
+    x = torch.randn(2, input_size, dtype=torch.float16)
+    output, _ = fc_no_quant(x)
+    assert output.shape == (2, output_size)
+
+
+def test_kv_cache_scale_name_handling():
+    # Mock a quant config that supports cache scales
+    mock_quant_config = Mock()
+    mock_quant_config.get_cache_scale = Mock(return_value="layers.0.self_attn.kv_scale")
+
+    # Condition check in load_weights
+    name = "layers.0.self_attn.k_proj.weight"
+    scale_name = mock_quant_config.get_cache_scale(name)
+
+    # Check if get_cache_scale is called and returns expected value
+    mock_quant_config.get_cache_scale.assert_called_once_with(name)
+    assert scale_name == "layers.0.self_attn.kv_scale"
+
+
+def test_kv_cache_scale_name_no_scale():
+    # Mock a quant config that returns None for get_cache_scale
+    mock_quant_config = Mock()
+    mock_quant_config.get_cache_scale = Mock(return_value=None)
+
+    name = "layers.0.mlp.gate_proj.weight"
+    scale_name = mock_quant_config.get_cache_scale(name)
+
+    # Should return None for weights that don't have cache scales
+    assert scale_name is None
+
+
+def test_maybe_remap_kv_scale_name():
+    from vllm.model_executor.model_loader.weight_utils import maybe_remap_kv_scale_name
+
+    params_dict = {
+        "layers.0.self_attn.kv_scale": Mock(),
+        "layers.1.self_attn.kv_scale": Mock(),
+    }
+
+    name = "layers.0.self_attn.some_scale"
+    remapped = maybe_remap_kv_scale_name(name, params_dict)
+
+    assert remapped in params_dict or remapped == name or remapped is None
+
+
+def test_load_weights_kv_scale_handling():
+    kv_scale_param = Mock()
+    kv_scale_param.weight_loader = Mock()
+
+    params_dict = {
+        "layers.0.self_attn.kv_scale": kv_scale_param,
+    }
+
+    mock_quant_config = Mock()
+    mock_quant_config.get_cache_scale = Mock(return_value="layers.0.self_attn.kv_scale")
+
+    # Load_weights logic for KV cache scales
+    name = "layers.0.self_attn.k_proj.weight"
+    loaded_weight_tensor = torch.tensor([1.0, 2.0])
+
+    if mock_quant_config is not None:
+        scale_name = mock_quant_config.get_cache_scale(name)
+        if scale_name:
+            param = params_dict[scale_name]
+            assert param is kv_scale_param
+            weight_to_load = (
+                loaded_weight_tensor
+                if loaded_weight_tensor.dim() == 0
+                else loaded_weight_tensor[0]
+            )
+
+            assert scale_name == "layers.0.self_attn.kv_scale"
+            assert weight_to_load == loaded_weight_tensor[0]

tests/model_executor/test_enabled_custom_ops.py

@@ -0,0 +1,169 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config import (
+    CompilationConfig,
+    VllmConfig,
+    get_cached_compilation_config,
+    set_current_vllm_config,
+)
+from vllm.model_executor.custom_op import CustomOp
+from vllm.model_executor.layers.activation import (
+    GeluAndMul,
+    ReLUSquaredActivation,
+    SiluAndMul,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+    dispatch_topk_func,
+    vllm_topk_softmax,
+)
+from vllm.model_executor.layers.layernorm import (
+    RMSNorm,
+    dispatch_rocm_rmsnorm_func,
+    fused_add_rms_norm,
+    rms_norm,
+)
+from vllm.platforms import current_platform
+
+RMS_NORM_SUPPORTED_DTYPES = [torch.float16, torch.bfloat16]
+
+
+# Registered subclass for test
+@CustomOp.register("relu3")
+class Relu3(ReLUSquaredActivation):
+    pass
+
+
+@pytest.mark.parametrize(
+    "env, compilation_mode, backend, ops_enabled, default_on",
+    [
+        # Default values based on compile level
+        # - All by default (no Inductor compilation)
+        (None, 0, "eager", [True] * 4, True),
+        (None, 1, "eager", [True] * 4, True),
+        (None, 2, "eager", [True] * 4, True),
+        (None, 3, "eager", [True] * 4, True),
+        # - None by default (with Inductor)
+        (None, 0, "inductor", [True] * 4, True),
+        # - None by default (with Inductor)
+        (None, 1, "inductor", [False] * 4, False),
+        (None, 2, "inductor", [False] * 4, False),
+        (None, 3, "inductor", [False] * 4, False),
+        # Explicitly enabling/disabling
+        #
+        # Default: all
+        #
+        # All but SiluAndMul
+        ("+rms_norm,-silu_and_mul", 0, "inductor", [1, 0, 1, 1], True),
+        # Only ReLU3
+        ("none,-rms_norm,+relu3", 1, "eager", [0, 0, 0, 1], False),
+        # All but SiluAndMul
+        ("all,-silu_and_mul", 2, "inductor", [1, 0, 1, 1], True),
+        # All but ReLU3 (even if ReLU2 is on)
+        ("-relu3,+relu2", 3, "eager", [1, 1, 1, 0], True),
+        # RMSNorm and SiluAndMul
+        ("none,-relu3,+rms_norm,+silu_and_mul", 3, "eager", [1, 1, 0, 0], False),
+        # All but RMSNorm
+        ("-rms_norm", 3, "eager", [0, 1, 1, 1], True),
+        #
+        # Default: none
+        #
+        # Only ReLU3
+        ("none,+relu3", 3, "inductor", [0, 0, 0, 1], False),
+        # All but RMSNorm
+        ("all,-rms_norm", 3, "inductor", [0, 1, 1, 1], True),
+    ],
+)
+def test_enabled_ops(
+    env: str | None,
+    compilation_mode: int,
+    backend: str,
+    ops_enabled: list[int],
+    default_on: bool,
+):
+    custom_ops = env.split(",") if env else []
+    vllm_config = VllmConfig(
+        compilation_config=CompilationConfig(
+            backend=backend, mode=compilation_mode, custom_ops=custom_ops
+        )
+    )
+    get_cached_compilation_config.cache_clear()
+    with set_current_vllm_config(vllm_config):
+        assert CustomOp.default_on() == default_on
+
+        ops_enabled = [bool(x) for x in ops_enabled]
+
+        assert RMSNorm(1024).enabled() == ops_enabled[0]
+        assert CustomOp.op_registry["rms_norm"].enabled() == ops_enabled[0]
+
+        assert SiluAndMul().enabled() == ops_enabled[1]
+        assert CustomOp.op_registry["silu_and_mul"].enabled() == ops_enabled[1]
+
+        assert GeluAndMul().enabled() == ops_enabled[2]
+        assert CustomOp.op_registry["gelu_and_mul"].enabled() == ops_enabled[2]
+
+        # If registered, subclasses should follow their own name
+        assert Relu3().enabled() == ops_enabled[3]
+        assert CustomOp.op_registry["relu3"].enabled() == ops_enabled[3]
+
+        # Unregistered subclass
+        class SiluAndMul2(SiluAndMul):
+            pass
+
+        # Subclasses should not require registration
+        assert SiluAndMul2().enabled() == SiluAndMul().enabled()
+
+
+@pytest.mark.parametrize(
+    "env", ["all,none", "all,+rms_norm,all", "+rms_norm,-rms_norm"]
+)
+def test_enabled_ops_invalid(env: str):
+    with pytest.raises(Exception):  # noqa
+        vllm_config = VllmConfig(
+            compilation_config=CompilationConfig(custom_ops=env.split(","))
+        )
+        with set_current_vllm_config(vllm_config):
+            RMSNorm(1024).enabled()
+
+
+@pytest.mark.parametrize(
+    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
+)
+def test_topk_dispatch(use_rocm_aiter: bool):
+    topk_func = dispatch_topk_func(use_rocm_aiter)
+
+    if current_platform.is_rocm() and use_rocm_aiter:
+        assert topk_func == rocm_aiter_ops.topk_softmax
+    else:
+        assert topk_func == vllm_topk_softmax
+
+
+@pytest.mark.parametrize("add_residual", [True, False])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("use_rocm_aiter", [True, False])
+@pytest.mark.skipif(
+    not current_platform.is_rocm(), reason="AITER is a feature exclusive for ROCm"
+)
+def test_rms_norm_dispatch(
+    add_residual: bool, dtype: torch.dtype, use_rocm_aiter: bool
+):
+    rms_norm_func = dispatch_rocm_rmsnorm_func(add_residual, dtype, use_rocm_aiter)
+
+    should_use_rocm_aiter = (
+        current_platform.is_rocm()
+        and use_rocm_aiter
+        and dtype in RMS_NORM_SUPPORTED_DTYPES
+    )
+
+    if add_residual and should_use_rocm_aiter:
+        assert rms_norm_func == rocm_aiter_ops.rms_norm2d_with_add
+    elif should_use_rocm_aiter:
+        assert rms_norm_func == rocm_aiter_ops.rms_norm
+    elif add_residual:
+        assert rms_norm_func == fused_add_rms_norm
+    else:
+        assert rms_norm_func == rms_norm

tests/model_executor/test_model_load_with_params.py

@@ -0,0 +1,141 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+
+import pytest
+
+from vllm.model_executor.layers.pooler import (
+    CLSPool,
+    DispatchPooler,
+    MeanPool,
+    PoolingType,
+)
+from vllm.model_executor.models.bert import BertEmbeddingModel
+from vllm.model_executor.models.roberta import RobertaEmbeddingModel
+from vllm.platforms import current_platform
+
+MAX_MODEL_LEN = 128
+MODEL_NAME = os.environ.get("MODEL_NAME", "BAAI/bge-base-en-v1.5")
+REVISION = os.environ.get("REVISION", "main")
+
+MODEL_NAME_ROBERTA = os.environ.get("MODEL_NAME", "intfloat/multilingual-e5-base")
+REVISION_ROBERTA = os.environ.get("REVISION", "main")
+
+
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
+)
+def test_model_loading_with_params(vllm_runner, monkeypatch):
+    """
+    Test parameter weight loading with tp>1.
+    """
+    # to use apply_model
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    with vllm_runner(
+        model_name=MODEL_NAME,
+        revision=REVISION,
+        dtype="float16",
+        max_model_len=MAX_MODEL_LEN,
+    ) as vllm_model:
+        output = vllm_model.embed(
+            "Write a short story about a robot that dreams for the first time.\n"
+        )
+
+        model_config = vllm_model.llm.llm_engine.model_config
+        model_tokenizer = vllm_model.llm.llm_engine.tokenizer
+
+        # asserts on the bert model config file
+        assert model_config.encoder_config["max_seq_length"] == 512
+        assert model_config.encoder_config["do_lower_case"]
+
+        # asserts on the pooling config files
+        assert model_config.pooler_config.pooling_type == PoolingType.CLS.name
+        assert model_config.pooler_config.normalize
+
+        # asserts on the tokenizer loaded
+        assert model_config.tokenizer == "BAAI/bge-base-en-v1.5"
+        assert model_tokenizer.model_max_length == 512
+
+        def check_model(model):
+            assert isinstance(model, BertEmbeddingModel)
+            assert isinstance(pooler := model.pooler, DispatchPooler)
+            assert isinstance(pooler.poolers_by_task["embed"].pooling, CLSPool)
+
+        vllm_model.apply_model(check_model)
+
+        assert output
+
+
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
+)
+def test_roberta_model_loading_with_params(vllm_runner, monkeypatch):
+    """
+    Test parameter weight loading with tp>1.
+    """
+    # to use apply_model
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    with vllm_runner(
+        model_name=MODEL_NAME_ROBERTA,
+        revision=REVISION_ROBERTA,
+        dtype="float16",
+        max_model_len=MAX_MODEL_LEN,
+    ) as vllm_model:
+        output = vllm_model.embed(
+            "Write a short story about a robot that dreams for the first time.\n"
+        )
+
+        model_config = vllm_model.llm.llm_engine.model_config
+        model_tokenizer = vllm_model.llm.llm_engine.tokenizer
+
+        # asserts on the bert model config file
+        assert model_config.encoder_config["max_seq_length"] == 512
+        assert not model_config.encoder_config["do_lower_case"]
+
+        # asserts on the pooling config files
+        assert model_config.pooler_config.pooling_type == PoolingType.MEAN.name
+        assert model_config.pooler_config.normalize
+
+        # asserts on the tokenizer loaded
+        assert model_config.tokenizer == "intfloat/multilingual-e5-base"
+        assert model_tokenizer.model_max_length == 512
+
+        def check_model(model):
+            assert isinstance(model, RobertaEmbeddingModel)
+            assert isinstance(pooler := model.pooler, DispatchPooler)
+            assert isinstance(pooler.poolers_by_task["embed"].pooling, MeanPool)
+
+        vllm_model.apply_model(check_model)
+
+        assert output
+
+
+@pytest.mark.skipif(
+    current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
+)
+def test_facebook_roberta_model_loading_with_params(vllm_runner, monkeypatch):
+    """
+    Test loading roberta-base model with no lm_head.
+    """
+    # to use apply_model
+    monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
+    model_name = "FacebookAI/roberta-base"
+    with vllm_runner(
+        model_name=model_name, dtype="float16", max_model_len=MAX_MODEL_LEN
+    ) as vllm_model:
+        output = vllm_model.embed(
+            "Write a short story about a robot that dreams for the first time.\n"
+        )
+
+        assert vllm_model.llm.llm_engine.model_config.tokenizer == model_name
+
+        def check_model(model):
+            assert isinstance(model, RobertaEmbeddingModel)
+            assert not hasattr(model, "lm_head")
+            assert isinstance(pooler := model.pooler, DispatchPooler)
+            assert isinstance(pooler.poolers_by_task["embed"].pooling, CLSPool)
+
+        vllm_model.apply_model(check_model)
+
+        assert output

tests/model_executor/test_qwen3_omni.py

@@ -0,0 +1,221 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock
+
+import pytest
+from transformers import PretrainedConfig
+
+from vllm.multimodal.processing import InputProcessingContext
+
+
+# Helper function to print input IDs with coalesced audio/video tokens.
+def print_input_ids(input_ids):
+    """
+    Print input IDs, compressing consecutive special tokens.
+    - 151675: <|audio_pad|>
+    - 151656: <|video_pad|>
+    """
+    if not input_ids:
+        print("[]")
+        return
+
+    result = []
+    i = 0
+
+    while i < len(input_ids):
+        current_id = input_ids[i]
+
+        # Check if it's a special token that should be compressed
+        if current_id in [151675, 151656]:
+            # Count consecutive occurrences
+            count = 1
+            while i + count < len(input_ids) and input_ids[i + count] == current_id:
+                count += 1
+
+            # Add compressed representation
+            token_name = "<|audio_pad|>" if current_id == 151675 else "<|video_pad|>"
+            result.append(f"{token_name} * {count}")
+            i += count
+        else:
+            # Regular token, just add it
+            result.append(str(current_id))
+            i += 1
+
+    print(", ".join(result))
+
+
+@pytest.fixture
+def mock_qwen3_omni_config():
+    """Create a mock Qwen3OmniMoeThinker config."""
+    config = Mock(spec=PretrainedConfig)
+    # Token IDs from https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
+    config.audio_token_id = 151675  # <|audio_pad|>
+    config.video_token_id = 151656  # <|video_pad|>
+    config.image_token_id = 151655  # <|image_pad|>
+    config.audio_start_token_id = 151669  # <|audio_start|>
+    config.audio_end_token_id = 151670  # <|audio_end|>
+    config.vision_start_token_id = 151652  # <|vision_start|>
+    config.position_id_per_seconds = 12.5
+
+    # Vision config
+    vision_config = Mock()
+    vision_config.spatial_merge_size = 2
+    config.vision_config = vision_config
+
+    return config
+
+
+@pytest.fixture
+def mock_processor():
+    """Create a mock HF processor."""
+    from transformers.models.whisper import WhisperFeatureExtractor
+
+    processor = Mock()
+    processor.audio_token = "<|audio_pad|>"
+    processor.image_token = "<|image_pad|>"
+    processor.video_token = "<|video_pad|>"
+
+    # Create a real WhisperFeatureExtractor instance for the feature_extractor attribute
+    feature_extractor = WhisperFeatureExtractor()
+    processor.feature_extractor = feature_extractor
+
+    return processor
+
+
+@pytest.fixture
+def mock_tokenizer():
+    """Create a mock tokenizer."""
+    tokenizer = Mock()
+    # Token IDs from https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json
+    tokenizer.get_vocab = Mock(
+        return_value={
+            "<|audio_pad|>": 151675,
+            "<|video_pad|>": 151656,
+            "<|image_pad|>": 151655,
+            "<|audio_start|>": 151669,
+            "<|audio_end|>": 151670,
+            "<|vision_start|>": 151652,
+            "<|vision_end|>": 151653,
+        }
+    )
+    tokenizer.encode = Mock(
+        side_effect=lambda x: {
+            "<|vision_start|>": [151652],
+            "<|vision_end|>": [151653],
+            "<|audio_start|>": [151669],
+            "<|audio_end|>": [151670],
+            "<|audio_pad|>": [151675],
+            "<|image_pad|>": [151655],
+            "<|video_pad|>": [151656],
+        }.get(x, [0])
+    )
+    tokenizer.vision_bos_token = "<|vision_start|>"
+    tokenizer.vision_eos_token = "<|vision_end|>"
+    tokenizer.audio_bos_token = "<|audio_start|>"
+    tokenizer.audio_eos_token = "<|audio_end|>"
+    return tokenizer
+
+
+@pytest.fixture
+def mock_image_processor():
+    """Create a mock image processor."""
+    image_processor = Mock()
+    image_processor.merge_size = 2
+    return image_processor
+
+
+def test_qwen3_omni_get_updates_use_audio_in_video(
+    mock_qwen3_omni_config,
+    mock_processor,
+    mock_tokenizer,
+    mock_image_processor,
+):
+    """Test the get_updates_use_audio_in_video method directly."""
+
+    from vllm.model_executor.models.qwen3_omni_moe_thinker import (
+        Qwen3OmniMoeThinkerMultiModalProcessor,
+        Qwen3OmniMoeThinkerProcessingInfo,
+    )
+
+    # Create a mock context
+    mock_ctx = Mock(spec=InputProcessingContext)
+
+    # Create processing info
+    info = Qwen3OmniMoeThinkerProcessingInfo(mock_ctx)
+    info.get_hf_config = Mock(return_value=mock_qwen3_omni_config)
+    info.get_hf_processor = Mock(return_value=mock_processor)
+    info.get_tokenizer = Mock(return_value=mock_tokenizer)
+    info.get_image_processor = Mock(return_value=mock_image_processor)
+
+    # Create a mock dummy_inputs builder
+    mock_dummy_inputs = Mock()
+
+    # Create the processor
+    processor = Qwen3OmniMoeThinkerMultiModalProcessor(info, mock_dummy_inputs)
+
+    # Test parameters from reference video
+    # https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-Omni/demo/draw.mp4
+    audio_len = 85
+    video_grid_thw = [6, 36, 64]
+    video_second_per_grid_t = 2.0
+
+    # Call the method
+    updates = processor.get_updates_use_audio_in_video(
+        thinker_config=mock_qwen3_omni_config,
+        audio_len=audio_len,
+        video_grid_thw=video_grid_thw,
+        video_second_per_grid_t=video_second_per_grid_t,
+    )
+
+    # Updated input ids should align with HF implementation.
+    # 151669,
+    # <|video_pad|> * 576, <|audio_pad|> * 25,
+    # <|video_pad|> * 576, <|audio_pad|> * 25,
+    # <|video_pad|> * 576, <|audio_pad|> * 25,
+    # <|video_pad|> * 576, <|audio_pad|> * 10,
+    # <|video_pad|> * 1152,
+    # 151670
+    print_input_ids(updates)
+
+    # Verify structure
+    assert isinstance(updates, list)
+    assert len(updates) > 0
+
+    # Verify start and end tokens
+    audio_start_token_id = mock_qwen3_omni_config.audio_start_token_id
+    audio_end_token_id = mock_qwen3_omni_config.audio_end_token_id
+
+    assert updates[0] == audio_start_token_id
+    assert updates[-1] == audio_end_token_id
+
+    # Verify both audio and video tokens are present
+    audio_token_id = mock_qwen3_omni_config.audio_token_id
+    video_token_id = mock_qwen3_omni_config.video_token_id
+
+    audio_count = updates.count(audio_token_id)
+    video_count = updates.count(video_token_id)
+
+    assert audio_count == audio_len, (
+        f"Expected {audio_len} audio tokens, got {audio_count}"
+    )
+
+    # Calculate expected video token count
+    spatial_merge_size = mock_qwen3_omni_config.vision_config.spatial_merge_size
+    height = video_grid_thw[1] // spatial_merge_size
+    width = video_grid_thw[2] // spatial_merge_size
+    expected_video_count = video_grid_thw[0] * height * width
+
+    assert video_count == expected_video_count, (
+        f"Expected {expected_video_count} video tokens, got {video_count}"
+    )
+
+    # Total tokens should be: 1 (start) + audio_len + video_count + 1 (end)
+    expected_total = 1 + audio_len + expected_video_count + 1
+    assert len(updates) == expected_total, (
+        f"Expected {expected_total} total tokens, got {len(updates)}"
+    )
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

tests/model_executor/test_weight_utils.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
 import os
 import tempfile
 
@@ -6,23 +9,24 @@ import pytest
 from huggingface_hub.utils import LocalEntryNotFoundError
 
 from vllm.model_executor.model_loader.weight_utils import (
-    download_weights_from_hf, enable_hf_transfer)
+    download_weights_from_hf,
+    enable_hf_transfer,
+)
 
 
 def test_hf_transfer_auto_activation():
     if "HF_HUB_ENABLE_HF_TRANSFER" in os.environ:
         # in case it is already set, we can't test the auto activation
-        pytest.skip(
-            "HF_HUB_ENABLE_HF_TRANSFER is set, can't test auto activation")
+        pytest.skip("HF_HUB_ENABLE_HF_TRANSFER is set, can't test auto activation")
     enable_hf_transfer()
     try:
         # enable hf hub transfer if available
         import hf_transfer  # type: ignore # noqa
-        HF_TRANFER_ACTIVE = True
+
+        HF_TRANSFER_ACTIVE = True
     except ImportError:
-        HF_TRANFER_ACTIVE = False
-    assert (huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER ==
-            HF_TRANFER_ACTIVE)
+        HF_TRANSFER_ACTIVE = False
+    assert huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER == HF_TRANSFER_ACTIVE
 
 
 def test_download_weights_from_hf():
@@ -31,22 +35,30 @@ def test_download_weights_from_hf():
         # if offline is set and model is not cached
         huggingface_hub.constants.HF_HUB_OFFLINE = True
         with pytest.raises(LocalEntryNotFoundError):
-            download_weights_from_hf("facebook/opt-125m",
-                                     allow_patterns=["*.safetensors", "*.bin"],
-                                     cache_dir=tmpdir)
+            download_weights_from_hf(
+                "facebook/opt-125m",
+                allow_patterns=["*.safetensors", "*.bin"],
+                cache_dir=tmpdir,
+            )
 
         # download the model
         huggingface_hub.constants.HF_HUB_OFFLINE = False
-        download_weights_from_hf("facebook/opt-125m",
-                                 allow_patterns=["*.safetensors", "*.bin"],
-                                 cache_dir=tmpdir)
+        download_weights_from_hf(
+            "facebook/opt-125m",
+            allow_patterns=["*.safetensors", "*.bin"],
+            cache_dir=tmpdir,
+        )
 
         # now it should work offline
         huggingface_hub.constants.HF_HUB_OFFLINE = True
-        assert download_weights_from_hf(
-            "facebook/opt-125m",
-            allow_patterns=["*.safetensors", "*.bin"],
-            cache_dir=tmpdir) is not None
+        assert (
+            download_weights_from_hf(
+                "facebook/opt-125m",
+                allow_patterns=["*.safetensors", "*.bin"],
+                cache_dir=tmpdir,
+            )
+            is not None
+        )
 
 
 if __name__ == "__main__":