Merge remote-tracking branch 'origin/main' into sync/lesj/qwen4exp-bugfixes

Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>

.buildkite/intel_jobs/test-intel.yaml

@@ -148,16 +148,10 @@ steps:
     source_file_dependencies:
       - vllm/
       - .buildkite/intel_jobs/test-intel.yaml 
-    # TODO: XPU fused MoE + model runner V2 currently hangs on some
-    # pre-quantized models, e.g.
-    # quantization/test_online.py::test_online_quantization_loads_real_weights[nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable]
-    # Remove the VLLM_USE_V2_MODEL_RUNNER=0 override below once the hang is
-    # root-caused and fixed.
     commands:
       - >-
         bash .buildkite/scripts/hardware_ci/run-intel-test.sh
         'cd tests &&
-        export VLLM_USE_V2_MODEL_RUNNER=0 &&
         pytest -v -s quantization/test_auto_round.py &&
         pytest -v -s quantization/test_online.py'
   - label: "XPU compressed tensors FP8 test"

.buildkite/scripts/ci-otel/README.md

@@ -27,9 +27,17 @@ These are device readings during each test's time interval, not process-level
 kernel attribution. Other processes and overlapping tests can contribute to
 the readings. Subsecond tests may have no sample. Missing readings never become
 zero utilization. MIG parent-device memory and utilization are omitted because
-they cannot be attributed to the job's partition. The initial implementation
-covers NVIDIA commands in the existing tracing rollout; CPU jobs and AMD
-mirrors do not collect GPU samples.
+they cannot be attributed to the job's partition. When the environment identifies
+an assigned MIG UUID, the collector queries that instance's used/total memory
+through vLLM's standalone bundled NVML binding in a bounded subprocess.
+The binding is loaded directly to avoid vLLM/Torch import side effects. Numeric
+`CUDA_VISIBLE_DEVICES` indices are resolved against explicit MIG UUIDs in
+`NVIDIA_VISIBLE_DEVICES`. For CDI assignments without an explicit UUID, a bounded
+helper initializes the CUDA driver and reads only CUDA-visible device UUIDs,
+without creating CUDA contexts or importing Torch. NVML verifies each MIG handle.
+Discovery or permission failures remain unknown; NVML sibling instances are never
+enumerated. MIG utilization stays absent.
+CPU jobs and AMD mirrors do not collect GPU samples.
 
 Run the helper tests without loading vLLM or requiring a GPU:
 

.buildkite/scripts/ci-otel/ci_gpu.py

@@ -6,7 +6,10 @@
 from __future__ import annotations
 
 import csv
+import ctypes
+import importlib.util
 import io
+import json
 import math
 import os
 import secrets
@@ -15,6 +18,8 @@ import subprocess
 import sys
 import threading
 import time
+import uuid as uuid_module
+from pathlib import Path
 
 from ci_otel import Span, export_spans, record_spans
 
@@ -67,6 +72,118 @@ def parse_samples(output: str, timestamp_ns: int) -> list[dict]:
     return events
 
 
+def visible_mig_devices() -> list[str]:
+    """Resolve explicit MIG assignments without enumerating sibling instances."""
+    nvidia = os.environ.get("NVIDIA_VISIBLE_DEVICES", "").split(",")
+    nvidia = [value.strip() for value in nvidia]
+    cuda = os.environ.get("CUDA_VISIBLE_DEVICES")
+    devices = nvidia if cuda is None else [value.strip() for value in cuda.split(",")]
+    if devices and all(value.isdigit() for value in devices):
+        if not all(value.startswith("MIG-") for value in nvidia):
+            return []
+        try:
+            devices = [nvidia[int(value)] for value in devices]
+        except IndexError:
+            return []
+    if not devices or not all(value.startswith("MIG-") for value in devices):
+        return []
+    return list(dict.fromkeys(devices))[:MAX_DEVICES]
+
+
+def load_nvml():
+    """Load the standalone bundled binding without importing vLLM or Torch."""
+    # Test images install vLLM as a wheel and keep only tests/helpers in the
+    # checkout. Finding the top-level package does not execute its __init__.
+    package = importlib.util.find_spec("vllm")
+    if package is None or package.origin is None:
+        raise ImportError("Installed vLLM package unavailable")
+    path = Path(package.origin).parent / "third_party/pynvml.py"
+    spec = importlib.util.spec_from_file_location("ci_nvml", path)
+    if spec is None or spec.loader is None:
+        raise ImportError("Bundled NVML binding unavailable")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def cuda_visible_mig_devices() -> list[str]:
+    """Discover CDI assignments in a subprocess without creating CUDA contexts."""
+    pynvml = load_nvml()
+
+    driver = ctypes.CDLL("libcuda.so.1")
+    count = ctypes.c_int()
+    if driver.cuInit(0) or driver.cuDeviceGetCount(ctypes.byref(count)):
+        return []
+    devices = []
+    pynvml.nvmlInit()
+    try:
+        for index in range(min(count.value, MAX_DEVICES)):
+            device = ctypes.c_int()
+            identifier = (ctypes.c_ubyte * 16)()
+            if driver.cuDeviceGet(ctypes.byref(device), index):
+                continue
+            # The v2 entry point returns the compute instance UUID under MIG.
+            if driver.cuDeviceGetUuid_v2(ctypes.byref(identifier), device):
+                continue
+            uuid = "MIG-" + str(uuid_module.UUID(bytes=bytes(identifier)))
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+                if pynvml.nvmlDeviceIsMigDeviceHandle(handle):
+                    devices.append(uuid)
+            except pynvml.NVMLError:
+                continue
+    finally:
+        pynvml.nvmlShutdown()
+    return devices
+
+
+def query_mig_samples(devices: list[str], timestamp_ns: int) -> list[dict]:
+    # Use the binding already shipped with vLLM; never initialize CUDA/Torch.
+    pynvml = load_nvml()
+
+    events = []
+    pynvml.nvmlInit()
+    try:
+        for index, uuid in enumerate(devices[:MAX_DEVICES]):
+            if not uuid.startswith("MIG-"):
+                continue
+            attributes: dict[str, str | int | bool] = {
+                "gpu.uuid": uuid,
+                "gpu.index": index,
+                "gpu.name": "NVIDIA MIG",
+                "gpu.status": "unavailable_mig_memory",
+            }
+            try:
+                handle = pynvml.nvmlDeviceGetHandleByUUID(uuid)
+                if not pynvml.nvmlDeviceIsMigDeviceHandle(handle):
+                    continue
+                memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
+                used = _number(str(memory.used))
+                total = _number(str(memory.total))
+                if used is not None and total and used <= total:
+                    attributes.update(
+                        {
+                            "gpu.memory.used": used,
+                            "gpu.memory.total": total,
+                            "gpu.status": "mig_memory_only",
+                        }
+                    )
+                # NVML does not support per-MIG utilization on these devices.
+            except pynvml.NVMLError:
+                pass
+            events.append(
+                {
+                    "time_ns": timestamp_ns,
+                    "name": "ci.gpu.sample",
+                    "attributes": attributes,
+                }
+            )
+    finally:
+        pynvml.nvmlShutdown()
+    return events
+
+
 def sample_span(events: list[dict]) -> Span:
     return Span(
         trace_id=os.environ["CI_INFRA_TRACE_ID"],
@@ -88,16 +205,23 @@ def collect(parent_pid: int, stop: threading.Event) -> None:
     events: list[dict] = []
     batch_started = time.monotonic()
     failures = 0
+    discovery_attempted = False
+    mig_devices = visible_mig_devices()
+    query = (
+        [sys.executable, __file__, "--query-mig", *mig_devices]
+        if mig_devices
+        else [
+            "nvidia-smi",
+            f"--query-gpu={QUERY}",
+            "--format=csv,noheader,nounits",
+        ]
+    )
     try:
         while not stop.is_set() and os.getppid() == parent_pid:
             started = time.monotonic()
             try:
                 result = subprocess.run(
-                    [
-                        "nvidia-smi",
-                        f"--query-gpu={QUERY}",
-                        "--format=csv,noheader,nounits",
-                    ],
+                    query,
                     capture_output=True,
                     text=True,
                     timeout=QUERY_TIMEOUT_SECONDS,
@@ -105,12 +229,41 @@ def collect(parent_pid: int, stop: threading.Event) -> None:
                 )
                 if stop.is_set():
                     break
-                samples = parse_samples(result.stdout, time.time_ns())
+                samples = (
+                    json.loads(result.stdout)
+                    if mig_devices
+                    else parse_samples(result.stdout, time.time_ns())
+                )
+                if (
+                    not mig_devices
+                    and not discovery_attempted
+                    and any(
+                        sample["attributes"].get("gpu.status") == "unsupported_mig"
+                        for sample in samples
+                    )
+                ):
+                    discovery_attempted = True
+                    # CDI can assign a MIG device without exporting its UUID.
+                    # Only CUDA-visible devices are inspected, never NVML siblings.
+                    try:
+                        discovered = subprocess.run(
+                            [sys.executable, __file__, "--discover-mig"],
+                            capture_output=True,
+                            text=True,
+                            timeout=QUERY_TIMEOUT_SECONDS,
+                            check=True,
+                        )
+                        mig_devices = json.loads(discovered.stdout)
+                    except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
+                        mig_devices = []
+                    if mig_devices:
+                        query = [sys.executable, __file__, "--query-mig", *mig_devices]
+                        continue  # Do not create an extra parent-device series.
                 if not samples:
                     break
                 events.extend(samples)
                 failures = 0
-            except (OSError, subprocess.SubprocessError):
+            except (OSError, subprocess.SubprocessError, json.JSONDecodeError):
                 failures += 1
                 if failures >= 3:
                     print(
@@ -142,4 +295,9 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    if len(sys.argv) > 1 and sys.argv[1] == "--discover-mig":
+        print(json.dumps(cuda_visible_mig_devices()))
+    elif len(sys.argv) > 1 and sys.argv[1] == "--query-mig":
+        print(json.dumps(query_mig_samples(sys.argv[2:], time.time_ns())))
+    else:
+        main()

.buildkite/scripts/ci-otel/tests/test_ci_otel.py

@@ -11,6 +11,7 @@ import sys
 import threading
 import time
 from pathlib import Path
+from types import SimpleNamespace
 
 import pytest
 
@@ -814,6 +815,85 @@ def test_gpu_samples_preserve_zero_and_omit_unsupported_metrics():
     assert "gpu.memory.total" not in events[2]["attributes"]
 
 
+@pytest.mark.parametrize(
+    "cuda,nvidia,expected",
+    [
+        ("MIG-a", "GPU-parent", ["MIG-a"]),
+        ("0", "MIG-a", ["MIG-a"]),
+        ("1,0", "MIG-a,MIG-b", ["MIG-b", "MIG-a"]),
+        (None, "MIG-a", ["MIG-a"]),
+        ("", "MIG-a", []),
+        ("-1", "MIG-a", []),
+        ("2", "MIG-a", []),
+        ("GPU-parent", "MIG-a", []),
+        ("0", "all", []),
+    ],
+)
+def test_mig_assignment_obeys_cuda_visibility(monkeypatch, cuda, nvidia, expected):
+    monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", nvidia)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    if cuda is not None:
+        monkeypatch.setenv("CUDA_VISIBLE_DEVICES", cuda)
+    assert ci_gpu.visible_mig_devices() == expected
+
+
+@pytest.mark.parametrize("used", [0, 3 * 1024**3])
+def test_mig_memory_uses_assigned_handle_without_parent_metrics(monkeypatch, used):
+    handles = []
+    closed = []
+
+    def get_handle(uuid):
+        handles.append(uuid)
+        return "slice-handle"
+
+    def get_memory(handle):
+        assert handle == "slice-handle"
+        return SimpleNamespace(used=used, total=16 * 1024**3)
+
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: closed.append(True),
+        nvmlDeviceGetHandleByUUID=get_handle,
+        nvmlDeviceIsMigDeviceHandle=lambda handle: handle == "slice-handle",
+        nvmlDeviceGetMemoryInfo=get_memory,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
+    assert handles == ["MIG-a"]
+    assert closed == [True]
+    assert samples[0]["time_ns"] == 1234
+    attributes = samples[0]["attributes"]
+    assert attributes["gpu.uuid"] == "MIG-a"
+    assert attributes["gpu.memory.used"] == used
+    assert attributes["gpu.memory.total"] == 16 * 1024**3
+    assert "gpu.utilization" not in attributes
+
+
+@pytest.mark.parametrize("is_mig", [True, False])
+def test_mig_memory_never_substitutes_parent_or_permission_failure(monkeypatch, is_mig):
+    def unavailable(handle):
+        if not is_mig:
+            pytest.fail("queried parent memory")
+        raise RuntimeError("Insufficient Permissions")
+
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: None,
+        nvmlDeviceGetHandleByUUID=lambda uuid: "handle",
+        nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
+        nvmlDeviceGetMemoryInfo=unavailable,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    samples = ci_gpu.query_mig_samples(["MIG-a"], 1234)
+    if is_mig:
+        assert samples[0]["attributes"]["gpu.status"] == "unavailable_mig_memory"
+        assert "gpu.memory.used" not in samples[0]["attributes"]
+    else:
+        assert samples == []
+
+
 def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
     monkeypatch, tmp_path
 ):
@@ -830,7 +910,11 @@ def test_gpu_batches_keep_command_identity_and_survive_spool_roundtrip(
     assert b"ci.gpu.sample" in ci_otel.encode_request([span])
 
 
-def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(monkeypatch):
+@pytest.mark.parametrize("mig", [False, True])
+def test_gpu_query_failures_are_bounded_and_never_synthesize_idle_samples(
+    monkeypatch, mig
+):
+    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "MIG-a" if mig else "0")
     calls = []
 
     def query(*args, **kwargs):
@@ -918,3 +1002,112 @@ def test_gpu_sampler_stops_with_command_and_preserves_failure_status(
             command["start_ns"] <= event["time_ns"] <= command["end_ns"]
             for event in batches[0]["events"]
         )
+
+
+@pytest.mark.parametrize("is_mig", [True, False])
+def test_cdi_discovery_uses_only_cuda_visible_uuids(monkeypatch, is_mig):
+    identifiers = []
+
+    def get_count(output):
+        output._obj.value = 1
+        return 0
+
+    def get_device(output, index):
+        assert index == 0
+        output._obj.value = 7
+        return 0
+
+    def get_uuid(output, device):
+        assert device.value == 7
+        output._obj[:] = bytes(range(16))
+        return 0
+
+    driver = SimpleNamespace(
+        cuInit=lambda flags: 0,
+        cuDeviceGetCount=get_count,
+        cuDeviceGet=get_device,
+        cuDeviceGetUuid_v2=get_uuid,
+    )
+    monkeypatch.setattr(ci_gpu.ctypes, "CDLL", lambda _: driver)
+    nvml = SimpleNamespace(
+        nvmlInit=lambda: None,
+        nvmlShutdown=lambda: None,
+        nvmlDeviceGetHandleByUUID=lambda value: identifiers.append(value) or "handle",
+        nvmlDeviceIsMigDeviceHandle=lambda handle: is_mig,
+        NVMLError=RuntimeError,
+    )
+    monkeypatch.setattr(ci_gpu, "load_nvml", lambda: nvml)
+    expected = "MIG-00010203-0405-0607-0809-0a0b0c0d0e0f"
+    assert ci_gpu.cuda_visible_mig_devices() == ([expected] if is_mig else [])
+    assert identifiers == [expected]
+
+
+@pytest.mark.parametrize("discovery_fails", [False, True])
+def test_cdi_sampling_switches_to_slice_or_keeps_unknown(monkeypatch, discovery_fails):
+    monkeypatch.setenv("CI_INFRA_TRACE_ID", "01" * 16)
+    monkeypatch.setenv("CI_INFRA_COMMAND_SPAN_ID", "02" * 8)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
+    monkeypatch.setenv("NVIDIA_VISIBLE_DEVICES", "void")
+    monkeypatch.setattr(ci_gpu, "INTERVAL_SECONDS", 0)
+    stop = threading.Event()
+    commands = []
+    records = []
+
+    def query(command, **kwargs):
+        assert kwargs["timeout"] == ci_gpu.QUERY_TIMEOUT_SECONDS
+        commands.append(command)
+        if "--discover-mig" in command:
+            if discovery_fails:
+                raise subprocess.TimeoutExpired(command, 2)
+            return SimpleNamespace(stdout='["MIG-assigned"]')
+        if len(commands) == 4:
+            stop.set()
+        if "--query-mig" in command:
+            assert command[-1] == "MIG-assigned"
+            return SimpleNamespace(
+                stdout=json.dumps(
+                    [
+                        {
+                            "time_ns": 2,
+                            "name": "ci.gpu.sample",
+                            "attributes": {
+                                "gpu.uuid": "MIG-assigned",
+                                "gpu.memory.used": 123,
+                            },
+                        }
+                    ]
+                )
+            )
+        return SimpleNamespace(stdout="0, GPU-parent, H200, 99, 100, 200, Enabled")
+
+    monkeypatch.setattr(ci_gpu.subprocess, "run", query)
+    monkeypatch.setattr(ci_gpu, "record_spans", lambda spans: records.extend(spans))
+    ci_gpu.collect(os.getppid(), stop)
+    assert sum("--discover-mig" in command for command in commands) == 1
+    attributes = [event["attributes"] for span in records for event in span.events]
+    if discovery_fails:
+        assert all("gpu.memory.used" not in event for event in attributes)
+    else:
+        assert any(event.get("gpu.memory.used") == 123 for event in attributes)
+        assert all(event["gpu.uuid"] == "MIG-assigned" for event in attributes)
+
+
+def test_bundled_nvml_load_does_not_import_vllm_or_torch(tmp_path):
+    # Match test images: the collector is outside the installed vLLM package.
+    helper = tmp_path / "ci_gpu_relocated.py"
+    helper.write_text((SCRIPTS_DIR / "ci_gpu.py").read_text())
+    result = subprocess.run(
+        [
+            sys.executable,
+            "-c",
+            "import sys; import ci_gpu_relocated as ci_gpu; "
+            "binding = ci_gpu.load_nvml(); "
+            "assert callable(binding.nvmlDeviceGetMemoryInfo); "
+            "assert 'vllm' not in sys.modules; assert 'torch' not in sys.modules",
+        ],
+        env={**os.environ, "PYTHONPATH": f"{tmp_path}{os.pathsep}{SCRIPTS_DIR}"},
+        capture_output=True,
+        text=True,
+        timeout=5,
+    )
+    assert result.returncode == 0, result.stderr

.buildkite/test_areas/cuda.yaml

@@ -30,8 +30,8 @@ steps:
       - vllm/
       - tests/cuda
 
-- label: ":nvidia: (H200 MIG 35GB) CUDAGraph"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) CUDAGraph"
+  device: h200_18gb
   key: cudagraph
   timeout_in_minutes: 30
   source_file_dependencies:

