chore: import upstream snapshot with attribution

tests/model_executor/layers/test_pooler_activations.py

@@ -0,0 +1,234 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for vllm.model_executor.layers.pooler.activations."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.model_executor.layers.pooler.activations import (
+    LambdaPoolerActivation,
+    PoolerClassify,
+    PoolerIdentity,
+    PoolerMultiLabelClassify,
+    PoolerNormalize,
+    get_act_fn,
+    resolve_classifier_act_fn,
+)
+
+
+# ---------------------------------------------------------------------------
+# PoolerIdentity
+# ---------------------------------------------------------------------------
+class TestPoolerIdentity:
+    def test_returns_input_unchanged(self):
+        pooler = PoolerIdentity()
+        x = torch.randn(4, 128)
+        out = pooler(x)
+        assert torch.equal(out, x)
+
+    def test_forward_list(self):
+        pooler = PoolerIdentity()
+        tensors = [torch.randn(128), torch.randn(256)]
+        out = pooler(tensors)
+        assert len(out) == 2
+        for orig, result in zip(tensors, out):
+            assert torch.equal(orig, result)
+
+
+# ---------------------------------------------------------------------------
+# PoolerNormalize
+# ---------------------------------------------------------------------------
+class TestPoolerNormalize:
+    def test_output_has_unit_norm(self):
+        pooler = PoolerNormalize()
+        x = torch.randn(4, 128)
+        out = pooler(x)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(4), atol=1e-5)
+
+    def test_single_vector(self):
+        pooler = PoolerNormalize()
+        x = torch.randn(1, 64)
+        out = pooler(x)
+        norm = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norm, torch.ones(1), atol=1e-5)
+
+    def test_forward_list(self):
+        pooler = PoolerNormalize()
+        tensors = [torch.randn(1, 64), torch.randn(1, 128)]
+        out = pooler(tensors)
+        for t in out:
+            norm = torch.linalg.norm(t, dim=-1)
+            assert torch.allclose(norm, torch.ones(1), atol=1e-5)
+
+
+# ---------------------------------------------------------------------------
+# PoolerMultiLabelClassify
+# ---------------------------------------------------------------------------
+class TestPoolerMultiLabelClassify:
+    def test_output_in_zero_one(self):
+        pooler = PoolerMultiLabelClassify()
+        x = torch.randn(4, 10)
+        out = pooler(x)
+        assert (out >= 0).all() and (out <= 1).all()
+
+    def test_large_positive_maps_near_one(self):
+        pooler = PoolerMultiLabelClassify()
+        x = torch.full((1, 3), 100.0)
+        out = pooler(x)
+        assert torch.allclose(out, torch.ones(1, 3), atol=1e-4)
+
+    def test_large_negative_maps_near_zero(self):
+        pooler = PoolerMultiLabelClassify()
+        x = torch.full((1, 3), -100.0)
+        out = pooler(x)
+        assert torch.allclose(out, torch.zeros(1, 3), atol=1e-4)
+
+
+# ---------------------------------------------------------------------------
+# PoolerClassify
+# ---------------------------------------------------------------------------
+class TestPoolerClassify:
+    def test_infers_from_shape_when_num_labels_none(self):
+        pooler = PoolerClassify(num_labels=None)
+        assert pooler.num_labels is None
+        x = torch.randn(2, 5)
+        out = pooler(x)
+        sums = out.sum(dim=-1)
+        assert torch.allclose(sums, torch.ones(2), atol=1e-5)
+
+    def test_sigmoid_when_num_labels_lt_2(self):
+        pooler = PoolerClassify(num_labels=1)
+        x = torch.zeros(1, 1)
+        out = pooler(x)
+        assert torch.allclose(out, torch.tensor([[0.5]]), atol=1e-5)
+
+    def test_num_labels_zero_uses_sigmoid(self):
+        pooler = PoolerClassify(num_labels=0)
+        assert pooler.num_labels == 0
+        x = torch.zeros(1, 3)
+        out = pooler(x)
+        assert torch.allclose(out, torch.full((1, 3), 0.5), atol=1e-5)
+
+    def test_num_labels_ge_2_uses_softmax(self):
+        pooler = PoolerClassify(num_labels=4)
+        assert pooler.num_labels == 4
+        x = torch.randn(2, 4)
+        out = pooler(x)
+        sums = out.sum(dim=-1)
+        assert torch.allclose(sums, torch.ones(2), atol=1e-5)
+
+    def test_default_num_labels_is_none(self):
+        pooler = PoolerClassify()
+        assert pooler.num_labels is None
+
+
+# ---------------------------------------------------------------------------
+# LambdaPoolerActivation
+# ---------------------------------------------------------------------------
+class TestLambdaPoolerActivation:
+    def test_applies_custom_fn(self):
+        pooler = LambdaPoolerActivation(nn.ReLU())
+        x = torch.tensor([[-1.0, 2.0, -3.0]])
+        out = pooler(x)
+        expected = torch.tensor([[0.0, 2.0, 0.0]])
+        assert torch.equal(out, expected)
+
+    def test_forward_list(self):
+        pooler = LambdaPoolerActivation(nn.ReLU())
+        tensors = [torch.tensor([-1.0, 2.0]), torch.tensor([3.0, -4.0])]
+        out = pooler(tensors)
+        assert torch.equal(out[0], torch.tensor([0.0, 2.0]))
+        assert torch.equal(out[1], torch.tensor([3.0, 0.0]))
+
+
+# ---------------------------------------------------------------------------
+# get_act_fn factory
+# ---------------------------------------------------------------------------
+class TestGetActFn:
+    @staticmethod
+    def _make_config(**kwargs):
+        return SimpleNamespace(**kwargs)
+
+    def test_regression(self):
+        cfg = self._make_config(problem_type="regression")
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerIdentity)
+
+    def test_single_label_classification(self):
+        cfg = self._make_config(
+            problem_type="single_label_classification", num_labels=3
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+        assert result.num_labels == 3
+
+    def test_multi_label_classification(self):
+        cfg = self._make_config(problem_type="multi_label_classification")
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerMultiLabelClassify)
+
+    def test_sentence_transformers_activation(self):
+        cfg = self._make_config(
+            problem_type="",
+            sentence_transformers={
+                "activation_fn": "torch.nn.modules.activation.Sigmoid"
+            },
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+
+    def test_sbert_activation(self):
+        cfg = self._make_config(
+            problem_type="",
+            sbert_ce_default_activation_function=(
+                "torch.nn.modules.activation.Sigmoid"
+            ),
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+
+    def test_default_fallback(self):
+        cfg = self._make_config(problem_type="")
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerClassify)
+
+    def test_sentence_transformers_takes_priority(self):
+        cfg = self._make_config(
+            problem_type="",
+            sentence_transformers={"activation_fn": "torch.nn.modules.linear.Identity"},
+            sbert_ce_default_activation_function=(
+                "torch.nn.modules.activation.Sigmoid"
+            ),
+        )
+        result = get_act_fn(cfg)
+        assert isinstance(result, PoolerIdentity)
+
+    def test_rejects_non_torch_activation(self):
+        cfg = self._make_config(
+            problem_type="",
+            sentence_transformers={"activation_fn": "os.system"},
+        )
+        with pytest.raises(ValueError, match="restricted"):
+            get_act_fn(cfg)
+
+
+# ---------------------------------------------------------------------------
+# resolve_classifier_act_fn
+# ---------------------------------------------------------------------------
+class TestResolveClassifierActFn:
+    def test_delegates_to_get_act_fn_when_none(self):
+        model_config = SimpleNamespace(
+            hf_config=SimpleNamespace(num_labels=3, problem_type="")
+        )
+        result = resolve_classifier_act_fn(model_config, act_fn=None)
+        assert isinstance(result, PoolerClassify)
+        assert result.num_labels == 3
+
+    def test_passes_through_provided_act_fn(self):
+        custom = PoolerIdentity()
+        result = resolve_classifier_act_fn(None, act_fn=custom)
+        assert result is custom

tests/model_executor/layers/test_pooler_heads.py

@@ -0,0 +1,481 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for sequence and token pooler head classes."""
+
+import torch
+import torch.nn as nn
+
+from vllm.model_executor.layers.pooler.activations import PoolerNormalize
+from vllm.model_executor.layers.pooler.seqwise.heads import (
+    ClassifierPoolerHead,
+    EmbeddingPoolerHead,
+)
+from vllm.model_executor.layers.pooler.tokwise.heads import (
+    TokenClassifierPoolerHead,
+    TokenEmbeddingPoolerHead,
+)
+from vllm.pooling_params import PoolingParams
+from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates
+
+_HIDDEN = 16
+_BATCH = 3
+
+
+def _make_params(
+    n: int,
+    *,
+    task: str = "embed",
+    dimensions: int | None = None,
+    use_activation: bool | None = None,
+) -> list[PoolingParams]:
+    return [
+        PoolingParams(task=task, dimensions=dimensions, use_activation=use_activation)
+        for _ in range(n)
+    ]
+
+
+def _make_metadata(pooling_params: list[PoolingParams]) -> PoolingMetadata:
+    n = len(pooling_params)
+    return PoolingMetadata(
+        prompt_lens=torch.ones(n, dtype=torch.long),
+        prompt_token_ids=None,
+        prompt_token_ids_cpu=None,
+        pooling_params=pooling_params,
+        pooling_states=[PoolingStates() for _ in range(n)],
+    )
+
+
+def _linear(in_f: int, out_f: int) -> nn.Linear:
+    torch.manual_seed(42)
+    return nn.Linear(in_f, out_f, bias=False)
+
+
+# ---------------------------------------------------------------------------
+# EmbeddingPoolerHead
+# ---------------------------------------------------------------------------
+class TestEmbeddingPoolerHead:
+    def test_supported_tasks(self):
+        head = EmbeddingPoolerHead()
+        assert head.get_supported_tasks() == {"embed"}
+
+    def test_passthrough(self):
+        head = EmbeddingPoolerHead()
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH))
+        out = head(x, meta)
+        assert torch.equal(out, x)
+
+    def test_head_dtype(self):
+        head = EmbeddingPoolerHead(head_dtype=torch.float16)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH))
+        out = head(x, meta)
+        assert out.dtype == torch.float16
+
+    def test_projector(self):
+        proj = _linear(_HIDDEN, 8)
+        head = EmbeddingPoolerHead(projector=proj)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH))
+        out = head(x, meta)
+        assert out.shape == (_BATCH, 8)
+        assert torch.allclose(out, proj(x))
+
+    def test_matryoshka_uniform(self):
+        head = EmbeddingPoolerHead()
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, dimensions=4)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert out.shape == (_BATCH, 4)
+        assert torch.equal(out, x[..., :4])
+
+    def test_matryoshka_mixed(self):
+        head = EmbeddingPoolerHead()
+        x = torch.randn(2, _HIDDEN)
+        params = [
+            PoolingParams(task="embed", dimensions=4),
+            PoolingParams(task="embed", dimensions=8),
+        ]
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert isinstance(out, list)
+        assert len(out) == 2
+        assert out[0].shape[-1] == 4
+        assert out[1].shape[-1] == 8
+
+    def test_matryoshka_mixed_with_none(self):
+        head = EmbeddingPoolerHead()
+        x = torch.randn(2, _HIDDEN)
+        params = [
+            PoolingParams(task="embed", dimensions=4),
+            PoolingParams(task="embed", dimensions=None),
+        ]
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert isinstance(out, list)
+        assert out[0].shape[-1] == 4
+        assert torch.equal(out[1], x[1])
+
+    def test_activation_uniform_true(self):
+        head = EmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, use_activation=True)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
+
+    def test_activation_uniform_false(self):
+        head = EmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, use_activation=False)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert torch.equal(out, x)
+
+    def test_activation_mixed_flags(self):
+        head = EmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(2, _HIDDEN)
+        params = [
+            PoolingParams(task="embed", use_activation=True),
+            PoolingParams(task="embed", use_activation=False),
+        ]
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert isinstance(out, list)
+        norm_0 = torch.linalg.norm(out[0], dim=-1)
+        assert torch.allclose(norm_0, torch.ones(1), atol=1e-5)
+        assert torch.equal(out[1], x[1])
+
+    def test_list_input_gets_stacked(self):
+        head = EmbeddingPoolerHead()
+        tensors = [torch.randn(_HIDDEN) for _ in range(_BATCH)]
+        meta = _make_metadata(_make_params(_BATCH))
+        out = head(tensors, meta)
+        assert out.shape == (_BATCH, _HIDDEN)
+        expected = torch.stack(tensors)
+        assert torch.equal(out, expected)
+
+    def test_projector_then_matryoshka(self):
+        proj = _linear(_HIDDEN, 8)
+        head = EmbeddingPoolerHead(projector=proj)
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, dimensions=4)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert out.shape == (_BATCH, 4)
+        assert torch.equal(out, proj(x)[..., :4])
+
+    def test_matryoshka_then_activation(self):
+        head = EmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, dimensions=4, use_activation=True)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert out.shape == (_BATCH, 4)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
+
+    def test_empty_batch(self):
+        head = EmbeddingPoolerHead()
+        x = torch.randn(0, _HIDDEN)
+        meta = _make_metadata([])
+        out = head(x, meta)
+        assert out.shape == (0, _HIDDEN)
+
+
+# ---------------------------------------------------------------------------
+# ClassifierPoolerHead
+# ---------------------------------------------------------------------------
+class TestClassifierPoolerHead:
+    def test_supported_tasks(self):
+        head = ClassifierPoolerHead()
+        assert head.get_supported_tasks() == {"classify"}
+
+    def test_passthrough(self):
+        head = ClassifierPoolerHead()
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        assert torch.equal(out, x)
+
+    def test_head_dtype(self):
+        head = ClassifierPoolerHead(head_dtype=torch.float16)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        assert out.dtype == torch.float16
+
+    def test_classifier(self):
+        clf = _linear(_HIDDEN, 3)
+        head = ClassifierPoolerHead(classifier=clf)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        assert out.shape == (_BATCH, 3)
+        assert torch.allclose(out, clf(x))
+
+    def test_logit_mean(self):
+        head = ClassifierPoolerHead(logit_mean=2.0)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        assert torch.allclose(out, x - 2.0)
+
+    def test_logit_sigma(self):
+        head = ClassifierPoolerHead(logit_sigma=0.5)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        assert torch.allclose(out, x / 0.5)
+
+    def test_platt_scaling_combined(self):
+        head = ClassifierPoolerHead(logit_mean=1.0, logit_sigma=2.0)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        assert torch.allclose(out, (x - 1.0) / 2.0)
+
+    def test_activation_uniform_true(self):
+        head = ClassifierPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, task="classify", use_activation=True)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(_BATCH), atol=1e-5)
+
+    def test_activation_uniform_false(self):
+        head = ClassifierPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(_BATCH, _HIDDEN)
+        params = _make_params(_BATCH, task="classify", use_activation=False)
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert torch.equal(out, x)
+
+    def test_activation_mixed_flags(self):
+        head = ClassifierPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(2, _HIDDEN)
+        params = [
+            PoolingParams(task="classify", use_activation=True),
+            PoolingParams(task="classify", use_activation=False),
+        ]
+        meta = _make_metadata(params)
+        out = head(x, meta)
+        assert isinstance(out, list)
+        norm_0 = torch.linalg.norm(out[0], dim=-1)
+        assert torch.allclose(norm_0, torch.ones(1), atol=1e-5)
+        assert torch.equal(out[1], x[1])
+
+    def test_list_input_gets_stacked(self):
+        head = ClassifierPoolerHead()
+        tensors = [torch.randn(_HIDDEN) for _ in range(_BATCH)]
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(tensors, meta)
+        assert out.shape == (_BATCH, _HIDDEN)
+        expected = torch.stack(tensors)
+        assert torch.equal(out, expected)
+
+    def test_classifier_then_platt_scaling(self):
+        clf = _linear(_HIDDEN, 3)
+        head = ClassifierPoolerHead(classifier=clf, logit_mean=1.0, logit_sigma=2.0)
+        x = torch.randn(_BATCH, _HIDDEN)
+        meta = _make_metadata(_make_params(_BATCH, task="classify"))
+        out = head(x, meta)
+        expected = (clf(x) - 1.0) / 2.0
+        assert torch.allclose(out, expected)
+
+    def test_empty_batch(self):
+        head = ClassifierPoolerHead()
+        x = torch.randn(0, _HIDDEN)
+        meta = _make_metadata([])
+        out = head(x, meta)
+        assert out.shape == (0, _HIDDEN)
+
+
+# ---------------------------------------------------------------------------
+# TokenEmbeddingPoolerHead
+# ---------------------------------------------------------------------------
+class TestTokenEmbeddingPoolerHead:
+    def test_supported_tasks(self):
+        head = TokenEmbeddingPoolerHead()
+        assert head.get_supported_tasks() == {"token_embed"}
+
+    def test_passthrough(self):
+        head = TokenEmbeddingPoolerHead()
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed")
+        out = head.forward_chunk(x, param)
+        assert torch.equal(out, x)
+
+    def test_none_chunked_prefill(self):
+        head = TokenEmbeddingPoolerHead()
+        param = PoolingParams(task="token_embed")
+        out = head.forward_chunk(None, param)
+        assert out is None
+
+    def test_head_dtype(self):
+        head = TokenEmbeddingPoolerHead(head_dtype=torch.float16)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed")
+        out = head.forward_chunk(x, param)
+        assert out.dtype == torch.float16
+
+    def test_projector(self):
+        proj = _linear(_HIDDEN, 8)
+        head = TokenEmbeddingPoolerHead(projector=proj)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed")
+        out = head.forward_chunk(x, param)
+        assert out.shape == (5, 8)
+        assert torch.allclose(out, proj(x))
+
+    def test_matryoshka_truncation(self):
+        head = TokenEmbeddingPoolerHead()
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed", dimensions=4)
+        out = head.forward_chunk(x, param)
+        assert out.shape == (5, 4)
+        assert torch.equal(out, x[..., :4])
+
+    def test_activation_true(self):
+        head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed", use_activation=True)
+        out = head.forward_chunk(x, param)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(5), atol=1e-5)
+
+    def test_activation_false(self):
+        head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed", use_activation=False)
+        out = head.forward_chunk(x, param)
+        assert torch.equal(out, x)
+
+    def test_projector_then_matryoshka(self):
+        proj = _linear(_HIDDEN, 8)
+        head = TokenEmbeddingPoolerHead(projector=proj)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed", dimensions=4)
+        out = head.forward_chunk(x, param)
+        assert out.shape == (5, 4)
+        assert torch.equal(out, proj(x)[..., :4])
+
+    def test_matryoshka_then_activation(self):
+        head = TokenEmbeddingPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_embed", dimensions=4, use_activation=True)
+        out = head.forward_chunk(x, param)
+        assert out.shape == (5, 4)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(5), atol=1e-5)
+
+    def test_forward_mixed_batch_chunked_prefill(self):
+        head = TokenEmbeddingPoolerHead()
+        pooled_data = [torch.randn(5, _HIDDEN), None, torch.randn(3, _HIDDEN)]
+        params = _make_params(3, task="token_embed")
+        meta = _make_metadata(params)
+        out = head(pooled_data, meta)
+        assert len(out) == 3
+        assert torch.equal(out[0], pooled_data[0])
+        assert out[1] is None
+        assert torch.equal(out[2], pooled_data[2])
+
+    def test_forward_empty_batch(self):
+        head = TokenEmbeddingPoolerHead()
+        meta = _make_metadata([])
+        out = head([], meta)
+        assert out == []
+
+
+# ---------------------------------------------------------------------------
+# TokenClassifierPoolerHead
+# ---------------------------------------------------------------------------
+class TestTokenClassifierPoolerHead:
+    def test_supported_tasks(self):
+        head = TokenClassifierPoolerHead()
+        assert head.get_supported_tasks() == {"token_classify"}
+
+    def test_passthrough(self):
+        head = TokenClassifierPoolerHead()
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(x, param)
+        assert torch.equal(out, x)
+
+    def test_none_chunked_prefill(self):
+        head = TokenClassifierPoolerHead()
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(None, param)
+        assert out is None
+
+    def test_head_dtype(self):
+        head = TokenClassifierPoolerHead(head_dtype=torch.float16)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(x, param)
+        assert out.dtype == torch.float16
+
+    def test_classifier(self):
+        clf = _linear(_HIDDEN, 3)
+        head = TokenClassifierPoolerHead(classifier=clf)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(x, param)
+        assert out.shape == (5, 3)
+        assert torch.allclose(out, clf(x))
+
+    def test_logit_mean(self):
+        head = TokenClassifierPoolerHead(logit_mean=2.0)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(x, param)
+        assert torch.allclose(out, x - 2.0)
+
+    def test_logit_sigma(self):
+        head = TokenClassifierPoolerHead(logit_sigma=0.5)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(x, param)
+        assert torch.allclose(out, x / 0.5)
+
+    def test_platt_scaling_combined(self):
+        head = TokenClassifierPoolerHead(logit_mean=1.0, logit_sigma=2.0)
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify")
+        out = head.forward_chunk(x, param)
+        assert torch.allclose(out, (x - 1.0) / 2.0)
+
+    def test_activation_true(self):
+        head = TokenClassifierPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify", use_activation=True)
+        out = head.forward_chunk(x, param)
+        norms = torch.linalg.norm(out, dim=-1)
+        assert torch.allclose(norms, torch.ones(5), atol=1e-5)
+
+    def test_activation_false(self):
+        head = TokenClassifierPoolerHead(activation=PoolerNormalize())
+        x = torch.randn(5, _HIDDEN)
+        param = PoolingParams(task="token_classify", use_activation=False)
+        out = head.forward_chunk(x, param)
+        assert torch.equal(out, x)
+
+    def test_forward_mixed_batch_chunked_prefill(self):
+        head = TokenClassifierPoolerHead()
+        pooled_data = [torch.randn(5, _HIDDEN), None, torch.randn(3, _HIDDEN)]
+        params = _make_params(3, task="token_classify")
+        meta = _make_metadata(params)
+        out = head(pooled_data, meta)
+        assert len(out) == 3
+        assert torch.equal(out[0], pooled_data[0])
+        assert out[1] is None
+        assert torch.equal(out[2], pooled_data[2])
+
+    def test_forward_empty_batch(self):
+        head = TokenClassifierPoolerHead()
+        meta = _make_metadata([])
+        out = head([], meta)
+        assert out == []

tests/model_executor/layers/test_pooler_methods.py

@@ -0,0 +1,499 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for sequence and token pooling methods and their factories."""
+
+from dataclasses import dataclass
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.pooler.seqwise.methods import (
+    CLSPool,
+    LastPool,
+    MeanPool,
+    get_seq_pooling_method,
+)
+from vllm.model_executor.layers.pooler.tokwise.methods import (
+    AllPool,
+    StepPool,
+    get_tok_pooling_method,
+)
+from vllm.pooling_params import PoolingParams
+from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates
+
+_CPU = torch.device("cpu")
+
+
+def _make_pooling_cursor(
+    prompt_lens: list[int],
+    *,
+    num_scheduled_tokens: list[int] | None = None,
+    seq_lens: list[int] | None = None,
+    device: torch.device = _CPU,
+) -> PoolingCursor:
+    """Build a PoolingCursor from a list of per-sequence prompt lengths."""
+    prompt_lens_cpu = torch.tensor(prompt_lens, dtype=torch.long)
+    if num_scheduled_tokens is None:
+        num_scheduled_tokens_cpu = prompt_lens_cpu.clone()
+    else:
+        num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, dtype=torch.long)
+    if seq_lens is None:
+        seq_lens_cpu = prompt_lens_cpu.clone()
+    else:
+        seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.long)
+
+    cumsum = torch.zeros(len(prompt_lens) + 1, dtype=torch.long, device=device)
+    torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])
+
+    return PoolingCursor(
+        first_token_indices_gpu=cumsum[: len(prompt_lens)].to(device),
+        last_token_indices_gpu=(cumsum[1:] - 1).to(device),
+        prompt_lens_cpu=prompt_lens_cpu,
+        seq_lens_cpu=seq_lens_cpu,
+        num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,
+    )
+
+
+def _make_metadata(
+    prompt_lens: list[int],
+    *,
+    tasks: list[str] | None = None,
+    token_ids: list[list[int]] | None = None,
+    pooling_params: list[PoolingParams] | None = None,
+    num_scheduled_tokens: list[int] | None = None,
+    seq_lens: list[int] | None = None,
+    device: torch.device = _CPU,
+) -> PoolingMetadata:
+    """Build a minimal PoolingMetadata for testing pooling methods."""
+    n_seqs = len(prompt_lens)
+    if tasks is None:
+        tasks = ["embed"] * n_seqs
+    if pooling_params is None:
+        pooling_params = [PoolingParams(task=t) for t in tasks]
+
+    prompt_lens_tensor = torch.tensor(prompt_lens, dtype=torch.long)
+
+    prompt_token_ids_cpu = None
+    prompt_token_ids = None
+    if token_ids is not None:
+        max_len = max(len(t) for t in token_ids)
+        padded = [t + [0] * (max_len - len(t)) for t in token_ids]
+        prompt_token_ids_cpu = torch.tensor(padded, dtype=torch.long)
+        prompt_token_ids = prompt_token_ids_cpu.to(device)
+
+    cursor = _make_pooling_cursor(
+        prompt_lens,
+        num_scheduled_tokens=num_scheduled_tokens,
+        seq_lens=seq_lens,
+        device=device,
+    )
+
+    pooling_states = [PoolingStates() for _ in range(n_seqs)]
+
+    return PoolingMetadata(
+        prompt_lens=prompt_lens_tensor,
+        prompt_token_ids=prompt_token_ids,
+        prompt_token_ids_cpu=prompt_token_ids_cpu,
+        pooling_params=pooling_params,
+        pooling_states=pooling_states,
+        pooling_cursor=cursor,
+    )
+
+
+# ---------------------------------------------------------------------------
+# CLSPool
+# ---------------------------------------------------------------------------
+class TestCLSPool:
+    def test_extracts_first_token(self):
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        pooler = CLSPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
+        assert torch.equal(out, expected)
+
+    def test_rejects_partial_prefill(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        metadata = _make_metadata([3], num_scheduled_tokens=[2])
+        pooler = CLSPool()
+        with pytest.raises(RuntimeError, match="partial prefill"):
+            pooler(hidden, metadata)
+
+
+# ---------------------------------------------------------------------------
+# LastPool
+# ---------------------------------------------------------------------------
+class TestLastPool:
+    def test_extracts_last_token(self):
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        pooler = LastPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[3.0, 4.0], [9.0, 10.0]])
+        assert torch.equal(out, expected)
+
+    def test_partial_prefill_extracts_last_scheduled(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        metadata = _make_metadata([4], num_scheduled_tokens=[2])
+        pooler = LastPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[3.0, 4.0]])
+        assert torch.equal(out, expected)
+
+
+# ---------------------------------------------------------------------------
+# MeanPool
+# ---------------------------------------------------------------------------
+class TestMeanPool:
+    def test_computes_mean(self):
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [10.0, 20.0]], dtype=torch.float32
+        )
+        metadata = _make_metadata([2, 1])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[2.0, 3.0], [10.0, 20.0]], dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_single_token_is_identity(self):
+        hidden = torch.tensor([[5.0, 10.0]], dtype=torch.float32)
+        metadata = _make_metadata([1])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        assert torch.allclose(out, hidden, atol=1e-5)
+
+    def test_uniform_values_return_same(self):
+        hidden = torch.full((4, 3), 7.0, dtype=torch.float32)
+        metadata = _make_metadata([4])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        expected = torch.full((1, 3), 7.0, dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_multiple_sequences(self):
+        hidden = torch.tensor(
+            [
+                [0.0, 0.0],
+                [2.0, 4.0],
+                [4.0, 8.0],
+                [10.0, 10.0],
+            ],
+            dtype=torch.float32,
+        )
+        metadata = _make_metadata([3, 1])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        expected = torch.tensor([[2.0, 4.0], [10.0, 10.0]], dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_empty_batch(self):
+        hidden = torch.empty((0, 8), dtype=torch.float32)
+        metadata = _make_metadata([])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        assert out.shape == (0, 8)
+
+    def test_rejects_partial_prefill(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)
+        metadata = _make_metadata([3], num_scheduled_tokens=[2])
+        pooler = MeanPool()
+        with pytest.raises(RuntimeError, match="partial prefill"):
+            pooler(hidden, metadata)
+
+    def test_chunked_accumulation(self):
+        hidden = torch.arange(20, dtype=torch.float32).reshape(5, 4)
+        metadata = _make_metadata([3, 2])
+        pooler = MeanPool()
+        with patch(
+            "vllm.model_executor.layers.pooler.seqwise.methods"
+            "._MEAN_POOL_ACCUMULATION_CHUNK_BYTES",
+            16,
+        ):
+            out = pooler(hidden, metadata)
+        expected_seq0 = hidden[:3].float().mean(dim=0, keepdim=True)
+        expected_seq1 = hidden[3:].float().mean(dim=0, keepdim=True)
+        expected = torch.cat([expected_seq0, expected_seq1], dim=0)
+        assert torch.allclose(out, expected, atol=1e-5)
+
+    def test_upcasts_to_float32(self):
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float16)
+        metadata = _make_metadata([2])
+        pooler = MeanPool()
+        out = pooler(hidden, metadata)
+        assert out.dtype == torch.float32
+        expected = torch.tensor([[2.0, 3.0]], dtype=torch.float32)
+        assert torch.allclose(out, expected, atol=1e-2)
+
+
+# ---------------------------------------------------------------------------
+# get_seq_pooling_method factory
+# ---------------------------------------------------------------------------
+class TestGetSeqPoolingMethod:
+    def test_cls(self):
+        assert isinstance(get_seq_pooling_method("CLS"), CLSPool)
+
+    def test_last(self):
+        assert isinstance(get_seq_pooling_method("LAST"), LastPool)
+
+    def test_mean(self):
+        assert isinstance(get_seq_pooling_method("MEAN"), MeanPool)
+
+    def test_unknown_raises(self):
+        with pytest.raises(NotImplementedError, match="UNKNOWN"):
+            get_seq_pooling_method("UNKNOWN")
+
+
+# ---------------------------------------------------------------------------
+# AllPool
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class _FakeSchedulerConfig:
+    enable_chunked_prefill: bool = False
+
+
+@dataclass
+class _FakeVllmConfig:
+    scheduler_config: _FakeSchedulerConfig
+
+
+class TestAllPool:
+    @staticmethod
+    def _make_all_pool(*, chunked: bool = False) -> AllPool:
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=chunked,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            return AllPool()
+
+    def test_splits_by_sequence(self):
+        pooler = self._make_all_pool()
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        out = pooler(hidden, metadata)
+        assert len(out) == 2
+        assert torch.equal(out[0], hidden[:2])
+        assert torch.equal(out[1], hidden[2:])
+
+    def test_single_sequence(self):
+        pooler = self._make_all_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        metadata = _make_metadata([3])
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert torch.equal(out[0], hidden)
+
+    def test_chunked_prefill_returns_none_for_unfinished(self):
+        pooler = self._make_all_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        metadata = _make_metadata(
+            [4],
+            num_scheduled_tokens=[2],
+            seq_lens=[2],
+        )
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert out[0] is None
+
+    def test_chunked_prefill_returns_concat_when_finished(self):
+        pooler = self._make_all_pool(chunked=True)
+
+        chunk1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        metadata1 = _make_metadata(
+            [4],
+            num_scheduled_tokens=[2],
+            seq_lens=[2],
+        )
+        out1 = pooler(chunk1, metadata1)
+        assert out1[0] is None
+
+        chunk2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]])
+        metadata2 = _make_metadata(
+            [4],
+            num_scheduled_tokens=[2],
+            seq_lens=[4],
+        )
+        metadata2.pooling_states = metadata1.pooling_states
+        out2 = pooler(chunk2, metadata2)
+        assert out2[0] is not None
+        expected = torch.cat([chunk1, chunk2], dim=0)
+        assert torch.equal(out2[0], expected)
+
+    def test_chunked_prefill_single_shot_matches_non_chunked(self):
+        pooler = self._make_all_pool(chunked=True)
+        hidden = torch.tensor(
+            [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]
+        )
+        metadata = _make_metadata([2, 3])
+        out = pooler(hidden, metadata)
+        assert len(out) == 2
+        assert torch.equal(out[0], hidden[:2])
+        assert torch.equal(out[1], hidden[2:])
+
+    def test_chunked_prefill_mixed_finished_unfinished(self):
+        pooler = self._make_all_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
+        metadata = _make_metadata(
+            [2, 4],
+            num_scheduled_tokens=[2, 1],
+            seq_lens=[2, 1],
+        )
+        out = pooler(hidden, metadata)
+        assert len(out) == 2
+        assert torch.equal(out[0], hidden[:2])
+        assert out[1] is None
+
+
+# ---------------------------------------------------------------------------
+# StepPool
+# ---------------------------------------------------------------------------
+class TestStepPool:
+    @staticmethod
+    def _make_step_pool(*, chunked: bool = False) -> StepPool:
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=chunked,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            return StepPool()
+
+    def test_filters_by_step_tag_id(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+        token_ids = [[10, 99, 10, 20]]
+        params = [PoolingParams(task="token_classify", step_tag_id=10)]
+        metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_filters_by_returned_token_ids(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
+        token_ids = [[10, 20]]
+        params = [PoolingParams(task="token_classify", returned_token_ids=[0, 2])]
+        metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[1.0, 3.0], [4.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_no_filtering_without_params(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        token_ids = [[10, 20]]
+        params = [PoolingParams(task="token_classify")]
+        metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert torch.equal(out[0], hidden)
+
+    def test_combined_step_tag_and_returned_token_ids(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])
+        token_ids = [[99, 10, 99]]
+        params = [
+            PoolingParams(
+                task="token_classify",
+                step_tag_id=10,
+                returned_token_ids=[0, 2],
+            )
+        ]
+        metadata = _make_metadata([3], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[4.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_step_tag_id_no_match_returns_empty(self):
+        pooler = self._make_step_pool()
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        token_ids = [[10, 20]]
+        params = [PoolingParams(task="token_classify", step_tag_id=999)]
+        metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert out[0].shape == (0, 2)
+
+    def test_chunked_prefill_propagates_none_for_unfinished(self):
+        pooler = self._make_step_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
+        token_ids = [[10, 20, 30, 40]]
+        params = [PoolingParams(task="token_classify", step_tag_id=10)]
+        metadata = _make_metadata(
+            [4],
+            token_ids=token_ids,
+            pooling_params=params,
+            num_scheduled_tokens=[2],
+            seq_lens=[2],
+        )
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        assert out[0] is None
+
+    def test_chunked_prefill_filters_when_finished(self):
+        pooler = self._make_step_pool(chunked=True)
+        hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+        token_ids = [[10, 99, 10, 20]]
+        params = [PoolingParams(task="token_classify", step_tag_id=10)]
+        metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)
+        out = pooler(hidden, metadata)
+        assert len(out) == 1
+        expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])
+        assert torch.equal(out[0], expected)
+
+    def test_requires_token_ids_update(self):
+        pooler = self._make_step_pool()
+        update = pooler.get_pooling_updates("token_classify")
+        assert update.requires_token_ids is True
+
+
+# ---------------------------------------------------------------------------
+# get_tok_pooling_method factory
+# ---------------------------------------------------------------------------
+class TestGetTokPoolingMethod:
+    def test_all(self):
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=False,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            assert isinstance(get_tok_pooling_method("ALL"), AllPool)
+
+    def test_step(self):
+        fake_config = _FakeVllmConfig(
+            scheduler_config=_FakeSchedulerConfig(
+                enable_chunked_prefill=False,
+            ),
+        )
+        with patch(
+            "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",
+            return_value=fake_config,
+        ):
+            assert isinstance(get_tok_pooling_method("STEP"), StepPool)
+
+    def test_unknown_raises(self):
+        with pytest.raises(NotImplementedError, match="UNKNOWN"):
+            get_tok_pooling_method("UNKNOWN")

tests/model_executor/layers/test_rocm_unquantized_gemm.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+
+from vllm.platforms import current_platform
+
+if current_platform.is_cuda():
+    pytest.skip(
+        "ROCm skinny GEMM tests are not supported on CUDA.",
+        allow_module_level=True,
+    )
+
+from vllm.model_executor.layers import utils
+
+
+def test_rocm_unquantized_gemm_gfx1x_wvsplitk_path(monkeypatch):
+    x = torch.randn(1, 64, dtype=torch.float16)
+    weight = torch.randn(128, 64, dtype=torch.float16)
+
+    monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
+    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
+
+    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+    llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitk_mock.assert_called_once()
+    llmm1_mock.assert_not_called()
+    assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
+
+
+def test_rocm_unquantized_gemm_gfx1x_n_gt_5_falls_back(monkeypatch):
+    # wvSplitK skinny GEMM handles n in [1, 5] (see PR #40687); n > 5 must
+    # fall back to torch.nn.functional.linear.
+    x = torch.randn(6, 64, dtype=torch.float16)
+    weight = torch.randn(128, 64, dtype=torch.float16)
+
+    monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: False)
+    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
+
+    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+    llmm1_mock = MagicMock(side_effect=lambda w, x_view, _: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "LLMM1", llmm1_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitk_mock.assert_not_called()
+    llmm1_mock.assert_not_called()
+    assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)
+
+
+def test_rocm_unquantized_gemm_gfx950_wvsplitkrc_path(monkeypatch):
+    x = torch.randn(16, 1024, dtype=torch.float16)
+    weight = torch.randn(256, 1024, dtype=torch.float16)
+
+    monkeypatch.setattr(utils, "use_aiter_triton_gemm", lambda *args: False)
+    monkeypatch.setattr(utils.envs, "VLLM_ROCM_USE_SKINNY_GEMM", True)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx1x", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx9", lambda: False)
+    monkeypatch.setattr("vllm.platforms.rocm.on_gfx950", lambda: True)
+    monkeypatch.setattr(utils, "num_compute_units", lambda: 120)
+
+    wvsplitkrc_mock = MagicMock(side_effect=lambda x_view, w, _, __: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "wvSplitKrc", wvsplitkrc_mock)
+    wvsplitk_mock = MagicMock(side_effect=lambda w, x_view, _, __: x_view @ w.t())
+    monkeypatch.setattr(utils.ops, "wvSplitK", wvsplitk_mock)
+
+    out = utils.rocm_unquantized_gemm_impl(x, weight, None)
+    ref = torch.nn.functional.linear(x, weight, None)
+
+    wvsplitkrc_mock.assert_called_once()
+    wvsplitk_mock.assert_not_called()
+    assert torch.allclose(out, ref, atol=1e-3, rtol=1e-3)

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

