fix test_spec_decode_acceptance_length for mrv2

Signed-off-by: yewentao256 <zhyanwentao@126.com>

.buildkite/scripts/coverage/aggregate-coverage.py

@@ -1,208 +0,0 @@
-#!/usr/bin/env python3
-"""Aggregate per-step coverage JSON files into a test-selection mapping.
-
-Downloads all coverage_*.json artifacts from the current Buildkite build,
-then produces two output files:
-
-1. coverage_map.json — inverted index: {source_file: [step_keys]}
-   Used by the pipeline generator to determine which steps to trigger.
-
-2. step_coverage.json — forward index: {step_key: [source_files]}
-   Useful for debugging and understanding test coverage.
-
-Usage:
-    # Run as a Buildkite step at the end of nightly CI
-    python3 .buildkite/scripts/coverage/aggregate-coverage.py
-
-    # Or locally with downloaded artifacts
-    python3 .buildkite/scripts/coverage/aggregate-coverage.py --local-dir ./artifacts/
-"""
-
-import argparse
-import json
-import os
-import subprocess
-import sys
-import tempfile
-from collections import defaultdict
-from pathlib import Path
-
-
-def download_artifacts(dest_dir: str) -> list[str]:
-    """Download all coverage_*.json artifacts from the current build."""
-    try:
-        subprocess.run(
-            ["buildkite-agent", "artifact", "download", "coverage_*.json", dest_dir],
-            check=True,
-            capture_output=True,
-            text=True,
-        )
-    except FileNotFoundError:
-        print("buildkite-agent not found, skipping download", file=sys.stderr)
-        return []
-    except subprocess.CalledProcessError as e:
-        print(f"Artifact download failed: {e.stderr}", file=sys.stderr)
-        return []
-
-    return list(Path(dest_dir).glob("coverage_*.json"))
-
-
-def load_coverage_files(files: list[Path]) -> dict[str, list[str]]:
-    """Load coverage JSON files and extract source files per step.
-
-    Returns: {step_key: [source_files]}
-    """
-    step_coverage = {}
-
-    for filepath in files:
-        filename = filepath.name
-        # coverage_<step_key>.json -> step_key
-        step_key = filename.removeprefix("coverage_").removesuffix(".json")
-
-        try:
-            with open(filepath) as f:
-                data = json.load(f)
-        except (json.JSONDecodeError, OSError) as e:
-            print(f"Warning: skipping {filename}: {e}", file=sys.stderr)
-            continue
-
-        source_files = []
-        for fpath, fdata in data.get("files", {}).items():
-            # Skip files with zero executed lines — coverage.py reports
-            # all files in the source tree, not just those actually run.
-            # Supports both full format (summary.covered_lines) and
-            # stripped format (covered_lines directly).
-            covered = fdata.get("covered_lines") or fdata.get("summary", {}).get("covered_lines", 0)
-            if covered == 0:
-                continue
-            # If function-level data is available, skip import-only files
-            # (files where only module-level code ran but no named functions
-            # were actually called).
-            funcs_called = fdata.get("functions_called")
-            if funcs_called is not None and funcs_called == 0:
-                continue
-            # Normalize paths to be relative to the vllm package root.
-            # coverage.py may report absolute paths or paths relative to
-            # the installed package location. We only care about files
-            # under the vllm/ directory.
-            normalized = _normalize_path(fpath)
-            if normalized:
-                source_files.append(normalized)
-
-        if source_files:
-            step_coverage[step_key] = sorted(set(source_files))
-            print(f"  {step_key}: {len(source_files)} source files")
-
-    return step_coverage
-
-
-def _normalize_path(path: str) -> str | None:
-    """Normalize a coverage path to a vllm-relative path.
-
-    Returns None for paths outside the vllm package (tests, third-party, etc).
-    """
-    # Strip common prefixes from installed package paths
-    markers = ["/site-packages/", "/dist-packages/", "/vllm-workspace/src/"]
-    for marker in markers:
-        idx = path.find(marker)
-        if idx != -1:
-            path = path[idx + len(marker):]
-            break
-
-    # Also handle paths that are already relative
-    if path.startswith("vllm/"):
-        return path
-
-    # Handle absolute paths that contain /vllm/
-    idx = path.find("/vllm/")
-    if idx != -1:
-        return path[idx + 1:]
-
-    return None
-
-
-def build_inverted_index(
-    step_coverage: dict[str, list[str]],
-) -> dict[str, list[str]]:
-    """Build {source_file: [step_keys]} from {step_key: [source_files]}."""
-    inverted = defaultdict(list)
-    for step_key, source_files in step_coverage.items():
-        for src_file in source_files:
-            inverted[src_file].append(step_key)
-
-    # Sort step lists for deterministic output
-    return {k: sorted(v) for k, v in sorted(inverted.items())}
-
-
-def main():
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument(
-        "--local-dir",
-        help="Directory containing coverage_*.json files (skip artifact download)",
-    )
-    parser.add_argument(
-        "--output-dir",
-        default=".",
-        help="Directory to write output files (default: cwd)",
-    )
-    args = parser.parse_args()
-
-    if args.local_dir:
-        artifact_dir = args.local_dir
-        files = list(Path(artifact_dir).glob("coverage_*.json"))
-    else:
-        artifact_dir = tempfile.mkdtemp(prefix="coverage_artifacts_")
-        files = download_artifacts(artifact_dir)
-
-    if not files:
-        print("No coverage files found. Nothing to aggregate.")
-        sys.exit(0)
-
-    print(f"Found {len(files)} coverage files:")
-
-    # Build the forward index: step -> source files
-    step_coverage = load_coverage_files(files)
-
-    if not step_coverage:
-        print("No valid coverage data found.")
-        sys.exit(0)
-
-    # Build the inverted index: source file -> steps
-    coverage_map = build_inverted_index(step_coverage)
-
-    # Write outputs
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    step_coverage_path = output_dir / "step_coverage.json"
-    with open(step_coverage_path, "w") as f:
-        json.dump(step_coverage, f, indent=2)
-    print(f"\nWrote {step_coverage_path} ({len(step_coverage)} steps)")
-
-    coverage_map_path = output_dir / "coverage_map.json"
-    with open(coverage_map_path, "w") as f:
-        json.dump(coverage_map, f, indent=2)
-    print(f"Wrote {coverage_map_path} ({len(coverage_map)} source files)")
-
-    # Summary stats
-    total_files = len(coverage_map)
-    total_mappings = sum(len(v) for v in coverage_map.values())
-    print(f"\nSummary: {total_files} source files mapped to "
-          f"{len(step_coverage)} steps ({total_mappings} total mappings)")
-
-    # Upload aggregated files as artifacts
-    for output_file in [step_coverage_path, coverage_map_path]:
-        try:
-            subprocess.run(
-                ["buildkite-agent", "artifact", "upload", str(output_file)],
-                check=True,
-                capture_output=True,
-                text=True,
-            )
-            print(f"Uploaded {output_file}")
-        except (FileNotFoundError, subprocess.CalledProcessError):
-            pass  # Not in Buildkite or upload failed — that's fine for local runs
-
-
-if __name__ == "__main__":
-    main()

.buildkite/scripts/coverage/upload-step-coverage.sh