.buildkite/test_areas/engine.yaml

@@ -74,18 +74,20 @@ steps:
       depends_on:
       - image-build-amd
 
-- label: ":nvidia: (H200 MIG 35GB) E2E Core"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) E2E Core"
+  device: h200_18gb
   key: e2e-core-1-gpu
   timeout_in_minutes: 50
   source_file_dependencies:
     - vllm/v1/
     - tests/v1/e2e/general/
   commands:
-    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
+    - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
   mirror:
     amd:
       label: ":amd: (MI250) E2E Core"
+      commands:
+      - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py
       device: mi250_1
       timeout_in_minutes: 50
       depends_on:
@@ -95,6 +97,18 @@ steps:
       - tests/v1/e2e/general/
       - vllm/platforms/rocm.py
 
+# Keep the large Gemma-3n and Mamba MRV2 warmup allocations on 35GB slices.
+- label: ":nvidia: (H200 MIG 35GB) E2E Core Large Memory"
+  device: h200_35gb
+  key: e2e-core-large-memory
+  timeout_in_minutes: 50
+  source_file_dependencies:
+    - vllm/v1/
+    - tests/v1/e2e/general/
+  commands:
+    - pytest -v -s v1/e2e/general/test_kv_sharing_fast_prefill.py
+    - pytest -v -s v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2_async
+
 - label: ":nvidia: (L4) V1 E2E"
   key: v1-e2e-2-gpus
   timeout_in_minutes: 25 # TODO: Fix timeout after we have more confidence in the test stability

.buildkite/test_areas/misc.yaml

@@ -103,8 +103,8 @@ steps:
       depends_on:
       - image-build-amd
 
-- label: ":nvidia: (H200 MIG 35GB) V1 Executor + Worker"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) V1 Executor + Worker"
+  device: h200_18gb
   key: v1-executor-worker
   timeout_in_minutes: 45
   source_file_dependencies: *v1-core-kv-metrics-deps

.buildkite/test_areas/samplers.yaml

@@ -2,8 +2,8 @@ group: Samplers
 depends_on: 
   - image-build
 steps:
-- label: ":nvidia: (H200 MIG 35GB) Samplers"
-  device: h200_35gb
+- label: ":nvidia: (H200 MIG 18GB) Samplers"
+  device: h200_18gb
   key: samplers-test
   timeout_in_minutes: 60
   source_file_dependencies:
@@ -15,8 +15,8 @@ steps:
   commands:
     # VLLM_USE_FLASHINFER_SAMPLER defaults to 1 now, so we need to pin both
     # values explicitly to still cover the PyTorch-native (Triton) path.
-    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers
-    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers
+    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
+    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers -k "not test_beam_search_passes_multimodal_data"
   mirror:
     amd:
       label: ":amd: (MI250) Samplers"
@@ -35,3 +35,18 @@ steps:
       - vllm/platforms/rocm.py
       commands:
       - pytest -v -s samplers
+
+# Qwen2-Audio-7B leaves no room for KV cache on a 16 GiB slice.
+- label: ":nvidia: (H200 MIG 35GB) Samplers Multimodal Beam Search"
+  device: h200_35gb
+  key: samplers-multimodal-beam-search
+  timeout_in_minutes: 30
+  source_file_dependencies:
+  - vllm/model_executor/layers
+  - vllm/sampling_metadata.py
+  - tests/samplers
+  - tests/conftest.py
+  - vllm/entrypoints/generate/beam_search
+  commands:
+    - VLLM_USE_FLASHINFER_SAMPLER=0 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data
+    - VLLM_USE_FLASHINFER_SAMPLER=1 pytest -v -s samplers/test_beam_search.py::test_beam_search_passes_multimodal_data

requirements/test/xpu.in

@@ -32,8 +32,8 @@ soundfile
 # --- Tool Parsing & Evaluation ---
 blobfile
 rapidfuzz
-gpt-oss[eval]==0.0.5
-schemathesis
+gpt-oss >= 0.0.7; python_version > '3.11'
+schemathesis>=4.0.0 # Required for openai schema test.
 jiwer
 bm25s
 pystemmer
@@ -42,12 +42,12 @@ num2words
 pqdm
 
 # --- Vision & Multimodal ---
-timm
+timm>=1.0.17 # required for internvl and gemma3n-mm test
 albumentations
 ftfy
 matplotlib
-decord==0.6.0
-mistral-common[image,audio]
+decord==0.6.0; platform_machine == "x86_64"
+mistral-common[image,audio]>=1.11.6 # required for voxtral test
 
 open_clip_torch==2.32.0 # Required for nemotron_vl test, Nemotron Parse in test_common.py
 

requirements/test/xpu.txt

@@ -67,9 +67,7 @@ azure-storage-blob==12.30.1
 blake3==1.0.9
     # via -r requirements/test/../common.txt
 blobfile==3.2.0
-    # via
-    #   -r requirements/test/xpu.in
-    #   gpt-oss
+    # via -r requirements/test/xpu.in
 bm25s==0.3.10
     # via
     #   -r requirements/test/xpu.in
@@ -259,7 +257,7 @@ googleapis-common-protos==1.75.1
     #   google-api-core
     #   opentelemetry-exporter-otlp-proto-grpc
     #   opentelemetry-exporter-otlp-proto-http
-gpt-oss==0.0.5
+gpt-oss==0.0.9
     # via -r requirements/test/xpu.in
 graphql-core==3.2.11
     # via hypothesis-graphql
@@ -379,7 +377,6 @@ jinja2==3.1.6
     # via
     #   -c requirements/xpu.txt
     #   fastapi
-    #   gpt-oss
     #   lm-eval
     #   torch
 jiter==0.16.0
@@ -409,6 +406,10 @@ jsonschema-rs==0.49.9
     # via schemathesis
 jsonschema-specifications==2025.9.1
     # via jsonschema
+jupyter-client==8.10.0
+    # via gpt-oss
+jupyter-core==5.9.1
+    # via jupyter-client
 kiwisolver==1.5.1
     # via matplotlib
 lark==1.2.2
@@ -529,7 +530,6 @@ numpy==2.2.6
     #   datasets
     #   decord
     #   evaluate
-    #   gpt-oss
     #   librosa
     #   lm-eval
     #   matplotlib
@@ -580,7 +580,6 @@ openai==3.2.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
-    #   gpt-oss
 openai-harmony==0.0.8
     # via
     #   -c requirements/common.txt
@@ -662,7 +661,6 @@ pandas==3.0.5
     # via
     #   datasets
     #   evaluate
-    #   gpt-oss
 partial-json-parser==0.2.1.1.post7
     # via -r requirements/test/../common.txt
 pathvalidate==3.3.1
@@ -674,7 +672,9 @@ pillow==12.3.0
     #   mistral-common
     #   torchvision
 platformdirs==4.11.3
-    # via pooch
+    # via
+    #   jupyter-core
+    #   pooch
 pluggy==1.6.0
     # via
     #   pytest
@@ -813,6 +813,7 @@ pytest-timeout==2.4.0
 python-dateutil==2.9.0.post0
     # via
     #   botocore
+    #   jupyter-client
     #   matplotlib
     #   pandas
     #   typepy
@@ -848,6 +849,7 @@ pyzmq==27.1.0
     # via
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
+    #   jupyter-client
 rapidfuzz==3.14.5
     # via
     #   -r requirements/test/xpu.in
@@ -1072,12 +1074,13 @@ torchvision==0.28.0+xpu
     # via
     #   open-clip-torch
     #   timm
+tornado==6.5.8
+    # via jupyter-client
 tqdm==4.70.0
     # via
     #   -r requirements/test/../common.txt
     #   datasets
     #   evaluate
-    #   gpt-oss
     #   huggingface-hub
     #   lm-eval
     #   modelscope
@@ -1088,6 +1091,10 @@ tqdm==4.70.0
     #   pqdm
     #   sentence-transformers
     #   transformers
+traitlets==5.16.1
+    # via
+    #   jupyter-client
+    #   jupyter-core
 transformers==5.16.1
     # via
     #   -c requirements/common.txt
@@ -1139,6 +1146,7 @@ typing-extensions==4.16.0
     #   grpcio
     #   httpx2
     #   huggingface-hub
+    #   jupyter-client
     #   lm-eval
     #   mcp
     #   mcp-types

rust/README.md

