Add vLLM v0.18.1 source tree with KV transfer abort fix

third_party/vllm/ now tracked in git for direct patch management.
Based on vLLM v0.18.1 release with one patch applied:

  vllm/v1/core/sched/scheduler.py:
    Replace fatal assert with graceful skip when KV transfer callback
    arrives for an already-aborted request during PD disaggregated serving.

Future vLLM modifications should be made directly in third_party/vllm/
and committed normally. The patches/ directory is kept as documentation
of what changed from upstream.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

third_party/vllm/tests/plugins_tests/test_bge_m3_sparse_io_processor_plugins.py

@@ -0,0 +1,212 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+
+import pytest
+import requests
+
+# Test configuration for BGE-M3 sparse plugin
+from tests.utils import RemoteOpenAIServer
+from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
+
+model_config = {
+    "model_name": "BAAI/bge-m3",
+    "plugin": "bge_m3_sparse_plugin",
+    "test_input": "What is the capital of France?",
+    "hf_overrides": json.dumps(
+        {"architectures": ["BgeM3EmbeddingModel"], "head_dtype": "float16"}
+    ),
+}
+
+
+def _float_close(expected: object, result: object):
+    assert isinstance(expected, float) and isinstance(result, float), (
+        f"{expected=}  or {result=} is not float"
+    )
+    return (expected - result) < 1e-3 or abs(expected / result - 1) < 1e-3
+
+
+def _get_attr_or_val(obj: object | dict, key: str):
+    if isinstance(obj, dict) and key in obj:
+        return obj[key]
+    return getattr(obj, key, None)
+
+
+def _check_sparse_embedding(data, check_tokens=False):
+    expected_weights = [
+        {"token_id": 32, "weight": 0.0552978515625, "token": "?"},
+        {"token_id": 70, "weight": 0.09808349609375, "token": "the"},
+        {"token_id": 83, "weight": 0.08154296875, "token": "is"},
+        {"token_id": 111, "weight": 0.11810302734375, "token": "of"},
+        {"token_id": 4865, "weight": 0.1171875, "token": "What"},
+        {"token_id": 9942, "weight": 0.292236328125, "token": "France"},
+        {"token_id": 10323, "weight": 0.2802734375, "token": "capital"},
+    ]
+    expected_embed = {x["token_id"]: x for x in expected_weights}
+
+    assert len(data) == len(expected_embed)
+    for entry in data:
+        expected_val = expected_embed[_get_attr_or_val(entry, "token_id")]
+        assert _float_close(
+            expected_val["weight"], _get_attr_or_val(entry, "weight")
+        ), f"actual embed {entry} not equal to {expected_val}"
+        if check_tokens:
+            assert expected_val["token"] == _get_attr_or_val(entry, "token"), (
+                f"actual embed {entry} not equal to {expected_val}"
+            )
+        else:
+            assert _get_attr_or_val(entry, "token") is None, (
+                f"{entry} should not return token"
+            )
+
+
+@pytest.fixture(scope="function")
+def server():
+    args = [
+        "--runner",
+        "pooling",
+        "--enforce-eager",
+        "--max-num-seqs",
+        "32",
+        "--hf_overrides",
+        model_config["hf_overrides"],
+        "--io-processor-plugin",
+        model_config["plugin"],
+    ]
+
+    with RemoteOpenAIServer(model_config["model_name"], args) as remote_server:
+        yield remote_server
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "return_tokens",
+    [True, False],
+)
+async def test_bge_m3_sparse_plugin_online(
+    server: RemoteOpenAIServer, return_tokens: bool
+):
+    """Test BGE-M3 sparse plugin in online mode via API."""
+    request_payload = {
+        "model": model_config["model_name"],
+        "task": "token_classify",
+        "data": {"input": model_config["test_input"], "return_tokens": return_tokens},
+    }
+
+    ret = requests.post(
+        server.url_for("pooling"),
+        json=request_payload,
+    )
+
+    response = ret.json()
+
+    # Verify the request response is in the correct format
+    assert (parsed_response := IOProcessorResponse(**response).data)
+
+    # Verify the output is formatted as expected for this plugin
+    assert _get_attr_or_val(parsed_response, "data")
+    assert len(_get_attr_or_val(parsed_response, "data")) > 0
+
+    data_entry = _get_attr_or_val(parsed_response, "data")[0]
+    assert _get_attr_or_val(data_entry, "object") == "sparse-embedding"
+    assert _get_attr_or_val(data_entry, "sparse_embedding")
+
+    # Verify sparse embedding format
+    sparse_embedding = _get_attr_or_val(data_entry, "sparse_embedding")
+    assert isinstance(sparse_embedding, list)
+    _check_sparse_embedding(sparse_embedding, return_tokens)
+
+    # Verify usage information
+    usage = _get_attr_or_val(parsed_response, "usage")
+    assert usage, f"usage not found for {parsed_response}"
+    assert _get_attr_or_val(usage, "prompt_tokens") > 0
+    assert _get_attr_or_val(usage, "total_tokens") == _get_attr_or_val(
+        usage, "prompt_tokens"
+    )
+
+
+@pytest.mark.parametrize(
+    "return_tokens",
+    [True, False],
+)
+def test_bge_m3_sparse_plugin_offline(vllm_runner, return_tokens: bool):
+    """Test BGE-M3 sparse plugin in offline mode."""
+    prompt = {
+        "data": {
+            "input": model_config["test_input"],
+            "return_tokens": return_tokens,
+        }
+    }
+
+    with vllm_runner(
+        model_config["model_name"],
+        runner="pooling",
+        enforce_eager=True,
+        max_num_seqs=32,
+        io_processor_plugin=model_config["plugin"],
+        hf_overrides=json.loads(model_config["hf_overrides"]),
+        default_torch_num_threads=1,
+    ) as llm_runner:
+        llm = llm_runner.get_llm()
+        pooler_output = llm.encode(prompt, pooling_task="token_classify")
+
+    outputs = pooler_output[0]
+
+    # Verify output structure
+    assert hasattr(outputs, "outputs")
+    response = outputs.outputs
+    assert hasattr(response, "data")
+    assert len(response.data) == 1
+    # Verify response data
+    for i, output in enumerate(response.data):
+        # Each output should have sparse embeddings
+        sparse_embedding = output.sparse_embedding
+        assert isinstance(sparse_embedding, list)
+        _check_sparse_embedding(sparse_embedding, return_tokens)
+
+    # Verify usage
+    assert response.usage.prompt_tokens > 0
+    assert response.usage.total_tokens == response.usage.prompt_tokens
+
+
+def test_bge_m3_sparse_plugin_offline_multiple_inputs(vllm_runner):
+    """Test BGE-M3 sparse plugin with multiple inputs in offline mode."""
+    prompts = {
+        "data": {
+            "input": [
+                "What is the capital of France?",
+                "What is the capital of Germany?",
+                "What is the capital of Spain?",
+            ],
+            "return_tokens": True,
+        }
+    }
+
+    with vllm_runner(
+        model_config["model_name"],
+        runner="pooling",
+        enforce_eager=True,
+        max_num_seqs=32,
+        io_processor_plugin=model_config["plugin"],
+        hf_overrides=json.loads(model_config["hf_overrides"]),
+        default_torch_num_threads=1,
+    ) as llm_runner:
+        llm = llm_runner.get_llm()
+        pooler_output = llm.encode(prompts, pooling_task="token_classify")
+
+    outputs = pooler_output[0]
+
+    # Verify output structure
+    assert hasattr(outputs, "outputs")
+    response = outputs.outputs
+    assert hasattr(response, "data")
+    assert len(response.data) == 3
+    for i, output in enumerate(response.data):
+        # Each output should have sparse embeddings
+        sparse_embedding = output.sparse_embedding
+        assert isinstance(sparse_embedding, list)
+
+    # Verify usage
+    assert response.usage.prompt_tokens > 0
+    assert response.usage.total_tokens == response.usage.prompt_tokens