@@ -1,42 +0,0 @@
-#!/bin/bash
-# Upload coverage data for the current Buildkite step.
-# Called automatically at the end of each step when COLLECT_COVERAGE=1.
-#
-# Expects:
-#   - .coverage.${BUILDKITE_STEP_KEY} data file from coverage run --append
-#   - BUILDKITE_STEP_KEY, BUILDKITE_BUILD_NUMBER env vars
-#
-# Produces:
-#   - coverage_${BUILDKITE_STEP_KEY}.json uploaded as a Buildkite artifact
-
-set -euo pipefail
-
-STEP_KEY="${BUILDKITE_STEP_KEY:-unknown}"
-DATA_FILE=".coverage.${STEP_KEY}"
-OUTPUT_JSON="coverage_${STEP_KEY}.json"
-
-if [ ! -f "$DATA_FILE" ]; then
-    echo "~~~ No coverage data file found ($DATA_FILE), skipping upload"
-    exit 0
-fi
-
-echo "~~~ :bar_chart: Exporting coverage data for step: ${STEP_KEY}"
-
-coverage json \
-    --data-file="$DATA_FILE" \
-    -o "$OUTPUT_JSON" \
-    --omit='*/tests/*,*/test_*,*/__pycache__/*' \
-    2>&1 || {
-        echo "Warning: coverage json export failed, skipping"
-        exit 0
-    }
-
-FILE_COUNT=$(python3 -c "import json; d=json.load(open('$OUTPUT_JSON')); print(len(d.get('files', {})))" 2>/dev/null || echo "?")
-echo "Coverage captured ${FILE_COUNT} source files for step ${STEP_KEY}"
-
-buildkite-agent artifact upload "$OUTPUT_JSON" 2>&1 || {
-    echo "Warning: artifact upload failed"
-    exit 0
-}
-
-echo "Uploaded $OUTPUT_JSON"

docs/benchmarking/cli.md

@@ -231,13 +231,21 @@ vllm bench serve \
 
 #### Custom Image Dataset
 
-If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and needs to have "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
+If the image dataset you want to benchmark is not supported yet in vLLM, then you can benchmark on it using `CustomImageDataset`. At inference time, use the option `--dataset-name custom_image`. Your data needs to be in the `.jsonl` format and can use "prompt" and "image_files" fields per entry, e.g., `image_data.jsonl`:
 
 ```json
 {"prompt": "How many animals are present in the given image?", "image_files": ["/path/to/image/folder/horsepony.jpg"]}
 {"prompt": "What colour is the bird shown in the image?", "image_files": ["/path/to/image/folder/flycatcher.jpeg"]}
 ```
 