@@ -43,6 +43,17 @@ For example:
 VLLM_USE_RUST_FRONTEND=1 vllm serve Qwen/Qwen3-0.6B
 ```
 
+### RL weight synchronization
+
+With `VLLM_SERVER_DEV_MODE=1`, the Rust frontend supports the HTTP weight-transfer
+lifecycle used by `HTTPVLLMWeightSyncClient`: initialization, starting an update,
+transferring weights, and finishing the update. It also supports draft-model
+updates and `/update_weight_version` and `/weight_info` for version tracking.
+
+For trainer-side usage, see the [IPC](../examples/rl/rlhf_http_ipc.py) and
+[NCCL](../examples/rl/rlhf_http_nccl.py) examples. Configure the weight-transfer
+backend on the server and pause generation during updates, as shown there.
+
 ### External Engine
 
 `vllm-rs serve` can be run standalone with `--data-parallel-size-local 0` when the Python engines

rust/src/server/src/error.rs

@@ -2,6 +2,7 @@
 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 use axum::Json;
+use axum::extract::rejection::JsonRejection;
 use axum::http::StatusCode;
 use axum::response::{IntoResponse, Response};
 use thiserror_ext::{AsReport as _, Construct, Macro};
@@ -70,6 +71,12 @@ impl ApiError {
     }
 }
 
+impl From<JsonRejection> for ApiError {
+    fn from(error: JsonRejection) -> Self {
+        Self::json_parse_error(error.body_text())
+    }
+}
+
 impl IntoResponse for ApiError {
     fn into_response(self) -> Response {
         (self.status_code(), Json(self.to_error_response())).into_response()

rust/src/server/src/routes.rs

@@ -17,6 +17,7 @@ mod server_info;
 mod sleep;
 mod tokenize;
 mod version;
+mod weight_transfer;
 mod world_size;
 
 use std::sync::Arc;
@@ -139,6 +140,28 @@ fn build_router_with_options(
             .route("/reset_mm_cache", post(cache::reset_mm_cache))
             .route("/reset_encoder_cache", post(cache::reset_encoder_cache))
             .route("/collective_rpc", post(collective_rpc::collective_rpc))
+            .route(
+                "/init_weight_transfer_engine",
+                post(weight_transfer::init_weight_transfer_engine),
+            )
+            .route(
+                "/start_weight_update",
+                post(weight_transfer::start_weight_update),
+            )
+            .route(
+                "/start_draft_weight_update",
+                post(weight_transfer::start_draft_weight_update),
+            )
+            .route("/update_weights", post(weight_transfer::update_weights))
+            .route(
+                "/finish_weight_update",
+                post(weight_transfer::finish_weight_update),
+            )
+            .route(
+                "/update_weight_version",
+                post(weight_transfer::update_weight_version),
+            )
+            .route("/weight_info", get(weight_transfer::weight_info))
             .route("/abort_requests", post(abort_requests::abort_requests))
             .route("/sleep", post(sleep::sleep))
             .route("/wake_up", post(sleep::wake_up))

rust/src/server/src/routes/abort_requests.rs

@@ -22,7 +22,7 @@ pub async fn abort_requests(
     State(state): State<Arc<AppState>>,
     body: Result<Json<AbortRequestsRequest>, JsonRejection>,
 ) -> Result<StatusCode, ApiError> {
-    let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
+    let Json(body) = body?;
     // Empty/missing `request_ids` aborts all in-flight requests.
     let request_ids = body.request_ids.unwrap_or_default();
 

rust/src/server/src/routes/collective_rpc.rs

@@ -36,7 +36,7 @@ pub async fn collective_rpc(
     State(state): State<Arc<AppState>>,
     body: Result<Json<CollectiveRpcRequest>, JsonRejection>,
 ) -> Result<Json<CollectiveRpcResponse>, ApiError> {
-    let Json(body) = body.map_err(|error| ApiError::json_parse_error(error.body_text()))?;
+    let Json(body) = body?;
     let method = body.method.ok_or_else(|| {
         ApiError::invalid_request(
             "Missing 'method' in request body".to_string(),

rust/src/server/src/routes/openai/utils/validated_json.rs

@@ -5,7 +5,6 @@
 //! Variation of https://github.com/lightseekorg/smg/blob/main/crates/protocols/src/validated.rs
 
 use axum::Json;
-use axum::extract::rejection::JsonRejection;
 use axum::extract::{FromRequest, Request};
 use serde::de::DeserializeOwned;
 use validator::Validate;
@@ -30,9 +29,7 @@ where
     type Rejection = ApiError;
 
     async fn from_request(req: Request, state: &S) -> Result<Self, Self::Rejection> {
-        let Json(mut data) = Json::<T>::from_request(req, state)
-            .await
-            .map_err(|err: JsonRejection| ApiError::json_parse_error(err.body_text()))?;
+        let Json(mut data) = Json::<T>::from_request(req, state).await?;
 
         data.normalize();
 

rust/src/server/src/routes/tests.rs

@@ -6113,6 +6113,13 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
         ("POST", "/reset_prefix_cache"),
         ("POST", "/reset_mm_cache"),
         ("POST", "/reset_encoder_cache"),
+        ("POST", "/init_weight_transfer_engine"),
+        ("POST", "/start_weight_update"),
+        ("POST", "/start_draft_weight_update"),
+        ("POST", "/update_weights"),
+        ("POST", "/finish_weight_update"),
+        ("POST", "/update_weight_version"),
+        ("GET", "/weight_info"),
     ] {
         let response = app
             .clone()
@@ -6132,6 +6139,248 @@ async fn admin_routes_are_hidden_when_dev_mode_is_disabled() {
     engine_task.abort_and_join().await;
 }
 
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn weight_transfer_routes_support_the_http_training_lifecycle() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let mut calls = Vec::new();
+            for _ in 0..10 {
+                let utility = recv_engine_message(dealer).await;
+                assert_eq!(utility[0].as_ref(), &[0x03]);
+                let payload = decode_value(&utility[1]).expect("decode utility");
+                let array = payload.as_array().expect("utility array");
+                let call_id = array[1].as_u64().expect("call id");
+                let args: serde_json::Value =
+                    rmpv::ext::from_value(array[3].clone()).expect("decode args");
+                calls.push(format!("{} {args}", array[2].as_str().expect("method")));
+                let result = if array[2] == Value::from("get_weight_version") {
+                    utility_result_value("step-2")
+                } else if array[2] == Value::from("collective_rpc") {
+                    utility_result_value(vec![(), ()])
+                } else {
+                    utility_none_result()
+                };
+                send_outputs(push, utility_outputs(call_id, result)).await;
+            }
+            expect_test::expect![[r#"
+                [
+                    "collective_rpc [\"init_weight_transfer_engine\",null,[],{\"init_info\":{\"master_address\":\"127.0.0.1\",\"master_port\":29500,\"rank_offset\":1,\"world_size\":3}}]",
+                    "collective_rpc [\"start_weight_update\",null,[],{}]",
+                    "collective_rpc [\"update_weights\",null,[],{\"update_info\":{\"names\":[\"model.embed_tokens.weight\"],\"dtype_names\":[\"bfloat16\"],\"shapes\":[[32,16]],\"is_checkpoint_format\":true}}]",
+                    "collective_rpc [\"update_weights\",null,[],{\"update_info\":[{\"rank\":0,\"ipc_handles_pickled\":\"opaque-handle-0\"},{\"rank\":1,\"ipc_handles_pickled\":\"opaque-handle-1\"}]}]",
+                    "collective_rpc [\"finish_weight_update\",null,[],{}]",
+                    "set_weight_version [\"step-1\"]",
+                    "collective_rpc [\"start_draft_weight_update\",null,[],{}]",
+                    "collective_rpc [\"finish_weight_update\",null,[],{}]",
+                    "set_weight_version [\"step-2\"]",
+                    "get_weight_version []",
+                ]
+            "#]].assert_debug_eq(&calls);
+        })
+    })
+    .await;
+
+    let mut responses = Vec::new();
+    for (method, path, body) in [
+        (
+            "POST",
+            "/init_weight_transfer_engine",
+            Some(json!({
+                "init_info": {"master_address": "127.0.0.1", "master_port": 29500,
+                    "rank_offset": 1, "world_size": 3}
+            })),
+        ),
+        ("POST", "/start_weight_update", None),
+        (
+            "POST",
+            "/update_weights",
+            Some(json!({
+                "update_info": {"names": ["model.embed_tokens.weight"],
+                    "dtype_names": ["bfloat16"], "shapes": [[32, 16]],
+                    "is_checkpoint_format": true}
+            })),
+        ),
+        (
+            "POST",
+            "/update_weights",
+            Some(json!({
+                "update_info": [{"rank": 0, "ipc_handles_pickled": "opaque-handle-0"},
+                    {"rank": 1, "ipc_handles_pickled": "opaque-handle-1"}]
+            })),
+        ),
+        (
+            "POST",
+            "/finish_weight_update",
+            Some(json!({"weight_version": "step-1"})),
+        ),
+        ("POST", "/start_draft_weight_update", None),
+        ("POST", "/finish_weight_update", None),
+        (
+            "POST",
+            "/update_weight_version",
+            Some(json!({"new_version": "step-2"})),
+        ),
+        ("GET", "/weight_info", None),
+    ] {
+        let mut request = Request::builder().method(method).uri(path);
+        let body = match body {
+            Some(body) => {
+                request = request.header("content-type", "application/json");
+                Body::from(body.to_string())
+            }
+            None => Body::empty(),
+        };
+        let response =
+            app.call(request.body(body).expect("build request")).await.expect("call app");
+        assert_eq!(response.status(), StatusCode::OK, "{path}");
+        let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+        responses.push(
+            serde_json::from_slice::<serde_json::Value>(&body)
+                .expect("decode json")
+                .to_string(),
+        );
+    }
+    engine_task.await.expect("mock engine task");
+    expect_test::expect![[r#"
+        [
+            "{\"message\":\"Weight transfer initialized\"}",
+            "{\"message\":\"Weight update started\"}",
+            "{\"message\":\"Weights updated\"}",
+            "{\"message\":\"Weights updated\"}",
+            "{\"message\":\"Weight update finished\"}",
+            "{\"message\":\"Draft weight update started\"}",
+            "{\"message\":\"Weight update finished\"}",
+            "{\"success\":true,\"new_version\":\"step-2\"}",
+            "{\"weight_version\":\"step-2\"}",
+        ]
+    "#]]
+    .assert_debug_eq(&responses);
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn weight_transfer_routes_reject_invalid_payloads_before_engine_calls() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, _push| {
+        boxed_test_future(async move {
+            let message = recv_engine_message(dealer).await;
+            panic!("invalid payload reached engine: {message:?}");
+        })
+    })
+    .await;
+
+    for (path, body) in [
+        ("/init_weight_transfer_engine", "{"),
+        ("/init_weight_transfer_engine", "{}"),
+        ("/init_weight_transfer_engine", r#"{"init_info":null}"#),
+        ("/init_weight_transfer_engine", r#"{"init_info":[]}"#),
+        ("/update_weights", "{"),
+        ("/update_weights", "{}"),
+        ("/update_weights", r#"{"update_info":null}"#),
+        ("/update_weights", r#"{"update_info":1}"#),
+        ("/update_weights", r#"{"update_info":[{},null]}"#),
+        ("/finish_weight_update", "{"),
+        ("/finish_weight_update", r#"{"weight_version":1}"#),
+        ("/update_weight_version", "{}"),
+        ("/update_weight_version", r#"{"new_version":null}"#),
+    ] {
+        let response = app
+            .call(
+                Request::builder()
+                    .method("POST")
+                    .uri(path)
+                    .header("content-type", "application/json")
+                    .body(Body::from(body))
+                    .expect("build request"),
+            )
+            .await
+            .expect("call app");
+        assert_eq!(response.status(), StatusCode::BAD_REQUEST, "{path}: {body}");
+    }
+    engine_task.abort_and_join().await;
+}
+
+#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
+#[serial]
+async fn weight_transfer_routes_do_not_commit_version_when_finish_fails() {
+    let (mut app, engine_task) = test_admin_app_with_engine_script(|dealer, push| {
+        boxed_test_future(async move {
+            let utility = recv_engine_message(dealer).await;
+            let payload = decode_value(&utility[1]).expect("decode utility");
+            let array = payload.as_array().expect("utility array");
+            assert_eq!(array[2], Value::from("collective_rpc"));
+            assert_eq!(
+                array[3].as_array().expect("args")[0],
+                Value::from("finish_weight_update")
+            );
+            send_outputs(
+                push,
+                UtilityCallOutput {
+                    output: UtilityOutput {
+                        call_id: array[1].as_u64().expect("call id").into(),
+                        failure_message: Some("weight reload failed".to_string()),
+                        result: None,
+                    },
+                    ..Default::default()
+                }
+                .into(),
+            )
+            .await;
+            let utility = recv_engine_message(dealer).await;
+            let payload = decode_value(&utility[1]).expect("decode utility");
+            let array = payload.as_array().expect("utility array");
+            assert_eq!(array[2], Value::from("get_weight_version"));
+            send_outputs(
+                push,
+                utility_outputs(
+                    array[1].as_u64().expect("call id"),
+                    utility_result_value("step-0"),
+                ),
+            )
+            .await;
+        })
+    })
+    .await;
+
+    let response = app
+        .call(
+            Request::builder()
+                .method("POST")
+                .uri("/finish_weight_update")
+                .header("content-type", "application/json")
+                .body(Body::from(r#"{"weight_version":"step-1"}"#))
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    let error: serde_json::Value = serde_json::from_slice(&body).expect("decode json");
+    assert!(
+        error["error"]["message"]
+            .as_str()
+            .expect("error message")
+            .contains("weight reload failed")
+    );
+
+    let response = app
+        .call(
+            Request::builder()
+                .uri("/weight_info")
+                .body(Body::empty())
+                .expect("build request"),
+        )
+        .await
+        .expect("call app");
+    assert_eq!(response.status(), StatusCode::OK);
+    let body = to_bytes(response.into_body(), usize::MAX).await.expect("read body");
+    assert_eq!(
+        serde_json::from_slice::<serde_json::Value>(&body).expect("decode json"),
+        json!({"weight_version": "step-0"})
+    );
+    engine_task.await.expect("mock engine task");
+}
+
 // ========================= Stop string tests =========================
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]

rust/src/server/src/routes/weight_transfer.rs

@@ -0,0 +1,190 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+use std::sync::Arc;
+
+use axum::Json;
+use axum::extract::State;
+use axum::extract::rejection::JsonRejection;
+use serde::{Deserialize, Serialize};
+use serde_json::Value as JsonValue;
+
+use crate::error::{ApiError, invalid_request};
+use crate::state::AppState;
+use crate::utils::utility_call_error;
+
+#[derive(Deserialize)]
+pub(crate) struct InitWeightTransferRequest {
+    init_info: Option<JsonValue>,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct UpdateWeightsRequest {
+    update_info: Option<JsonValue>,
+}
+
+#[derive(Default, Deserialize)]
+pub(crate) struct FinishWeightUpdateRequest {
+    weight_version: Option<String>,
+}
+
+#[derive(Deserialize)]
+pub(crate) struct UpdateWeightVersionRequest {
+    new_version: String,
+}
+
+#[derive(Serialize)]
+pub(crate) struct MessageResponse {
+    message: &'static str,
+}
+
+#[derive(Serialize)]
+pub(crate) struct UpdateWeightVersionResponse {
+    success: bool,
+    new_version: String,
+}
+
+#[derive(Serialize)]
+pub(crate) struct WeightInfoResponse {
+    weight_version: String,
+}
+
+/// Initialize weight transfer for RL training.
+pub async fn init_weight_transfer_engine(
+    State(state): State<Arc<AppState>>,
+    body: Result<Json<InitWeightTransferRequest>, JsonRejection>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    let Json(body) = body?;
+    let init_info = body.init_info.filter(JsonValue::is_object).ok_or_else(|| {
+        invalid_request!(
+            param = Some("init_info"),
+            "'init_info' must be a JSON object"
+        )
+    })?;
+
+    state
+        .engine_core_client()
+        .init_weight_transfer_engine(init_info)
+        .await
+        .map_err(|error| utility_call_error("init_weight_transfer_engine", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weight transfer initialized",
+    }))
+}
+
+/// Start a weight update transaction.
+pub async fn start_weight_update(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    state
+        .engine_core_client()
+        .start_weight_update()
+        .await
+        .map_err(|error| utility_call_error("start_weight_update", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weight update started",
+    }))
+}
+
+/// Start a draft-model weight update transaction.
+pub async fn start_draft_weight_update(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    state
+        .engine_core_client()
+        .start_draft_weight_update()
+        .await
+        .map_err(|error| utility_call_error("start_draft_weight_update", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Draft weight update started",
+    }))
+}
+
+/// Update model weights with backend-specific transfer metadata.
+pub async fn update_weights(
+    State(state): State<Arc<AppState>>,
+    body: Result<Json<UpdateWeightsRequest>, JsonRejection>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    let Json(body) = body?;
+    let update_info = body
+        .update_info
+        .filter(|info| {
+            info.is_object()
+                || info.as_array().is_some_and(|items| items.iter().all(JsonValue::is_object))
+        })
+        .ok_or_else(|| {
+            invalid_request!(
+                param = Some("update_info"),
+                "'update_info' must be a JSON object or a list of per-worker JSON objects",
+            )
+        })?;
+
+    state
+        .engine_core_client()
+        .update_weights(update_info)
+        .await
+        .map_err(|error| utility_call_error("update_weights", error))?;
+
+    Ok(Json(MessageResponse {
+        message: "Weights updated",
+    }))
+}
+
+/// Finish a weight update transaction and optionally set the weight version.
+pub async fn finish_weight_update(
+    State(state): State<Arc<AppState>>,
+    body: Result<Option<Json<FinishWeightUpdateRequest>>, JsonRejection>,
+) -> Result<Json<MessageResponse>, ApiError> {
+    // HTTPVLLMWeightSyncClient omits the body when no version is supplied.
+    let request = body?.map(|Json(request)| request).unwrap_or_default();
+
+    let client = state.engine_core_client();
+    client
+        .finish_weight_update()
+        .await
+        .map_err(|error| utility_call_error("finish_weight_update", error))?;
+    if let Some(version) = request.weight_version {
+        client
+            .set_weight_version(&version)
+            .await
+            .map_err(|error| utility_call_error("set_weight_version", error))?;
+    }
+
+    Ok(Json(MessageResponse {
+        message: "Weight update finished",
+    }))
+}
+
+/// Update the weight version identifier.
+pub async fn update_weight_version(
+    State(state): State<Arc<AppState>>,
+    body: Result<Json<UpdateWeightVersionRequest>, JsonRejection>,
+) -> Result<Json<UpdateWeightVersionResponse>, ApiError> {
+    let Json(body) = body?;
+    state
+        .engine_core_client()
+        .set_weight_version(&body.new_version)
+        .await
+        .map_err(|error| utility_call_error("set_weight_version", error))?;
+
+    Ok(Json(UpdateWeightVersionResponse {
+        success: true,
+        new_version: body.new_version,
+    }))
+}
+
+/// Get current weight metadata.
+pub async fn weight_info(
+    State(state): State<Arc<AppState>>,
+) -> Result<Json<WeightInfoResponse>, ApiError> {
+    let weight_version = state
+        .engine_core_client()
+        .get_weight_version()
+        .await
+        .map_err(|error| utility_call_error("get_weight_version", error))?;
+
+    Ok(Json(WeightInfoResponse { weight_version }))
+}

tests/kernels/attention/test_minimax_m3.py

@@ -1025,17 +1025,20 @@ def test_decode_index_topk_correctness(
         "expected",
     ),
     [
-        (0, 24, 1, 16, True, (2, False, False)),
-        (1, 1, 1, 16, True, (16, True, True)),
-        (512, 4, 1, 16, True, (16, True, True)),
-        (513, 24, 4, 16, True, (16, True, True)),
-        (2760, 5, 2, 16, True, (16, True, True)),
-        (8192, 64, 4, 16, True, (16, True, True)),
-        (8193, 4, 1, 16, True, (16, False, False)),
-        (8193, 16, 1, 16, True, (4, False, False)),
-        (65534, 64, 1, 16, True, (1, False, False)),
-        (513, 16, 1, 13, True, (4, False, False)),
-        (513, 16, 4, 16, False, (1, False, False)),
+        (0, 24, 1, 16, True, (2, 512, 8, 2, False, False)),
+        (1, 1, 1, 16, True, (1, 128, 2, 1, True, True)),
+        (74, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
+        (128, 4, 1, 16, True, (1, 128, 2, 1, True, True)),
+        (129, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
+        (512, 4, 1, 16, True, (16, 512, 4, 2, True, True)),
+        (513, 24, 4, 16, True, (16, 512, 4, 2, True, True)),
+        (2760, 5, 2, 16, True, (16, 512, 4, 2, True, True)),
+        (8192, 64, 4, 16, True, (16, 512, 4, 2, True, True)),
+        (8193, 4, 1, 16, True, (16, 512, 8, 2, False, False)),
+        (8193, 16, 1, 16, True, (4, 512, 8, 2, False, False)),
+        (65534, 64, 1, 16, True, (1, 512, 8, 2, False, False)),
+        (513, 16, 1, 13, True, (4, 512, 8, 2, False, False)),
+        (513, 16, 4, 16, False, (1, 512, 8, 2, False, False)),
     ],
 )
 def test_amd_decode_topk_launch_policy(
@@ -1044,9 +1047,9 @@ def test_amd_decode_topk_launch_policy(
     num_idx_heads: int,
     topk: int,
     is_gfx950: bool,
-    expected: tuple[int, bool, bool],
+    expected: tuple[int, int, int, int, bool, bool],
 ):
-    """All supported head and query layouts use adaptive selection."""
+    """The short-context specialization has strict dispatch boundaries."""
     from vllm.models.minimax_m3.amd.ops.index_topk import (
         _decode_topk_launch_policy,
     )
@@ -1073,19 +1076,23 @@ def test_amd_decode_topk_launch_policy(
         "single_tile_guaranteed",
         "adaptive_final_merge",
         "max_blocks",
+        "block_size_k",
+        "num_warps",
+        "num_stages",
         "topk",
         "num_idx_heads",
     ),
     [
-        (1, False, False, 65534, 16, 1),
-        (2, False, False, 8193, 16, 1),
-        (4, False, False, 8193, 16, 1),
-        (8, False, False, 8193, 16, 1),
-        (16, False, False, 8193, 16, 1),
-        (16, False, False, 65534, 16, 1),
-        (2, True, False, 513, 13, 1),
-        (4, False, True, 8192, 16, 2),
-        (16, True, True, 8192, 16, 4),
+        (1, False, False, 65534, 512, 8, 2, 16, 1),
+        (2, False, False, 8193, 512, 8, 2, 16, 1),
+        (4, False, False, 8193, 512, 8, 2, 16, 1),
+        (8, False, False, 8193, 512, 8, 2, 16, 1),
+        (16, False, False, 8193, 512, 8, 2, 16, 1),
+        (16, False, False, 65534, 512, 8, 2, 16, 1),
+        (2, True, False, 513, 512, 8, 2, 13, 1),
+        (4, False, True, 8192, 512, 4, 2, 16, 2),
+        (16, True, True, 8192, 512, 4, 2, 16, 4),
+        (1, True, True, 74, 128, 2, 1, 16, 1),
     ],
 )
 def test_amd_decode_fused_topk_total_order_and_replay(
@@ -1093,10 +1100,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
     single_tile_guaranteed: bool,
     adaptive_final_merge: bool,
     max_blocks: int,
+    block_size_k: int,
+    num_warps: int,
+    num_stages: int,
     topk: int,
     num_idx_heads: int,
 ):
-    """Atomic multi-chunk selection is ordered and resets graph state."""
+    """Selector specializations preserve ordering and graph state."""
     from vllm.models.minimax_m3.amd.ops.index_topk import (
         SPARSE_BLOCK_SIZE,
         _decode_topk_fused_kernel,
@@ -1104,18 +1114,25 @@ def test_amd_decode_fused_topk_total_order_and_replay(
 
     set_random_seed(0)
     block_size_t = 1 << (topk - 1).bit_length()
+    block_counts = tuple(
+        count
+        for count in (
+            15,
+            16,
+            17,
+            block_size_k // 2 - 1,
+            block_size_k // 2,
+            block_size_k // 2 + 1,
+            block_size_k - 1,
+            block_size_k,
+            block_size_k + 1,
+        )
+        if count <= max_blocks
+    )
     long_seq_lens_list = [
         0,
         1,
-        15 * SPARSE_BLOCK_SIZE,
-        16 * SPARSE_BLOCK_SIZE,
-        17 * SPARSE_BLOCK_SIZE,
-        255 * SPARSE_BLOCK_SIZE,
-        256 * SPARSE_BLOCK_SIZE,
-        257 * SPARSE_BLOCK_SIZE,
-        511 * SPARSE_BLOCK_SIZE,
-        512 * SPARSE_BLOCK_SIZE,
-        513 * SPARSE_BLOCK_SIZE,
+        *(blocks * SPARSE_BLOCK_SIZE for blocks in block_counts),
         max_blocks * SPARSE_BLOCK_SIZE,
     ]
     short_seq_lens_list = [0, 1, 127, 128, 129, 0, 1, 128, 129, 0, 1, 129]
@@ -1126,12 +1143,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
         boundary_limit = 4
     if boundary_limit:
         boundary_blocks = [
-            (active_chunks - 1) * 512 + 1
+            (active_chunks - 1) * block_size_k + 1
             for active_chunks in range(3, boundary_limit + 1)
         ]
-        assert [(num_blocks + 511) // 512 for num_blocks in boundary_blocks] == list(
-            range(3, boundary_limit + 1)
-        )
+        assert [
+            (num_blocks + block_size_k - 1) // block_size_k
+            for num_blocks in boundary_blocks
+        ] == list(range(3, boundary_limit + 1))
         long_seq_lens_list.extend(
             num_blocks * SPARSE_BLOCK_SIZE for num_blocks in boundary_blocks
         )
@@ -1139,6 +1157,8 @@ def test_amd_decode_fused_topk_total_order_and_replay(
             num_blocks * SPARSE_BLOCK_SIZE for num_blocks in reversed(boundary_blocks)
         )
     batch = len(long_seq_lens_list)
+    repeats = (batch + len(short_seq_lens_list) - 1) // len(short_seq_lens_list)
+    short_seq_lens_list = (short_seq_lens_list * repeats)[:batch]
     score = torch.randn((num_idx_heads, batch, max_blocks), device="cuda")
     score[:, :, 0] = float("nan")
     score[:, :, 1] = float("inf")
@@ -1147,8 +1167,9 @@ def test_amd_decode_fused_topk_total_order_and_replay(
     score[:, :, 4] = -0.0
     score[:, :, 5] = 7.0
     score[:, :, 6] = -1e30
-    score[:, :, 256] = 7.0
-    score[:, :, 512] = 7.0
+    for index in (block_size_k // 2, block_size_k):
+        if index < max_blocks:
+            score[:, :, index] = 7.0
 
     seq_lens = torch.empty(batch, device="cuda", dtype=torch.int32)
     partial = torch.full(
@@ -1198,13 +1219,13 @@ def test_amd_decode_fused_topk_total_order_and_replay(
             pages_per_sparse_block=8,
             block_page_stride=8,
             NUM_TOPK_CHUNKS=num_topk_chunks,
-            BLOCK_SIZE_K=512,
+            BLOCK_SIZE_K=block_size_k,
             BLOCK_SIZE_T=block_size_t,
             EMIT_SPARSE_TABLE=False,
             SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
             ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
-            num_warps=4 if adaptive_final_merge else 8,
-            num_stages=2,
+            num_warps=num_warps,
+            num_stages=num_stages,
         )
 
     score_cpu = score.cpu()
@@ -1251,6 +1272,7 @@ def test_amd_decode_fused_topk_total_order_and_replay(
         (16, False, False, 8193),
         (4, False, True, 8192),
         (16, True, True, 8192),
+        (1, True, True, 74),
     ):
         graph = torch.cuda.CUDAGraph()
         torch.accelerator.synchronize()

tests/model_executor/layers/test_fused_shared_expert.py

@@ -832,6 +832,7 @@ def test_models_fse_init(
         runner_type="generate",
         is_moe=True,
         logits_processors=None,
+        rswa_window=None,
     )
     vllm_config.parallel_config.enable_expert_parallel = False
     if model_type == "deepseek_v4":

tests/tokenizers_/test_deepseek_v41.py

@@ -99,6 +99,24 @@ def test_raw_text_parts_preserve_reference_separator():
     assert messages == original
 
 
+@pytest.mark.parametrize(
+    ("role", "responses_type"),
+    [("user", "input_text"), ("assistant", "output_text")],
+)
+def test_responses_text_parts_match_chat_text_parts(role, responses_type):
+    responses_message = {
+        "role": role,
+        "content": [{"type": responses_type, "text": "hello"}],
+    }
+    chat_message = {
+        "role": role,
+        "content": [{"type": "text", "text": "hello"}],
+    }
+    assert render([responses_message], thinking=False) == render(
+        [chat_message], thinking=False
+    )
+
+
 def test_mid_system_gets_its_own_marker_and_generation_header():
     assert render(
         [

tests/v1/core/test_kv_cache_utils.py

@@ -35,7 +35,7 @@ from vllm.multimodal.inputs import (
 from vllm.sampling_params import SamplingParams
 from vllm.utils.hashing import sha256, sha256_cbor, xxhash, xxhash_cbor
 from vllm.utils.mem_constants import GiB_bytes
-from vllm.v1.core.kv_cache_manager import KVCacheManager
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
 from vllm.v1.core.kv_cache_utils import (
     BlockHash,
     FreeKVCacheBlockQueue,
@@ -61,6 +61,7 @@ from vllm.v1.hisparse.layout import (
 )
 from vllm.v1.kv_cache_interface import (
     ChunkedLocalAttentionSpec,
+    CircularBufferSpec,
     FullAttentionSpec,
     HiddenStateCacheSpec,
     HiSparseHotSpec,
@@ -324,6 +325,45 @@ def test_kv_cache_config_selects_only_transferable_groups():
     )
 
 
+def test_kv_cache_config_selects_prefix_cacheable_groups():
+    """Prefix stores exclude scratch state without changing transfer groups."""
+    full_group = KVCacheGroupSpec(["full"], new_kv_cache_spec())
+    qsa_group = KVCacheGroupSpec(
+        ["qsa"],
+        CircularBufferSpec(
+            block_size=4,
+            num_kv_heads=1,
+            head_size=64,
+            head_size_v=0,
+            dtype=torch.float16,
+        ),
+    )
+    disabled_group = KVCacheGroupSpec(
+        ["disabled"], new_kv_cache_spec(), enable_kv_transfer=False
+    )
+    config = KVCacheConfig(
+        num_blocks=1,
+        kv_cache_tensors=[],
+        kv_cache_groups=[full_group, qsa_group, disabled_group],
+    )
+    assert config.transfer_group_ids == (0, 1)
+    assert config.select_transfer_block_ids(([1], [2], [3])) == ([1], [2])
+    assert config.prefix_cacheable_group_ids == (0,)
+    assert config.prefix_cacheable_groups == (full_group,)
+
+
+def test_kv_cache_blocks_selects_requested_groups():
+    blocks = KVCacheBlocks(
+        (
+            [KVCacheBlock(1)],
+            [KVCacheBlock(2)],
+            [KVCacheBlock(3)],
+        )
+    )
+
+    assert blocks.get_block_ids(group_ids=(0, 2)) == ([1], [3])
+
+
 def new_sliding_window_spec(
     block_size=16,
     num_kv_heads=2,

tests/v1/kv_connector/unit/test_mooncake_store_connector.py

@@ -30,6 +30,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
     MooncakeStoreConnectorStats,
 )
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
     KVCacheConfig,
     KVCacheGroupSpec,
@@ -91,6 +92,53 @@ def _make_kv_cache_config() -> KVCacheConfig:
     )
 
 
+def _make_qsa_hybrid_kv_cache_config(mamba_mode: str):
+    full_spec = FullAttentionSpec(
+        block_size=800, num_kv_heads=8, head_size=64, dtype=None
+    )
+    qsa_spec = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=64,
+        head_size_v=0,
+        dtype=torch.float16,
+    )
+    mamba_spec = MambaSpec(
+        block_size=800,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode=mamba_mode,
+    )
+    return KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full_spec),
+            KVCacheGroupSpec(["qsa"], qsa_spec),
+            KVCacheGroupSpec(["mamba"], mamba_spec),
+        ],
+    )
+
+
+def test_validation_accepts_aligned_mamba_after_block_size_rewrite():
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 8
+
+    mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
+        vllm_config, _make_qsa_hybrid_kv_cache_config("align")
+    )
+
+
+def test_validation_rejects_mamba_mode_directly():
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 800
+
+    with pytest.raises(ValueError, match="mamba_cache_mode=\x27none\x27"):
+        mooncake_store_connector.MooncakeStoreConnector._validate_kv_cache_config(
+            vllm_config, _make_qsa_hybrid_kv_cache_config("none")
+        )
+
+
 def test_scheduler_requires_align_mode_for_mamba():
     vllm_config = _make_vllm_config()
     mamba_spec = MambaSpec(

tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py

@@ -759,8 +759,9 @@ def test_worker_setup_tolerates_finer_scratch_group():
     worker.pp_size = 1
     worker.num_kv_head = 8
     assert worker.coord.enable_partial_hash_hits
-    # The scratch DB is keyed at its own block size and never probed.
-    assert worker.token_dbs[2].hash_block_size == 4
+    # Only the participating full-attention and Mamba groups are registered.
+    assert len(worker.token_dbs) == 2
+    assert all(db.hash_block_size == 8 for db in worker.token_dbs)
 
     for g_idx, db in enumerate(worker.token_dbs):
         db.set_kv_caches_base_addr([g_idx * 10_000])
@@ -772,11 +773,11 @@ def test_worker_setup_tolerates_finer_scratch_group():
         block_size=worker.block_size,
         coord=worker.coord,
         tp_rank=0,
-        group_put_steps=[1, 1, 1],
+        group_put_steps=[1, 1],
         kv_role="kv_both",
         ready_event=threading.Event(),
         replicate_config=MagicMock(),
-        group_participates=[True, True, False],
+        group_participates=[True, True],
     )
 
     # Persist the sub-block partial tail at boundary 12 (keyed by hs[12//8-1]).
@@ -784,7 +785,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
     req = ReqMeta(
         req_id="r0",
         token_len_chunk=0,
-        block_ids=([1], [2], [3]),
+        block_ids=([1], [2]),
         block_hashes=hs,
         can_save=True,
         num_prompt_tokens=20,
@@ -797,8 +798,7 @@ def test_worker_setup_tolerates_finer_scratch_group():
     # A 13-token prompt sharing the prefix must hit the first hash unit.
     assert worker.lookup(num_tokens=13, block_hashes=hs).hit_length == 8
     # The scratch group's namespace never enters the store.
-    scratch_prefix = worker.token_dbs[2].key_for(hs[0]).rsplit("@", 1)[0]
-    assert not any(key.startswith(scratch_prefix) for key in store._data)
+    assert not any("@group:2" in key for key in store._data)
 
 
 def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
@@ -838,7 +838,8 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
     worker.tp_size = 1
     worker.pp_size = 1
     worker.num_kv_head = 8
-    assert worker.token_dbs[1].hash_block_size == 8
+    assert len(worker.token_dbs) == 1
+    assert worker.token_dbs[0].hash_block_size == 16
 
     raw_full = torch.zeros(4 * full.page_size_bytes, dtype=torch.int8)
     raw_ring = torch.zeros(4 * ring.page_size_bytes, dtype=torch.int8)
@@ -879,14 +880,14 @@ def test_ring_scratch_group_is_never_stored_and_does_not_block_hits():
         kv_role=worker.kv_role,
         ready_event=threading.Event(),
         enable_kv_event=False,
-        group_participates=[True, False],
+        group_participates=[True],
     )
     hs = [BlockHash(bytes([i + 1]) * 4) for i in range(4)]
     save_req = ReqMeta(
         req_id="r0",
         token_len_chunk=64,
-        # The ring group holds one block for the request's lifetime.
-        block_ids=([1, 2, 3, 4], [1]),
+        # Scheduler metadata is projected to the participating store groups.
+        block_ids=([1, 2, 3, 4],),
         block_hashes=hs,
         can_save=True,
         store_job_id=1,

tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py

@@ -2,12 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from types import SimpleNamespace
+from unittest.mock import patch
 
 import pytest
+import torch
 
 from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.data import (
     LoadSpec,
     MooncakeLookupResult,
+    MooncakeStoreConnectorMetadata,
     MooncakeStoreWorkerMetadata,
     ReqMeta,
     RequestTracker,
@@ -16,7 +19,16 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.scheduler impor
     MooncakeStoreScheduler,
 )
 from vllm.v1.core.block_pool import BlockPool
+from vllm.v1.core.kv_cache_manager import KVCacheBlocks
+from vllm.v1.core.kv_cache_utils import KVCacheBlock
 from vllm.v1.core.sched.output import KVConnectorBlockState
+from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
+    FullAttentionSpec,
+    KVCacheConfig,
+    KVCacheGroupSpec,
+    MambaSpec,
+)
 
 
 def _make_bare_scheduler(
@@ -35,9 +47,9 @@ def _make_bare_scheduler(
     scheduler._block_size = 16
     scheduler._hash_block_size = hash_block_size
     scheduler.enable_partial_hash_hits = enable_partial_hash_hits
-    scheduler.kv_cache_config = SimpleNamespace(
-        select_transfer_block_ids=lambda block_ids: tuple(block_ids)
-    )
+    scheduler.kv_cache_config = SimpleNamespace()
+    scheduler._store_group_ids = (0,)
+    scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 1: 1}
     scheduler.load_specs = {}
     scheduler._unfinished_request_ids = {"req-0"}
     scheduler._unfinished_requests = {}
@@ -228,14 +240,82 @@ def test_pending_load_for_non_chosen_connector_is_dropped():
     assert "req-0" not in scheduler._request_trackers
 
 
+def _make_qsa_hybrid_cache_config():
+    full = FullAttentionSpec(block_size=800, num_kv_heads=8, head_size=64, dtype=None)
+    circular = CircularBufferSpec(
+        block_size=8,
+        num_kv_heads=1,
+        head_size=64,
+        head_size_v=0,
+        dtype=torch.float16,
+    )
+    mamba = MambaSpec(
+        block_size=800,
+        shapes=((1, 1),),
+        dtypes=(torch.float32,),
+        mamba_cache_mode="align",
+    )
+    return KVCacheConfig(
+        num_blocks=4,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(["full"], full),
+            KVCacheGroupSpec(["qsa"], circular),
+            KVCacheGroupSpec(["mamba"], mamba),
+        ],
+    )
+
+
+def test_scheduler_projects_nonprefix_groups_and_mamba_ids():
+    vllm_config = SimpleNamespace(
+        kv_transfer_config=SimpleNamespace(
+            kv_role="kv_both", kv_connector_extra_config={}
+        ),
+        kv_events_config=None,
+        cache_config=SimpleNamespace(
+            block_size=800, enable_prefix_caching=True, prefix_match_unit=None
+        ),
+        parallel_config=SimpleNamespace(decode_context_parallel_size=1, world_size=1),
+    )
+
+    with patch(
+        "vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store."
+        "scheduler.LookupKeyClient"
+    ):
+        scheduler = MooncakeStoreScheduler(vllm_config, _make_qsa_hybrid_cache_config())
+
+    assert scheduler._store_group_ids == (0, 2)
+    assert scheduler._boundary_state_group_ids == frozenset({1})
+
+
+def test_current_save_block_ids_use_store_group_projection():
+    scheduler = _make_bare_scheduler()
+    scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
+    scheduler._store_group_ids = (0, 2)
+    meta = MooncakeStoreConnectorMetadata(set(), set())
+    req_meta = ReqMeta(
+        req_id="req-0",
+        token_len_chunk=800,
+        block_ids=(),
+        block_hashes=[b"h0"],
+        can_save=True,
+    )
+    meta.add_request(req_meta)
+    output = SimpleNamespace(
+        kv_connector_block_state=_make_connector_block_state(([10], [80], [30]))
+    )
+
+    scheduler._apply_current_save_block_ids(meta, output)
+
+    assert req_meta.block_ids == ([10], [30])
+
+
 def test_update_state_excludes_nontransfer_groups():
     """Store metadata must match the worker's registered cache groups."""
     scheduler = _make_bare_scheduler()
-    scheduler.kv_cache_config = SimpleNamespace(
-        select_transfer_block_ids=lambda block_ids: (block_ids[0],)
-    )
+    scheduler._store_group_ids = (0,)
     request = SimpleNamespace(request_id="req-1")
-    blocks = SimpleNamespace(get_block_ids=lambda: ([1, 2], [9]))
+    blocks = KVCacheBlocks(([KVCacheBlock(1), KVCacheBlock(2)], [KVCacheBlock(9)]))
 
     scheduler.update_state_after_alloc(request, blocks, num_external_tokens=32)
 
@@ -1003,6 +1083,7 @@ def test_pending_partial_tail_emits_offload_only_reqmeta():
 
 def test_finished_partial_tail_is_pre_pinned_as_store_job():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler._store_group_ids = (0, 1)
     scheduler.client = SimpleNamespace(discard=lambda *_: None)
     request = SimpleNamespace(
         request_id="req-0",
@@ -1143,6 +1224,20 @@ def _make_offload_only_output(entries, block_ids=([0],)):
     )
 
 
+def test_boundary_state_group_ids_are_remapped_to_store_projection():
+    scheduler = _make_bare_scheduler(hash_block_size=800, enable_partial_hash_hits=True)
+    scheduler.kv_cache_config = _make_qsa_hybrid_cache_config()
+    scheduler._store_group_ids = (0, 2)
+    scheduler._store_group_id_by_kv_cache_group_id = {0: 0, 2: 1}
+    scheduler._boundary_state_group_ids = frozenset({1})
+    _register_offload_request(scheduler, prefill_end_tokens=800, num_prompt_tokens=800)
+    meta = MooncakeStoreConnectorMetadata(set(), set())
+
+    scheduler._handle_boundary_state_offloads({"req-0": [(2, 7, 800)]}, meta)
+
+    assert meta.requests[0].boundary_state_offloads == [(1, 7, 800)]
+
+
 def test_resumed_prefill_claims_boundaries_past_prompt_length():
     # A resumed request re-prefills its previously generated tokens, so its
     # save window (`prefill_end_tokens`) extends past `num_prompt_tokens`.
@@ -1195,6 +1290,7 @@ def test_boundary_state_job_pins_exact_blocks_once():
 
 def test_store_job_pins_current_non_null_non_mamba_blocks():
     scheduler = _make_bare_scheduler(hash_block_size=4, enable_partial_hash_hits=True)
+    scheduler._store_group_ids = (0, 1)
     request = SimpleNamespace(
         all_token_ids=list(range(48)),
         block_hashes=[bytes([i]) for i in range(12)],

tests/v1/kv_connector/unit/test_mooncake_store_worker.py

@@ -46,8 +46,11 @@ from vllm.distributed.kv_transfer.kv_connector.v1.mooncake.store.metrics import
 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
 from vllm.v1.core.kv_cache_utils import BlockHash, maybe_convert_block_hash
 from vllm.v1.kv_cache_interface import (
+    CircularBufferSpec,
     FullAttentionSpec,
+    KVCacheConfig,
     KVCacheGroupSpec,
+    MambaSpec,
     RSWASpec,
 )
 from vllm.v1.kv_cache_layout import KVCacheLayout
@@ -342,6 +345,43 @@ def _make_kv_cache_config(
     )
 
 
+def _make_qsa_hybrid_kv_cache_config():
+    return KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                ["full"],
+                FullAttentionSpec(
+                    block_size=800,
+                    num_kv_heads=8,
+                    head_size=64,
+                    dtype=None,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["qsa"],
+                CircularBufferSpec(
+                    block_size=8,
+                    num_kv_heads=1,
+                    head_size=64,
+                    head_size_v=0,
+                    dtype=torch.float16,
+                ),
+            ),
+            KVCacheGroupSpec(
+                ["mamba"],
+                MambaSpec(
+                    block_size=800,
+                    shapes=((1, 1),),
+                    dtypes=(torch.float32,),
+                    mamba_cache_mode="align",
+                ),
+            ),
+        ],
+    )
+
+
 def _write_mooncake_config(tmp_path, config: dict[str, object]) -> str:
     config_path = tmp_path / "mooncake_config.json"
     config_path.write_text(json.dumps(config), encoding="utf-8")
@@ -2107,6 +2147,31 @@ def test_recv_thread_reports_unsplittable_key_larger_than_budget():
     assert thread.get_and_clear_block_ids_with_load_errors() == {0, 1, 2}
 
 
+def test_worker_init_excludes_nonprefix_cache_groups(monkeypatch):
+    store = MagicMock()
+    store.setup.return_value = 0
+    _install_fake_mooncake(monkeypatch, store)
+    _patch_worker_runtime(monkeypatch)
+    monkeypatch.setattr(
+        worker.MooncakeStoreConfig,
+        "load_from_config",
+        staticmethod(lambda: _make_config()),
+    )
+    vllm_config = _make_vllm_config()
+    vllm_config.cache_config.block_size = 800
+    vllm_config.cache_config.enable_prefix_caching = True
+    vllm_config.cache_config.prefix_match_unit = None
+
+    store_worker = worker.MooncakeStoreWorker(
+        vllm_config, _make_qsa_hybrid_kv_cache_config()
+    )
+
+    assert [
+        group.kv_cache_spec.block_size for group in store_worker._kv_cache_groups
+    ] == [800, 800]
+    assert [db.block_size for db in store_worker.token_dbs] == [800, 800]
+
+
 def test_requester_worker_init_uses_positional_setup(tmp_path, monkeypatch):
     store = MagicMock()
     store.setup.return_value = 0
@@ -4144,6 +4209,31 @@ def test_register_kv_caches_uses_transfer_group_memory_domain():
     )
 
 
+def test_register_kv_caches_excludes_nonprefix_cacheable_layers():
+    num_blocks = 2
+    worker = _make_bare_worker(num_gpu_blocks=num_blocks)
+    store_layout = MagicMock()
+    worker.token_dbs[0].store_layout = store_layout
+    full_cache = torch.zeros(num_blocks, 4, dtype=torch.float16)
+    qsa_cache = torch.zeros(num_blocks, 4, dtype=torch.float16)
+
+    _register_with_mocked_threads(
+        worker,
+        {"layer0": full_cache, "qsa": qsa_cache},
+    )
+
+    store_layout.register_kv_caches.assert_called_once()
+    registered_caches, registered_num_blocks = (
+        store_layout.register_kv_caches.call_args.args
+    )
+    assert registered_num_blocks == num_blocks
+    assert len(registered_caches) == 1
+    assert registered_caches[0] is full_cache
+    worker.store.register_buffer.assert_called_once_with(
+        full_cache.data_ptr(), full_cache.untyped_storage().nbytes()
+    )
+
+
 def test_register_kv_caches_separate_head_groups():
     # LHBNC gives each head group its own region; the K/V split doubles heads.
     layout = KVCacheLayout.LHBNC

vllm/config/speculative.py

@@ -1746,6 +1746,7 @@ class SpeculativeConfig:
         draft_parallel_config = ParallelConfig(
             pipeline_parallel_size=1,
             tensor_parallel_size=speculative_draft_tensor_parallel_size,
+            enable_expert_parallel=target_parallel_config.enable_expert_parallel,
             distributed_executor_backend=target_parallel_config.distributed_executor_backend,
             max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers,
             disable_custom_all_reduce=target_parallel_config.disable_custom_all_reduce,

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py

@@ -132,20 +132,22 @@ class MooncakeStoreConnector(KVConnectorBase_V1, SupportsHMA):
         from vllm.v1.kv_cache_interface import CrossAttentionSpec, MambaSpec
 
         unsupported: list[str] = []
-        cache_block_size = vllm_config.cache_config.block_size
-        for g_idx, g in enumerate(kv_cache_config.transfer_groups):
-            spec = g.kv_cache_spec
+        store_group_ids = kv_cache_config.prefix_cacheable_group_ids
+        if not store_group_ids:
+            raise ValueError(
+                "MooncakeStore requires at least one prefix-cacheable KV cache group"
+            )
+        for group_id in store_group_ids:
+            spec = kv_cache_config.kv_cache_groups[group_id].kv_cache_spec
             if isinstance(spec, CrossAttentionSpec):
-                unsupported.append(f"group {g_idx}: CrossAttentionSpec")
-            # Enforce Mamba align mode
-            if isinstance(spec, MambaSpec) and spec.block_size != cache_block_size:
+                unsupported.append(f"group {group_id}: CrossAttentionSpec")
+            if isinstance(spec, MambaSpec) and spec.mamba_cache_mode != "align":
                 unsupported.append(
-                    f"group {g_idx}: MambaSpec with block_size="
-                    f"{spec.block_size} != cache_config.block_size="
-                    f"{cache_block_size} (mamba_cache_mode != 'align')"
+                    f"group {group_id}: mamba_cache_mode="
+                    f"{spec.mamba_cache_mode!r} != 'align'"
                 )
         pcp = vllm_config.parallel_config.prefill_context_parallel_size
-        if len(kv_cache_config.transfer_groups) > 1 and pcp > 1:
+        if len(store_group_ids) > 1 and pcp > 1:
             unsupported.append(f"PCP > 1 (pcp={pcp}) with hybrid attention")
         if unsupported:
             raise ValueError(

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py

@@ -452,7 +452,7 @@ def _unwrap_spec(spec: KVCacheSpec) -> KVCacheSpec:
 
 
 def partial_hash_hits_enabled(
-    kv_cache_groups: list[KVCacheGroupSpec],
+    kv_cache_groups: Sequence[KVCacheGroupSpec],
     hash_block_size: int,
     dcp_world_size: int = 1,
 ) -> bool:

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py

@@ -71,19 +71,27 @@ class MooncakeStoreScheduler:
         )
         self.client = LookupKeyClient(vllm_config)
         self.kv_cache_config = kv_cache_config
+        self._store_group_ids = kv_cache_config.prefix_cacheable_group_ids
+        # Map scheduler group IDs to store group indices. Groups outside the
+        # store projection, such as the QSA ring, map to None and are skipped.
+        self._store_group_id_by_kv_cache_group_id = {
+            group_id: store_group_id
+            for store_group_id, group_id in enumerate(self._store_group_ids)
+        }
+        store_groups = kv_cache_config.prefix_cacheable_groups
 
         # Align with the engine's own scheduler_block_size and hash_block_size.
         self._block_size, self._hash_block_size = resolve_kv_cache_block_sizes(
             kv_cache_config, vllm_config
         )
         self.enable_partial_hash_hits = partial_hash_hits_enabled(
-            kv_cache_config.kv_cache_groups,
+            store_groups,
             self._hash_block_size,
             vllm_config.parallel_config.decode_context_parallel_size,
         )
         mamba_groups = {
             group_id: group.kv_cache_spec
-            for group_id, group in enumerate(kv_cache_config.kv_cache_groups)
+            for group_id, group in enumerate(store_groups)
             if isinstance(group.kv_cache_spec, MambaSpec)
         }
         assert all(
@@ -173,8 +181,8 @@ class MooncakeStoreScheduler:
         """Update state after block allocation."""
         local_block_ids: tuple[list[int], ...] = ()
         if num_external_tokens > 0:
-            local_block_ids = self.kv_cache_config.select_transfer_block_ids(
-                blocks.get_block_ids()
+            local_block_ids = blocks.get_block_ids(
+                group_ids=self._store_group_ids,
             )
 
         self._unfinished_requests[request.request_id] = (request, local_block_ids)
@@ -239,10 +247,7 @@ class MooncakeStoreScheduler:
             request_real = request_tuple[0]  # type: ignore[index]
 
             unfolded_block_ids = tuple(
-                blocks.copy()
-                for blocks in self.kv_cache_config.select_transfer_block_ids(
-                    request.block_ids
-                )
+                request.block_ids[group_id].copy() for group_id in self._store_group_ids
             )
 
             prefill_tokens = _new_req_prefill_tokens(request)
@@ -278,8 +283,8 @@ class MooncakeStoreScheduler:
             for i, req_id in enumerate(cached_reqs.req_ids):
                 new_block_ids = cached_reqs.new_block_ids[i]
                 if new_block_ids:
-                    new_block_ids = self.kv_cache_config.select_transfer_block_ids(
-                        new_block_ids
+                    new_block_ids = tuple(
+                        new_block_ids[group_id] for group_id in self._store_group_ids
                     )
 
                 req_meta = None
@@ -438,7 +443,9 @@ class MooncakeStoreScheduler:
             assert block_ids is not None, (
                 f"Missing current block table for store request {req_meta.req_id}"
             )
-            req_meta.block_ids = block_ids
+            req_meta.block_ids = tuple(
+                block_ids[group_id] for group_id in self._store_group_ids
+            )
 
     def _reference_save_blocks(self, meta: MooncakeStoreConnectorMetadata) -> None:
         """Take a GPU block reference for every store job this step emits.
@@ -509,12 +516,15 @@ class MooncakeStoreScheduler:
             return False
 
         pinned_block_ids: list[int] = []
-        for group_id, block_id, _ in partial_tail_offloads:
-            if group_id not in self._boundary_state_group_ids:
+        remapped_offloads: list[tuple[int, int, int]] = []
+        for group_id, block_id, boundary in partial_tail_offloads:
+            store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id)
+            if store_group_id not in self._boundary_state_group_ids:
                 return False
             if block_id == NULL_BLOCK_ID:
                 return False
             pinned_block_ids.append(block_id)
+            remapped_offloads.append((store_group_id, block_id, boundary))
         pinned_block_ids = list(dict.fromkeys(pinned_block_ids))
 
         pool = self._gpu_block_pool
@@ -530,12 +540,14 @@ class MooncakeStoreScheduler:
         self._finished_partial_tail_metas[request.request_id] = ReqMeta(
             req_id=request.request_id,
             token_len_chunk=0,
-            block_ids=tuple(group.copy() for group in block_ids),
+            block_ids=tuple(
+                block_ids[group_id].copy() for group_id in self._store_group_ids
+            ),
             block_hashes=list(request.block_hashes),
             can_save=True,
             num_prompt_tokens=tracker.prefill_end_tokens,
             store_job_id=store_job_id,
-            boundary_state_offloads=partial_tail_offloads,
+            boundary_state_offloads=remapped_offloads,
         )
         tracker.has_pending_offload = True
         # The store job owns exact block refs, so request cleanup need not wait.
@@ -574,9 +586,10 @@ class MooncakeStoreScheduler:
                     continue
                 if block_id == NULL_BLOCK_ID:
                     continue
-                if group_id not in self._boundary_state_group_ids:
+                store_group_id = self._store_group_id_by_kv_cache_group_id.get(group_id)
+                if store_group_id not in self._boundary_state_group_ids:
                     continue
-                accepted.append((group_id, block_id, boundary_tokens))
+                accepted.append((store_group_id, block_id, boundary_tokens))
             if not accepted:
                 continue
             tracker.has_pending_offload = True

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py

@@ -1672,10 +1672,11 @@ class MooncakeStoreWorker:
             dataclasses.replace(
                 group,
                 kv_cache_spec=resolve_dcp_kv_cache_spec(
-                    group.kv_cache_spec, self.dcp_size
+                    group.kv_cache_spec,
+                    self.dcp_size,
                 ),
             )
-            for group in kv_cache_config.transfer_groups
+            for group in kv_cache_config.prefix_cacheable_groups
         ]
         spec_cfg = getattr(vllm_config, "speculative_config", None)
         use_eagle_block_drop = bool(
@@ -1923,10 +1924,16 @@ class MooncakeStoreWorker:
         use_group_regions = self._kv_cache_config.hisparse_host_num_blocks is not None
 
         if not use_group_regions:
+            store_layer_names = {
+                layer_name
+                for group in self._kv_cache_groups
+                for layer_name in group.layer_names
+            }
             seen_storage_ptrs: set[int] = set()
             cache_tensors = [
                 group_kernel_blocks(_repr_tensor(cache), self.num_blocks)
-                for cache in kv_caches.values()
+                for layer_name, cache in kv_caches.items()
+                if layer_name in store_layer_names
             ]
             for cache in cache_tensors:
                 cache_storage = cache.untyped_storage()

vllm/entrypoints/cli/serve.py

@@ -54,7 +54,7 @@ class ServeSubcommand(CLISubcommand):
             args.model = args.model_tag
 
         if getattr(args, "grpc", False):
-            from vllm.entrypoints.grpc_server import serve_grpc
+            from vllm.entrypoints.launchers.grpc_server import serve_grpc
 
             uvloop.run(serve_grpc(args))
             return

vllm/entrypoints/grpc_server.py

@@ -1,200 +1,27 @@
-#!/usr/bin/env python3
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-# mypy: ignore-errors
-"""
-vLLM gRPC Server
+import warnings
 
-Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer
-from the smg-grpc-servicer package.
+from vllm.entrypoints.launchers.grpc_server import main, serve_grpc
 
-Usage:
-    python -m vllm.entrypoints.grpc_server --model <model_path>
+warnings.warn(
+    "`vllm.entrypoints.grpc_server is deprecated and will likely be "
+    "unsupported in a future version. Use the corresponding function from "
+    "`vllm.entrypoints.launchers.grpc_server` instead.",
+    DeprecationWarning,
+    stacklevel=1,
+)
 