third_party/vllm/tests/plugins_tests/test_io_processor_plugins.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections.abc import Sequence
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from vllm.config import VllmConfig
+from vllm.inputs.data import PromptType
+from vllm.outputs import PoolingRequestOutput
+from vllm.plugins.io_processors import get_io_processor
+from vllm.plugins.io_processors.interface import IOProcessor
+from vllm.renderers import BaseRenderer
+
+
+class DummyIOProcessor(IOProcessor):
+    """Minimal IOProcessor used as the target of the mocked plugin entry point."""
+
+    def pre_process(
+        self,
+        prompt: object,
+        request_id: str | None = None,
+        **kwargs,
+    ) -> PromptType | Sequence[PromptType]:
+        raise NotImplementedError
+
+    def post_process(
+        self,
+        model_output: Sequence[PoolingRequestOutput],
+        request_id: str | None = None,
+        **kwargs,
+    ) -> object:
+        raise NotImplementedError
+
+
+@pytest.fixture
+def my_plugin_entry_points():
+    """Patch importlib.metadata.entry_points to expose a single 'my_plugin'
+    entry point backed by DummyIOProcessor, exercising the full plugin-loading
+    code path: entry_points → plugin.load() → func() →
+    resolve_obj_by_qualname → IOProcessor.__init__."""
+    qualname = f"{DummyIOProcessor.__module__}.{DummyIOProcessor.__qualname__}"
+    ep = MagicMock()
+    ep.name = "my_plugin"
+    ep.value = qualname
+    ep.load.return_value = lambda: qualname
+    with patch("importlib.metadata.entry_points", return_value=[ep]):
+        yield
+
+
+def test_loading_missing_plugin():
+    vllm_config = VllmConfig()
+    renderer = MagicMock(spec=BaseRenderer)
+    with pytest.raises(ValueError):
+        get_io_processor(
+            vllm_config, renderer=renderer, plugin_from_init="wrong_plugin"
+        )
+
+
+def test_loading_plugin(my_plugin_entry_points):
+    # Plugin name supplied via plugin_from_init.
+    vllm_config = MagicMock(spec=VllmConfig)
+    renderer = MagicMock(spec=BaseRenderer)
+
+    result = get_io_processor(
+        vllm_config, renderer=renderer, plugin_from_init="my_plugin"
+    )
+
+    assert isinstance(result, DummyIOProcessor)
+
+
+def test_loading_missing_plugin_from_model_config():
+    # Build a mock VllmConfig whose hf_config advertises a plugin name,
+    # exercising the model-config code path without loading a real model.
+    mock_hf_config = MagicMock()
+    mock_hf_config.to_dict.return_value = {"io_processor_plugin": "wrong_plugin"}
+
+    vllm_config = MagicMock(spec=VllmConfig)
+    vllm_config.model_config.hf_config = mock_hf_config
+
+    renderer = MagicMock(spec=BaseRenderer)
+    with pytest.raises(ValueError):
+        get_io_processor(vllm_config, renderer=renderer)
+
+
+def test_loading_plugin_from_model_config(my_plugin_entry_points):
+    # Plugin name supplied via the model's hf_config.
+    mock_hf_config = MagicMock()
+    mock_hf_config.to_dict.return_value = {"io_processor_plugin": "my_plugin"}
+
+    vllm_config = MagicMock(spec=VllmConfig)
+    vllm_config.model_config.hf_config = mock_hf_config
+
+    renderer = MagicMock(spec=BaseRenderer)
+
+    result = get_io_processor(vllm_config, renderer=renderer)
+
+    assert isinstance(result, DummyIOProcessor)