tests/model_executor/test_cpu_unquantized_gemm_dispatch.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for CPU unquantized GEMM dispatch behavior."""
+
+import pytest
+import torch
+
+from vllm.model_executor.layers import utils
+from vllm.platforms import current_platform
+
+
+@pytest.fixture(scope="module")
+def _mock_zentorch_linear_unary():
+    """Register a mock zentorch_linear_unary op when zentorch is not installed.
+
+    Allows the dispatch tests to run in CI without a real zentorch build.
+    Skips registration when zentorch is already available.
+    """
+    if hasattr(torch.ops.zentorch, "zentorch_linear_unary"):
+        yield
+        return
+
+    lib_def = torch.library.Library("zentorch", "DEF")
+    lib_def.define(
+        "zentorch_linear_unary("
+        "Tensor input, "
+        "Tensor weight, "
+        "Tensor? bias, "
+        "bool is_weight_prepacked=False"
+        ") -> Tensor"
+    )
+
+    lib_impl = torch.library.Library("zentorch", "IMPL", "CPU")
+    lib_impl.impl(
+        "zentorch_linear_unary",
+        lambda input, weight, bias, is_weight_prepacked=False: (
+            torch.nn.functional.linear(input, weight, bias)
+        ),
+    )
+
+    yield
+
+    lib_impl._destroy()
+    lib_def._destroy()
+
+
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
+def test_dispatch_cpu_unquantized_gemm_uses_zentorch_on_zen(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
+
+    layer = torch.nn.Linear(16, 8, bias=True)
+    x = torch.randn(4, 16)
+    expected = torch.nn.functional.linear(x, layer.weight, layer.bias)
+
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
+    output = layer.cpu_linear(x, layer.weight, layer.bias)
+
+    torch.testing.assert_close(output, expected)
+
+
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
+def test_dispatch_cpu_unquantized_gemm_zen_remove_weight(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
+
+    layer = torch.nn.Linear(16, 8, bias=True)
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=True)
+
+    assert layer.weight.numel() == 0
+
+
+@pytest.mark.usefixtures("_mock_zentorch_linear_unary")
+def test_dispatch_cpu_unquantized_gemm_logs_zentorch_dispatch(monkeypatch):
+    monkeypatch.setattr(current_platform, "is_zen_cpu", lambda: True)
+    expected_prepacked = bool(utils.envs.VLLM_ZENTORCH_WEIGHT_PREPACK) and hasattr(
+        torch.ops.zentorch, "zentorch_weight_prepack_for_linear"
+    )
+
+    log_calls = []
+    monkeypatch.setattr(
+        utils.logger, "debug_once", lambda *args: log_calls.append(args)
+    )
+
+    layer = torch.nn.Linear(16, 8, bias=True)
+    utils.dispatch_cpu_unquantized_gemm(layer, remove_weight=False)
+
+    assert log_calls == [
+        (
+            "CPU unquantized GEMM dispatch: using zentorch_linear_unary (prepacked=%s)",
+            expected_prepacked,
+        )
+    ]

tests/model_executor/test_eagle_quantization.py

@@ -0,0 +1,159 @@
+# 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
+
+DEVICE_TYPE = current_platform.device_type
+DEVICES = (
+    [f"{DEVICE_TYPE}:{i}" for i in range(min(torch.accelerator.device_count(), 2))]
+    if not current_platform.is_cpu()
+    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(default_vllm_config, dist_init, device) -> None:
+    import torch
+
+    from vllm.model_executor.layers.linear import ReplicatedLinear
+
+    if current_platform.is_cuda_alike():
+        torch.accelerator.set_device_index(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_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_eagle3_lm_head_receives_quant_config():
+    """Eagle3LlamaForCausalLM must pass quant_config to ParallelLMHead.
+
+    Without quant_config, quantized lm_head weights (e.g. INT8 per-channel)
+    in Eagle3 drafter checkpoints fail to load because ParallelLMHead doesn't
+    expect weight_packed tensors.
+    """
+    from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.draft_vocab_size = 1000
+    mock_hf_config.hidden_size = 256
+    mock_hf_config.vocab_size = 32000
+    mock_hf_config.logit_scale = 1.0
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.speculative_config.draft_model_config.hf_config = mock_hf_config
+    mock_vllm_config.model_config.get_num_layers.return_value = 32
+    mock_vllm_config.speculative_config.parallel_drafting = False
+
+    with (
+        patch("vllm.model_executor.models.llama_eagle3.LlamaModel") as MockModel,
+        patch("vllm.model_executor.models.llama_eagle3.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.llama_eagle3.LogitsProcessor"),
+        patch(
+            "vllm.model_executor.models.llama_eagle3.get_draft_quant_config",
+            return_value=mock_quant_config,
+        ),
+    ):
+        MockModel.return_value.use_aux_hidden_state = True
+
+        Eagle3LlamaForCausalLM(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert "quant_config" in call_kwargs, (
+            "ParallelLMHead must receive quant_config for quantized lm_head weights"
+        )
+        assert call_kwargs["quant_config"] is mock_quant_config, (
+            "ParallelLMHead must receive the draft model's quant_config"
+        )

tests/model_executor/test_enabled_custom_ops.py

@@ -0,0 +1,151 @@
+# 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, op_registry
+from vllm.model_executor.layers.activation import (
+    GeluAndMul,
+    ReLUSquaredActivation,
+    SiluAndMul,
+)
+from vllm.model_executor.layers.fused_moe.router.fused_topk_router import (
+    dispatch_topk_sigmoid_func,
+    dispatch_topk_softmax_func,
+    vllm_topk_sigmoid,
+    vllm_topk_softmax,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+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 op_registry["rms_norm"].enabled() == ops_enabled[0]
+
+        assert SiluAndMul().enabled() == ops_enabled[1]
+        assert op_registry["silu_and_mul"].enabled() == ops_enabled[1]
+
+        assert GeluAndMul().enabled() == ops_enabled[2]
+        assert op_registry["gelu_and_mul"].enabled() == ops_enabled[2]
+
+        # If registered, subclasses should follow their own name
+        assert Relu3().enabled() == ops_enabled[3]
+        assert 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_softmax_dispatch(use_rocm_aiter: bool):
+    topk_func = dispatch_topk_softmax_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(
+    "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
+)
+def test_topk_sigmoid_dispatch(use_rocm_aiter: bool):
+    topk_func = dispatch_topk_sigmoid_func(use_rocm_aiter)
+
+    if current_platform.is_rocm() and use_rocm_aiter:
+        assert topk_func == rocm_aiter_ops.topk_sigmoid
+    else:
+        assert topk_func == vllm_topk_sigmoid

tests/model_executor/test_ernie45_vl_mrope.py

@@ -0,0 +1,143 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm.model_executor.models.ernie45_vl import (
+    Ernie4_5_VLMoeForConditionalGeneration,
+)
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+@dataclass
+class DummyConfig:
+    spatial_conv_size: int = 2
+    temporal_conv_size: int = 2
+
+
+def make_model(config: DummyConfig) -> Ernie4_5_VLMoeForConditionalGeneration:
+    model = object.__new__(Ernie4_5_VLMoeForConditionalGeneration)
+    model.config = config
+    return model
+
+
+def make_mm_feature(
+    *,
+    modality: str,
+    offset: int,
+    length: int,
+    grid_thw: tuple[int, int, int],
+) -> MultiModalFeatureSpec:
+    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
+    return MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                field_name: MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality=modality,
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def test_get_mrope_input_positions_text_only():
+    model = make_model(DummyConfig())
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[11, 12, 13, 14, 15],
+        mm_features=[],
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == 0
+
+
+def test_get_mrope_input_positions_single_image():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4],
+            [0, 1, 1, 2, 2, 3, 4],
+            [0, 1, 2, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2
+
+
+def test_get_mrope_input_positions_interleaved_image_and_video():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        ),
+        make_mm_feature(
+            modality="video",
+            offset=7,
+            length=2,
+            grid_thw=(2, 4, 2),
+        ),
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 50, 51],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 8],
+            [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8],
+            [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 8],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2

tests/model_executor/test_gemma_hidden_act.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.model_executor.layers.activation import (
+    GeluAndMul,
+    SiluAndMul,
+    get_act_and_mul_fn,
+    get_act_fn,
+)
+from vllm.model_executor.models.gemma3 import Gemma3MLP
+from vllm.model_executor.models.gemma4 import Gemma4MLP
+
+
+@pytest.mark.parametrize(
+    ("activation_name", "expected_type"),
+    [
+        ("gelu_pytorch_tanh", GeluAndMul),
+        ("silu", SiluAndMul),
+        ("swish", SiluAndMul),
+    ],
+)
+def test_get_act_and_mul_fn_supports_gemma_hidden_act_aliases(
+    activation_name: str,
+    expected_type: type[torch.nn.Module],
+    default_vllm_config,
+) -> None:
+    assert isinstance(get_act_and_mul_fn(activation_name), expected_type)
+
+
+def test_get_act_fn_supports_swish_alias() -> None:
+    assert isinstance(get_act_fn("swish"), torch.nn.SiLU)
+
+
+@pytest.mark.parametrize("mlp_cls", [Gemma3MLP, Gemma4MLP])
+@pytest.mark.parametrize(
+    ("activation_name", "expected_type"),
+    [
+        ("gelu_pytorch_tanh", GeluAndMul),
+        ("silu", SiluAndMul),
+        ("swish", SiluAndMul),
+    ],
+)
+def test_gemma_mlp_supports_hidden_act_variants(
+    mlp_cls: type[torch.nn.Module],
+    activation_name: str,
+    expected_type: type[torch.nn.Module],
+    default_vllm_config,
+    dist_init,
+) -> None:
+    mlp = mlp_cls(
+        hidden_size=16,
+        intermediate_size=32,
+        hidden_activation=activation_name,
+    )
+
+    assert isinstance(mlp.act_fn, expected_type)
+    assert mlp(torch.randn(3, 16)).shape == (3, 16)

tests/model_executor/test_keye_mrope.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+
+import pytest
+import torch
+
+from vllm.model_executor.models.keye import KeyeForConditionalGeneration
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+@dataclass
+class DummyVisionConfig:
+    spatial_merge_size: int = 2
+
+
+@dataclass
+class DummyConfig:
+    vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
+
+
+def make_model(config: DummyConfig) -> KeyeForConditionalGeneration:
+    model = object.__new__(KeyeForConditionalGeneration)
+    model.config = config
+    return model
+
+
+def make_mm_feature(
+    *,
+    modality: str,
+    offset: int,
+    length: int,
+    grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
+) -> MultiModalFeatureSpec:
+    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
+    return MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                field_name: MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality=modality,
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def test_get_mrope_input_positions_text_only():
+    model = make_model(DummyConfig())
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[11, 12, 13, 14, 15],
+        mm_features=[],
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == 0
+
+
+def test_get_mrope_input_positions_single_image():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4],
+            [0, 1, 1, 2, 2, 3, 4],
+            [0, 1, 2, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2
+
+
+def test_get_mrope_input_positions_interleaved_image_and_video():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        ),
+        make_mm_feature(
+            modality="video",
+            offset=7,
+            length=4,
+            grid_thw=[(2, 4, 2)],
+        ),
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31, 40, 41, 42, 43, 50, 51],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4, 5, 5, 7, 7, 9, 10],
+            [0, 1, 1, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            [0, 1, 2, 1, 2, 3, 4, 5, 5, 7, 7, 9, 10],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2

tests/model_executor/test_keye_vl1_5_mrope.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+
+import pytest
+import torch
+
+from vllm.model_executor.models.keye_vl1_5 import KeyeVL1_5ForConditionalGeneration
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+@dataclass
+class DummyVisionConfig:
+    spatial_merge_size: int = 2
+
+
+@dataclass
+class DummyConfig:
+    vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
+
+
+def make_model(config: DummyConfig) -> KeyeVL1_5ForConditionalGeneration:
+    model = object.__new__(KeyeVL1_5ForConditionalGeneration)
+    model.config = config
+    return model
+
+
+def make_mm_feature(
+    *,
+    modality: str,
+    offset: int,
+    length: int,
+    grid_thw: tuple[int, int, int] | list[tuple[int, int, int]],
+    is_embed: list[bool] | None = None,
+) -> MultiModalFeatureSpec:
+    field_name = "image_grid_thw" if modality == "image" else "video_grid_thw"
+    return MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                field_name: MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality=modality,
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(
+            offset=offset,
+            length=length,
+            is_embed=None if is_embed is None else torch.tensor(is_embed),
+        ),
+    )
+
+
+def test_get_mrope_input_positions_text_only():
+    model = make_model(DummyConfig())
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[11, 12, 13, 14, 15],
+        mm_features=[],
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == 0
+
+
+def test_get_mrope_input_positions_single_image():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="image",
+            offset=1,
+            length=4,
+            grid_thw=(1, 4, 4),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 20, 21, 22, 23, 30, 31],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4],
+            [0, 1, 1, 2, 2, 3, 4],
+            [0, 1, 2, 1, 2, 3, 4],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == -2
+
+
+def test_get_mrope_input_positions_video_uses_embed_ranges():
+    model = make_model(DummyConfig())
+    mm_features = [
+        make_mm_feature(
+            modality="video",
+            offset=1,
+            length=8,
+            grid_thw=[(2, 4, 2)],
+            is_embed=[False, False, True, True, False, False, True, True],
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=[10, 101, 102, 20, 21, 103, 104, 30, 31, 40, 41],
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
+            [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
+            [0, 1, 2, 3, 3, 5, 6, 7, 7, 9, 10],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    assert delta == 0

tests/model_executor/test_mistral_large_3_eagle.py

@@ -0,0 +1,146 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+
+from vllm.config.compilation import CompilationMode
+from vllm.model_executor.models import deepseek_v2 as deepseek_mod
+from vllm.model_executor.models import mistral_large_3_eagle as eagle_mod
+
+
+class DummyPPGroup:
+    world_size = 1
+    is_first_rank = True
+    is_last_rank = True
+
+
+class DummyEmbedding(nn.Module):
+    def __init__(self, vocab_size, hidden_size, *args, **kwargs):
+        super().__init__()
+        self.hidden_size = hidden_size
+
+    def forward(self, input_ids):
+        return torch.zeros(
+            (*input_ids.shape, self.hidden_size),
+            dtype=torch.float32,
+            device=input_ids.device,
+        )
+
+
+class DummyLinear(nn.Module):
+    def __init__(self, in_features, out_features, *args, **kwargs):
+        super().__init__()
+        self.out_features = out_features
+
+    def forward(self, x):
+        return torch.zeros(
+            (*x.shape[:-1], self.out_features),
+            dtype=x.dtype,
+            device=x.device,
+        )
+
+
+class DummyNorm(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, hidden_states, residual=None):
+        return hidden_states, residual
+
+
+class DummyDecoderLayer(nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+
+    def forward(self, positions, hidden_states, residual, llama_4_scaling=None):
+        return hidden_states, residual
+
+
+def make_vllm_config(
+    *, model_type="mistral3", qk_nope_head_dim=128, qk_rope_head_dim=64
+):
+    hf_config = SimpleNamespace(
+        model_type=model_type,
+        first_k_dense_replace=0,
+        vocab_size=32000,
+        hidden_size=16,
+        num_hidden_layers=1,
+        rms_norm_eps=1e-5,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+    )
+
+    return SimpleNamespace(
+        model_config=SimpleNamespace(hf_config=hf_config),
+        quant_config=None,
+        parallel_config=SimpleNamespace(
+            eplb_config=SimpleNamespace(num_redundant_experts=0),
+        ),
+        scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
+        cache_config=None,
+        compilation_config=SimpleNamespace(mode=CompilationMode.NONE),
+    )
+
+
+@pytest.fixture(autouse=True)
+def patch_heavy_modules(monkeypatch):
+    monkeypatch.setattr(eagle_mod, "get_pp_group", lambda: DummyPPGroup())
+    monkeypatch.setattr(deepseek_mod, "get_pp_group", lambda: DummyPPGroup())
+
+    monkeypatch.setattr(eagle_mod, "VocabParallelEmbedding", DummyEmbedding)
+    monkeypatch.setattr(eagle_mod, "RowParallelLinear", DummyLinear)
+    monkeypatch.setattr(eagle_mod, "RMSNorm", DummyNorm)
+    monkeypatch.setattr(eagle_mod, "DeepseekV2DecoderLayer", DummyDecoderLayer)
+
+
+@pytest.mark.cpu_test
+@pytest.mark.parametrize(
+    ("model_type", "qk_nope_head_dim", "qk_rope_head_dim", "expected_use_mha"),
+    [
+        # MLA-style config: should not use MHA.
+        ("mistral3", 128, 64, False),
+        # No MLA dims: should use MHA, matching DeepseekV2Model.__init__ logic.
+        ("mistral3", 0, 0, True),
+        # DeepSeek model type always uses MHA by the parent logic.
+        ("deepseek", 128, 64, True),
+    ],
+)
+def test_eagle_mistral_large3_initializes_deepseek_runtime_attrs(
+    model_type,
+    qk_nope_head_dim,
+    qk_rope_head_dim,
+    expected_use_mha,
+):
+    vllm_config = make_vllm_config(
+        model_type=model_type,
+        qk_nope_head_dim=qk_nope_head_dim,
+        qk_rope_head_dim=qk_rope_head_dim,
+    )
+
+    model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config)
+
+    assert model.aux_hidden_state_layers == ()
+    assert model.use_mha is expected_use_mha
+
+    # Add this if your fix also copies num_redundant_experts from
+    # DeepseekV2Model.__init__.
+    assert model.num_redundant_experts == 0
+
+
+@pytest.mark.cpu_test
+def test_eagle_mistral_large3_forward_reuses_deepseek_parent_forward():
+    vllm_config = make_vllm_config()
+    model = eagle_mod.EagleMistralLarge3Model(vllm_config=vllm_config)
+
+    input_ids = torch.tensor([[1, 2, 3]])
+    positions = torch.tensor([[0, 1, 2]])
+    hidden_states = torch.zeros((1, 3, 16))
+
+    output = model(input_ids, positions, hidden_states)
+
+    assert isinstance(output, torch.Tensor)
+    assert output.shape == hidden_states.shape

tests/model_executor/test_model_load_with_params.py

@@ -0,0 +1,139 @@
+# 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 DispatchPooler
+from vllm.model_executor.layers.pooler.seqwise import CLSPool, MeanPool
+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.seq_pooling_type == "CLS"
+        assert model_config.pooler_config.tok_pooling_type == "ALL"
+        assert model_config.pooler_config.use_activation
+
+        # 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.seq_pooling_type == "MEAN"
+        assert model_config.pooler_config.tok_pooling_type == "ALL"
+        assert model_config.pooler_config.use_activation
+
+        # 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_nemotron_h_quantization.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock, patch
+
+
+def test_nemotron_h_lm_head_receives_quant_config():
+    from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.vocab_size = 128
+    mock_hf_config.hidden_size = 64
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.model_config.hf_config = mock_hf_config
+    mock_vllm_config.model_config.dtype = None
+    mock_vllm_config.scheduler_config = Mock()
+    mock_vllm_config.quant_config = mock_quant_config
+
+    with (
+        patch("vllm.model_executor.models.nemotron_h.NemotronHModel") as MockModel,
+        patch("vllm.model_executor.models.nemotron_h.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.nemotron_h.LogitsProcessor"),
+    ):
+        MockModel.return_value.make_empty_intermediate_tensors = Mock()
+        MockModel.return_value.has_moe = False
+
+        NemotronHForCausalLM(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert call_kwargs["quant_config"] is mock_quant_config

tests/model_executor/test_oink_integration.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import multiprocessing
+import types
+
+import pytest
+
+from vllm.platforms import current_platform
+
+
+def _test_oink_availability_impl(
+    device_capability: tuple[int, int],
+    has_rmsnorm: bool,
+    has_fused_add_rms_norm: bool,
+    expected_available: bool,
+    expected_fused: bool,
+) -> None:
+    """Test OINK support detection with mocked state."""
+    import torch
+
+    from vllm import platforms
+
+    # Mock device capability (class method, override on class)
+    dc = platforms.interface.DeviceCapability(*device_capability)
+    platforms.current_platform.__class__.get_device_capability = lambda device_id=0: dc
+
+    # Mock oink ops
+    oink_ops = types.SimpleNamespace()
+    if has_rmsnorm:
+        oink_ops.rmsnorm = lambda x, w, eps: x
+    if has_fused_add_rms_norm:
+        oink_ops.fused_add_rms_norm = lambda x, residual, w, eps: None
+
+    torch.ops.oink = oink_ops
+
+    # Now import vllm modules with mocks in place (fresh import with mocked platform)
+    import vllm.kernels.oink_ops  # noqa: F401
+    from vllm.ir.ops import fused_add_rms_norm, rms_norm
+
+    # Verify support checks
+    assert rms_norm.impls["oink"].supported is expected_available
+    assert fused_add_rms_norm.impls["oink"].supported is expected_fused
+
+
+@pytest.mark.parametrize(
+    "device_capability,has_rmsnorm,has_fused_add_rms_norm,expected_available,expected_fused",
+    [
+        # Case 1: < SM100, ops not supported
+        ((9, 0), True, False, False, False),
+        # Case 2: CUDA available and SM100, rmsnorm op registered
+        ((10, 0), True, False, True, False),
+        # Case 3: SM100 with both rmsnorm and fused_add_rms_norm
+        ((10, 0), True, True, True, True),
+    ],
+)
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
+def test_oink_availability_checks(
+    device_capability: tuple[int, int],
+    has_rmsnorm: bool,
+    has_fused_add_rms_norm: bool,
+    expected_available: bool,
+    expected_fused: bool,
+):
+    """Test OINK support detection with clean import state for each parameter set."""
+
+    # Use spawn to run function in fresh process with clean imports
+    # TODO migrate to spawn utility:
+    # https://github.com/vllm-project/vllm/issues/41415
+    ctx = multiprocessing.get_context("spawn")
+    process = ctx.Process(
+        target=_test_oink_availability_impl,
+        args=(
+            device_capability,
+            has_rmsnorm,
+            has_fused_add_rms_norm,
+            expected_available,
+            expected_fused,
+        ),
+    )
+    process.start()
+    process.join()
+
+    if process.exitcode != 0:
+        raise AssertionError(
+            f"Subprocess test failed with exit code {process.exitcode}"
+        )
+
+
+def test_can_view_as_2d_stride_guard():
+    # No global import
+    import torch
+
+    # Import the helper from the kernels module.
+    from vllm.kernels.oink_ops import _can_view_as_2d
+
+    x = torch.zeros((2, 3, 4))
+    assert _can_view_as_2d(x) is True
+
+    # Size-1 dims should be ignored by the viewability check.
+    # Create a tensor where stride(0) != stride(1) * size(1) due to padding,
+    # but view(-1, H) is still valid because dim 1 has size 1.
+    base = torch.zeros((2, 10, 4))
+    x_singleton = base[:, :1, :]
+    x_singleton.view(-1, x_singleton.shape[-1])
+    assert _can_view_as_2d(x_singleton) is True
+
+    # Middle-dimension stride break: view(-1, hidden) should be invalid.
+    x2 = x[:, ::2, :]
+    with pytest.raises(RuntimeError):
+        x2.view(-1, x2.shape[-1])
+    assert _can_view_as_2d(x2) is False

tests/model_executor/test_paddleocr_vl_mrope.py

@@ -0,0 +1,210 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from dataclasses import dataclass, field
+
+import pytest
+import torch
+
+from vllm.model_executor.models.paddleocr_vl import (
+    PaddleOCRVLForConditionalGeneration,
+)
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+@dataclass
+class DummyVisionConfig:
+    spatial_merge_size: int = 2
+    patch_size: int = 14
+
+
+@dataclass
+class DummyConfig:
+    image_token_id: int = 151655
+    video_token_id: int = 151654
+    vision_start_token_id: int = 151652
+    vision_end_token_id: int = 151653
+    vision_config: DummyVisionConfig = field(default_factory=DummyVisionConfig)
+
+
+def make_model(config: DummyConfig) -> PaddleOCRVLForConditionalGeneration:
+    model = object.__new__(PaddleOCRVLForConditionalGeneration)
+    model.config = config
+    return model
+
+
+def make_mm_feature(
+    *,
+    offset: int,
+    length: int,
+    image_grid_thw: tuple[int, int, int],
+) -> MultiModalFeatureSpec:
+    return MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                "image_grid_thw": MultiModalFieldElem(
+                    data=torch.tensor(image_grid_thw),
+                    field=None,
+                ),
+            }
+        ),
+        modality="image",
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=offset, length=length),
+    )
+
+
+def test_get_mrope_input_positions_text_only():
+    model = make_model(DummyConfig())
+    input_tokens = [11, 12, 13, 14, 15]
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=input_tokens,
+        mm_features=[],
+    )
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+            [0, 1, 2, 3, 4],
+        ]
+    )
+    assert torch.equal(positions, expected)
+    assert delta == 0
+
+
+def test_get_mrope_input_positions_single_image():
+    model = make_model(DummyConfig())
+    spatial_merge_size = model.config.vision_config.spatial_merge_size
+
+    t, h, w = 1, 2, 2
+    num_image_tokens = t * h * w
+
+    input_tokens = (
+        [10]
+        + [model.config.vision_start_token_id]
+        + [model.config.image_token_id] * num_image_tokens
+        + [model.config.vision_end_token_id]
+        + [30, 31]
+    )
+
+    mm_features = [
+        make_mm_feature(
+            offset=2,  # 1 (text) + 1 (vision_start)
+            length=num_image_tokens,
+            image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=input_tokens,
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 2, 2, 2, 2, 4, 5, 6],
+            [0, 1, 2, 2, 3, 3, 4, 5, 6],
+            [0, 1, 2, 3, 2, 3, 4, 5, 6],
+        ]
+    )
+
+    assert torch.equal(positions, expected)
+    expected_delta = (positions.max().item() + 1) - len(input_tokens)
+    assert delta == expected_delta
+
+
+def test_get_mrope_input_positions_multiple_images():
+    model = make_model(DummyConfig())
+    spatial_merge_size = model.config.vision_config.spatial_merge_size
+
+    t1, h1, w1 = 1, 2, 2
+    num1 = t1 * h1 * w1
+
+    t2, h2, w2 = 1, 1, 3
+    num2 = t2 * h2 * w2
+
+    input_tokens = (
+        [10]
+        + [model.config.vision_start_token_id]
+        + [model.config.image_token_id] * num1
+        + [model.config.vision_end_token_id]
+        + [20, 21]
+        + [model.config.vision_start_token_id]
+        + [model.config.image_token_id] * num2
+        + [model.config.vision_end_token_id]
+        + [30]
+    )
+
+    mm_features = [
+        make_mm_feature(
+            offset=2,
+            length=num1,
+            image_grid_thw=(t1, h1 * spatial_merge_size, w1 * spatial_merge_size),
+        ),
+        make_mm_feature(
+            offset=2 + num1 + 1 + 2 + 1,
+            length=num2,
+            image_grid_thw=(t2, h2 * spatial_merge_size, w2 * spatial_merge_size),
+        ),
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=input_tokens,
+        mm_features=mm_features,
+    )
+
+    assert positions.shape == (3, 15)
+    assert not torch.equal(positions[:, 2:6], torch.arange(4).expand(3, 4) + 2)
+    assert not torch.equal(positions[:, 10:13], torch.arange(3).expand(3, 3) + 10)
+
+
+def test_get_mrope_input_positions_image_at_start():
+    model = make_model(DummyConfig())
+    spatial_merge_size = model.config.vision_config.spatial_merge_size
+
+    t, h, w = 1, 2, 2
+    num_tokens = t * h * w
+
+    input_tokens = (
+        [model.config.vision_start_token_id]
+        + [model.config.image_token_id] * num_tokens
+        + [model.config.vision_end_token_id]
+        + [10, 11]
+    )
+
+    mm_features = [
+        make_mm_feature(
+            offset=1,  # start token at index 0
+            length=num_tokens,
+            image_grid_thw=(t, h * spatial_merge_size, w * spatial_merge_size),
+        )
+    ]
+
+    positions, delta = model.get_mrope_input_positions(
+        input_tokens=input_tokens,
+        mm_features=mm_features,
+    )
+
+    expected = torch.tensor(
+        [
+            [0, 1, 1, 1, 1, 3, 4, 5],
+            [0, 1, 1, 2, 2, 3, 4, 5],
+            [0, 1, 2, 1, 2, 3, 4, 5],
+        ]
+    )
+
+    assert torch.equal(positions, expected)

tests/model_executor/test_qwen3_5_quantization.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from unittest.mock import Mock, patch
+
+
+def test_qwen3_5_lm_head_receives_quant_config():
+    from vllm.model_executor.models.qwen3_5 import Qwen3_5ForCausalLMBase
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.tie_word_embeddings = False
+    mock_hf_config.vocab_size = 128
+    mock_hf_config.hidden_size = 64
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.model_config.hf_text_config = mock_hf_config
+    mock_vllm_config.cache_config.mamba_cache_mode = "align"
+    mock_vllm_config.scheduler_config = Mock()
+    mock_vllm_config.quant_config = mock_quant_config
+    mock_vllm_config.lora_config = None
+
+    mock_pp_group = Mock()
+    mock_pp_group.is_last_rank = True
+
+    with (
+        patch("vllm.model_executor.models.qwen3_5.Qwen3_5Model") as MockModel,
+        patch("vllm.model_executor.models.qwen3_5.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.qwen3_5.LogitsProcessor"),
+        patch(
+            "vllm.model_executor.models.qwen3_5.get_pp_group",
+            return_value=mock_pp_group,
+        ),
+    ):
+        MockModel.return_value.make_empty_intermediate_tensors = Mock()
+
+        Qwen3_5ForCausalLMBase(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert call_kwargs["quant_config"] is mock_quant_config
+
+
+def test_qwen3_5_mtp_lm_head_receives_quant_config():
+    from vllm.config import CompilationMode
+    from vllm.model_executor.models.qwen3_5_mtp import Qwen3_5MTP
+
+    mock_quant_config = Mock()
+
+    mock_hf_config = Mock()
+    mock_hf_config.tie_word_embeddings = False
+    mock_hf_config.vocab_size = 128
+    mock_hf_config.hidden_size = 64
+
+    mock_vllm_config = Mock()
+    mock_vllm_config.model_config.hf_text_config = mock_hf_config
+    mock_vllm_config.cache_config.mamba_cache_mode = "align"
+    mock_vllm_config.compilation_config.mode = CompilationMode.NONE
+    mock_vllm_config.quant_config = mock_quant_config
+
+    mock_pp_group = Mock()
+    mock_pp_group.is_last_rank = True
+
+    with (
+        patch("vllm.model_executor.models.qwen3_5_mtp.Qwen3_5MultiTokenPredictor"),
+        patch("vllm.model_executor.models.qwen3_5_mtp.ParallelLMHead") as MockLMHead,
+        patch("vllm.model_executor.models.qwen3_5_mtp.LogitsProcessor"),
+        patch(
+            "vllm.model_executor.models.qwen3_5_mtp.get_pp_group",
+            return_value=mock_pp_group,
+        ),
+    ):
+        Qwen3_5MTP(vllm_config=mock_vllm_config)
+
+        MockLMHead.assert_called_once()
+        call_kwargs = MockLMHead.call_args.kwargs
+        assert call_kwargs["quant_config"] is mock_quant_config

tests/model_executor/test_qwen3_omni.py

@@ -0,0 +1,222 @@
+# 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_expected_hidden_size = lambda: 100
+    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_qwen3_vl_mrope.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import dataclasses
+import random
+from dataclasses import dataclass
+
+import pytest
+import torch
+
+from vllm.model_executor.models.qwen3_vl import Qwen3VLForConditionalGeneration
+from vllm.multimodal.inputs import (
+    MultiModalFeatureSpec,
+    MultiModalFieldElem,
+    MultiModalKwargsItem,
+    PlaceholderRange,
+)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _force_cpu_default_device():
+    # _get_mrope_input_positions returns CPU tensors (via torch.from_numpy).
+    # Ensure the default device is CPU so the rest of the test tensors match.
+    original = torch.get_default_device()
+    torch.set_default_device("cpu")
+    yield
+    torch.set_default_device(original)
+
+
+IMAGE_TOKEN_ID = 999
+VIDEO_TOKEN_ID = 888
+VISION_START_TOKEN_ID = 777
+VISION_END_TOKEN_ID = 778
+
+
+@dataclass
+class DummyVisionConfig:
+    spatial_merge_size: int = 1
+
+
+@dataclass
+class DummyConfig:
+    image_token_id: int = IMAGE_TOKEN_ID
+    video_token_id: int = VIDEO_TOKEN_ID
+    vision_start_token_id: int = VISION_START_TOKEN_ID
+    vision_end_token_id: int = VISION_END_TOKEN_ID
+    vision_config: DummyVisionConfig = dataclasses.field(
+        default_factory=DummyVisionConfig
+    )
+
+
+def make_video_embedding(
+    t, h, w, interleave_text_tokens: tuple[int, int], video_pruning_rate: float = 0.0
+):
+    """
+    Helper function to make a video embedding for a given video size and pruning rate.
+
+    Args:
+        t: Number of frames.
+        h: Number of rows.
+        w: Number of columns.
+        interleave_text_tokens: Tuple of minimum and maximum number of text tokens to
+            interleave with the video.
+        video_pruning_rate: Pruning rate for the video.
+
+    Returns:
+        Tuple of (unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask)
+    """
+    unpruned_tokens_sequence = []
+    population = list(range(1, 100))
+
+    for _ in range(t):
+        num_prefix_tokens = random.randint(
+            interleave_text_tokens[0], interleave_text_tokens[1]
+        )
+
+        prefix_tokens = random.choices(population, k=num_prefix_tokens)
+        vision_tokens = (
+            [VISION_START_TOKEN_ID] + [VIDEO_TOKEN_ID] * h * w + [VISION_END_TOKEN_ID]
+        )
+
+        unpruned_tokens_sequence.extend(prefix_tokens)
+        unpruned_tokens_sequence.extend(vision_tokens)
+
+    unpruned_tokens_sequence = torch.tensor(unpruned_tokens_sequence, dtype=torch.long)
+    video_token_mask = unpruned_tokens_sequence == VIDEO_TOKEN_ID
+
+    pruning_mask = torch.bernoulli(video_token_mask.float() * video_pruning_rate).bool()  # type: ignore[attr-defined]
+    # Sanity check that we don't prune what should not be pruned.
+    assert not pruning_mask[~video_token_mask].any()
+
+    retention_mask = ~pruning_mask
+    pruned_tokens_sequence = unpruned_tokens_sequence[retention_mask]
+    return unpruned_tokens_sequence, pruned_tokens_sequence, retention_mask
+
+
+@pytest.mark.parametrize("spatial_merge_size", [1, 2])
+@pytest.mark.parametrize("grid_thw", [[3, 8, 7], [128, 10, 12]])
+@pytest.mark.parametrize("num_prefix_tokens", [1, 11])
+@pytest.mark.parametrize("num_suffix_tokens", [0, 7])
+@pytest.mark.parametrize("video_pruning_rate", [0, 0.25, 0.75])
+@pytest.mark.parametrize("interleave_text_tokens", [(0, 0), (1, 4)])
+def test_match_qwen3vl_mrope_evs_on(
+    spatial_merge_size: int,
+    num_prefix_tokens: int,
+    grid_thw: tuple[int, int, int],
+    num_suffix_tokens: int,
+    video_pruning_rate: float,
+    interleave_text_tokens: tuple[int, int],
+):
+    hf_config = DummyConfig()
+    hf_config.vision_config.spatial_merge_size = spatial_merge_size
+
+    t, h, w = grid_thw
+    population = list(range(1, 100))
+    prefix_tokens = random.choices(population, k=num_prefix_tokens)
+    suffix_tokens = random.choices(population, k=num_suffix_tokens)
+
+    video_tokens, video_tokens_pruned, retention_mask = make_video_embedding(
+        t,
+        h // spatial_merge_size,
+        w // spatial_merge_size,
+        interleave_text_tokens=interleave_text_tokens,
+        video_pruning_rate=video_pruning_rate,
+    )
+    assert len(video_tokens) == len(retention_mask)
+
+    input_tokens = prefix_tokens + video_tokens.tolist() + suffix_tokens
+    input_tokens_pruned = prefix_tokens + video_tokens_pruned.tolist() + suffix_tokens
+
+    whole_sequence_retention_mask = torch.cat(
+        [
+            torch.ones(len(prefix_tokens), dtype=torch.bool),
+            retention_mask,
+            torch.ones(len(suffix_tokens), dtype=torch.bool),
+        ],
+        dim=0,
+    )
+
+    # Build the GT mrope for unpruned input.
+    mm_feature = MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                "video_grid_thw": MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality="video",
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=0, length=len(input_tokens)),
+    )
+    expected_mrope, _ = Qwen3VLForConditionalGeneration._get_mrope_input_positions(
+        input_tokens=input_tokens,
+        mm_features=[mm_feature],
+        config=hf_config,
+    )
+
+    # Compute mrope for a video-only media (unpruned).
+    mm_feature = MultiModalFeatureSpec(
+        data=MultiModalKwargsItem(
+            {
+                "video_grid_thw": MultiModalFieldElem(
+                    data=torch.tensor(grid_thw),
+                    field=None,  # HACK.
+                ),
+            }
+        ),
+        modality="video",
+        identifier="DUMMY",
+        mm_position=PlaceholderRange(offset=0, length=video_tokens.numel()),
+    )
+    video_mrope, _ = Qwen3VLForConditionalGeneration._get_mrope_input_positions(
+        input_tokens=video_tokens.tolist(),
+        mm_features=[mm_feature],
+        config=hf_config,
+    )
+    video_mrope = video_mrope.permute(1, 0)  # [N, 3]
+    hidden_size = 16
+
+    is_video_embed = torch.isin(
+        video_tokens_pruned, torch.tensor([VIDEO_TOKEN_ID], dtype=torch.long)
+    )
+
+    expanded_positions = torch.full(
+        (len(video_tokens_pruned), 5),
+        fill_value=-100,
+        device=video_mrope.device,
+        dtype=torch.long,
+    )
+    expanded_positions[is_video_embed, :3] = video_mrope[retention_mask][is_video_embed]
+    expanded_positions[~is_video_embed, :3] = video_mrope[retention_mask][
+        ~is_video_embed
+    ]
+
+    is_vision_start = video_tokens_pruned == VISION_START_TOKEN_ID
+    expanded_positions[..., 3] = is_vision_start
+    expanded_positions[..., 4] = is_video_embed
+
+    # Check that all positions were filled, since we initialized them as negative.
+    assert (expanded_positions >= 0).all()
+
+    video_embeddings = torch.empty(
+        (len(video_tokens_pruned), hidden_size), device=video_mrope.device
+    )
+
+    video_embeddings = torch.cat(
+        [
+            video_embeddings,
+            expanded_positions.float(),
+        ],
+        dim=1,
+    )
+    multimodal_embeddings = [video_embeddings]
+
+    expected_mrope_masked = expected_mrope[:, whole_sequence_retention_mask]
+
+    # Initialize computed_mrope with sequential positions for all prefix tokens
+    computed_mrope = torch.empty((3, len(input_tokens_pruned)), dtype=torch.long)
+    computed_mrope[:, 0 : len(prefix_tokens)] = expected_mrope[
+        :, 0 : len(prefix_tokens)
+    ]
+
+    # Paranoia check that computed_mrope is wrong.
+    assert not torch.equal(computed_mrope, expected_mrope_masked)
+
+    _, actual_mrope, _ = Qwen3VLForConditionalGeneration._recompute_mrope_positions(
+        input_ids=input_tokens_pruned,
+        multimodal_embeddings=multimodal_embeddings,
+        mrope_positions=computed_mrope,
+        num_computed_tokens=len(prefix_tokens),
+        vision_start_token_id=hf_config.vision_start_token_id,
+        image_token_id=hf_config.image_token_id,
+        video_token_id=hf_config.video_token_id,
+    )
+
+    assert torch.equal(actual_mrope, expected_mrope_masked)

tests/model_executor/test_routed_experts_capture.py

@@ -0,0 +1,252 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import types
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from vllm.distributed.eplb.eplb_state import EplbLayerState
+from vllm.model_executor.layers.fused_moe.config import RoutingMethodType
+from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
+    RoutedExpertsCapturer,
+)
+from vllm.model_executor.layers.fused_moe.router.base_router import BaseRouter
+
+pytestmark = pytest.mark.cpu_test
+
+_REC_MODULE = "vllm.model_executor.layers.fused_moe.routed_experts_capturer"
+
+
+def _capturer_with_buffer(
+    *,
+    max_tokens: int = 8,
+    num_layers: int = 4,
+    num_experts_per_tok: int = 2,
+    dp_rank: int = 0,
+    tp_size: int = 1,
+) -> RoutedExpertsCapturer:
+    # Bypass __init__ so the test can use a CPU buffer and skip the
+    # VllmConfig dependency. The CUDA device-tensor allocation in the
+    # real constructor is not what we are exercising here.
+    c = RoutedExpertsCapturer.__new__(RoutedExpertsCapturer)
+    c.dp_rank = dp_rank
+    c.tp_size = tp_size
+    c.device_buffer = torch.full(
+        (max_tokens, num_layers, num_experts_per_tok),
+        -1,
+        dtype=torch.int32,
+    )
+    return c
+
+
+class DummyRouter(BaseRouter):
+    @property
+    def routing_method_type(self) -> RoutingMethodType:
+        return RoutingMethodType.FUSED_TOPK
+
+    def _compute_routing(
+        self, hidden_states, router_logits, indices_type, *, input_ids=None
+    ):
+        topk_ids = torch.tensor([[1, 2], [3, 4]], dtype=torch.int64)
+        topk_weights = torch.ones_like(topk_ids, dtype=torch.float32)
+        return topk_weights, topk_ids
+
+    def _apply_eplb_mapping(self, topk_ids: torch.Tensor) -> torch.Tensor:
+        # Make mapping observable without requiring CUDA EPLB path.
+        return topk_ids + 10
+
+
+def _make_router(eplb_state: EplbLayerState | None = None) -> DummyRouter:
+    return DummyRouter(
+        top_k=2,
+        global_num_experts=16,
+        eplb_state=eplb_state,
+    )
+
+
+def test_base_router_capture_pre_eplb_mapping():
+    router = _make_router()
+    captured = []
+
+    def capture_fn(ids):
+        captured.append(ids.clone())
+
+    router.set_capture_fn(capture_fn)
+    topk_weights, topk_ids = router.select_experts(
+        hidden_states=torch.empty(1),
+        router_logits=torch.empty(1),
+    )
+
+    assert topk_weights.shape == topk_ids.shape
+    assert len(captured) == 1
+    assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
+    assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
+
+
+def test_base_router_capture_with_eplb_enabled():
+    eplb_state = EplbLayerState()
+    eplb_state.expert_load_view = torch.zeros(32, dtype=torch.int64)
+    eplb_state.logical_to_physical_map = torch.arange(32).view(32, 1)
+    eplb_state.logical_replica_count = torch.ones(32, dtype=torch.int64)
+    eplb_state.should_record_tensor = torch.ones((), dtype=torch.bool)
+    eplb_state.num_unpadded_tokens_tensors = [torch.tensor(0, dtype=torch.int32)]
+    router = _make_router(eplb_state=eplb_state)
+
+    captured = []
+
+    def capture_fn(ids):
+        captured.append(ids.clone())
+
+    router.set_capture_fn(capture_fn)
+    _, topk_ids = router.select_experts(
+        hidden_states=torch.empty(1),
+        router_logits=torch.empty(1),
+    )
+
+    assert len(captured) == 1
+    # Capture should see logical ids pre-EPLB mapping.
+    assert torch.equal(captured[0], torch.tensor([[1, 2], [3, 4]]))
+    # Our DummyRouter mapping adds +10.
+    assert torch.equal(topk_ids, torch.tensor([[11, 12], [13, 14]]))
+
+
+def test_gpu_model_runner_binds_router_capture(monkeypatch):
+    from vllm.v1.worker import gpu_model_runner as gmr
+
+    class _DummyRouter:
+        _routing_replay_out: torch.Tensor | None = None
+
+    class DummyFusedMoE:
+        def __init__(self):
+            self.layer_id = 7
+            self.router = _make_router()
+
+    class DummyCapturer:
+        def __init__(self):
+            self.calls = []
+
+        def capture(self, layer_id, topk_ids):
+            self.calls.append((layer_id, topk_ids))
+
+    dummy_module = DummyFusedMoE()
+
+    # Patch the runtime import inside _bind_routed_experts_capturer.
+    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
+
+    monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
+
+    dummy_self = types.SimpleNamespace(
+        compilation_config=types.SimpleNamespace(
+            static_forward_context={"dummy": dummy_module}
+        )
+    )
+
+    capturer = DummyCapturer()
+    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
+
+    assert dummy_module.router.capture_fn is not None
+    dummy_module.router.capture_fn(torch.tensor([[5, 6]]))
+
+    assert len(capturer.calls) == 1
+    layer_id, topk_ids = capturer.calls[0]
+    assert layer_id == 7
+    assert torch.equal(topk_ids, torch.tensor([[5, 6]]))
+
+
+def test_gpu_model_runner_binding_stage(monkeypatch):
+    from vllm.v1.worker import gpu_model_runner as gmr
+
+    class DummyFusedMoE:
+        def __init__(self):
+            self.layer_id = 11
+            self.router = _make_router()
+
+    class DummyCapturer:
+        def __init__(self):
+            self.calls = []
+
+        def capture(self, layer_id, topk_ids):
+            self.calls.append((layer_id, topk_ids))
+
+    dummy_module = DummyFusedMoE()
+
+    import vllm.model_executor.layers.fused_moe.layer as fused_moe_layer
+
+    monkeypatch.setattr(fused_moe_layer, "MoERunner", DummyFusedMoE)
+
+    dummy_self = types.SimpleNamespace(
+        compilation_config=types.SimpleNamespace(
+            static_forward_context={"dummy": dummy_module}
+        )
+    )
+
+    # Before binding, no capture hook.
+    assert dummy_module.router.capture_fn is None
+
+    capturer = DummyCapturer()
+    gmr.GPUModelRunner._bind_routed_experts_capturer(dummy_self, capturer)
+
+    # After binding, hook should exist and be callable.
+    assert callable(dummy_module.router.capture_fn)
+    dummy_module.router.capture_fn(torch.tensor([[9, 10]]))
+    assert len(capturer.calls) == 1
+
+
+def test_routed_experts_capturer_single_dp_no_metadata():
+    """dp_metadata is None: capture writes the full topk_ids rows."""
+    capturer = _capturer_with_buffer(dp_rank=0)
+    topk = torch.tensor([[1, 2], [3, 4], [5, 6]], dtype=torch.int32)
+    ctx = SimpleNamespace(dp_metadata=None)
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
+    assert capturer.device_buffer[3, 0, 0].item() == -1
+
+
+def test_routed_experts_capturer_dp_naive_concatenated_all_ranks():
+    """n == sum(num_tokens_dp): slice this rank's segment from concatenated topk."""
+    capturer = _capturer_with_buffer(dp_rank=1)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    # Concatenated order: rank0 rows then rank1 rows.
+    topk = torch.tensor(
+        [[0, 1], [2, 3], [10, 11], [12, 13], [14, 15]], dtype=torch.int32
+    )
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    want = topk[2:5]
+    assert torch.equal(capturer.device_buffer[:3, 0, :], want)
+
+
+def test_routed_experts_capturer_dp_modular_local_tokens():
+    """n == token_num_per_dp: topk is already local to this DP rank."""
+    capturer = _capturer_with_buffer(dp_rank=1)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    topk = torch.tensor([[10, 11], [12, 13], [14, 15]], dtype=torch.int32)
+    with patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    assert torch.equal(capturer.device_buffer[:3, 0, :], topk)
+
+
+def test_routed_experts_capturer_dp_unexpected_batch_raises():
+    """Mismatch between topk batch dim and DP layout: fail fast."""
+    capturer = _capturer_with_buffer(dp_rank=0)
+    num_tokens_dp = torch.tensor([2, 3], dtype=torch.int32)
+    ctx = SimpleNamespace(
+        dp_metadata=SimpleNamespace(num_tokens_across_dp_cpu=num_tokens_dp)
+    )
+    # total=5, local=2: n=1 matches neither naive (5) nor modular (2).
+    topk = torch.tensor([[1, 2]], dtype=torch.int32)
+    with (
+        patch(f"{_REC_MODULE}.get_forward_context", return_value=ctx),
+        pytest.raises(AssertionError, match="unexpected topk_ids batch dim"),
+    ):
+        capturer.capture(layer_id=0, topk_ids=topk)
+    assert capturer.device_buffer[0, 0, 0].item() == -1