-Example:
-    python -m vllm.entrypoints.grpc_server \
-        --model meta-llama/Llama-2-7b-hf \
-        --host 0.0.0.0 \
-        --port 50051
-"""
-
-import argparse
-import asyncio
-import signal
-import sys
-import time
-
-try:
-    import grpc
-    from grpc_health.v1 import health_pb2_grpc
-    from grpc_reflection.v1alpha import reflection
-    from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
-    from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer
-    from smg_grpc_servicer.vllm.servicer import VllmEngineServicer
-except ImportError as e:
-    raise ImportError(
-        "gRPC mode requires smg-grpc-servicer. "
-        "If not installed, run: pip install vllm[grpc]. "
-        "If already installed, there may be a broken import due to a "
-        "version mismatch — see the chained exception above for details."
-    ) from e
-
-import uvloop
-
-from vllm import envs
-from vllm.engine.arg_utils import AsyncEngineArgs
-from vllm.entrypoints.serve.utils.api_utils import log_version_and_model
-from vllm.logger import init_logger
-from vllm.usage.usage_lib import UsageContext
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-from vllm.v1.engine.async_llm import AsyncLLM
-from vllm.version import __version__ as VLLM_VERSION
-
-logger = init_logger(__name__)
-
-
-async def serve_grpc(args: argparse.Namespace):
-    """
-    Main gRPC serving function.
-
-    Args:
-        args: Parsed command line arguments
-    """
-    log_version_and_model(logger, VLLM_VERSION, args.model)
-    logger.info("vLLM gRPC server args: %s", args)
-
-    start_time = time.time()
-
-    # Create engine args
-    engine_args = AsyncEngineArgs.from_cli_args(args)
-
-    # Build vLLM config
-    vllm_config = engine_args.create_engine_config(
-        usage_context=UsageContext.OPENAI_API_SERVER,
-    )
-
-    # Create AsyncLLM
-    async_llm = AsyncLLM.from_vllm_config(
-        vllm_config=vllm_config,
-        usage_context=UsageContext.OPENAI_API_SERVER,
-        enable_log_requests=args.enable_log_requests,
-        disable_log_stats=args.disable_log_stats,
-    )
-
-    # Create servicer
-    servicer = VllmEngineServicer(async_llm, start_time)
-
-    # Create gRPC server
-    server = grpc.aio.server(
-        options=[
-            ("grpc.max_send_message_length", -1),
-            ("grpc.max_receive_message_length", -1),
-            # Tolerate client keepalive pings every 10s (default 300s is too
-            # strict for non-streaming requests where no DATA frames flow
-            # during generation)
-            ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000),
-            ("grpc.keepalive_permit_without_calls", True),
-        ],
-    )
-
-    # Add servicer to server
-    vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server)
-
-    # Add standard gRPC health service for Kubernetes probes
-    health_servicer = VllmHealthServicer(async_llm)
-    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
-
-    # Enable reflection for grpcurl and other tools
-    service_names = (
-        vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name,
-        "grpc.health.v1.Health",
-        reflection.SERVICE_NAME,
-    )
-    reflection.enable_server_reflection(service_names, server)
-
-    # Bind to address
-    host = args.host or "0.0.0.0"
-    address = f"{host}:{args.port}"
-    server.add_insecure_port(address)
-
-    try:
-        # Start server
-        await server.start()
-        logger.info("vLLM gRPC server started on %s", address)
-        logger.info("Server is ready to accept requests")
-
-        # Start periodic stats logging (mirrors the HTTP server's lifespan task)
-        if not args.disable_log_stats:
-
-            async def _force_log():
-                while True:
-                    await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL)
-                    await async_llm.do_log_stats()
-
-            stats_task = asyncio.create_task(_force_log())
-        else:
-            stats_task = None
-
-        # Handle shutdown signals
-        loop = asyncio.get_running_loop()
-        stop_event = asyncio.Event()
-
-        def signal_handler():
-            logger.info("Received shutdown signal")
-            stop_event.set()
-
-        for sig in (signal.SIGTERM, signal.SIGINT):
-            loop.add_signal_handler(sig, signal_handler)
-
-        try:
-            await stop_event.wait()
-        except KeyboardInterrupt:
-            logger.info("Interrupted by user")
-    finally:
-        logger.info("Shutting down vLLM gRPC server...")
-        if stats_task is not None:
-            stats_task.cancel()
-        try:
-            health_servicer.set_not_serving()
-        except Exception:  # broad: must not prevent server.stop() / shutdown()
-            logger.warning("Failed to set health status to NOT_SERVING", exc_info=True)
-        await server.stop(grace=5.0)
-        logger.info("gRPC server stopped")
-        async_llm.shutdown()
-        logger.info("AsyncLLM engine stopped")
-        logger.info("Shutdown complete")
-
-
-def main():
-    """Main entry point for python -m vllm.entrypoints.grpc_server."""
-    parser = FlexibleArgumentParser(
-        description="vLLM gRPC Server",
-    )
-
-    # Server args
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="0.0.0.0",
-        help="Host to bind gRPC server to",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=50051,
-        help="Port to bind gRPC server to",
-    )
-    parser = AsyncEngineArgs.add_cli_args(parser)
-
-    args = parser.parse_args()
-
-    # Run server
-    try:
-        uvloop.run(serve_grpc(args))
-    except Exception as e:
-        logger.exception("Server failed: %s", e)
-        sys.exit(1)
+__all__ = ["main", "serve_grpc"]
 
 
 if __name__ == "__main__":
+    warnings.warn(
+        "`python -m vllm.entrypoints.grpc_server is deprecated and "
+        "will likely be unsupported in a future version. "
+        "Please use `vllm serve <model_path> --grpc` instead.",
+        DeprecationWarning,
+        stacklevel=1,
+    )
+
     main()

vllm/entrypoints/launchers/grpc_server.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# mypy: ignore-errors
+"""
+vLLM gRPC Server
+
+Starts a gRPC server backed by AsyncLLM, using the VllmEngineServicer
+from the smg-grpc-servicer package.
+
+Usage:
+    python -m vllm.entrypoints.launchers.grpc_server --model <model_path>
+
+Example:
+    python -m vllm.entrypoints.launchers.grpc_server \
+        --model meta-llama/Llama-2-7b-hf \
+        --host 0.0.0.0 \
+        --port 50051
+"""
+
+import argparse
+import asyncio
+import signal
+import sys
+import time
+
+try:
+    import grpc
+    from grpc_health.v1 import health_pb2_grpc
+    from grpc_reflection.v1alpha import reflection
+    from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
+    from smg_grpc_servicer.vllm.health_servicer import VllmHealthServicer
+    from smg_grpc_servicer.vllm.servicer import VllmEngineServicer
+except ImportError as e:
+    raise ImportError(
+        "gRPC mode requires smg-grpc-servicer. "
+        "If not installed, run: pip install vllm[grpc]. "
+        "If already installed, there may be a broken import due to a "
+        "version mismatch — see the chained exception above for details."
+    ) from e
+
+import uvloop
+
+from vllm import envs
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.entrypoints.serve.utils.api_utils import log_version_and_model
+from vllm.logger import init_logger
+from vllm.usage.usage_lib import UsageContext
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.version import __version__ as VLLM_VERSION
+
+logger = init_logger(__name__)
+
+
+async def serve_grpc(args: argparse.Namespace):
+    """
+    Main gRPC serving function.
+
+    Args:
+        args: Parsed command line arguments
+    """
+    log_version_and_model(logger, VLLM_VERSION, args.model)
+    logger.info("vLLM gRPC server args: %s", args)
+
+    start_time = time.time()
+
+    # Create engine args
+    engine_args = AsyncEngineArgs.from_cli_args(args)
+
+    # Build vLLM config
+    vllm_config = engine_args.create_engine_config(
+        usage_context=UsageContext.OPENAI_API_SERVER,
+    )
+
+    # Create AsyncLLM
+    async_llm = AsyncLLM.from_vllm_config(
+        vllm_config=vllm_config,
+        usage_context=UsageContext.OPENAI_API_SERVER,
+        enable_log_requests=args.enable_log_requests,
+        disable_log_stats=args.disable_log_stats,
+    )
+
+    # Create servicer
+    servicer = VllmEngineServicer(async_llm, start_time)
+
+    # Create gRPC server
+    server = grpc.aio.server(
+        options=[
+            ("grpc.max_send_message_length", -1),
+            ("grpc.max_receive_message_length", -1),
+            # Tolerate client keepalive pings every 10s (default 300s is too
+            # strict for non-streaming requests where no DATA frames flow
+            # during generation)
+            ("grpc.http2.min_recv_ping_interval_without_data_ms", 10000),
+            ("grpc.keepalive_permit_without_calls", True),
+        ],
+    )
+
+    # Add servicer to server
+    vllm_engine_pb2_grpc.add_VllmEngineServicer_to_server(servicer, server)
+
+    # Add standard gRPC health service for Kubernetes probes
+    health_servicer = VllmHealthServicer(async_llm)
+    health_pb2_grpc.add_HealthServicer_to_server(health_servicer, server)
+
+    # Enable reflection for grpcurl and other tools
+    service_names = (
+        vllm_engine_pb2.DESCRIPTOR.services_by_name["VllmEngine"].full_name,
+        "grpc.health.v1.Health",
+        reflection.SERVICE_NAME,
+    )
+    reflection.enable_server_reflection(service_names, server)
+
+    # Bind to address
+    host = args.host or "0.0.0.0"
+    address = f"{host}:{args.port}"
+    server.add_insecure_port(address)
+
+    try:
+        # Start server
+        await server.start()
+        logger.info("vLLM gRPC server started on %s", address)
+        logger.info("Server is ready to accept requests")
+
+        # Start periodic stats logging (mirrors the HTTP server's lifespan task)
+        if not args.disable_log_stats:
+
+            async def _force_log():
+                while True:
+                    await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL)
+                    await async_llm.do_log_stats()
+
+            stats_task = asyncio.create_task(_force_log())
+        else:
+            stats_task = None
+
+        # Handle shutdown signals
+        loop = asyncio.get_running_loop()
+        stop_event = asyncio.Event()
+
+        def signal_handler():
+            logger.info("Received shutdown signal")
+            stop_event.set()
+
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            loop.add_signal_handler(sig, signal_handler)
+
+        try:
+            await stop_event.wait()
+        except KeyboardInterrupt:
+            logger.info("Interrupted by user")
+    finally:
+        logger.info("Shutting down vLLM gRPC server...")
+        if stats_task is not None:
+            stats_task.cancel()
+        try:
+            health_servicer.set_not_serving()
+        except Exception:  # broad: must not prevent server.stop() / shutdown()
+            logger.warning("Failed to set health status to NOT_SERVING", exc_info=True)
+        await server.stop(grace=5.0)
+        logger.info("gRPC server stopped")
+        async_llm.shutdown()
+        logger.info("AsyncLLM engine stopped")
+        logger.info("Shutdown complete")
+
+
+def main():
+    """Main entry point for python -m vllm.entrypoints.launchers.grpc_server."""
+    parser = FlexibleArgumentParser(
+        description="vLLM gRPC Server",
+    )
+
+    # Server args
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="0.0.0.0",
+        help="Host to bind gRPC server to",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=50051,
+        help="Port to bind gRPC server to",
+    )
+    parser = AsyncEngineArgs.add_cli_args(parser)
+
+    args = parser.parse_args()
+
+    # Run server
+    try:
+        uvloop.run(serve_grpc(args))
+    except Exception as e:
+        logger.exception("Server failed: %s", e)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

vllm/model_executor/kernels/mhc/warmup.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+from collections.abc import Iterable
 from dataclasses import asdict, dataclass
 from typing import Any
 
@@ -61,6 +62,21 @@ class MHCPreNormKernel(VllmJitKernel["MHCPreNormKernel.CompileKey"]):
                 **asdict(compile_key)
             )
 
+    def compile_many(self, compile_keys: Iterable[CompileKey]) -> None:
+        missing = list(
+            dict.fromkeys(
+                key for key in compile_keys if key not in self._compiled_cache
+            )
+        )
+        if len(missing) < 2:
+            return super().compile_many(missing)
+
+        # TileLang elaborates serially before compiling independent TIR modules.
+        compiled = self.kernel.par_compile(
+            [asdict(key) for key in missing], num_workers=min(4, len(missing))
+        )
+        self._compiled_cache.update(zip(missing, compiled, strict=True))
+
     def __call__(self, *tensors, **fields):
         compile_key = self.dispatch(
             n_splits=tensors[0].shape[0],

vllm/model_executor/model_loader/weight_cache/daemon.py

@@ -17,8 +17,30 @@ Engines then load from the daemons with:
     vllm serve /path/to/model --tensor-parallel-size 4 \\
         --load-format ipc_cache
 
-Only tensor and expert parallelism are supported; pipeline and data
-parallelism are rejected at launch.
+Only tensor and expert parallelism are supported; pipeline and data parallelism
+are rejected at launch.
+
+For multi-node tensor parallelism, run one launcher per node with a shared
+rendezvous so the global TP group forms across nodes (CUDA IPC handles are
+node-local, so each node serves only its local GPUs' shards). Reuse the same
+``--nnodes``/``--node-rank``/``--master-addr`` flags you pass the engine, plus a
+``--weight-cache-master-port`` distinct from the engine's ``--master-port``:
+
+    # node 0 (8 local GPUs)
+    python -m vllm.model_executor.model_loader.weight_cache.daemon \\
+        --model /path/to/model --tensor-parallel-size 16 \\
+        --nnodes 2 --node-rank 0 --master-addr 10.0.0.1 \\
+        --weight-cache-master-port 29600
+    # node 1 (8 local GPUs)
+    python -m vllm.model_executor.model_loader.weight_cache.daemon \\
+        --model /path/to/model --tensor-parallel-size 16 \\
+        --nnodes 2 --node-rank 1 --master-addr 10.0.0.1 \\
+        --weight-cache-master-port 29600
+
+The global TP rank of local GPU ``i`` on node ``r`` is
+``r * (tp_size // nnodes) + i``, matching vLLM's contiguous per-node rank
+assignment, so each engine worker maps its shard from the daemon on its own
+node.
 """
 
 import contextlib