+Every image listed in "image_files" is added to the request in the listed order after the prompt text. To preserve an interleaved order of text and images, use a "content" field with OpenAI-compatible content parts:
+
+```json
+{"content": [{"type": "text", "text": "Compare "}, {"type": "image", "image": "/path/to/image/folder/chart_a.png"}, {"type": "text", "text": " with "}, {"type": "image_url", "image_url": {"url": "/path/to/image/folder/chart_b.png"}}]}
+```
+
+The "image" shorthand accepts the same values as "image_files". The "image_url" field accepts either an OpenAI-style object with a "url" field or a URL string.
+
 ```bash
 # need a model with vision capability here
 vllm serve Qwen/Qwen2-VL-7B-Instruct

tests/benchmarks/test_custom_image_dataset.py

@@ -0,0 +1,246 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+from argparse import Namespace
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from vllm.benchmarks.datasets import CustomImageDataset, get_samples
+from vllm.benchmarks.lib.endpoint_request_func import (
+    RequestFuncInput,
+    _get_chat_content,
+    _get_chat_messages,
+)
+
+pytestmark = pytest.mark.skip_global_cleanup
+
+
+class _TokenizedPrompt:
+    def __init__(self, prompt: str) -> None:
+        self.input_ids = prompt.split()
+
+
+class _Tokenizer:
+    def __call__(self, prompt: str) -> _TokenizedPrompt:
+        return _TokenizedPrompt(prompt)
+
+
+def _write_jsonl(path: Path, rows: list[dict[str, Any]]) -> None:
+    with path.open("w") as f:
+        for row in rows:
+            f.write(json.dumps(row) + "\n")
+
+
+def _args_for_custom_image(dataset_path: Path) -> Namespace:
+    return Namespace(
+        dataset_name="custom_image",
+        dataset_path=str(dataset_path),
+        disable_shuffle=True,
+        seed=0,
+        num_prompts=2,
+        custom_output_len=32,
+        enable_multimodal_chat=False,
+        request_id_prefix="req-",
+        no_oversample=False,
+    )
+
+
+@pytest.mark.benchmark
+def test_get_samples_custom_image_cli_path_supports_multi_image_and_content(
+    tmp_path: Path,
+) -> None:
+    image_a = tmp_path / "chart_a.png"
+    image_b = tmp_path / "chart_b.png"
+    image_c = tmp_path / "chart_c.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "prompt": "Compare the first two charts.",
+                "image_files": [str(image_a), str(image_b)],
+            },
+            {
+                "content": [
+                    {"type": "text", "text": "Now compare "},
+                    {"type": "image", "image": str(image_c)},
+                ],
+            },
+        ],
+    )
+
+    samples = get_samples(_args_for_custom_image(jsonl), _Tokenizer())
+
+    assert len(samples) == 2
+    assert samples[0].request_id == "req-0"
+    assert isinstance(samples[0].multi_modal_data, list)
+    assert [part["image_url"]["url"] for part in samples[0].multi_modal_data] == [
+        f"file://{image_a}",
+        f"file://{image_b}",
+    ]
+
+    assert samples[1].request_id == "req-1"
+    assert samples[1].multi_modal_data is None
+    assert isinstance(samples[1].prompt, list)
+    assert samples[1].prompt[0] == {"type": "text", "text": "Now compare "}
+    assert samples[1].prompt[1]["image_url"]["url"] == f"file://{image_c}"
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_uses_all_image_files(tmp_path: Path) -> None:
+    image_a = tmp_path / "chart_a.png"
+    image_b = tmp_path / "chart_b.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "prompt": "Compare the charts.",
+                "image_files": [str(image_a), str(image_b)],
+            }
+        ],
+    )
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+    )
+
+    assert len(samples) == 1
+    sample = samples[0]
+    assert sample.prompt == "Compare the charts."
+    assert sample.prompt_len == 3
+    assert isinstance(sample.multi_modal_data, list)
+    assert [part["image_url"]["url"] for part in sample.multi_modal_data] == [
+        f"file://{image_a}",
+        f"file://{image_b}",
+    ]
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_preserves_interleaved_content_order(
+    tmp_path: Path,
+) -> None:
+    image_a = tmp_path / "chart_a.png"
+    image_b = tmp_path / "chart_b.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "content": [
+                    {"type": "text", "text": "Compare "},
+                    {"type": "image", "image": str(image_a)},
+                    {"type": "text", "text": " with "},
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": str(image_b),
+                            "detail": "low",
+                        },
+                    },
+                ],
+            }
+        ],
+    )
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+    )
+
+    assert len(samples) == 1
+    sample = samples[0]
+    assert sample.multi_modal_data is None
+    assert sample.prompt_len == 2
+    assert isinstance(sample.prompt, list)
+    assert [part["type"] for part in sample.prompt] == [
+        "text",
+        "image_url",
+        "text",
+        "image_url",
+    ]
+    assert sample.prompt[1]["image_url"]["url"] == f"file://{image_a}"
+    assert sample.prompt[3]["image_url"] == {
+        "url": f"file://{image_b}",
+        "detail": "low",
+    }
+
+    request_input = RequestFuncInput(
+        prompt=sample.prompt,
+        api_url="http://localhost:8000/v1/chat/completions",
+        prompt_len=sample.prompt_len,
+        output_len=32,
+        model="test-model",
+    )
+    assert _get_chat_content(request_input) == sample.prompt
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_wraps_interleaved_content_for_multimodal_chat(
+    tmp_path: Path,
+) -> None:
+    image = tmp_path / "chart.png"
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(
+        jsonl,
+        [
+            {
+                "content": [
+                    {"type": "text", "text": "Describe "},
+                    {"type": "image", "image": str(image)},
+                ],
+            }
+        ],
+    )
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    samples = dataset.sample(
+        tokenizer=_Tokenizer(),
+        num_requests=1,
+        output_len=32,
+        enable_multimodal_chat=True,
+    )
+
+    sample = samples[0]
+    assert sample.multi_modal_data is None
+    assert sample.prompt == [
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe "},
+                {"type": "image_url", "image_url": {"url": f"file://{image}"}},
+            ],
+        }
+    ]
+
+    request_input = RequestFuncInput(
+        prompt=sample.prompt,
+        api_url="http://localhost:8000/v1/chat/completions",
+        prompt_len=sample.prompt_len,
+        output_len=32,
+        model="test-model",
+    )
+    assert _get_chat_messages(request_input) == sample.prompt
+
+
+@pytest.mark.benchmark
+def test_custom_image_dataset_rejects_invalid_content_part(
+    tmp_path: Path,
+) -> None:
+    jsonl = tmp_path / "images.jsonl"
+    _write_jsonl(jsonl, [{"content": [{"type": "audio", "audio": "clip.wav"}]}])
+
+    dataset = CustomImageDataset(dataset_path=str(jsonl), disable_shuffle=True)
+    with pytest.raises(ValueError, match="type 'text', 'image', or 'image_url'"):
+        dataset.sample(
+            tokenizer=_Tokenizer(),
+            num_requests=1,
+            output_len=32,
+        )

tests/entrypoints/openai/chat_completion/test_thinking_token_budget_validation.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import pytest
+from pydantic import ValidationError
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
+from vllm.entrypoints.openai.completion.protocol import CompletionRequest
+
+
+@pytest.mark.parametrize("raw_value", [-2, 0.6, 10.5])
+def test_chat_completion_request_rejects_invalid_thinking_token_budget(raw_value):
+    with pytest.raises(ValidationError, match="thinking_token_budget"):
+        ChatCompletionRequest.model_validate(
+            {
+                "model": "qwen",
+                "messages": [{"role": "user", "content": "hello"}],
+                "thinking_token_budget": raw_value,
+            }
+        )
+
+
+def test_chat_completion_request_accepts_valid_thinking_token_budget():
+    request = ChatCompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking_token_budget": 10,
+        }
+    )
+    assert request.thinking_token_budget == 10
+
+
+def test_chat_completion_request_accepts_minus_one_as_unlimited():
+    request = ChatCompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "messages": [{"role": "user", "content": "hello"}],
+            "thinking_token_budget": -1,
+        }
+    )
+    assert request.thinking_token_budget is None
+
+
+@pytest.mark.parametrize("raw_value", [0.6, 3.14, -2])
+def test_completion_request_rejects_invalid_thinking_token_budget(raw_value):
+    with pytest.raises(ValidationError, match="thinking_token_budget"):
+        CompletionRequest.model_validate(
+            {
+                "model": "qwen",
+                "prompt": "hello",
+                "thinking_token_budget": raw_value,
+            }
+        )
+
+
+def test_completion_request_accepts_valid_thinking_token_budget():
+    request = CompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "prompt": "hello",
+            "thinking_token_budget": 5,
+        }
+    )
+    assert request.thinking_token_budget == 5
+
+
+def test_completion_request_accepts_minus_one_as_unlimited():
+    request = CompletionRequest.model_validate(
+        {
+            "model": "qwen",
+            "prompt": "hello",
+            "thinking_token_budget": -1,
+        }
+    )
+    assert request.thinking_token_budget is None

tests/v1/kv_connector/unit/test_hma_auto_config.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Regression tests for HMA auto-disable with KV transfer connectors."""
+
+import pytest
+
+from vllm.config import DeviceConfig, KVTransferConfig, SchedulerConfig, VllmConfig
+from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
+from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorRole
+from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import KVCacheConfig
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.fixture(autouse=True)
+def mock_hybrid_kv_cache_supported(monkeypatch):
+    monkeypatch.setattr(current_platform, "support_hybrid_kv_cache", lambda: True)
+
+
+@pytest.mark.parametrize(
+    "kv_transfer_config,expect_disabled",
+    [
+        (  # HMA-supporting connector → HMA stays enabled
+            KVTransferConfig(
+                kv_connector="SimpleCPUOffloadConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={"cpu_bytes_to_use": 1 << 30},
+            ),
+            False,
+        ),
+        (  # Non-HMA connector → HMA is auto-disabled
+            KVTransferConfig(kv_connector="ExampleConnector", kv_role="kv_both"),
+            True,
+        ),
+        (  # MultiConnector: all HMA children → HMA stays enabled
+            KVTransferConfig(
+                kv_connector="MultiConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={
+                    "connectors": [
+                        {
+                            "kv_connector": "SimpleCPUOffloadConnector",
+                            "kv_role": "kv_both",
+                            "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                        },
+                        {
+                            "kv_connector": "OffloadingConnector",
+                            "kv_role": "kv_both",
+                            "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                        },
+                    ]
+                },
+            ),
+            False,
+        ),
+        (  # MultiConnector: mixed children → HMA is auto-disabled
+            KVTransferConfig(
+                kv_connector="MultiConnector",
+                kv_role="kv_both",
+                kv_connector_extra_config={
+                    "connectors": [
+                        {
+                            "kv_connector": "SimpleCPUOffloadConnector",
+                            "kv_role": "kv_both",
+                            "kv_connector_extra_config": {"cpu_bytes_to_use": 1 << 30},
+                        },
+                        {"kv_connector": "ExampleConnector", "kv_role": "kv_both"},
+                    ]
+                },
+            ),
+            True,
+        ),
+    ],
+    ids=["hma_connector", "non_hma_connector", "multi_all_hma", "multi_mixed"],
+)
+def test_hma_auto_config(kv_transfer_config, expect_disabled):
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig("cpu"),
+        kv_transfer_config=kv_transfer_config,
+    )
+    assert (
+        vllm_config.scheduler_config.disable_hybrid_kv_cache_manager is expect_disabled
+    )
+
+
+def test_explicit_hma_with_non_hma_connector_errors_at_factory():
+    vllm_config = VllmConfig(
+        device_config=DeviceConfig("cpu"),
+        scheduler_config=SchedulerConfig(
+            max_model_len=16,
+            is_encoder_decoder=False,
+            disable_hybrid_kv_cache_manager=False,
+        ),
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="ExampleConnector",
+            kv_role="kv_both",
+        ),
+    )
+    kv_cache_config = KVCacheConfig(
+        num_blocks=0, kv_cache_tensors=[], kv_cache_groups=[]
+    )
+    with pytest.raises(ValueError, match="does not support HMA but HMA is enabled"):
+        KVConnectorFactory.create_connector(
+            vllm_config, KVConnectorRole.SCHEDULER, kv_cache_config
+        )

tests/v1/kv_connector/unit/test_moriio_connector.py

@@ -1,7 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import importlib.util
-import os
 import subprocess
 import uuid
 from unittest.mock import MagicMock, patch
@@ -202,6 +201,7 @@ def create_vllm_config(
     enable_chunked_prefill: bool = True,
     enable_permute_local_kv: bool = False,
     role="kv_consumer",
+    read_mode: bool = False,
 ) -> VllmConfig:
     """Initialize VllmConfig for testing."""
     scheduler_config = SchedulerConfig(
@@ -228,6 +228,7 @@ def create_vllm_config(
         kv_connector="MoRIIOConnector",
         kv_role=role,
         enable_permute_local_kv=enable_permute_local_kv,
+        kv_connector_extra_config={"read_mode": read_mode},
     )
     return VllmConfig(
         scheduler_config=scheduler_config,
@@ -238,15 +239,6 @@ def create_vllm_config(
     )
 
 
-@pytest.fixture
-def moriio_read_mode():
-    """Force the connector into read mode via env for tests."""
-    os.environ["VLLM_MORIIO_CONNECTOR_READ_MODE"] = "True"
-    yield
-    # Cleanup after test
-    os.environ.pop("VLLM_MORIIO_CONNECTOR_READ_MODE", None)
-
-
 def test_write_mode_saves_local_block_ids():
     """Write mode records local block ids in MoRIIOConnectorMetadata.reqs_to_save."""
 
@@ -358,11 +350,11 @@ def test_write_mode_with_chunked_prefill_saves_local_block_ids():
         assert block_id == block.block_id, f"{block_id} != {block.block_id}"
 
 
-def test_read_mode_loads_remote_block_ids(moriio_read_mode):
+def test_read_mode_loads_remote_block_ids():
     """Read mode loads remote block ids into local cache mapping."""
 
     # Setup Scheduler and Request
-    vllm_config = create_vllm_config(role="kv_consumer")
+    vllm_config = create_vllm_config(role="kv_consumer", read_mode=True)
     scheduler = create_scheduler(vllm_config)
 
     # 2 Full Blocks and 1 Half Block.

tests/v1/kv_connector/unit/test_multi_connector.py

@@ -1000,11 +1000,8 @@ def _make_multi_connector(connector_names: list[str]) -> MultiConnector:
     )
 
 
-def test_multi_connector_hma_opt_in():
+def test_multi_connector_hma_support_detection():
     """
-    MultiConnector currently assumes HMA is opt-in: it needs
-    --no-disable-hybrid-kv-cache-manager to be enabled.
-
     At runtime, _all_support_hma is True only when every sub-connector
     implements SupportsHMA. Test all combinations of HMA / non-HMA
     sub-connectors.

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -723,8 +723,7 @@ def test_has_mamba_init(
 
     block_size = 16
     vllm_config = create_vllm_config(block_size=block_size)
-    # VllmConfig.__post_init__ auto-disables HMA when kv_transfer_config
-    # is set; override so we can test the scheduler's own derivation.
+    # Explicitly enable HMA so we can test the scheduler's own derivation.
     vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False
     kv_cache_config = make_kv_cache_config(
         block_size=block_size,

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -280,7 +280,7 @@ def test_cpu_offloading(
         kv_events_config=kv_events_config,
         kv_transfer_config=kv_transfer_config,
         **({"attention_config": {"backend": attn_backend}} if attn_backend else {}),
-        # HMA models need explicit opt-in when kv_transfer_config is set
+        # Keep HMA explicitly enabled for HMA model coverage.
         **({"disable_hybrid_kv_cache_manager": False} if uses_hma else {}),
         **({"enable_prefix_caching": True} if force_prefix_caching else {}),
         # ROCm: batch size 1 to reduce variability

tests/v1/logits_processors/test_correctness.py

@@ -20,7 +20,7 @@ from tests.v1.sample.utils import (
 )
 from vllm.config import VllmConfig
 from vllm.platforms import current_platform
-from vllm.sampling_params import SamplingParams
+from vllm.sampling_params import SamplingParams, validate_thinking_token_budget
 from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.v1.sample.logits_processor import (
     BatchUpdate,
@@ -1194,3 +1194,37 @@ def test_thinking_budget_enforced_without_penalties():
         "Budget exceeded: in_end should be True so that apply_to_logits "
         "forces the end token"
     )
+
+
+@pytest.mark.parametrize(
+    ("raw_value", "expected"),
+    [
+        (None, None),
+        (-1, None),
+        (10, 10),
+        (0, 0),
+    ],
+)
+def test_validate_thinking_token_budget(raw_value, expected):
+    assert validate_thinking_token_budget(raw_value) == expected
+
+
+def test_sampling_params_minus_one_normalizes_to_none():
+    params = SamplingParams(thinking_token_budget=-1)
+    assert params.thinking_token_budget is None
+
+
+@pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5, True])
+def test_validate_thinking_token_budget_rejects_invalid(invalid_budget):
+    from vllm.exceptions import VLLMValidationError
+
+    with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
+        validate_thinking_token_budget(invalid_budget)
+
+
+@pytest.mark.parametrize("invalid_budget", [-2, 0.6, 10.5])
+def test_thinking_budget_invalid_budget_rejected(invalid_budget):
+    from vllm.exceptions import VLLMValidationError
+
+    with pytest.raises(VLLMValidationError, match="thinking_token_budget"):
+        SamplingParams(thinking_token_budget=invalid_budget)

vllm/benchmarks/datasets/datasets.py

@@ -2323,15 +2323,152 @@ class CustomImageDataset(CustomDataset):
         "prompt": "Which country has the most pokemons based on the given graphs?",
         "image_files": ["path/to/image.png"],
     }
+    {
+        "content": [
+            {"type": "text", "text": "Compare these images: "},
+            {"type": "image", "image": "path/to/image1.png"},
+            {"type": "text", "text": " and "},
+            {"type": "image_url", "image_url": {"url": "path/to/image2.png"}},
+        ],
+    }
     ```
 