third_party/vllm/tests/plugins_tests/test_platform_plugins.py

@@ -0,0 +1,47 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+import torch
+
+from vllm.plugins import load_general_plugins
+
+
+def test_platform_plugins():
+    # simulate workload by running an example
+    import runpy
+
+    current_file = __file__
+    import os
+
+    example_file = os.path.join(
+        os.path.dirname(os.path.dirname(os.path.dirname(current_file))),
+        "examples",
+        "basic/offline_inference/basic.py",
+    )
+    runpy.run_path(example_file)
+
+    # check if the plugin is loaded correctly
+    from vllm.platforms import _init_trace, current_platform
+
+    assert current_platform.device_name == "DummyDevice", (
+        f"Expected DummyDevice, got {current_platform.device_name}, "
+        "possibly because current_platform is imported before the plugin"
+        f" is loaded. The first import:\n{_init_trace}"
+    )
+
+
+def test_oot_custom_op(default_vllm_config, monkeypatch: pytest.MonkeyPatch):
+    # simulate workload by running an example
+    load_general_plugins()
+    from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
+
+    layer = RotaryEmbedding(16, 16, 16, 16, True, torch.float16)
+    assert layer.__class__.__name__ == "DummyRotaryEmbedding", (
+        f"Expected DummyRotaryEmbedding, got {layer.__class__.__name__}, "
+        "possibly because the custom op is not registered correctly."
+    )
+    assert hasattr(layer, "addition_config"), (
+        "Expected DummyRotaryEmbedding to have an 'addition_config' attribute, "
+        "which is set by the custom op."
+    )