@@ -132,11 +154,13 @@ class WeightCacheDaemon:
         self,
         vllm_config: VllmConfig,
         tp_rank: int,
+        local_rank: int,
         distributed_init_method: str,
         socket_dir: str | None = None,
     ):
         self.vllm_config = vllm_config
         self.tp_rank = tp_rank
+        self.local_rank = local_rank
         self.distributed_init_method = distributed_init_method
         self.socket_dir = socket_dir
         self.model: torch.nn.Module | None = None
@@ -152,12 +176,12 @@ class WeightCacheDaemon:
 
     def load_model(self) -> None:
         tp_size = self.cache_config.tp_size
-        torch.accelerator.set_device_index(self.tp_rank)
+        torch.accelerator.set_device_index(self.local_rank)
         init_distributed_environment(
             world_size=tp_size,
             rank=self.tp_rank,
             distributed_init_method=self.distributed_init_method,
-            local_rank=self.tp_rank,
+            local_rank=self.local_rank,
             backend=current_platform.dist_backend,
         )
         with set_current_vllm_config(self.vllm_config):
@@ -307,13 +331,14 @@ class WeightCacheDaemon:
 
 def _run_daemon(
     tp_rank: int,
+    local_rank: int,
     vllm_config: VllmConfig,
     distributed_init_method: str,
     socket_dir: str | None,
     ready_queue: "multiprocessing.Queue[int]",
 ) -> None:
     daemon = WeightCacheDaemon(
-        vllm_config, tp_rank, distributed_init_method, socket_dir
+        vllm_config, tp_rank, local_rank, distributed_init_method, socket_dir
     )
     daemon.load_model()
     daemon.serve_forever(ready_callback=lambda: ready_queue.put(tp_rank))