-    NOTE: Only the first image file in "image_files" is used for each sample request.
-
     This is used to benchmark multimodal LLMs on arbitrary datasets.
     """
 
     IS_MULTIMODAL = True
 
+    def load_data(self) -> None:
+        if self.dataset_path is None:
+            raise ValueError("dataset_path must be provided for loading data.")
+
+        self.data: list[dict] = []
+
+        if not self.dataset_path.endswith(".jsonl"):
+            raise NotImplementedError(
+                "Only JSONL format is supported for CustomImageDataset."
+            )
+
+        with open(self.dataset_path, encoding="utf-8") as f:
+            for line_number, line in enumerate(f, start=1):
+                line = line.strip()
+                if not line:
+                    continue
+
+                try:
+                    item = json.loads(line)
+                except json.JSONDecodeError as e:
+                    raise ValueError(
+                        f"Invalid JSON in custom image dataset line {line_number}: {e}"
+                    ) from e
+
+                if not isinstance(item, dict):
+                    raise ValueError(
+                        "Each custom image dataset line must contain a JSON object. "
+                        f"Found {type(item)} on line {line_number}."
+                    )
+
+                has_legacy_fields = "prompt" in item and "image_files" in item
+                has_interleaved_content = "content" in item
+                if not has_legacy_fields and not has_interleaved_content:
+                    raise ValueError(
+                        "Each custom image dataset line must contain either "
+                        "'prompt' and 'image_files' fields, or a 'content' field. "
+                        f"Invalid line: {line_number}."
+                    )
+
+                self.data.append(item)
+
+        random.seed(self.random_seed)
+        if not getattr(self, "disable_shuffle", False):
+            random.shuffle(self.data)
+
+    @staticmethod
+    def _validate_content_parts(content: Any) -> list[dict[str, Any]]:
+        if not isinstance(content, list):
+            raise ValueError(
+                "'content' must be a list of text and image content dictionaries."
+            )
+
+        if not content:
+            raise ValueError("'content' must contain at least one item.")
+
+        parts: list[dict[str, Any]] = []
+        for part in content:
+            if not isinstance(part, dict):
+                raise ValueError(
+                    f"Each item in 'content' must be a dictionary. Found {type(part)}."
+                )
+            parts.append(part)
+
+        return parts
+
+    @classmethod
+    def _process_content_part(cls, part: dict[str, Any]) -> dict[str, Any]:
+        content_type = part.get("type")
+        if content_type == "text":
+            text = part.get("text")
+            if not isinstance(text, str):
+                raise ValueError("Text content parts must contain a string 'text'.")
+            return {"type": "text", "text": text}
+
+        if content_type == "image":
+            if "image" not in part:
+                raise ValueError("Image content parts must contain an 'image' field.")
+            return dict(process_image(part["image"]))
+
+        if content_type == "image_url":
+            image_url = part.get("image_url")
+            if isinstance(image_url, str):
+                return dict(process_image(image_url))
+
+            if isinstance(image_url, dict):
+                url = image_url.get("url")
+                if not isinstance(url, str):
+                    raise ValueError(
+                        "Image URL content parts must contain a string 'image_url.url'."
+                    )
+
+                processed_part = dict(process_image(url))
+                processed_image_url = dict(processed_part["image_url"])
+                processed_image_url.update(
+                    {key: value for key, value in image_url.items() if key != "url"}
+                )
+                processed_part["image_url"] = processed_image_url
+                return processed_part
+
+            raise ValueError(
+                "Image URL content parts must contain an 'image_url' string "
+                "or dictionary."
+            )
+
+        raise ValueError(
+            "Content parts must have type 'text', 'image', or 'image_url'. "
+            f"Found: {content_type!r}."
+        )
+
+    @classmethod
+    def _process_interleaved_content(cls, content: Any) -> list[dict[str, Any]]:
+        return [
+            cls._process_content_part(part)
+            for part in cls._validate_content_parts(content)
+        ]
+
+    @staticmethod
+    def _get_text_from_content(content: list[dict[str, Any]]) -> str:
+        return "".join(part["text"] for part in content if part.get("type") == "text")
+
+    @staticmethod
+    def _process_image_files(images: Any) -> dict[str, Any] | list[dict[str, Any]]:
+        if not isinstance(images, list) or not images:
+            raise ValueError("'image_files' must be a non-empty list.")
+
+        mm_content = [dict(process_image(image)) for image in images]
+        if len(mm_content) == 1:
+            return mm_content[0]
+
+        return mm_content
+
     def sample(
         self,
         tokenizer: TokenizerLike,
@@ -2356,17 +2493,33 @@ class CustomImageDataset(CustomDataset):
         for i, item in enumerate(self.data):
             if len(sampled_requests) >= num_requests:
                 break
+
+            if "content" in item:
+                content = self._process_interleaved_content(item["content"])
+                text_prompt = self._get_text_from_content(content)
+                prompt_len = len(tokenizer(text_prompt).input_ids)
+                prompt = (
+                    [{"role": "user", "content": content}]
+                    if enable_multimodal_chat
+                    else content
+                )
+                sampled_requests.append(
+                    SampleRequest(
+                        prompt=prompt,
+                        prompt_len=prompt_len,
+                        expected_output_len=output_len,
+                        multi_modal_data=None,
+                        request_id=request_id_prefix + str(i),
+                    )
+                )
+                continue
+
             prompt = item["prompt"]
+            if not isinstance(prompt, str):
+                raise ValueError("'prompt' must be a string.")
 
             prompt_len = len(tokenizer(prompt).input_ids)
-            images = item["image_files"]
-            if len(images) > 1:
-                logger.warning(
-                    "Multiple image files found for sample %d. "
-                    "Only the first image will be used.",
-                    i,
-                )
-            mm_content = process_image(images[0])
+            mm_content = self._process_image_files(item["image_files"])
             if enable_multimodal_chat:
                 # Note: when chat is enabled the request prompt_len is no longer
                 # accurate and we will be using request output to count the

vllm/benchmarks/lib/endpoint_request_func.py

@@ -66,7 +66,7 @@ class StreamedResponseHandler:
 class RequestFuncInput:
     """The input for the request function."""
 
-    prompt: str | list[str]
+    prompt: str | list[str] | list[dict[str, Any]]
     api_url: str
     prompt_len: int
     output_len: int
@@ -268,8 +268,6 @@ def _get_chat_content(
     request_func_input: RequestFuncInput,
     mm_position: Literal["first", "last"] = "last",
 ) -> list[dict[str, Any]]:
-    text_contents = [{"type": "text", "text": request_func_input.prompt}]
-
     mm_contents = []
     if request_func_input.multi_modal_content:
         mm_content = request_func_input.multi_modal_content
@@ -282,12 +280,60 @@ def _get_chat_content(
                 "multi_modal_content must be a dict or list[dict] for openai-chat"
             )
 
+    prompt = request_func_input.prompt
+    if (
+        isinstance(prompt, list)
+        and prompt
+        and all(
+            isinstance(item, dict) and isinstance(item.get("type"), str)
+            for item in prompt
+        )
+    ):
+        if mm_position == "first":
+            return mm_contents + prompt
+
+        return prompt + mm_contents
+
+    text_contents = [{"type": "text", "text": prompt}]
+
     if mm_position == "first":
         return mm_contents + text_contents
 
     return text_contents + mm_contents
 
 
+def _is_chat_messages(prompt: Any) -> bool:
+    return (
+        isinstance(prompt, list)
+        and prompt
+        and all(
+            isinstance(item, dict)
+            and isinstance(item.get("role"), str)
+            and isinstance(item.get("content"), (str, list))
+            for item in prompt
+        )
+    )
+
+
+def _get_chat_messages(
+    request_func_input: RequestFuncInput,
+    mm_position: Literal["first", "last"] = "last",
+) -> list[dict[str, Any]]:
+    prompt = request_func_input.prompt
+    if _is_chat_messages(prompt):
+        return prompt
+
+    return [
+        {
+            "role": "user",
+            "content": _get_chat_content(
+                request_func_input,
+                mm_position=mm_position,
+            ),
+        }
+    ]
+
+
 async def async_request_openai_chat_completions(
     request_func_input: RequestFuncInput,
     session: aiohttp.ClientSession,
@@ -297,15 +343,13 @@ async def async_request_openai_chat_completions(
     api_url = request_func_input.api_url
     _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions")
 
-    content = _get_chat_content(request_func_input, mm_position=mm_position)
+    messages = _get_chat_messages(request_func_input, mm_position=mm_position)
 
     payload = {
         "model": request_func_input.model_name
         if request_func_input.model_name
         else request_func_input.model,
-        "messages": [
-            {"role": "user", "content": content},
-        ],
+        "messages": messages,
         "max_completion_tokens": request_func_input.output_len,
         "stream": True,
         "stream_options": {
@@ -608,15 +652,13 @@ async def async_request_openai_embeddings_chat(
     api_url = request_func_input.api_url
     _validate_api_url(api_url, "OpenAI Embeddings API", "embeddings")
 
-    content = _get_chat_content(request_func_input, mm_position=mm_position)
+    messages = _get_chat_messages(request_func_input, mm_position=mm_position)
 
     payload = {
         "model": request_func_input.model_name
         if request_func_input.model_name
         else request_func_input.model,
-        "messages": [
-            {"role": "user", "content": content},
-        ],
+        "messages": messages,
         # Many embedding models have short context length,
         # this is to avoid dropping some of the requests.
         "truncate_prompt_tokens": -1,

vllm/config/vllm.py

@@ -1406,7 +1406,7 @@ class VllmConfig:
         # Hybrid KV cache manager (HMA) runtime rules:
         # - Explicit enable (--no-disable-kv-cache-manager): error if runtime
         #   disables it
-        # - No preference: auto-disable for unsupported features (e.g. kv connector)
+        # - No preference: auto-disable for unsupported features or connector configs
         # - Explicit disable (--disable-kv-cache-manager): always respect it
         need_disable_hybrid_kv_cache_manager = False
         # logger should only print warning message for hybrid models. As we
@@ -1438,43 +1438,25 @@ class VllmConfig:
                 need_disable_hybrid_kv_cache_manager = True
 
         if self.scheduler_config.disable_hybrid_kv_cache_manager is None:
-            # Default to disable HMA, but only if the user didn't express a preference.
+            # Auto-disable HMA only when the connector config does not support it.
             if self.kv_transfer_config is not None:
-                from vllm.config.kv_transfer import KVTransferConfig
                 from vllm.distributed.kv_transfer.kv_connector.factory import (
                     KVConnectorFactory,
                 )
-                from vllm.distributed.kv_transfer.kv_connector.v1.base import (
-                    supports_hma,
-                )
 
-                connector_cls = KVConnectorFactory.get_connector_class(
-                    self.kv_transfer_config
-                )
-                all_support_hma = supports_hma(connector_cls)
-                # MultiConnector subclasses SupportsHMA; only effectively
-                # supports HMA when every sub-connector does.
-                if all_support_hma and connector_cls.__name__ == "MultiConnector":
-                    sub_ktcs = self.kv_transfer_config.kv_connector_extra_config.get(
-                        "connectors", []
-                    )
-                    all_support_hma = all(
-                        supports_hma(
-                            KVConnectorFactory.get_connector_class(
-                                KVTransferConfig(**sub)
-                            )
-                        )
-                        for sub in sub_ktcs
-                    )
-                if not all_support_hma:
+                if not KVConnectorFactory.supports_hma_config(self.kv_transfer_config):
                     need_disable_hybrid_kv_cache_manager = True
                     logger.warning(
                         "Turning off hybrid kv cache manager because "
-                        "connector %s does not subclass `SupportsHMA`. "
-                        "This will reduce performance on models with "
-                        "sliding window or Mamba attention. See "
-                        "kv_connector/v1/base.py for details.",
-                        connector_cls.__name__,
+                        "`--kv-transfer-config` selects a KV connector that "
+                        "does not support it. Impact: hybrid SSM models "
+                        "(e.g. Jamba, Bamba) require HMA and will fail at "
+                        "startup without it; models with sliding window "
+                        "attention will run with reduced performance. "
+                        "To add HMA support to a KV connector, subclass "
+                        "`SupportsHMA` defined in kv_connector/v1/base.py "
+                        "(for MultiConnector, all child connectors must "
+                        "support HMA)."
                     )
             self.scheduler_config.disable_hybrid_kv_cache_manager = (
                 need_disable_hybrid_kv_cache_manager

vllm/distributed/kv_transfer/kv_connector/factory.py

@@ -5,6 +5,7 @@ import importlib
 from collections.abc import Callable
 from typing import TYPE_CHECKING, cast
 
+from vllm.config.kv_transfer import KVTransferConfig
 from vllm.distributed.kv_transfer.kv_connector.base import (
     KVConnectorBase,
     KVConnectorBaseType,
@@ -18,7 +19,6 @@ from vllm.utils.func_utils import supports_kw
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
-    from vllm.config.kv_transfer import KVTransferConfig
     from vllm.v1.kv_cache_interface import KVCacheConfig
 
 logger = init_logger(__name__)
@@ -53,7 +53,7 @@ class KVConnectorFactory:
 
         # check if the connector supports HMA
         hma_enabled = not config.scheduler_config.disable_hybrid_kv_cache_manager
-        if hma_enabled and not supports_hma(connector_cls):
+        if hma_enabled and not cls.supports_hma_config(kv_transfer_config):
             raise ValueError(
                 f"Connector {connector_cls.__name__} does not support HMA but "
                 f"HMA is enabled. Please set `--disable-hybrid-kv-cache-manager`."
@@ -127,6 +127,23 @@ class KVConnectorFactory:
             raise ValueError(f"Unsupported connector type: {connector_name}")
         return connector_cls
 
+    @classmethod
+    def supports_hma_config(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Return whether this KV transfer config supports HMA.
+
+        MultiConnector is a special case: the wrapper class implements
+        SupportsHMA, but effective support depends on every configured child.
+        """
+        connector_cls = cls.get_connector_class(kv_transfer_config)
+        if kv_transfer_config.kv_connector != "MultiConnector":
+            return supports_hma(connector_cls)
+
+        from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
+            MultiConnector,
+        )
+
+        return MultiConnector.all_children_support_hma(kv_transfer_config)
+
 
 # Register various connectors here.
 # The registration should not be done in each individual file, as we want to

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
+import os
 import threading
 import time
 from collections.abc import Iterator
