diff --git a/.buildkite/intel_jobs/test-intel.yaml b/.buildkite/intel_jobs/test-intel.yaml index 8f4e83eee8..a76331f4d5 100644 --- a/.buildkite/intel_jobs/test-intel.yaml +++ b/.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" diff --git a/.buildkite/scripts/ci-otel/README.md b/.buildkite/scripts/ci-otel/README.md index 3ad552b5db..318cec5f21 100644 --- a/.buildkite/scripts/ci-otel/README.md +++ b/.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: diff --git a/.buildkite/scripts/ci-otel/ci_gpu.py b/.buildkite/scripts/ci-otel/ci_gpu.py index 304d738b88..0c77f5170b 100644 --- a/.buildkite/scripts/ci-otel/ci_gpu.py +++ b/.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() diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index f636e7ed18..ba54f9f273 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.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 diff --git a/.buildkite/test_areas/cuda.yaml b/.buildkite/test_areas/cuda.yaml index f0b1f4af2b..230781cd5d 100644 --- a/.buildkite/test_areas/cuda.yaml +++ b/.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: diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index a6b98bf931..bb848cd9e5 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.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 diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index d24ec5b57f..b45b26b2f2 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.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 diff --git a/.buildkite/test_areas/samplers.yaml b/.buildkite/test_areas/samplers.yaml index dae7e96468..18f857f65a 100644 --- a/.buildkite/test_areas/samplers.yaml +++ b/.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 diff --git a/requirements/test/xpu.in b/requirements/test/xpu.in index e93b418e9a..3aab5bdb00 100644 --- a/requirements/test/xpu.in +++ b/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 diff --git a/requirements/test/xpu.txt b/requirements/test/xpu.txt index 6c8bc1fcdd..b7cf49649b 100644 --- a/requirements/test/xpu.txt +++ b/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 diff --git a/rust/README.md b/rust/README.md index 06b50dbbfe..0cc33d64f8 100644 --- a/rust/README.md +++ b/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 diff --git a/rust/src/server/src/error.rs b/rust/src/server/src/error.rs index b658388cbd..9f083a1a61 100644 --- a/rust/src/server/src/error.rs +++ b/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 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() diff --git a/rust/src/server/src/routes.rs b/rust/src/server/src/routes.rs index 393a08450f..e7e91938bb 100644 --- a/rust/src/server/src/routes.rs +++ b/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)) diff --git a/rust/src/server/src/routes/abort_requests.rs b/rust/src/server/src/routes/abort_requests.rs index 2987703160..d08e5af3fa 100644 --- a/rust/src/server/src/routes/abort_requests.rs +++ b/rust/src/server/src/routes/abort_requests.rs @@ -22,7 +22,7 @@ pub async fn abort_requests( State(state): State>, body: Result, JsonRejection>, ) -> Result { - 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(); diff --git a/rust/src/server/src/routes/collective_rpc.rs b/rust/src/server/src/routes/collective_rpc.rs index bd769bee43..fa3782cbe9 100644 --- a/rust/src/server/src/routes/collective_rpc.rs +++ b/rust/src/server/src/routes/collective_rpc.rs @@ -36,7 +36,7 @@ pub async fn collective_rpc( State(state): State>, body: Result, JsonRejection>, ) -> Result, 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(), diff --git a/rust/src/server/src/routes/openai/utils/validated_json.rs b/rust/src/server/src/routes/openai/utils/validated_json.rs index 1ddec3a9f5..a9b68bc3b4 100644 --- a/rust/src/server/src/routes/openai/utils/validated_json.rs +++ b/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 { - let Json(mut data) = Json::::from_request(req, state) - .await - .map_err(|err: JsonRejection| ApiError::json_parse_error(err.body_text()))?; + let Json(mut data) = Json::::from_request(req, state).await?; data.normalize(); diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 210ae044a0..2851ec1e88 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/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::(&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::(&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)] diff --git a/rust/src/server/src/routes/weight_transfer.rs b/rust/src/server/src/routes/weight_transfer.rs new file mode 100644 index 0000000000..df55e25593 --- /dev/null +++ b/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, +} + +#[derive(Deserialize)] +pub(crate) struct UpdateWeightsRequest { + update_info: Option, +} + +#[derive(Default, Deserialize)] +pub(crate) struct FinishWeightUpdateRequest { + weight_version: Option, +} + +#[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>, + body: Result, JsonRejection>, +) -> Result, 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>, +) -> Result, 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>, +) -> Result, 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>, + body: Result, JsonRejection>, +) -> Result, 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>, + body: Result>, JsonRejection>, +) -> Result, 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>, + body: Result, JsonRejection>, +) -> Result, 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>, +) -> Result, 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 })) +} diff --git a/tests/entrypoints/unit_tests/test_grpc_health.py b/tests/entrypoints/launchers/test_grpc_health.py similarity index 100% rename from tests/entrypoints/unit_tests/test_grpc_health.py rename to tests/entrypoints/launchers/test_grpc_health.py diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index c45b87d04b..a7e6d499c1 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/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() diff --git a/tests/model_executor/layers/test_fused_shared_expert.py b/tests/model_executor/layers/test_fused_shared_expert.py index 37693f18c4..507fd93e9d 100644 --- a/tests/model_executor/layers/test_fused_shared_expert.py +++ b/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": diff --git a/tests/tokenizers_/test_deepseek_v41.py b/tests/tokenizers_/test_deepseek_v41.py index 493fdd88ba..13c0375fbb 100644 --- a/tests/tokenizers_/test_deepseek_v41.py +++ b/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( [ diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 3b29799d3f..dbddc556c9 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py b/tests/v1/kv_connector/unit/test_mooncake_store_connector.py index 96315b026d..6afb529d14 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_connector.py +++ b/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( diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py b/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py index e6b7302429..a93698bba9 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_hma_e2e.py +++ b/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, diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py b/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py index 67fa0c88aa..7ced38d65a 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_scheduler.py +++ b/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)], diff --git a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py b/tests/v1/kv_connector/unit/test_mooncake_store_worker.py index f613ba720e..9baf16e261 100644 --- a/tests/v1/kv_connector/unit/test_mooncake_store_worker.py +++ b/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 diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index cfecff72b6..1861266820 100644 --- a/vllm/config/speculative.py +++ b/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, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py index 2c8d118cdd..134d5055b6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/connector.py +++ b/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( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py index 375730807d..d05b553b64 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/coordinator.py +++ b/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: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py index ae61710e25..8f277d7eea 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/scheduler.py +++ b/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 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py index ebcf5708c8..6f154ea991 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/mooncake/store/worker.py +++ b/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() diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index a284b4bf39..e43b273e5e 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/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 diff --git a/vllm/entrypoints/grpc_server.py b/vllm/entrypoints/grpc_server.py index 59269dd180..913d084e4a 100644 --- a/vllm/entrypoints/grpc_server.py +++ b/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 +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 --grpc` instead.", + DeprecationWarning, + stacklevel=1, + ) + main() diff --git a/vllm/entrypoints/launchers/grpc_server.py b/vllm/entrypoints/launchers/grpc_server.py new file mode 100644 index 0000000000..3dc51f2e77 --- /dev/null +++ b/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 + +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() diff --git a/vllm/model_executor/kernels/mhc/warmup.py b/vllm/model_executor/kernels/mhc/warmup.py index 6a42ca6222..0abe1a2161 100644 --- a/vllm/model_executor/kernels/mhc/warmup.py +++ b/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], diff --git a/vllm/model_executor/model_loader/weight_cache/daemon.py b/vllm/model_executor/model_loader/weight_cache/daemon.py index 31f59a6ac5..96cedbe1e9 100644 --- a/vllm/model_executor/model_loader/weight_cache/daemon.py +++ b/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)) diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 72626964ae..84f713aa64 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/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) diff --git a/vllm/models/minimax_m3/amd/ops/index_topk.py b/vllm/models/minimax_m3/amd/ops/index_topk.py index b5c75371f0..ce27ea1220 100644 --- a/vllm/models/minimax_m3/amd/ops/index_topk.py +++ b/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 diff --git a/vllm/multimodal/media/video.py b/vllm/multimodal/media/video.py index f1f5783b0e..7528331cdf 100644 --- a/vllm/multimodal/media/video.py +++ b/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 diff --git a/vllm/tokenizers/deepseek_v41.py b/vllm/tokenizers/deepseek_v41.py index 3a57717270..a1449e59e2 100644 --- a/vllm/tokenizers/deepseek_v41.py +++ b/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) diff --git a/vllm/v1/core/kv_cache_manager.py b/vllm/v1/core/kv_cache_manager.py index 20e037f60e..226da3a1ed 100644 --- a/vllm/v1/core/kv_cache_manager.py +++ b/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.""" diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index e2869e047f..8cd867eb9c 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/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."""