@@ -348,6 +373,15 @@ def main() -> None:
         default=None,
         help="Directory for the daemon Unix sockets (default: tempdir).",
     )
+    parser.add_argument(
+        "--weight-cache-master-port",
+        type=int,
+        default=None,
+        help="Rendezvous port for the daemon's own TP group. Must differ from "
+        "the engine's --master-port (the daemon holds its group open while "
+        "serving) and match across nodes. Required when --nnodes > 1; defaults "
+        "to a free port for single-node.",
+    )
     args = parser.parse_args()
     engine_args = EngineArgs.from_cli_args(args)
     vllm_config = engine_args.create_engine_config()
@@ -363,22 +397,41 @@ def main() -> None:
     _reject_unsupported_parallelism(parallel_config)
     tp_size = parallel_config.tensor_parallel_size
 
-    distributed_init_method = get_distributed_init_method("127.0.0.1", get_open_port())
+    nnodes = parallel_config.nnodes
+    node_rank = parallel_config.node_rank
+    local_world_size = tp_size // nnodes
+
+    # The daemon forms its own TP group and holds it open while serving, so it
+    # needs a rendezvous port distinct from the engine's --master-port. All
+    # nodes must agree on it; single-node can auto-pick a free port. The master
+    # address is reused from the engine's --master-addr (node_rank 0).
+    master_addr = parallel_config.master_addr
+    if nnodes > 1 and args.weight_cache_master_port is None:
+        raise ValueError("--weight-cache-master-port is required when --nnodes > 1")
+    master_port = args.weight_cache_master_port or get_open_port()
+    distributed_init_method = get_distributed_init_method(master_addr, master_port)
+
     ctx = multiprocessing.get_context("spawn")
     ready_queue: multiprocessing.Queue[int] = ctx.Queue()