@@ -12,8 +13,7 @@ import regex as re
 import torch
 import zmq
 
-from vllm import envs
-from vllm.config import VllmConfig
+from vllm.config import KVTransferConfig, VllmConfig
 from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorMetadata,
 )
@@ -162,8 +162,10 @@ class TransferError(MoRIIOError):
     pass
 
 
-def get_moriio_mode() -> MoRIIOMode:
-    read_mode = envs.VLLM_MORIIO_CONNECTOR_READ_MODE
+def get_moriio_mode(kv_transfer_config: KVTransferConfig) -> MoRIIOMode:
+    read_mode = str(
+        kv_transfer_config.kv_connector_extra_config.get("read_mode", "false")
+    ).lower().strip() in ("true", "1")
     logger.debug("MoRIIO Connector read_mode: %s", read_mode)
     if read_mode:
         return MoRIIOMode.READ
@@ -175,6 +177,26 @@ def get_port_offset(dp_rank: int, tp_rank: int, tp_size: int = 1) -> int:
     return (dp_rank) * tp_size + tp_rank
 
 
+_DEPRECATED_ENV_VARS: dict[str, str] = {
+    "VLLM_MORIIO_CONNECTOR_READ_MODE": "read_mode",
+    "VLLM_MORIIO_QP_PER_TRANSFER": "qp_per_transfer",
+    "VLLM_MORIIO_POST_BATCH_SIZE": "post_batch_size",
+    "VLLM_MORIIO_NUM_WORKERS": "num_workers",
+}
+
+
+def _warn_deprecated_env_vars() -> None:
+    for env_var, new_key in _DEPRECATED_ENV_VARS.items():
+        if env_var in os.environ:
+            logger.warning_once(
+                "The environment variable %s is deprecated and ignored. "
+                "Set %r inside kv_transfer_config.kv_connector_extra_config "
+                "instead.",
+                env_var,
+                new_key,
+            )
+
+
 @dataclass
 class MoRIIOConfig:
     local_ip: str