tests/model_executor/test_weight_utils.py

@@ -0,0 +1,285 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import tempfile
+
+import huggingface_hub.constants
+import pytest
+from huggingface_hub.utils import LocalEntryNotFoundError
+
+from vllm.model_executor.model_loader.weight_utils import (
+    download_weights_from_hf,
+    maybe_remap_kv_scale_name,
+)
+
+
+def test_download_weights_from_hf():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        # assert LocalEntryNotFoundError error is thrown
+        # 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 the model
+        huggingface_hub.constants.HF_HUB_OFFLINE = False
+        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
+        )
+
+
+class TestMaybeRemapKvScaleName:
+    """Tests for maybe_remap_kv_scale_name covering all checkpoint formats."""
+
+    PARAMS_DICT = {
+        "model.layers.0.self_attn.attn.k_scale": None,
+        "model.layers.0.self_attn.attn.v_scale": None,
+        "model.layers.0.self_attn.attn.q_scale": None,
+        "model.layers.0.self_attn.qkv_proj.weight": None,
+    }
+
+    def test_qkv_proj_k_scale(self):
+        """Qwen3-MoE / llm-compressor format: qkv_proj.k_scale -> attn.k_scale
+        Regression test for https://github.com/vllm-project/vllm/issues/25047"""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.k_scale"
+
+    def test_qkv_proj_v_scale(self):
+        """Qwen3-MoE / llm-compressor format: qkv_proj.v_scale -> attn.v_scale
+        Regression test for https://github.com/vllm-project/vllm/issues/25047"""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.v_scale"
+
+    def test_modelopt_k_proj_k_scale(self):
+        """ModelOpt format: k_proj.k_scale -> attn.k_scale"""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.k_scale"
+
+    def test_modelopt_v_proj_v_scale(self):
+        """ModelOpt format: v_proj.v_scale -> attn.v_scale"""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.v_scale"
+
+    def test_deprecated_kv_scale(self):
+        """Old format: kv_scale -> attn.k_scale (deprecated)"""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.kv_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.k_scale"
+
+    def test_default_bare_k_scale(self):
+        """Default format: .k_scale -> .attn.k_scale"""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.k_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.k_scale"
+
+    def test_non_scale_name_unchanged(self):
+        """Non-scale names should be returned unchanged."""
+        name = "model.layers.0.self_attn.qkv_proj.weight"
+        result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
+        assert result == name
+
+    def test_nvfp4_modelopt_k_proj_k_scale(self):
+        """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4):
+        k_proj.k_scale -> attn.k_scale.
+        Validates that NVFP4 checkpoints are not broken by this change."""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.k_proj.k_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.k_scale"
+
+    def test_nvfp4_modelopt_v_proj_v_scale(self):
+        """ModelOpt NVFP4 format (e.g. nvidia/Qwen3-30B-A3B-NVFP4):
+        v_proj.v_scale -> attn.v_scale.
+        Validates that NVFP4 checkpoints are not broken by this change."""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.v_proj.v_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.v_scale"
+
+    def test_qwen3_vl_moe_qkv_proj_k_scale(self):
+        """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE.
+        Regression test for qwen3_vl_moe.py fix (same bug as #25047)."""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.qkv_proj.k_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.k_scale"
+
+    def test_qwen3_vl_moe_qkv_proj_v_scale(self):
+        """Qwen3-VL-MoE uses the same fused qkv_proj naming as Qwen3-MoE.
+        Regression test for qwen3_vl_moe.py fix (same bug as #25047)."""
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.qkv_proj.v_scale", self.PARAMS_DICT
+        )
+        assert result == "model.layers.0.self_attn.attn.v_scale"
+
+    def test_nvfp4_weight_scale_not_remapped(self):
+        """NVFP4 weight_scale should not be touched by remap (not a kv scale)."""
+        name = "model.layers.0.self_attn.k_proj.weight_scale"
+        result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
+        assert result == name
+
+    def test_nvfp4_input_scale_not_remapped(self):
+        """NVFP4 input_scale should not be touched by remap (not a kv scale)."""
+        name = "model.layers.0.self_attn.k_proj.input_scale"
+        result = maybe_remap_kv_scale_name(name, self.PARAMS_DICT)
+        assert result == name
+
+    def test_missing_target_returns_none(self):
+        """If remapped name not in params_dict, return None."""
+        empty_params: dict[str, None] = {}
+        result = maybe_remap_kv_scale_name(
+            "model.layers.0.self_attn.qkv_proj.k_scale", empty_params
+        )
+        assert result is None
+
+
+class TestKvCacheScaleMapper:
+    """The `WeightsMapper` returned by `get_cache_scale_mapper` replaces the
+    per-model `maybe_remap_kv_scale_name` calls. It must remap the same set of
+    checkpoint formats (the non-`params_dict`-dependent ones) and be idempotent
+    so it composes safely with a model's own qkv/gate_up `hf_to_vllm_mapper`."""
+
+    def _mapper(self):
+        # `get_cache_scale_mapper` does not use `self`; call it on the base
+        # class to get the default (non-config-specific) mapper.
+        from vllm.model_executor.layers.quantization.base_config import (
+            QuantizationConfig,
+        )
+
+        return QuantizationConfig.get_cache_scale_mapper()
+
+    def _map(self, name: str) -> str | None:
+        return self._mapper()._map_name(name)
+
+    @pytest.mark.parametrize(
+        "name,expected",
+        [
+            # Qwen3-MoE / llm-compressor fused qkv_proj
+            (
+                "model.layers.0.self_attn.qkv_proj.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            (
+                "model.layers.0.self_attn.qkv_proj.v_scale",
+                "model.layers.0.self_attn.attn.v_scale",
+            ),
+            # ModelOpt / NVFP4 k_proj/v_proj
+            (
+                "model.layers.0.self_attn.k_proj.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            (
+                "model.layers.0.self_attn.v_proj.v_scale",
+                "model.layers.0.self_attn.attn.v_scale",
+            ),
+            # deprecated fused kv_scale and bare scales
+            (
+                "model.layers.0.self_attn.kv_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            (
+                "model.layers.0.self_attn.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            # NemotronH mixer
+            (
+                "model.layers.0.mixer.k_proj.k_scale",
+                "model.layers.0.mixer.attn.k_scale",
+            ),
+            # already in vLLM form -> unchanged (idempotent)
+            (
+                "model.layers.0.self_attn.attn.k_scale",
+                "model.layers.0.self_attn.attn.k_scale",
+            ),
+            # non-kv scales must not be touched
+            (
+                "model.layers.0.self_attn.k_proj.weight_scale",
+                "model.layers.0.self_attn.k_proj.weight_scale",
+            ),
+            (
+                "model.layers.0.self_attn.k_proj.input_scale",
+                "model.layers.0.self_attn.k_proj.input_scale",
+            ),
+            # regular weights untouched
+            (
+                "model.layers.0.self_attn.q_proj.weight",
+                "model.layers.0.self_attn.q_proj.weight",
+            ),
+        ],
+    )
+    def test_remap(self, name, expected):
+        assert self._map(name) == expected
+
+    @pytest.mark.parametrize(
+        "name",
+        [
+            "model.layers.0.self_attn.k_scale",
+            "model.layers.0.self_attn.k_proj.k_scale",
+            "model.layers.0.self_attn.qkv_proj.v_scale",
+            "model.layers.0.mixer.k_proj.k_scale",
+        ],
+    )
+    def test_idempotent(self, name):
+        once = self._map(name)
+        assert once is not None
+        assert self._map(once) == once
+
+    def test_composes_with_qkv_mapper(self):
+        """Applied together with a model's qkv/gate_up mapper, the regex scale
+        rules run before the substr rename, so scales are normalized to `.attn.`
+        and regular projections are still fused correctly."""
+        from vllm.model_executor.models.utils import WeightsMapper
+
+        model_mapper = WeightsMapper(
+            orig_to_new_substr={
+                ".q_proj": ".qkv_proj.q",
+                ".k_proj": ".qkv_proj.k",
+                ".v_proj": ".qkv_proj.v",
+            }
+        )
+        # AutoWeightsLoader does `mapper |= cache_scale_mapper`
+        combined = model_mapper | self._mapper()
+
+        assert (
+            combined._map_name("model.layers.0.self_attn.q_proj.weight")
+            == "model.layers.0.self_attn.qkv_proj.q.weight"
+        )
+        assert (
+            combined._map_name("model.layers.0.self_attn.k_proj.k_scale")
+            == "model.layers.0.self_attn.attn.k_scale"
+        )
+        assert (
+            combined._map_name("model.layers.0.self_attn.k_scale")
+            == "model.layers.0.self_attn.attn.k_scale"
+        )
+
+
+if __name__ == "__main__":
+    test_download_weights_from_hf()