+    # Global TP rank of local GPU i on this node; local index == device index.
+    global_ranks = [
+        node_rank * local_world_size + local_rank
+        for local_rank in range(local_world_size)
+    ]
     procs = [
         ctx.Process(
             target=_run_daemon,
             args=(
-                rank,
+                global_rank,
+                local_rank,
                 vllm_config,
                 distributed_init_method,
                 args.weight_cache_socket_dir,
                 ready_queue,
             ),
-            name=f"vllm-weight-cache-daemon-{rank}",
+            name=f"vllm-weight-cache-daemon-{global_rank}",
         )
-        for rank in range(tp_size)
+        for local_rank, global_rank in enumerate(global_ranks)
     ]
     for proc in procs:
         proc.start()
@@ -391,7 +444,7 @@ def main() -> None:
     signal.signal(signal.SIGTERM, _shutdown)
 
     ready_ranks: set[int] = set()
-    while len(ready_ranks) < tp_size:
+    while len(ready_ranks) < local_world_size:
         try:
             ready_ranks.add(ready_queue.get(timeout=1.0))
         except queue.Empty:
@@ -407,11 +460,16 @@ def main() -> None:
                 for proc in procs:
                     proc.join()
                 sys.exit(max((p.exitcode or 0) for p in procs))
-    logger.info(
-        "Weight cache daemon ready: all %d ranks serving in %s",
-        tp_size,
-        args.weight_cache_socket_dir or "the default socket dir",
+    socket_dir_msg = args.weight_cache_socket_dir or "the default socket dir"
+    logger.info_once(
+        "===== Weight cache daemon READY: node %d/%d serving %d local rank(s) "
+        "in %s =====",
+        node_rank,
+        nnodes,
+        local_world_size,
+        socket_dir_msg,
     )
+
     for proc in procs:
         proc.join()
     sys.exit(max(proc.exitcode or 0 for proc in procs))

vllm/model_executor/warmup/jit_warmup.py

@@ -963,14 +963,18 @@ class VllmJitKernel(Generic[CompileKeyT], ABC):
         """Compile one warmup key."""
         raise NotImplementedError
 
+    def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None:
+        """Compile a batch of warmup keys, allowing backend-specific scheduling."""
+        for compile_key in compile_keys:
+            self.compile(compile_key)
+
     def register_warmup(self, *args: Any, **kwargs: Any) -> None:
         """Register this kernel with the active runner's warmup registry."""
         JitWarmupRegistry.register(self, *args, **kwargs)
 
     def warmup(self, *args: Any, **kwargs: Any) -> None:
         """Compile this kernel's warmup keys."""
-        for compile_key in self.get_warmup_keys(*args, **kwargs):
-            self.compile(compile_key)
+        self.compile_many(self.get_warmup_keys(*args, **kwargs))
 
 
 def _same_value(left: Any, right: Any) -> bool:
@@ -1111,5 +1115,4 @@ class JitWarmupRegistry:
                     f"{kernel.__class__.__name__} ({len(compile_keys)} keys)",
                     refresh=False,
                 )