@@ -189,6 +211,10 @@ class MoRIIOConfig:
     dp_rank: int
     dp_size: int
     tp_size: int
+    read_mode: bool = False
+    qp_per_transfer: int = 1
+    post_batch_size: int = -1
+    num_workers: int = 1
     backend: str = "rdma"
 
     @classmethod
@@ -201,11 +227,24 @@ class MoRIIOConfig:
         # notify_port       -> For synchronizing stages between prefill and decode
         # handshake_port    -> For initial handshake between mori engine
 
+        # Optional tuning knobs
+        # read_mode        -> If true, run the connector in READ mode (consumer
+        #                     pulls KV from producer) instead of the default
+        #                     WRITE mode.
+
+        # Knobs for RDMA transfers, ignored if on xgmi backend
+        # qp_per_transfer  -> Number of RDMA Queue Pairs per KV transfer.
+        # post_batch_size  -> Batch size for posting transfer work requests
+        #                     (-1 lets the MoRI backend choose).
+        # num_workers      -> Number of background worker threads the MoRI
+        #                     engine uses for transfer processing.
+
         # TODO : merge notify_port and handshake_port to simplify port management
         #        supports non-contiguous ports
         assert vllm_config.kv_transfer_config is not None, (
             "kv_transfer_config must be set for MoRIIOConnector"
         )
+        _warn_deprecated_env_vars()
         kv_transfer_config = vllm_config.kv_transfer_config
         extra_config = kv_transfer_config.kv_connector_extra_config
         tp_rank = get_tensor_model_parallel_rank()
@@ -234,6 +273,10 @@ class MoRIIOConfig:
             dp_rank=dp_rank,
             dp_size=dp_size,
             tp_size=tp_size,
+            read_mode=get_moriio_mode(kv_transfer_config) == MoRIIOMode.READ,
+            qp_per_transfer=int(extra_config.get("qp_per_transfer", 1)),
+            post_batch_size=int(extra_config.get("post_batch_size", -1)),
+            num_workers=int(extra_config.get("num_workers", 1)),
             backend=backend,
         )
 

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

@@ -108,7 +108,7 @@ class MoRIIOConnector(KVConnectorBase_V1):
             + ":"
             + str(self.kv_transfer_config.kv_connector_extra_config["handshake_port"])
         )