third_party/vllm/tests/plugins_tests/test_scheduler_plugins.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+
+from vllm.engine.arg_utils import EngineArgs
+from vllm.sampling_params import SamplingParams
+from vllm.v1.core.sched.scheduler import Scheduler
+from vllm.v1.engine.llm_engine import LLMEngine
+
+
+class DummyV1Scheduler(Scheduler):
+    def schedule(self):
+        raise Exception("Exception raised by DummyV1Scheduler")
+
+
+def test_scheduler_plugins_v1(monkeypatch: pytest.MonkeyPatch):
+    with monkeypatch.context() as m:
+        # Explicitly turn off engine multiprocessing so
+        # that the scheduler runs in this process
+        m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+        with pytest.raises(Exception) as exception_info:
+            engine_args = EngineArgs(
+                model="facebook/opt-125m",
+                enforce_eager=True,  # reduce test time
+                scheduler_cls=DummyV1Scheduler,
+            )
+
+            engine = LLMEngine.from_engine_args(engine_args=engine_args)
+
+            sampling_params = SamplingParams(max_tokens=1)
+            engine.add_request("0", "foo", sampling_params)
+            engine.step()
+
+        assert str(exception_info.value) == "Exception raised by DummyV1Scheduler"

third_party/vllm/tests/plugins_tests/test_stats_logger_plugins.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from dummy_stat_logger.dummy_stat_logger import DummyStatLogger
+
+from vllm.config import VllmConfig
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.metrics.loggers import load_stat_logger_plugin_factories
+
+
+def test_stat_logger_plugin_is_discovered(monkeypatch: pytest.MonkeyPatch):
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
+
+        factories = load_stat_logger_plugin_factories()
+        assert len(factories) == 1, f"Expected 1 factory, got {len(factories)}"
+        assert factories[0] is DummyStatLogger, (
+            f"Expected DummyStatLogger class, got {factories[0]}"
+        )
+
+        # instantiate and confirm the right type
+        vllm_config = VllmConfig()
+        instance = factories[0](vllm_config)
+        assert isinstance(instance, DummyStatLogger)
+
+
+def test_no_plugins_loaded_if_env_empty(monkeypatch: pytest.MonkeyPatch):
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_PLUGINS", "")
+
+        factories = load_stat_logger_plugin_factories()
+        assert factories == []
+
+
+def test_invalid_stat_logger_plugin_raises(monkeypatch: pytest.MonkeyPatch):
+    def fake_plugin_loader(group: str):
+        assert group == "vllm.stat_logger_plugins"
+        return {"bad": object()}
+
+    with monkeypatch.context() as m:
+        m.setattr(
+            "vllm.v1.metrics.loggers.load_plugins_by_group",
+            fake_plugin_loader,
+        )
+        with pytest.raises(
+            TypeError,
+            match="Stat logger plugin 'bad' must be a subclass of StatLoggerBase",
+        ):
+            load_stat_logger_plugin_factories()
+
+
+@pytest.mark.asyncio
+async def test_stat_logger_plugin_integration_with_engine(
+    monkeypatch: pytest.MonkeyPatch,
+):
+    with monkeypatch.context() as m:
+        m.setenv("VLLM_PLUGINS", "dummy_stat_logger")
+
+        engine_args = AsyncEngineArgs(
+            model="facebook/opt-125m",
+            enforce_eager=True,  # reduce test time
+            disable_log_stats=True,  # disable default loggers
+        )
+
+        engine = AsyncLLM.from_engine_args(engine_args=engine_args)
+
+        assert len(engine.logger_manager.stat_loggers) == 2
+        assert len(engine.logger_manager.stat_loggers[0].per_engine_stat_loggers) == 1
+        assert isinstance(
+            engine.logger_manager.stat_loggers[0].per_engine_stat_loggers[0],
+            DummyStatLogger,
+        )
+
+        engine.shutdown()