-                for compile_key in compile_keys:
-                    kernel.compile(compile_key)
+                kernel.compile_many(compile_keys)

vllm/models/minimax_m3/amd/ops/index_topk.py

@@ -37,6 +37,7 @@ MIN_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 20
 MAX_DECODE_SCORE_GFX950_HIGH_BATCH_REQUESTS = 64
 MAX_DECODE_SCORE_KV_CHUNKS = 256
 DECODE_TOPK_BLOCKS_PER_CHUNK = 512
+DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK = 128
 MAX_DECODE_TOPK_FAST_CHUNKS = 16
 DECODE_TOPK_TARGET_GRID = 64
 
@@ -103,14 +104,23 @@ def _decode_topk_launch_policy(
     topk: int,
     *,
     is_gfx950: bool,
-) -> tuple[int, bool, bool]:
-    """Choose the graph grid and bounded-context selector specialization."""
+) -> tuple[int, int, int, int, bool, bool]:
+    """Choose the selector grid and compile-time launch configuration."""
     if (
         is_gfx950
         and topk == 16
         and 0 < max_block <= MAX_DECODE_TOPK_FAST_CHUNKS * DECODE_TOPK_BLOCKS_PER_CHUNK
     ):
-        return MAX_DECODE_TOPK_FAST_CHUNKS, True, True
+        if max_block <= DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK:
+            return 1, DECODE_TOPK_SHORT_BLOCKS_PER_CHUNK, 2, 1, True, True
+        return (
+            MAX_DECODE_TOPK_FAST_CHUNKS,
+            DECODE_TOPK_BLOCKS_PER_CHUNK,
+            4,
+            2,
+            True,
+            True,
+        )
 
     target = max(
         1,
@@ -119,7 +129,14 @@ def _decode_topk_launch_policy(
             DECODE_TOPK_TARGET_GRID // max(1, total_q * num_idx_heads),
         ),
     )
-    return 1 << (target.bit_length() - 1), False, False
+    return (
+        1 << (target.bit_length() - 1),
+        DECODE_TOPK_BLOCKS_PER_CHUNK,
+        8,
+        2,
+        False,
+        False,
+    )
 
 
 # ---------------------------------------------------------------------------
@@ -1525,6 +1542,9 @@ def minimax_m3_index_decode(
     # only the number of chunks needed for its live context.
     (
         num_topk_chunks,
+        topk_block_size,
+        selector_num_warps,
+        selector_num_stages,
         single_tile_guaranteed,
         adaptive_final_merge,
     ) = _decode_topk_launch_policy(
@@ -1608,12 +1628,12 @@ def minimax_m3_index_decode(
         pages_per_sparse_block=PAGES_PER_SPARSE_BLOCK,
         block_page_stride=selector_block_page_stride,
         NUM_TOPK_CHUNKS=num_topk_chunks,
-        BLOCK_SIZE_K=512,
+        BLOCK_SIZE_K=topk_block_size,
         BLOCK_SIZE_T=block_size_t,
         EMIT_SPARSE_TABLE=emit_sparse_table,
         SINGLE_TILE_GUARANTEED=single_tile_guaranteed,
         ADAPTIVE_FINAL_MERGE=adaptive_final_merge,
-        num_warps=4 if adaptive_final_merge else 8,
-        num_stages=2,
+        num_warps=selector_num_warps,
+        num_stages=selector_num_stages,
     )
     return topk_idx

vllm/multimodal/media/video.py

@@ -190,7 +190,7 @@ class VideoMediaIO(MediaIO[MediaWithBytes[tuple[DecodedFrames, dict[str, Any]]]]
             }
             return MediaWithBytes((frames, metadata), data.encode())
 
-        return self.load_bytes(pybase64.b64decode(data))
+        return self.load_bytes(pybase64.b64decode(data, validate=True))
 
     def load_file(
         self, filepath: Path

vllm/tokenizers/deepseek_v41.py

@@ -32,7 +32,7 @@ def _normalize_messages(
             parts = []
             for block in content:
                 part_type = block.get("type")
-                if part_type == "text":
+                if part_type in ("text", "input_text", "output_text"):
                     parts.append(block.get("text", ""))
                 elif part_type in ("image_url", "input_image", "image_pil"):
                     parts.append(IMAGE_PLACEHOLDER)

vllm/v1/core/kv_cache_manager.py

@@ -67,30 +67,45 @@ class KVCacheBlocks:
     def get_block_ids(
         self,
         allow_none: Literal[False] = False,
+        *,
+        group_ids: Sequence[int] | None = None,
     ) -> tuple[list[int], ...]: ...
 
     @overload
     def get_block_ids(
         self,
         allow_none: Literal[True] = True,
+        *,
+        group_ids: Sequence[int] | None = None,
     ) -> tuple[list[int], ...] | None: ...
 
     def get_block_ids(
         self,
         allow_none: bool = False,
+        *,
+        group_ids: Sequence[int] | None = None,
     ) -> tuple[list[int], ...] | None:
         """
         Converts the KVCacheBlocks instance to block_ids.
 
+        Args:
+            allow_none: Return None when every selected group is empty.
+            group_ids: KV cache groups to include. Includes all groups by default.
+
         Returns:
             tuple[list[int], ...]: A tuple of lists where:
                 - the outer tuple corresponds to KV cache groups
                 - each inner list contains the block_ids of the blocks in that
                   group
         """
-        if allow_none and all(len(group) == 0 for group in self.blocks):
+        groups = (
+            self.blocks
+            if group_ids is None
+            else tuple(self.blocks[group_id] for group_id in group_ids)
+        )
+        if allow_none and all(len(group) == 0 for group in groups):
             return None
-        return tuple([blk.block_id for blk in group] for group in self.blocks)
+        return tuple([blk.block_id for blk in group] for group in groups)
 
     def get_unhashed_block_ids(self) -> list[int]:
         """Get block_ids of unhashed blocks from KVCacheBlocks instance."""

vllm/v1/kv_cache_interface.py

@@ -1450,6 +1450,23 @@ class KVCacheConfig:
             if group.enable_kv_transfer
         )
 
+    @cached_property
+    def prefix_cacheable_group_ids(self) -> tuple[int, ...]:
+        """IDs of transferable groups eligible for hash-addressed stores."""
+        return tuple(
+            group_id
+            for group_id in self.transfer_group_ids
+            if self.kv_cache_groups[group_id].kv_cache_spec.prefix_cacheable
+        )
+
+    @cached_property
+    def prefix_cacheable_groups(self) -> tuple[KVCacheGroupSpec, ...]:
+        """Transferable groups eligible for hash-addressed stores."""
+        return tuple(
+            self.kv_cache_groups[group_id]
+            for group_id in self.prefix_cacheable_group_ids
+        )
+
     @cached_property
     def transfer_groups(self) -> tuple[KVCacheGroupSpec, ...]:
         """Cache groups that participate in external KV transfer."""