-        self.mode = get_moriio_mode()
+        self.mode = get_moriio_mode(self.kv_transfer_config)
         if role == KVConnectorRole.SCHEDULER:
             self.connector_scheduler: MoRIIOConnectorScheduler | None = (
                 MoRIIOConnectorScheduler(vllm_config, self.engine_id)
@@ -250,7 +250,7 @@ class MoRIIOConnectorScheduler:
         self.kv_transfer_config = vllm_config.kv_transfer_config
         self.block_size = vllm_config.cache_config.block_size
         self.engine_id: EngineId = engine_id
-        self.mode = get_moriio_mode()
+        self.mode = get_moriio_mode(self.kv_transfer_config)
         self.host_ip = get_ip()
         self.handshake_port = self.kv_transfer_config.kv_connector_extra_config[
             "handshake_port"
@@ -615,8 +615,11 @@ class MoRIIOConnectorWorker:
                 "is installed and properly configured."
             )
 
+        assert vllm_config.kv_transfer_config is not None
         self.moriio_config = MoRIIOConfig.from_vllm_config(vllm_config)
-        self.mode = get_moriio_mode()
+        self.mode = (
+            MoRIIOMode.READ if self.moriio_config.read_mode else MoRIIOMode.WRITE
+        )
 
         logger.info("Initializing MoRIIO worker %s", engine_id)
 
@@ -700,7 +703,12 @@ class MoRIIOConnectorWorker:
             if self.moriio_config.backend == "xgmi"
             else BackendType.RDMA
         )
-        self.moriio_wrapper.set_backend_type(backend)
+        self.moriio_wrapper.set_backend_type(
+            backend,
+            qp_per_transfer=self.moriio_config.qp_per_transfer,
+            post_batch_size=self.moriio_config.post_batch_size,
+            num_workers=self.moriio_config.num_workers,
+        )
         self.moriio_wrapper.notify_port = self.moriio_config.notify_port
         self.local_kv_cache_metadata: list[bytes] = []
         self.local_kv_cache_size: list[int] = []

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py

@@ -8,7 +8,6 @@ import msgpack
 import torch
 import zmq
 
-from vllm import envs
 from vllm.logger import init_logger
 from vllm.utils.network_utils import (
     make_zmq_path,
@@ -16,7 +15,7 @@ from vllm.utils.network_utils import (
 )
 
 if TYPE_CHECKING:
-    pass
+    from mori.io import BackendType
 
 from queue import Empty, Queue
 
@@ -376,7 +375,13 @@ class MoRIIOWrapper:
         )
         self.moriio_engine = moriio_engine
 
-    def set_backend_type(self, backend_type):
+    def set_backend_type(
+        self,
+        backend_type: "BackendType",
+        qp_per_transfer: int = 1,
+        post_batch_size: int = -1,
+        num_workers: int = 1,
+    ) -> None:
         assert self.moriio_engine is not None, "MoRIIO engine must be set first"
         if backend_type == BackendType.XGMI:
             logger.info("Using MoRIIO backend: XGMI")
@@ -385,14 +390,14 @@ class MoRIIOWrapper:
             logger.info(
                 "Using MoRIIO backend: RDMA "
                 "(qp_per_transfer=%d, post_batch_size=%d, num_workers=%d)",
-                envs.VLLM_MORIIO_QP_PER_TRANSFER,
-                envs.VLLM_MORIIO_POST_BATCH_SIZE,
-                envs.VLLM_MORIIO_NUM_WORKERS,
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
             )
             rdma_cfg = RdmaBackendConfig(
-                envs.VLLM_MORIIO_QP_PER_TRANSFER,
-                envs.VLLM_MORIIO_POST_BATCH_SIZE,
-                envs.VLLM_MORIIO_NUM_WORKERS,
+                qp_per_transfer,
+                post_batch_size,
+                num_workers,
                 PollCqMode.POLLING,
             )
             self.moriio_engine.create_backend(backend_type, rdma_cfg)

vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py

@@ -19,7 +19,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
     KVConnectorRole,
     KVConnectorWorkerMetadata,
     SupportsHMA,
-    supports_hma,
 )
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
     KVConnectorPromMetrics,
@@ -151,6 +150,22 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
                 return True
         return False
 
+    @classmethod
+    def all_children_support_hma(cls, kv_transfer_config: "KVTransferConfig") -> bool:
+        """Return True only if every configured child connector supports HMA."""
+        connectors_config = kv_transfer_config.kv_connector_extra_config.get(
+            "connectors", []
+        )
+        if not connectors_config:
+            return False
+        for conn_config in connectors_config:
+            child_config = KVTransferConfig(
+                **{"engine_id": kv_transfer_config.engine_id, **conn_config}
+            )
+            if not KVConnectorFactory.supports_hma_config(child_config):
+                return False
+        return True
+
     def __init__(
         self,
         vllm_config: "VllmConfig",
@@ -169,7 +184,10 @@ class MultiConnector(KVConnectorBase_V1, SupportsHMA):
             self._connectors.append(connector_cls(temp_config, role, kv_cache_config))
             self._ktc_kv_transfer_config.append(temp_config.kv_transfer_config)
 
-        self._all_support_hma = all(supports_hma(c) for c in self._connectors)
+        assert vllm_config.kv_transfer_config is not None
+        self._all_support_hma = MultiConnector.all_children_support_hma(
+            vllm_config.kv_transfer_config
+        )
         assert (
             vllm_config.scheduler_config.disable_hybrid_kv_cache_manager
             or self._all_support_hma

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -41,6 +41,7 @@ from vllm.sampling_params import (
     RequestOutputKind,
     SamplingParams,
     StructuredOutputsParams,
+    ThinkingTokenBudget,
 )
 from vllm.utils import random_uuid
 
@@ -225,7 +226,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
             "part of the standard OpenAI API specification."
         ),
     )
-    thinking_token_budget: int | None = None
+    thinking_token_budget: ThinkingTokenBudget = None
     include_reasoning: bool = True
     parallel_tool_calls: bool | None = True
 

vllm/entrypoints/openai/completion/protocol.py

@@ -29,6 +29,7 @@ from vllm.sampling_params import (
     RequestOutputKind,
     SamplingParams,
     StructuredOutputsParams,
+    ThinkingTokenBudget,
 )
 from vllm.utils import random_uuid
 
@@ -185,11 +186,12 @@ class CompletionRequest(OpenAIBaseModel):
         "can detect such behavior and terminate early, saving time and tokens.",
     )
 