third_party/vllm/tests/plugins_tests/test_terratorch_io_processor_plugins.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import base64
+import io
+
+import imagehash
+import pytest
+import requests
+from PIL import Image
+
+from tests.utils import RemoteOpenAIServer
+from vllm.entrypoints.pooling.pooling.protocol import IOProcessorResponse
+
+models_config = {
+    "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-TL-Sen1Floods11": {
+        "image_url": "https://huggingface.co/christian-pinto/Prithvi-EO-2.0-300M-TL-VLLM/resolve/main/valencia_example_2024-10-26.tiff",  # noqa: E501
+        "out_hash": "aa6d92ad25926a5e",
+        "plugin": "prithvi_to_tiff",
+    },
+    "ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars": {
+        "image_url": "https://huggingface.co/ibm-nasa-geospatial/Prithvi-EO-2.0-300M-BurnScars/resolve/main/examples/subsetted_512x512_HLS.S30.T10SEH.2018190.v1.4_merged.tif",  # noqa: E501
+        "out_hash": "c07f4f602da73552",
+        "plugin": "prithvi_to_tiff",
+    },
+}
+
+
+def _compute_image_hash(base64_data: str) -> str:
+    # Decode the base64 output and create image from byte stream
+    decoded_image = base64.b64decode(base64_data)
+    image = Image.open(io.BytesIO(decoded_image))
+
+    # Compute perceptual hash of the output image
+    return str(imagehash.phash(image))
+
+
+@pytest.fixture(scope="function")
+def server(model_name, plugin):
+    args = [
+        "--runner",
+        "pooling",
+        "--enforce-eager",
+        "--skip-tokenizer-init",
+        # Limit the maximum number of parallel requests
+        # to avoid the model going OOM in CI.
+        "--max-num-seqs",
+        "32",
+        "--io-processor-plugin",
+        plugin,
+        "--enable-mm-embeds",
+    ]
+
+    with RemoteOpenAIServer(model_name, args) as remote_server:
+        yield remote_server
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "model_name, image_url, plugin, expected_hash",
+    [
+        (model_name, config["image_url"], config["plugin"], config["out_hash"])
+        for model_name, config in models_config.items()
+    ],
+)
+async def test_prithvi_mae_plugin_online(
+    server: RemoteOpenAIServer,
+    model_name: str,
+    image_url: str | dict,
+    plugin: str,
+    expected_hash: str,
+):
+    request_payload_url = {
+        "data": {
+            "data": image_url,
+            "data_format": "url",
+            "image_format": "tiff",
+            "out_data_format": "b64_json",
+        },
+        "priority": 0,
+        "model": model_name,
+        "softmax": False,
+    }
+
+    ret = requests.post(
+        server.url_for("pooling"),
+        json=request_payload_url,
+    )
+
+    response = ret.json()
+
+    # verify the request response is in the correct format
+    assert (parsed_response := IOProcessorResponse(**response))
+
+    # verify the output is formatted as expected for this plugin
+    plugin_data = parsed_response.data
+    assert all(plugin_data.get(attr) for attr in ["type", "format", "data"])
+
+    # Compute the output image hash and compare it against the expected hash
+    image_hash = _compute_image_hash(plugin_data["data"])
+    assert image_hash == expected_hash, (
+        f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
+    )
+
+
+@pytest.mark.parametrize(
+    "model_name, image_url, plugin, expected_hash",
+    [
+        (model_name, config["image_url"], config["plugin"], config["out_hash"])
+        for model_name, config in models_config.items()
+    ],
+)
+def test_prithvi_mae_plugin_offline(
+    vllm_runner, model_name: str, image_url: str | dict, plugin: str, expected_hash: str
+):
+    img_data = dict(
+        data=image_url,
+        data_format="url",
+        image_format="tiff",
+        out_data_format="b64_json",
+    )
+
+    prompt = dict(data=img_data)
+
+    with vllm_runner(
+        model_name,
+        runner="pooling",
+        skip_tokenizer_init=True,
+        enable_mm_embeds=True,
+        enforce_eager=True,
+        # Limit the maximum number of parallel requests
+        # to avoid the model going OOM in CI.
+        max_num_seqs=32,
+        io_processor_plugin=plugin,
+        default_torch_num_threads=1,
+    ) as llm_runner:
+        pooler_output = llm_runner.get_llm().encode(prompt, pooling_task="plugin")
+
+    output = pooler_output[0].outputs
+
+    # verify the output is formatted as expected for this plugin
+    assert all(hasattr(output, attr) for attr in ["type", "format", "data"])
+
+    # Compute the output image hash and compare it against the expected hash
+    image_hash = _compute_image_hash(output.data)
+    assert image_hash == expected_hash, (
+        f"Image hash mismatch: expected {expected_hash}, got {image_hash}"
+    )