-    thinking_token_budget: int | None = Field(
+    thinking_token_budget: ThinkingTokenBudget = Field(
         default=None,
         description=(
             "Maximum number of tokens allowed for thinking operations "
-            "(reasoning models). -1 = unlimited."
+            "(reasoning models). Non-negative integer sets the limit; "
+            "-1 means unlimited (treated as unset)."
         ),
     )
 

vllm/envs.py

@@ -216,10 +216,6 @@ if TYPE_CHECKING:
     VLLM_ROCM_QUICK_REDUCE_MAX_SIZE_BYTES_MB: int | None = None
     VLLM_ROCM_QUICK_REDUCE_MIN_SIZE_BYTES_MB: int | None = None
     VLLM_ROCM_QUICK_REDUCE_QUANTIZATION_MIN_SIZE_KB: int | None = None
-    VLLM_MORIIO_CONNECTOR_READ_MODE: bool = False
-    VLLM_MORIIO_QP_PER_TRANSFER: int = 1
-    VLLM_MORIIO_POST_BATCH_SIZE: int = -1
-    VLLM_MORIIO_NUM_WORKERS: int = 1
     VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: int = 480
     VLLM_ENABLE_CUDAGRAPH_GC: bool = False
     VLLM_LOOPBACK_IP: str = ""
@@ -1642,20 +1638,6 @@ environment_variables: dict[str, Callable[[], Any]] = {
         "Use --linear-backend emulation.",
         lambda: bool(int(os.getenv("VLLM_USE_NVFP4_CT_EMULATIONS", "0"))),
     ),
-    # Controls the read mode for the Mori-IO connector
-    "VLLM_MORIIO_CONNECTOR_READ_MODE": lambda: (
-        os.getenv("VLLM_MORIIO_CONNECTOR_READ_MODE", "False").lower() in ("true", "1")
-    ),
-    # Controls the QP (Queue Pair) per transfer configuration for the Mori-IO connector
-    "VLLM_MORIIO_QP_PER_TRANSFER": lambda: int(
-        os.getenv("VLLM_MORIIO_QP_PER_TRANSFER", "1")
-    ),
-    # Controls the post-processing batch size for the Mori-IO connector
-    "VLLM_MORIIO_POST_BATCH_SIZE": lambda: int(
-        os.getenv("VLLM_MORIIO_POST_BATCH_SIZE", "-1")
-    ),
-    # Controls the number of workers for Mori operations for the Mori-IO connector
-    "VLLM_MORIIO_NUM_WORKERS": lambda: int(os.getenv("VLLM_MORIIO_NUM_WORKERS", "1")),
     # Timeout (in seconds) for MooncakeConnector in PD disaggregated setup.
     "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT": lambda: int(
         os.getenv("VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", "480")

vllm/lora/ops/triton_ops/fused_moe_lora_op.py

@@ -7,6 +7,7 @@ from vllm.distributed import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 from vllm.triton_utils.allocation import set_triton_allocator
 from vllm.utils.torch_utils import direct_register_custom_op
@@ -406,7 +407,7 @@ def _run_fused_moe_lora_one_shot(
 
     # NPID_FACTOR heuristic: scale N-axis parallelism when base CTA count is
     # short of saturating the SM array. Cap by the cost of redundant shrink.
-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    sm_count = current_platform.num_compute_units(device.index)
     base_programs = max(M_blocks * num_slices * grid_lora_dim, 1)
     shrink_ratio = K / max(K + N_per_slice, 1)
     max_npid_by_budget = max(1, int(1.5 / max(shrink_ratio, 1e-3)) + 1)
@@ -786,7 +787,7 @@ def _run_fused_moe_lora_small_batch(
     N_tiles = triton.cdiv(N_per_slice, BLOCK_N)
     pair_slices = M_grid * num_slices
 
-    sm_count = torch.cuda.get_device_properties(device).multi_processor_count
+    sm_count = current_platform.num_compute_units(device.index)
     n_tiles_per_program = _pick_small_batch_chunk(pair_slices, N_tiles, sm_count)
     n_chunks = triton.cdiv(N_tiles, n_tiles_per_program)
     work_total = pair_slices * n_chunks

vllm/model_executor/models/internlm2.py

@@ -41,6 +41,7 @@ from vllm.sequence import IntermediateTensors
 from .interfaces import SupportsLoRA, SupportsPP
 from .interfaces_base import default_pooling_type
 from .utils import (
+    AutoWeightsLoader,
     StageMissingLayer,
     is_pp_missing_parameter,
     make_empty_intermediate_tensors_factory,
@@ -308,6 +309,42 @@ class InternLM2Model(nn.Module):
         hidden_states, _ = self.norm(hidden_states, residual)
         return hidden_states
 
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        stacked_params_mapping = [
+            # (param_name, shard_name, shard_id)
+            ("gate_up_proj", "w1", 0),
+            ("gate_up_proj", "w3", 1),
+        ]
+        params_dict = dict(self.named_parameters())
+        loaded_params: set[str] = set()
+        for name, loaded_weight in weights:
+            if "rotary_emb.inv_freq" in name:
+                continue
+            for param_name, weight_name, shard_id in stacked_params_mapping:
+                if weight_name not in name:
+                    continue
+                name = name.replace(weight_name, param_name)
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = param.weight_loader
+                weight_loader(param, loaded_weight, shard_id)
+                break
+            else:
+                # Skip loading extra bias for GPTQ models.
+                if name.endswith(".bias") and name not in params_dict:
+                    continue
+                if is_pp_missing_parameter(name, self):
+                    continue
+                param = params_dict[name]
+                weight_loader = getattr(param, "weight_loader", default_weight_loader)
+                weight_loader(param, loaded_weight)
+            loaded_params.add(name)
+        return loaded_params
+
 
 class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
     packed_modules_mapping = {
@@ -368,40 +405,11 @@ class InternLM2ForCausalLM(nn.Module, SupportsPP, SupportsLoRA):
         return logits
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
-        stacked_params_mapping = [
-            # (param_name, shard_name, shard_id)
-            ("gate_up_proj", "w1", 0),
-            ("gate_up_proj", "w3", 1),
-        ]
-        params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
-        for name, loaded_weight in weights:
-            if "rotary_emb.inv_freq" in name:
-                continue
-            for param_name, weight_name, shard_id in stacked_params_mapping:
-                if weight_name not in name:
-                    continue
-                name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = param.weight_loader
-                weight_loader(param, loaded_weight, shard_id)
-                break
-            else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                if is_pp_missing_parameter(name, self):
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader", default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
-        return loaded_params
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["output."] if self.config.tie_word_embeddings else None),
+        )
+        return loader.load_weights(weights)
 
 
 @default_pooling_type(tok_pooling_type="ALL")

vllm/sampling_params.py

@@ -7,9 +7,10 @@ import json as json_mod
 from dataclasses import field
 from enum import Enum, IntEnum
 from functools import cached_property
-from typing import Any
+from typing import Annotated, Any
 
 import msgspec
+from pydantic import BeforeValidator
 from pydantic.dataclasses import dataclass
 
 import vllm.envs as envs
@@ -30,6 +31,35 @@ MAX_LOGPROB_TOKEN_IDS = 128
 the per-request row width allocated by the sampler's `LogprobTokenIdsState`."""
 
 
+def validate_thinking_token_budget(value: int | float | bool | None) -> int | None:
+    """Validate ``thinking_token_budget``; return ``None`` if unset."""
+    if value is None:
+        return None
+    if isinstance(value, (bool, float)) or not isinstance(value, int):
+        raise VLLMValidationError(
+            "`thinking_token_budget` must be a non-negative integer "
+            "or -1 for unlimited.",
+            parameter="thinking_token_budget",
+            value=value,
+        )
+    if value == -1:
+        return None
+    if value < 0:
+        raise VLLMValidationError(
+            "`thinking_token_budget` must be a non-negative integer "
+            "or -1 for unlimited.",
+            parameter="thinking_token_budget",
+            value=value,
+        )
+    return value
+
+
+ThinkingTokenBudget = Annotated[
+    int | None,
+    BeforeValidator(validate_thinking_token_budget),
+]
+
+
 class SamplingType(IntEnum):
     GREEDY = 0
     RANDOM = 1
@@ -409,6 +439,10 @@ class SamplingParams(
         if self.seed == -1:
             self.seed = None
 
+        self.thinking_token_budget = validate_thinking_token_budget(
+            self.thinking_token_budget
+        )
+
         if self.stop is None:
             self.stop = []
         elif isinstance(self.stop, str):

vllm/v1/worker/gpu/model_runner.py

@@ -244,6 +244,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         # For transferring state from execute_model to subsequent sample_tokens call.
         self.execute_model_state: ExecuteModelState | None = None
+        self._deferred_kv_connector_scheduler_output: SchedulerOutput | None = None
 
         # Expert parallelism load balancer.
         self.eplb = EPLBController(self.parallel_config, self.device)
@@ -1207,7 +1208,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             aux_hidden_states = None
             output_intermediate_tensors = model_output
 
-        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
+        kv_connector_output = None
+        self._deferred_kv_connector_scheduler_output = None
+        if self.is_last_pp_rank and self.speculator is not None and not dummy_run:
+            self._deferred_kv_connector_scheduler_output = scheduler_output
+        else:
+            kv_connector_output = self.kv_connector.post_forward(scheduler_output)
         self.execute_model_state = ExecuteModelState(
             input_batch=input_batch,
             attn_metadata=attn_metadata,
@@ -1344,6 +1350,18 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
             self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
 
+            if kv_connector_output is None:
+                assert self._deferred_kv_connector_scheduler_output is not None
+                # delay KV connector finalization until the drafter has written
+                # its KV cache.
+                kv_connector_output = self.kv_connector.post_forward(
+                    self._deferred_kv_connector_scheduler_output
+                )
+                async_output.model_runner_output.kv_connector_output = (
+                    kv_connector_output
+                )
+                self._deferred_kv_connector_scheduler_output = None
+
         if self.use_async_scheduling:
             return async_output
         return async_output.get_output()