From 41b8c158097c57ccf61626680d0fd13790bd120b Mon Sep 17 00:00:00 2001 From: YukioZzz Date: Sat, 17 Oct 2026 12:57:58 +0800 Subject: [PATCH] [Feature][MoRIIO] Support hybrid WRITE transfers with DSpark (#60544) Signed-off-by: Yichao Zhu Signed-off-by: tjtanaa Co-authored-by: OpenAI Codex Co-authored-by: tjtanaa Co-authored-by: Claude Opus 5.5 (1M context) --- .buildkite/test_areas/misc.yaml | 53 ++ docs/features/moriio_connector_usage.md | 8 + .../moriio_toy_proxy_server.py | 4 +- .../moriio_integration/__init__.py | 2 + .../moriio_integration/moriio_wide_dummy.py | 55 ++ .../test_write_dummy_e2e_rocm.py | 736 ++++++++++++++++++ .../unit/test_moriio_connector.py | 6 + .../unit/test_moriio_hma_scheduler.py | 321 +++++++- .../unit/test_moriio_kv_layout.py | 94 ++- .../unit/test_moriio_transfer_completion.py | 61 ++ .../unit/test_moriio_write_done_routing.py | 84 +- .../kv_connector/v1/moriio/moriio_common.py | 7 +- .../v1/moriio/moriio_connector.py | 145 ++-- .../kv_connector/v1/moriio/moriio_engine.py | 101 ++- 14 files changed, 1589 insertions(+), 88 deletions(-) create mode 100644 tests/v1/kv_connector/moriio_integration/__init__.py create mode 100644 tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py create mode 100644 tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index c79b451749..53ce97f6be 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -164,6 +164,59 @@ steps: depends_on: - image-build-amd +# The shared KV Connectors steps above exclude vllm/distributed/kv_transfer/ from +# their triggers, and CUDA skips MoRIIO's mori/aiter-gated tests. Run the whole +# MoRIIO suite on ROCm whenever the connector or its tests change. +- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" + key: amd-moriio-connector-unit-mi355 + dind: false + device: mi355_dpx + num_devices: 1 + timeout_in_minutes: 30 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/unit/test_moriio_ + - tests/v1/kv_connector/unit/utils.py + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py + +# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate +# GPUs over xGMI, compared exactly against a direct serve. +- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" + key: amd-moriio-write-dummy-e2e-mi355 + dind: false + device: mi355_2 + num_devices: 2 + # ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the + # largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports + # for its servers and proxy and its own compile-cache dir. + parallelism: 6 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - vllm/model_executor/model_loader/dummy_loader.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/moriio_integration/ + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval" device: h200_35gb key: v1-metrics-lmeval diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md index 6c906e7db1..1c67908a58 100644 --- a/docs/features/moriio_connector_usage.md +++ b/docs/features/moriio_connector_usage.md @@ -136,6 +136,14 @@ WRITE mode is used by default. READ mode can be configured by setting `--kv-tran !!! note `notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host. +### Hybrid attention and recurrent state + +WRITE mode supports hybrid attention plus Mamba/KDA models, including DSpark, with equal producer and consumer TP sizes. Use matching cache layouts on both instances, `--mamba-cache-mode align`, and `read_mode: false` in each instance's `kv_connector_extra_config`. Both peers must support hybrid WRITE. Sliding-window WRITE and hybrid speculative methods other than DSpark remain unsupported. + +For DSpark, configure the same draft model on both instances through `--speculative-config`. Target and draft attention groups retain their own block IDs and token spans; only the running recurrent state is transferred, not speculative scratch slots. The producer transfers the state after the first N-1 prompt tokens, and the consumer recomputes the last prompt token. Convolution and SSM transfers must both finish before the layer is complete. + +Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format. + ### Transport configuration MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments. diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index ce4d558fb9..e611b2e690 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -473,9 +473,11 @@ async def stop_profile(): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=10001) + # Port the prefill/decode servers register on (their proxy_ping_port). + parser.add_argument("--service-discovery-port", type=int, default=36367) args = parser.parse_args() - t = start_service_discovery("0.0.0.0", 36367) + t = start_service_discovery("0.0.0.0", args.service_discovery_port) # High-concurrency hardening. Quart's app.run() uses a shallow listen # backlog (100) and, with app.debug=True, adds per-request overhead that # starves the single accept loop. Under a burst of ~512 simultaneous client diff --git a/tests/v1/kv_connector/moriio_integration/__init__.py b/tests/v1/kv_connector/moriio_integration/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/v1/kv_connector/moriio_integration/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py b/tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py new file mode 100644 index 0000000000..ac80a6bc9b --- /dev/null +++ b/tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Widen the dummy weight init for the MoRIIO dummy-weight P/D test. + +The dummy loader fills float weights in +-1e-3, which keeps every activation so +small that outputs barely depend on the input, so a wrong KV prefix would go +unnoticed. The P/D test names ``WideDummyWeights`` as each server's +``--worker-extension-cls``. Workers import it before they load the model, which +re-registers the ``dummy`` load format in every worker process with a loader +that fills 16/32-bit float weights in +-``MORIIO_DUMMY_WEIGHT_RANGE``; 8-bit +weights and scales keep the default init. The format keeps the name ``dummy`` +so every config path that special-cases dummy weights still applies. + +The test process itself never imports this module. +""" + +import os + +import torch +from torch import nn + +from vllm.config import ModelConfig +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info +from vllm.model_executor.model_loader.reload.utils import get_layer_tensors +from vllm.model_executor.model_loader.weight_utils import ( + initialize_single_dummy_weight, +) + +_RANGE = float(os.environ["MORIIO_DUMMY_WEIGHT_RANGE"]) + + +@register_model_loader("dummy") +class WideDummyModelLoader(DummyModelLoader): + def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: + # DummyModelLoader.load_weights with a wider init for 16/32-bit floats. + for layer in model.modules(): + info = get_layerwise_info(layer) + if info.can_load(): + self._process_online_quant_layer(layer, info) + continue + for name, tensor in get_layer_tensors(layer).items(): + if name in layer._non_persistent_buffers_set: + continue + if tensor.is_floating_point() and torch.finfo(tensor.dtype).bits >= 16: + initialize_single_dummy_weight( + tensor, low=-_RANGE, high=_RANGE, name=name + ) + else: + initialize_single_dummy_weight(tensor, name=name) + + +class WideDummyWeights: + """Worker extension that only exists to import this module in workers.""" diff --git a/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py new file mode 100644 index 0000000000..74c9a95698 --- /dev/null +++ b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py @@ -0,0 +1,736 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Single-node MoRIIO WRITE P/D checks on dummy weights. + +Each case serves a real model architecture with ``--load-format dummy``, widened +by ``moriio_wide_dummy`` (see ``DUMMY_WEIGHT_RANGE``): the +config keeps the real hidden size, head geometry, MLA ranks and recurrent-state +shapes, so the connector registers and transfers exactly the production cache +layout, but the layer stack is cut to a few layers that still contain every +cache-group kind and the MoE is cut to a handful of experts. Dummy weights are +seeded, so the prefill, decode and reference servers hold identical weights. +Cases marked ``HMA_off`` run every server with +``--disable-hybrid-kv-cache-manager``, so all layers, sliding-window ones +included, share a single cache group. + +Prefill runs on one GPU, decode on another, and MoRI moves the KV over xGMI. +Either way decode computes the last prompt token itself. For hybrid models +prefill computes only the first ``n - 1`` tokens and transfers the KV and the +recurrent state after them. For attention-only models prefill computes and +transfers all ``n``, and decode's scheduler, seeing the whole prompt cached, +recomputes the last token. A direct serve splits the same prompt the same way +when the first ``n - 1`` tokens end on a chunk boundary that both servers share, +so the prompts are chosen to make that happen: one is a single chunk plus one +token, the others are a whole number of hybrid cache blocks plus one token +(prefix caching ends chunks on block boundaries) or, for attention-only models, +a whole number of chunks plus one token. With the split matched, a correct +transfer reproduces the direct serve's first-token log-probabilities exactly, +so the comparison needs no noise allowance. The +servers compile with Inductor's deterministic mode, since otherwise each one +benchmarks its own reduction launch configs and separately started servers +can drift apart by a few bf16 logit steps. Guards: + +* decode must report an external prefix-cache hit for every transferred token + of every prompt, so the match is not a local recompute; +* changing the first half of a prompt must move the direct serve's log-probs by + at least ``MIN_CONTEXT_EFFECT``, so a wrong prefix would be visible; +* ``test_corrupted_prefill_state_is_detected`` makes prefill compute its KV with + a different norm epsilon (same shapes, different values) and requires every + prompt to fail the comparison. +""" + +import atexit +import contextlib +import json +import os +import random +import shutil +import signal +import subprocess +import sys +import tempfile +import time +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest +import requests + +from tests.utils import RemoteOpenAIServer +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port + +_ROOT = Path(__file__).resolve().parents[4] +_PROXY = ( + _ROOT / "examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py" +) + +# Max allowed difference between centered first-token log-probabilities of the +# P/D and direct outputs, over the direct serve's top-k tokens. The split is +# matched, so correct transfers produce identical logits; this only absorbs the +# fp32 rounding of the log-softmax (~1e-7), far below one bf16 logit step. +TOLERANCE = float(os.getenv("MORIIO_DUMMY_E2E_TOL", "1e-5")) +# Several bf16 steps of the logits at DUMMY_WEIGHT_RANGE. +MIN_CONTEXT_EFFECT = 0.05 +# Half-width of the uniform dummy init for 16/32-bit float weights, applied in +# every server worker by moriio_wide_dummy.py. The loader default (1e-3) shrinks +# every activation so far that a context change moves the first-token log-probs +# by ~1e-4; at 0.1 it moves them by >0.1 and a corrupted prefill changes the +# top-k set entirely. +DUMMY_WEIGHT_RANGE = 0.1 +_WIDE_DUMMY_DIR = Path(__file__).resolve().parent +TOP_K = 20 +MAX_TOKENS = 4 +MAX_MODEL_LEN = 4096 + +# Prefill token budget of every server. Prompts longer than it are chunked, so +# the final-chunk save path runs; the short prompt fills exactly one chunk. +PREFILL_CHUNK = 299 +SHORT_PROMPT_LEN = PREFILL_CHUNK + 1 +# Hybrid blocks per aligned prompt, before the trailing token. +ALIGNED_BLOCKS = (1, 2) +# Prefill chunks per long attention-only prompt, before the trailing token. +ALIGNED_CHUNKS = (2, 4) + + +def _mori_available() -> bool: + try: + import mori.io # noqa: F401 + except ImportError: + return False + return True + + +pytestmark = [ + pytest.mark.skipif( + not current_platform.is_rocm() or not _mori_available(), + reason="MoRIIO P/D needs ROCm and MoRI", + ), + pytest.mark.skipif( + current_platform.is_rocm() and current_platform.device_count() < 2, + reason="needs one GPU for prefill and one for decode", + ), +] + + +# --------------------------------------------------------------------------- +# Model architectures + + +def _kimi_linear_overrides(cfg) -> dict[str, Any]: + # 3 KDA layers then 1 MLA layer; layer 0 stays dense (first_k_dense_replace). + linear = dict(cfg.linear_attn_config) + linear["kda_layers"] = [1, 2, 3] + linear["full_attn_layers"] = [4] + groups = max(1, cfg.num_expert_group) + return { + "num_hidden_layers": 4, + "linear_attn_config": linear, + "num_experts": 2 * groups, + "num_experts_per_token": 2, + "topk_group": min(cfg.topk_group, groups), + } + + +def _qwen3_5_moe_overrides(cfg) -> dict[str, Any]: + # 3 Gated DeltaNet layers then 1 full-attention layer. + return { + "text_config": { + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["full_attention"], + "num_experts": 4, + "num_experts_per_tok": 2, + } + } + + +def _qwen3_dense_overrides(cfg) -> dict[str, Any]: + # Attention-only: decode gets the KV for the whole prompt. + return {"num_hidden_layers": 4} + + +def _deepseek_v3_overrides(cfg) -> dict[str, Any]: + # MLA attention-only: layer 0 dense, layer 1 MoE, no MTP layer. + return { + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "n_routed_experts": 8, + "num_experts_per_tok": 2, + "n_group": 1, + "topk_group": 1, + "num_nextn_predict_layers": 0, + } + + +def _gemma3_overrides(cfg) -> dict[str, Any]: + # Five sliding-window layers then one global layer, as in the real + # 5:1 pattern. Without the hybrid KV cache manager vLLM gives the sliding + # layers full-attention cache specs, so all six share one group. + return { + "num_hidden_layers": 6, + "layer_types": ["sliding_attention"] * 5 + ["full_attention"], + } + + +@dataclass(frozen=True) +class ArchCase: + model: str + overrides: Callable[[Any], dict[str, Any]] + # Norm-epsilon override that perturbs every computed KV/state value + # without changing a single shape. Used by the corruption control. + corrupt_overrides: dict[str, Any] + env: dict[str, str] = field(default_factory=dict) + args: tuple[str, ...] = () + # Has recurrent-state (mamba-style) layers, so prefill stops at n - 1. + hybrid: bool = True + + +ARCHS: dict[str, ArchCase] = { + "kimi_linear": ArchCase( + model="moonshotai/Kimi-Linear-48B-A3B-Instruct", + overrides=_kimi_linear_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "1"}, + args=("--trust-remote-code", "--mamba-ssm-cache-dtype", "float32"), + ), + "qwen3_5_moe": ArchCase( + model="Qwen/Qwen3.5-35B-A3B", + overrides=_qwen3_5_moe_overrides, + corrupt_overrides={"text_config": {"rms_norm_eps": 1.0}}, + # AITER's GDN conv-update kernel rejects the DS conv-state layout that + # MoRIIO's recurrent-state transfer requires. + env={"VLLM_ROCM_USE_AITER": "0"}, + args=("--language-model-only",), + ), + "qwen3_dense": ArchCase( + model="Qwen/Qwen3-32B", + overrides=_qwen3_dense_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "0"}, + hybrid=False, + ), + "deepseek_v3": ArchCase( + model="deepseek-ai/DeepSeek-V3", + overrides=_deepseek_v3_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "1"}, + hybrid=False, + ), + # Text-only Gemma 3. The multimodal Gemma3ForConditionalGeneration takes + # the inputs_embeds path on KV consumers (enable_mm_embeds is forced on + # there) and computes differently from a direct serve without any transfer. + "gemma3": ArchCase( + model="google/gemma-3-1b-it", + overrides=_gemma3_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "0"}, + hybrid=False, + ), +} + + +def _arch_overrides(arch: ArchCase) -> dict[str, Any]: + from vllm.transformers_utils.config import get_config + + return arch.overrides(get_config(arch.model, trust_remote_code=True)) + + +def _merge(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key, value in extra.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _merge(merged[key], value) + else: + merged[key] = value + return merged + + +# --------------------------------------------------------------------------- +# Serving + + +@dataclass(frozen=True) +class Case: + arch: str + attention_backend: str + prefill_graph: str + decode_graph: str + # Run every server of the case without the hybrid KV cache manager. + hma_off: bool = False + + @property + def extra_args(self) -> tuple[str, ...]: + return ("--disable-hybrid-kv-cache-manager",) if self.hma_off else () + + @property + def id(self) -> str: + return ( + f"{self.arch}-{self.attention_backend}-" + f"P_{self.prefill_graph}-D_{self.decode_graph}" + + ("-HMA_off" if self.hma_off else "") + ) + + +_DEFAULT = "FULL_AND_PIECEWISE" +CASES = [ + # Kimi-Linear: KDA + MLA groups. + Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT), + Case("kimi_linear", "ROCM_AITER_MLA", "NONE", "NONE"), + Case("kimi_linear", "ROCM_AITER_MLA", "PIECEWISE", "PIECEWISE"), + Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, "FULL_DECODE_ONLY"), + Case("kimi_linear", "TRITON_MLA", _DEFAULT, _DEFAULT), + # Qwen3.5: Gated DeltaNet + full-attention groups. + Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, _DEFAULT), + Case("qwen3_5_moe", "TRITON_ATTN", "NONE", "NONE"), + Case("qwen3_5_moe", "TRITON_ATTN", "PIECEWISE", "PIECEWISE"), + Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, "FULL_DECODE_ONLY"), + # ROCM_ATTN keeps a (2, num_blocks, ...) KV layout and opts out of KV + # connectors, so the second ROCm attention backend is the unified one. + Case("qwen3_5_moe", "ROCM_AITER_UNIFIED_ATTN", _DEFAULT, _DEFAULT), + # Attention-only: prefill transfers the whole prompt. + Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT), + Case("qwen3_dense", "ROCM_AITER_UNIFIED_ATTN", _DEFAULT, _DEFAULT), + # Without the hybrid KV cache manager: one cache group per model. + Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT, hma_off=True), + Case("deepseek_v3", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT, hma_off=True), + Case("gemma3", "TRITON_ATTN", _DEFAULT, _DEFAULT, hma_off=True), +] + + +def _graph_args(mode: str) -> list[str]: + if mode == "NONE": + return ["--enforce-eager"] + # Inductor otherwise picks reduction launch configs by timing them, so each + # separately compiled server may sum in a different order and drift by a + # few bf16 logit steps; deterministic mode fixes the choice. + config = { + "cudagraph_mode": mode, + "inductor_compile_config": {"deterministic": True}, + } + return ["--compilation-config", json.dumps(config)] + + +def _visible_devices() -> tuple[str, list[str]]: + selectors = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES") + selector = next((k for k in selectors if k in os.environ), selectors[0]) + visible = os.getenv(selector) + devices = ( + visible.split(",") + if visible + else [str(i) for i in range(current_platform.device_count())] + ) + return selector, devices + + +def _device_env(selector: str, device: str) -> dict[str, str]: + # vLLM mirrors HIP_VISIBLE_DEVICES into CUDA_VISIBLE_DEVICES at import, so + # the test process carries both; pin every selector or the server refuses + # the mismatch. ROCR narrows the set first, so HIP/CUDA then index into it. + env = {selector: device} + for key in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + if key != selector: + env[key] = "0" if selector == "ROCR_VISIBLE_DEVICES" else device + return env + + +# KV consumers force enable_mm_embeds on, which moves multimodal-capable models +# (Qwen3.5) from input_ids to inputs_embeds, but the AOT compile-cache key does +# not cover it, so a consumer can load a producer's graph and crash in +# profile_run. Keep each role's compile cache apart until +# vllm-project/vllm#54642 lands (vllm-project/vllm#50891). +_CACHE_ROOT = Path(tempfile.mkdtemp(prefix="moriio_dummy_e2e_cache_")) +atexit.register(shutil.rmtree, _CACHE_ROOT, ignore_errors=True) + + +def _server( + arch: ArchCase, + overrides: dict[str, Any], + backend: str, + graph: str, + device: str, + selector: str, + *, + port: int | None = None, + extra_args: tuple[str, ...] = (), + kv_transfer_config: dict[str, Any] | None = None, +) -> RemoteOpenAIServer: + port = port or get_open_port() + args = [ + *arch.args, + *extra_args, + "--port", + str(port), + "--load-format", + "dummy", + "--worker-extension-cls", + "moriio_wide_dummy.WideDummyWeights", + "--hf-overrides", + json.dumps(overrides), + "--dtype", + "bfloat16", + "--attention-backend", + backend, + "--max-model-len", + str(MAX_MODEL_LEN), + "--max-num-batched-tokens", + str(PREFILL_CHUNK), + "--max-num-seqs", + "16", + "--kv-cache-memory-bytes", + str(2 * 1024**3), + "--gpu-memory-utilization", + "0.5", + "--enable-prefix-caching", + *_graph_args(graph), + ] + role = "direct" + if kv_transfer_config is not None: + args += ["--kv-transfer-config", json.dumps(kv_transfer_config)] + role = kv_transfer_config["kv_role"] + env = { + "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"], + **_device_env(selector, device), + "VLLM_HOST_IP": "127.0.0.1", + "VLLM_PORT": str(get_open_port()), + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "VLLM_SSM_CONV_STATE_LAYOUT": "DS", + "VLLM_CACHE_ROOT": str(_CACHE_ROOT / role), + # Only this directory, so the source tree cannot shadow installed vllm. + "PYTHONPATH": os.pathsep.join( + filter(None, [str(_WIDE_DUMMY_DIR), os.environ.get("PYTHONPATH")]) + ), + "MORIIO_DUMMY_WEIGHT_RANGE": str(DUMMY_WEIGHT_RANGE), + **arch.env, + } + return RemoteOpenAIServer( + arch.model, args, env_dict=env, auto_port=False, max_wait_seconds=900 + ) + + +def _moriio_config(role: str, http_port: int, ping_port: int) -> dict[str, Any]: + return { + "kv_connector": "MoRIIOConnector", + "kv_role": role, + "kv_connector_extra_config": { + "backend": "xgmi", + "read_mode": False, + "host_ip": "127.0.0.1", + "proxy_ip": "127.0.0.1", + "proxy_ping_port": str(ping_port), + "http_port": str(http_port), + "handshake_port": str(get_open_port()), + "notify_port": str(get_open_port()), + }, + } + + +@contextlib.contextmanager +def _proxy(ping_port: int) -> Iterator[str]: + # Free ports for both sockets, so cases from parallel shards on one host + # do not collide. + port = get_open_port() + process = subprocess.Popen( + [ + sys.executable, + str(_PROXY), + "--port", + str(port), + "--service-discovery-port", + str(ping_port), + ], + start_new_session=True, + ) + try: + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + assert process.poll() is None, "MoRIIO proxy exited during startup" + try: + requests.get(f"http://127.0.0.1:{port}/", timeout=1) + break + except requests.ConnectionError: + time.sleep(0.2) + else: + raise TimeoutError("MoRIIO proxy did not start") + yield f"http://127.0.0.1:{port}" + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=10) + + +@contextlib.contextmanager +def _pd_pair( + arch: ArchCase, + prefill_overrides: dict[str, Any], + decode_overrides: dict[str, Any], + case: Case, +) -> Iterator[tuple[str, RemoteOpenAIServer]]: + selector, devices = _visible_devices() + assert len(devices) >= 2, "needs one GPU for prefill and one for decode" + prefill_port, decode_port, ping_port = (get_open_port() for _ in range(3)) + with _proxy(ping_port) as proxy_url, contextlib.ExitStack() as stack: + servers: list[RemoteOpenAIServer] = [] + stack.callback(RemoteOpenAIServer.shutdown_many, servers) + with ThreadPoolExecutor(2) as pool: + futures = [ + pool.submit( + _server, + arch, + prefill_overrides, + case.attention_backend, + case.prefill_graph, + devices[0], + selector, + port=prefill_port, + extra_args=case.extra_args, + kv_transfer_config=_moriio_config( + "kv_producer", prefill_port, ping_port + ), + ), + pool.submit( + _server, + arch, + decode_overrides, + case.attention_backend, + case.decode_graph, + devices[1], + selector, + port=decode_port, + extra_args=case.extra_args, + kv_transfer_config=_moriio_config( + "kv_consumer", decode_port, ping_port + ), + ), + ] + for future in futures: + servers.append(future.result()) + _wait_for_registration(proxy_url, arch.model) + yield proxy_url, servers[1] + + +def _wait_for_registration(proxy_url: str, model: str) -> None: + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + # The proxy hands decode max_tokens - 1 (prefill emits the first + # token), so the probe needs at least 2 for decode to accept it. + try: + response = requests.post( + f"{proxy_url}/v1/completions", + json={"model": model, "prompt": [1000, 1001], "max_tokens": 2}, + timeout=300, + ) + # The proxy streams the decode body, so a decode-side error can + # surface as a truncated response rather than a status code. + if response.status_code == 200 and response.json()["choices"]: + return + except requests.RequestException: + pass + time.sleep(1) + raise TimeoutError("prefill/decode never registered with the MoRIIO proxy") + + +# --------------------------------------------------------------------------- +# Measurement + + +def _metric(server: RemoteOpenAIServer, name: str) -> list[Any]: + from prometheus_client.parser import text_string_to_metric_families + + response = requests.get(server.url_for("metrics"), timeout=10) + response.raise_for_status() + return [ + sample + for family in text_string_to_metric_families(response.text) + for sample in family.samples + if sample.name == name + ] + + +def _block_size(server: RemoteOpenAIServer) -> int: + # Hybrid models pad the attention block to the recurrent-state page, so the + # effective size is only known once the server has started. + sizes = { + int(sample.labels["block_size"]) + for sample in _metric(server, "vllm:cache_config_info") + } + assert len(sizes) == 1, f"expected one cache block size, got {sizes}" + return sizes.pop() + + +def _prompt_lengths(arch: ArchCase, server: RemoteOpenAIServer) -> list[int]: + # Hybrid chunks end on cache-block boundaries; attention-only chunks end on + # the prefill budget, so the last token is computed on its own either way. + if arch.hybrid: + unit, counts = _block_size(server), ALIGNED_BLOCKS + else: + unit, counts = PREFILL_CHUNK, ALIGNED_CHUNKS + aligned = [ + k * unit + 1 for k in counts if k * unit + 1 + MAX_TOKENS <= MAX_MODEL_LEN + ] + assert aligned, f"alignment unit {unit} leaves no aligned prompt" + return [SHORT_PROMPT_LEN, *aligned] + + +def _prompts(lengths: list[int]) -> list[list[int]]: + rng = random.Random(0) + return [[rng.randrange(1000, 30000) for _ in range(n)] for n in lengths] + + +def _first_step(url: str, model: str, prompt: list[int]) -> dict[str, float]: + response = requests.post( + f"{url}/v1/completions", + json={ + "model": model, + "prompt": prompt, + "max_tokens": MAX_TOKENS, + "temperature": 0, + "logprobs": TOP_K, + "return_tokens_as_token_ids": True, + }, + timeout=600, + ) + response.raise_for_status() + top = response.json()["choices"][0]["logprobs"]["top_logprobs"] + assert top, "server returned no logprobs" + return top[0] + + +def _distance(ref: dict[str, float], got: dict[str, float]) -> float: + """Max difference of centered log-probs over the shared top-k tokens. + + Centering removes the log-sum-exp offset, leaving the logit differences the + KV actually determines. + """ + common = sorted(ref.keys() & got.keys()) + if len(common) < len(ref) // 2: + return float("inf") + ref_mean = sum(ref[t] for t in common) / len(common) + got_mean = sum(got[t] for t in common) / len(common) + return max(abs((ref[t] - ref_mean) - (got[t] - got_mean)) for t in common) + + +def _external_hits(server: RemoteOpenAIServer) -> float: + samples = _metric(server, "vllm:external_prefix_cache_hits_total") + return sum(sample.value for sample in samples) + + +@dataclass +class Reference: + prompts: list[list[int]] + logprobs: list[dict[str, float]] + hybrid: bool + + @property + def expected_external_hits(self) -> int: + # Hybrid prefill stops one token short; attention-only prefill hands + # decode the whole prompt. + return sum(len(p) - int(self.hybrid) for p in self.prompts) + + +def _reference(arch: ArchCase, overrides: dict[str, Any], case: Case) -> Reference: + """Direct serve; also checks that the comparison can see a wrong prefix.""" + selector, devices = _visible_devices() + # Prefill graphs run the chunks and the decode-only last step runs the same + # captured graph in every decode mode used here, so the prefill mode + # reproduces both halves of the P/D computation. + with _server( + arch, + overrides, + case.attention_backend, + case.prefill_graph, + devices[0], + selector, + extra_args=case.extra_args, + ) as server: + url = server.url_root + lengths = _prompt_lengths(arch, server) + prompts = _prompts(lengths) + refs = [_first_step(url, arch.model, p) for p in prompts] + # Same tail, different early context: the comparison has to see it. + for prompt, ref in zip(prompts, refs): + altered = [t + 1 for t in prompt[: len(prompt) // 2]] + altered += prompt[len(prompt) // 2 :] + moved = _distance(ref, _first_step(url, arch.model, altered)) + print(f"[sensitivity] len={len(prompt)} context-change distance={moved}") + assert moved >= MIN_CONTEXT_EFFECT, ( + f"first-token log-probs barely depend on the context (moved " + f"{moved:.4g} for len {len(prompt)}); a wrong prefix would go " + "unnoticed" + ) + # Determinism of the reference itself (served from its prefix cache). + again = _first_step(url, arch.model, prompts[-1]) + print(f"[noise] repeat distance={_distance(refs[-1], again)}") + return Reference(prompts, refs, arch.hybrid) + + +def _pd_distances( + arch: ArchCase, + prefill_overrides: dict[str, Any], + decode_overrides: dict[str, Any], + case: Case, + ref: Reference, +) -> tuple[list[float], float]: + with _pd_pair(arch, prefill_overrides, decode_overrides, case) as ( + proxy_url, + decode, + ): + hits_before = _external_hits(decode) + distances = [ + _distance(expected, _first_step(proxy_url, arch.model, prompt)) + for prompt, expected in zip(ref.prompts, ref.logprobs) + ] + hits = _external_hits(decode) - hits_before + lengths = [len(p) for p in ref.prompts] + print(f"[pd] {case.id} lens={lengths} distances={distances} external_hits={hits}") + return distances, hits + + +# --------------------------------------------------------------------------- +# Tests + + +@pytest.mark.parametrize("case", CASES, ids=lambda c: c.id) +def test_write_matches_direct_serve(case: Case): + arch = ARCHS[case.arch] + overrides = _arch_overrides(arch) + ref = _reference(arch, overrides, case) + distances, hits = _pd_distances(arch, overrides, overrides, case, ref) + + # Every transferred token of every prompt must come from prefill. + assert hits == ref.expected_external_hits, ( + f"decode reported {hits} external prefix-cache hits, expected " + f"{ref.expected_external_hits}; it did not use all the transferred KV" + ) + lengths = [len(p) for p in ref.prompts] + bad = {n: d for n, d in zip(lengths, distances) if not d <= TOLERANCE} + assert not bad, f"P/D first-token log-probs differ from direct serve: {bad}" + + +@pytest.mark.parametrize("arch_name", sorted(ARCHS)) +def test_corrupted_prefill_state_is_detected(arch_name: str): + """Negative control: KV computed with different values must fail.""" + arch = ARCHS[arch_name] + default = next(c for c in CASES if c.arch == arch_name) + case = Case( + arch_name, default.attention_backend, _DEFAULT, _DEFAULT, default.hma_off + ) + overrides = _arch_overrides(arch) + ref = _reference(arch, overrides, case) + corrupted = _merge(overrides, arch.corrupt_overrides) + distances, hits = _pd_distances(arch, corrupted, overrides, case, ref) + + assert hits == ref.expected_external_hits, ( + "decode did not use the transferred KV; the control proves nothing" + ) + assert all(d > TOLERANCE for d in distances), ( + f"P/D outputs with wrongly computed prefill KV still matched the " + f"reference (distances {distances}); the comparison cannot detect " + "transfer corruption" + ) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 73cc5f0583..ffd78969a1 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -146,6 +146,7 @@ def _write_consumer_scheduler_for_finished_request(tp_size: int = 2): scheduler.is_producer = False scheduler.mode = MoRIIOMode.WRITE scheduler.tp_size = tp_size + scheduler._has_mamba = False scheduler._reqs_need_recv = {} scheduler.unmap_request_id = MagicMock() return scheduler @@ -158,6 +159,10 @@ def _write_producer_scheduler(block_size: int = 1) -> Any: scheduler.mode = MoRIIOMode.WRITE scheduler.is_producer = True scheduler.block_size = block_size + scheduler._has_mamba = False + scheduler.kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=object())] + ) scheduler.blocks_per_sw = [0] scheduler._group_block_sizes = [block_size] scheduler.transfer_id_to_request_id = {} @@ -944,6 +949,7 @@ def test_clamp_to_prompt_blocks_clamps_each_group_by_its_own_block_size(): # their own prompt-block count; a sliding-window group is left untouched. scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) scheduler.blocks_per_sw = [0, 0, 3] + scheduler._has_mamba = False scheduler._group_block_sizes = [16, 32, 16] req = SimpleNamespace(num_prompt_tokens=40) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 6d8dc24ea2..1120062e54 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -12,12 +12,13 @@ wiring. import importlib from dataclasses import asdict from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from vllm.v1.core.block_pool import BlockPool -from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.output import KVConnectorBlockState, SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs @@ -52,6 +53,7 @@ class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore attributes read by the method under test are set.""" def __init__(self, **attrs): + self.mode = MoRIIOMode.READ self._mamba_group_ids: list[int] = [] self._attn_group_ids: list[int] = [0] self._num_ssm_scratch_blocks = 0 @@ -287,7 +289,8 @@ def test_scheduler_rejects_mamba_group_without_two_states(): moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) -def test_scheduler_rejects_non_dspark_speculative_hybrid_read(): +@pytest.mark.parametrize("read_mode", [True, False]) +def test_scheduler_rejects_non_dspark_speculative_hybrid(read_mode): config = SimpleNamespace( kv_cache_groups=[ SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), @@ -298,7 +301,9 @@ def test_scheduler_rejects_non_dspark_speculative_hybrid_read(): with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"): moriio_connector.MoRIIOConnectorScheduler( - _gate_vllm_config(speculative_config=_spec_config("ngram")), + _gate_vllm_config( + read_mode=read_mode, speculative_config=_spec_config("ngram") + ), "engine", config, ) @@ -336,14 +341,16 @@ def _dspark_wrapped_attention_config(*, dcp_sharded: bool = True): @pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize("read_mode", [True, False]) @pytest.mark.parametrize( "dcp_size,dcp_sharded,expected_tail", [(1, True, 4), (4, True, 1), (4, False, 4)] ) def test_scheduler_accepts_dspark_with_wrapped_full_attention( - dcp_size, dcp_sharded, expected_tail + dcp_size, dcp_sharded, expected_tail, read_mode ): config = _dspark_wrapped_attention_config(dcp_sharded=dcp_sharded) vllm_config = _gate_vllm_config( + read_mode=read_mode, speculative_config=_spec_config("dspark"), block_size=4, num_lookahead_tokens=7, @@ -384,19 +391,234 @@ def test_scheduler_accepts_dspark_with_wrapped_full_attention( ) -def test_scheduler_rejects_hybrid_write(): - config = SimpleNamespace( - kv_cache_groups=[ - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), - ], +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_clamps_attention_and_excludes_scratch_slots(): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config ) - config.transfer_groups = tuple(config.kv_cache_groups) + assert sched._clamp_to_prompt_blocks( + SimpleNamespace(num_prompt_tokens=9), [[10, 11, 12, 13], [1, 2, 3, 4]] + ) == [[11], [1, 2]] - with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"): - moriio_connector.MoRIIOConnectorScheduler( - _gate_vllm_config(read_mode=False), "engine", config + +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_resolves_current_state_instead_of_cached_block_ids(): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config + ) + state = KVConnectorBlockState( + req_ids={"req"}, + resolve_block_ids=lambda _: ([0, 51, 52, 53], [31, 32, 33]), + boundary_state_offloads={}, + ) + assert sched._get_write_blocks( + SimpleNamespace(request_id="req", num_prompt_tokens=9), + [[10, 11, 12], [1]], + SimpleNamespace(kv_connector_block_state=state), + ) == [[51], [31, 32]] + + +@pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize("external_tokens", [0, 9]) +def test_hybrid_write_notifies_grouped_slots_or_releases_a_complete_local_hit( + external_tokens, +): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False), "engine", config + ) + sched.send_notify_block = Mock() + sched._send_transfer_release = Mock() + req = _mk_request( + list(range(10)), + params={ + "do_remote_prefill": True, + "transfer_id": "tx", + "remote_host": "127.0.0.2", + "remote_notify_port": 7000, + }, + ) + req.request_id = "req" + sched.update_state_after_alloc( + req, _FakeBlocks(([10, 11, 12], [1, 2, 3])), external_tokens + ) + if external_tokens: + sched.send_notify_block.assert_called_once_with( + req_id="req", + transfer_id="tx", + block_notify_list=[[10], [1, 2, 3]], + host="127.0.0.2", + port=7000, ) + sched._send_transfer_release.assert_not_called() + else: + sched.send_notify_block.assert_not_called() + sched._send_transfer_release.assert_called_once_with("tx", "127.0.0.2", 7000) + assert not req.kv_transfer_params["do_remote_prefill"] + + +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_keeps_separate_draft_groups_and_their_block_spans(): + config = _dspark_wrapped_attention_config() + # A disabled cache is absent from the wire. Target and draft may have + # different token spans, and neither can borrow the other's block IDs. + draft = SimpleNamespace( + enable_kv_transfer=True, + kv_cache_spec=FullAttentionSpec( + block_size=4, + num_kv_heads=1, + head_size=2, + dtype=torch.float32, + dcp_sharded=False, + ), + ) + disabled = SimpleNamespace(enable_kv_transfer=False, kv_cache_spec=object()) + config.kv_cache_groups.insert(1, disabled) + config.kv_cache_groups.append(draft) + config.transfer_groups = tuple( + g for g in config.kv_cache_groups if g.enable_kv_transfer + ) + config.select_transfer_block_ids = lambda blocks: (blocks[0], blocks[2], blocks[3]) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config( + read_mode=False, dcp_size=2, speculative_config=_spec_config("dspark") + ), + "engine", + config, + ) + req = _mk_request( + list(range(9)), + params={ + "do_remote_decode": True, + "transfer_id": "tx", + }, + ) + req.request_id = "req" + req.status = moriio_connector.RequestStatus.FINISHED_LENGTH_CAPPED + blocks = ([10, 11, 12, 13], [99], [20, 21, 22], [30, 31, 32, 33]) + expected = [[11], [20, 21], [30, 31, 32]] + assert sched._clamp_to_prompt_blocks(req, blocks) == expected + # Exercise the connector entry point too: completion must not flatten + # a multi-group WRITE payload through READ's single-attention helper. + delay_free, params = _FakeConnector(sched).request_finished_all_groups(req, blocks) + assert delay_free + assert params["remote_block_ids"] == expected + + +@pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize( + "blocks,delay_free", + [(([0, 51, 52, 53], [31, 32, 33]), True), (([], []), False)], +) +def test_hybrid_write_request_finished_defers_free_only_with_blocks_to_send( + blocks, delay_free +): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config + ) + sched.is_producer = True + req = _mk_request( + list(range(9)), params={"do_remote_decode": True, "transfer_id": "tx"} + ) + req.request_id = "req" + req.status = moriio_connector.RequestStatus.FINISHED_LENGTH_CAPPED + + assert sched.request_finished(req, blocks)[0] is delay_free + # The WRITE ACK, not the READ watchdog, releases the deferred blocks. + if delay_free: + assert sched._deferred_send_deadlines["req"][1] == "tx" + else: + assert "req" not in sched._deferred_send_deadlines + + +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_build_meta_saves_current_state_on_the_final_chunk(): + """Chunked prefill must not ship block ids cached at admission: align-mode + state slots move between chunks without appending a block.""" + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config + ) + sched.is_producer = True + set_role(moriio_common.ROLE.PRODUCER) + params = { + "do_remote_decode": True, + "transfer_id": "tx", + "remote_engine_id": "decode", + "remote_block_ids": [], + "remote_host": "127.0.0.1", + "remote_handshake_port": 6000, + "remote_notify_port": 6001, + } + req = _mk_request(list(range(9)), params=params) + req.request_id = "req" + req.num_computed_tokens = 0 + sched._reqs_need_save["req"] = (req, [[10, 11, 12], [1]]) + sched._req_kv_params["req"] = dict(params) + + def _output(cached_req_ids, new_block_ids, num_scheduled): + return SimpleNamespace( + scheduled_cached_reqs=SimpleNamespace( + req_ids=cached_req_ids, new_block_ids=new_block_ids + ), + num_scheduled_tokens={"req": num_scheduled}, + new_block_ids_to_zero=None, + kv_connector_block_state=KVConnectorBlockState( + req_ids={"req"}, + resolve_block_ids=lambda _: ([0, 51, 52, 53], [31, 32, 33]), + boundary_state_offloads={}, + ), + ) + + # First chunk: not final, so nothing is saved yet. + meta = sched.build_connector_meta(_output([], [], 4)) + assert not meta.reqs_to_save + assert "req" in sched._reqs_need_pending_save + + # Final chunk: stale cached ids plus new ids are ignored in favour of + # the scheduler's current table, clamped to the prompt. + req.num_computed_tokens = 4 + meta = sched.build_connector_meta(_output(["req"], [([77], [78])], 5)) + assert meta.reqs_to_save["req"].local_block_ids == [[51], [31, 32]] + assert "req" not in sched._reqs_need_pending_save + + +@pytest.mark.parametrize( + "mode,deadlines,expected", + [ + (MoRIIOMode.WRITE, {"req": (0.0, "tx")}, True), + (MoRIIOMode.WRITE, {}, False), + (MoRIIOMode.READ, {"req": (0.0, "tx")}, False), + ], +) +def test_has_pending_push_work_tracks_write_deferred_frees(mode, deadlines, expected): + connector = _FakeConnector(_FakeScheduler(_deferred_send_deadlines=deadlines)) + connector.mode = mode + assert connector.has_pending_push_work() is expected + + +def test_has_pending_push_work_is_false_on_worker_role(): + connector = _FakeConnector(None) + connector.mode = MoRIIOMode.WRITE + assert connector.has_pending_push_work() is False + + +def test_loaded_kv_cache_groups_are_every_transfer_group(): + # The connector restores the request's own blocks, not a hash-addressed + # prefix, so a transfer-enabled Mamba group counts as loaded too. + connector = _FakeConnector(None) + connector._kv_cache_config = SimpleNamespace( + transfer_group_ids=(0, 2), prefix_cacheable_group_ids=(0,) + ) + assert connector.get_loaded_kv_cache_group_ids(SimpleNamespace()) == (0, 2) def test_split_block_groups_rejects_non_hybrid_use(): @@ -728,11 +950,15 @@ def test_session_build_rejects_per_layer_region_count_mismatch(): worker._get_built_session("prefill") -def test_register_kv_caches_rejects_hybrid_write_before_registration(): - worker = _FakeWorker(mode=MoRIIOMode.WRITE, _transfer_layer_names={"kda.0"}) +def test_register_kv_caches_rejects_invalid_hybrid_write_state_layout(): + worker = _FakeWorker( + mode=MoRIIOMode.WRITE, + _transfer_layer_names={"kda.0"}, + layer_to_spec={"kda.0": _mamba_spec(1)}, + ) worker._is_mamba_layer = lambda _layer_name: True - with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"): + with pytest.raises(moriio_common.MoRIIOError, match="exactly two Mamba states"): worker.register_kv_caches({"kda.0": object()}) @@ -933,6 +1159,40 @@ def test_read_matches_only_requests_with_pending_remote_prefill(params, expected assert scheduler.get_num_new_matched_tokens(request, 0) == expected +@pytest.mark.parametrize( + "has_mamba,expected", [(False, (32, True)), (True, (31, True))] +) +def test_write_consumer_matches_full_prompt_unless_hybrid(has_mamba, expected): + """Without a mamba group (attention-only, or HMA disabled) prefill pushed + KV for all N tokens; with one it pushed h(N-1).""" + scheduler = _FakeScheduler( + is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=has_mamba + ) + request = SimpleNamespace( + num_prompt_tokens=32, kv_transfer_params={"do_remote_prefill": True} + ) + + assert scheduler.get_num_new_matched_tokens(request, 0) == expected + + +@pytest.mark.parametrize( + "is_producer,has_mamba,expected_prompt_len", + [(True, True, 4), (True, False, 5), (False, True, 5)], +) +def test_on_new_request_truncates_only_hybrid_producers( + is_producer, has_mamba, expected_prompt_len +): + """Attention-only prefill (including models run with HMA disabled) + computes all N prompt tokens; only a hybrid producer drops the last.""" + scheduler = _FakeScheduler(is_producer=is_producer, _has_mamba=has_mamba) + request = _mk_request([10, 11, 12, 13, 14], params={"do_remote_decode": True}) + + scheduler.on_new_request(request) + + assert request.num_prompt_tokens == expected_prompt_len + assert len(request.prompt_token_ids) == expected_prompt_len + + def test_read_rematch_after_alloc_keeps_queued_transfer(): """Clearing the flag blocks another external hit, not the queued READ.""" sched = _make_read_scheduler() @@ -981,12 +1241,37 @@ def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens(): assert is_async is True +@pytest.mark.parametrize("params", [None, {}, {"do_remote_prefill": False}]) +def test_write_does_not_wait_again_after_receive_or_for_local_requests(params): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True) + req = _mk_request(list(range(10))) + req.kv_transfer_params = params + assert sched.get_num_new_matched_tokens(req, 0) == (0, False) + + +def test_write_never_declares_an_empty_async_load(): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=False) + req = _mk_request(list(range(10)), params={"do_remote_prefill": True}) + assert sched.get_num_new_matched_tokens(req, 10) == (0, False) + + +@pytest.mark.parametrize( + "prompt_len,computed,expected", [(16, 0, 15), (17, 16, 0), (1, 0, 0)] +) +def test_hybrid_write_advertises_only_the_state_computed_by_prefill( + prompt_len, computed, expected +): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True) + req = _mk_request(list(range(prompt_len)), params={"do_remote_prefill": True}) + assert sched.get_num_new_matched_tokens(req, computed) == (expected, expected > 0) + + @pytest.mark.parametrize( ("mode", "num_computed_tokens", "expected", "is_async"), [ (MoRIIOMode.READ, 0, 9, False), (MoRIIOMode.READ, 10, 0, False), - (MoRIIOMode.WRITE, 2, 8, True), + (MoRIIOMode.WRITE, 2, 7, True), ], ) def test_get_num_new_matched_tokens_supports_embeds_only_prompts( diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 5ec6fb5755..f0559f259d 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -7,6 +7,7 @@ import threading from collections import OrderedDict from types import SimpleNamespace from typing import Any +from unittest.mock import Mock import pytest import torch @@ -106,7 +107,7 @@ def _write_task(layer_name: str, transfer_id: str = "xfer") -> Any: request_id="req", transfer_id=transfer_id, dst_engine_id="remote-engine", - local_block_ids=[1, 3], + local_block_ids=[[1, 3]], remote_block_ids_hint=None, layer_name=layer_name, event=None, @@ -345,6 +346,8 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry(): calls: list[str] = [] class FakeWorker: + _has_mamba: bool + layer_to_group: dict[str, int] kv_caches: dict[str, torch.Tensor] layer_name_to_local_kv_cache_metadata: dict[str, list[Any]] @@ -364,11 +367,13 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry(): return ([call_id], [call_id + 10], [call_id + 20]) fake_worker = FakeWorker() + fake_worker._has_mamba = False + fake_worker.layer_to_group = dict.fromkeys(kv_caches, 0) fake_worker.kv_caches = kv_caches fake_worker.layer_name_to_local_kv_cache_metadata = {name: [] for name in kv_caches} writer = MoRIIOWriter.__new__(MoRIIOWriter) writer._worker_ref = lambda: fake_worker - request_info = RemoteAllocInfo(block_ids=[4, 5]) + request_info = RemoteAllocInfo(block_ids=[[4, 5]]) remote_meta = _remote_meta() dense0_plan = writer._prepare_transfer_plan( @@ -395,7 +400,7 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry(): def test_write_scheduler_deduplicates_layers_and_seals_expected_count(): - request_info = RemoteAllocInfo(block_ids=[4, 5]) + request_info = RemoteAllocInfo(block_ids=[[4, 5]]) wrapper = _wrapper_for_messages() wrapper.done_remote_allocate_req_dict["xfer"] = request_info writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) @@ -411,6 +416,80 @@ def test_write_scheduler_deduplicates_layers_and_seals_expected_count(): assert writer._sealed_writes["xfer"] == 2 +def test_hybrid_write_attention_uses_its_own_group_not_the_first_group(): + worker = SimpleNamespace( + _has_mamba=True, + kv_caches={"draft": torch.empty((8, 1, 4, 3))}, + layer_to_group={"draft": 2}, + _region_session_indices=lambda _: [5], + _compute_block_transfer_offsets=Mock(return_value=([1], [2], [3])), + ) + writer = make_moriio_writer(worker) + task = _write_task("draft") + task.local_block_ids = [[10], [20], [30, 31]] + remote = RemoteAllocInfo(block_ids=[[40], [50], [60, 61, 62]]) + meta = _remote_meta() + plan = writer._prepare_transfer_plan(task, remote, meta) + worker._compute_block_transfer_offsets.assert_called_once_with( + "draft", [30, 31], [60, 61], meta, remote_engine_id=task.dst_engine_id + ) + assert plan.sess_idx == 5 + + +def test_hybrid_write_posts_both_regions_before_layer_completion(): + from unittest.mock import Mock + + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorWorker, + ) + + wrapper = _wrapper_for_messages() + info = RemoteAllocInfo(block_ids=[[], [7], [9]]) + wrapper.done_remote_allocate_req_dict["xfer"] = info + wrapper.write_remote_data = Mock(side_effect=["conv-status", "ssm-status"]) + worker = SimpleNamespace( + moriio_wrapper=wrapper, + tp_rank=0, + get_engine_name_with_dp=lambda name, rank: name, + _get_built_session=lambda name: (["conv-session", "ssm-session"], None), + _is_mamba_layer=lambda name: True, + _mamba_payload_index_by_layer={"kda": 2}, + _region_session_indices=lambda name: [0, 1], + _compute_mamba_transfer_offsets=Mock( + return_value=([10, 20, 30, 40], [50, 60, 70, 80], [2, 4, 8, 16], 3) + ), + ) + worker._mamba_blocks_for_layer = lambda name, blocks: ( + MoRIIOConnectorWorker._mamba_blocks_for_layer(worker, name, blocks) + ) + writer = make_moriio_writer(worker) + task = _write_task("kda") + task.local_block_ids = [[], [1], [3]] + task.remote_tp_size = 2 + task.event = Mock() + writer._execute_write_task(task) + + task.event.synchronize.assert_called_once() + worker._compute_mamba_transfer_offsets.assert_called_once_with( + "kda", [3], [9], remote_tp_size=2 + ) + assert wrapper.write_remote_data.call_args_list[0].args == ( + [2, 4, 8], + [10, 20, 30], + [50, 60, 70], + "conv-session", + ) + assert wrapper.write_remote_data.call_args_list[1].args == ( + [16], + [40], + [80], + "ssm-session", + ) + assert info.transfer_statuses == ["conv-status", "ssm-status"] + assert info.writes_done == 1 + assert not info.completion_notified + + def test_write_completion_notifies_once_after_all_sealed_writes_finish(): class FakeWrapper: def __init__(self): @@ -445,7 +524,7 @@ def test_write_completion_notifies_once_after_all_sealed_writes_finish(): ) wrapper = FakeWrapper() - request_info = RemoteAllocInfo(block_ids=[4, 5], writes_expected=2) + request_info = RemoteAllocInfo(block_ids=[[4, 5]], writes_expected=2) request_info.transfer_statuses.extend(["status-a", "status-b"]) request_info.completion_request_id = "req" # Final notify port is resolved in _execute_write_task; finalize reads it @@ -487,6 +566,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain(): wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) wrapper.lock = threading.Lock() wrapper._transfer_timeout = 1 + wrapper.moriio_engine = object() global_status = FakeStatus() scoped_status = FakeStatus() wrapper.transfer_status = [global_status] @@ -500,7 +580,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain(): def test_write_failure_marks_terminal_and_clears_scheduled_state(): wrapper = _wrapper_for_messages() - wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5]) + wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[[4, 5]]) writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) writer._scheduled_writes["xfer"] = 2 writer._scheduled_layers["xfer"] = {"dense0", "indexer"} @@ -518,7 +598,7 @@ def test_write_failure_marks_terminal_and_clears_scheduled_state(): def test_schedule_write_rejects_terminal_transfer_without_recreating_state(): wrapper = _wrapper_for_messages() - wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5]) + wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[[4, 5]]) writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) writer._scheduled_writes["xfer"] = 1 writer._scheduled_layers["xfer"] = {"dense0"} @@ -610,7 +690,7 @@ def test_moriio_wrapper_routes_valid_messages(role, payload, expected): if expected == "remote_blocks": request_info = wrapper.done_remote_allocate_req_dict["xfer"] - assert request_info.block_ids == [4, 5] + assert request_info.block_ids == [[4, 5]] assert request_info.decode_dp_rank == 3 elif expected == "write_done": assert wrapper.done_write_cache_req_ids == ["xfer"] diff --git a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py index 2a3092ad5a..8fa59e5fb5 100644 --- a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py +++ b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py @@ -11,6 +11,8 @@ import importlib import importlib.util import threading import time +from types import SimpleNamespace +from unittest.mock import Mock import pytest @@ -127,3 +129,62 @@ def test_wait_blocks_until_terminal(): assert time.monotonic() - start >= 0.15 finally: timer.cancel() + + +@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 50), (0.0001, 1)]) +def test_native_wait_receives_all_statuses_and_a_positive_deadline( + timeout, expected_ms +): + statuses = [_pending(), _ok()] + engine = SimpleNamespace( + wait_all=Mock(return_value=moriio_engine.StatusCode.SUCCESS) + ) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=timeout) + wrapper.waiting_for_transfer_complete(statuses) + engine.wait_all.assert_called_once_with(statuses, timeout_ms=expected_ms) + + +@pytest.mark.parametrize( + "result,statuses,message", + [ + ("ERR_RDMA_OP", [_pending(), _bad()], "SQ full"), + ("IN_PROGRESS", [_pending(), _ok()], "timed out"), + ("ERR_INVALID_ARGS", [_pending()], "ERR_INVALID_ARGS"), + ], +) +def test_native_wait_reports_failure_or_timeout(result, statuses, message): + engine = SimpleNamespace( + wait_all=Mock(return_value=getattr(moriio_engine.StatusCode, result)) + ) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + with pytest.raises(TransferError, match=message): + wrapper.waiting_for_transfer_complete(statuses) + + +def test_wait_falls_back_without_native_status_enum(monkeypatch): + monkeypatch.setattr(moriio_engine, "StatusCode", None) + engine = SimpleNamespace(wait_all=Mock()) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + with pytest.raises(TransferError, match="SQ full"): + wrapper.waiting_for_transfer_complete([_bad(), _ok()]) + engine.wait_all.assert_not_called() + + +def test_nonblocking_poll_does_not_call_native_wait(): + engine = SimpleNamespace(wait_all=Mock()) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + assert wrapper.poll_transfer_batch([_pending()]) is TransferBatchState.PENDING + wrapper.waiting_for_transfer_complete([]) + engine.wait_all.assert_not_called() + + +def test_native_wait_drains_only_the_unscoped_status_snapshot(): + statuses = [_pending(), _ok()] + engine = SimpleNamespace( + wait_all=Mock(return_value=moriio_engine.StatusCode.SUCCESS) + ) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + wrapper.transfer_status.extend(statuses) + wrapper.waiting_for_transfer_complete() + engine.wait_all.assert_called_once_with(statuses, timeout_ms=50) + assert wrapper.transfer_status == [] diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index 57b6893869..ee04c2cabe 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -14,6 +14,8 @@ import pytest from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( ROLE, + LayerTransferPlan, + MoRIIOError, MoRIIOMode, RemoteAllocInfo, ReqMeta, @@ -96,7 +98,7 @@ def _write_task( request_id=request_id, transfer_id=transfer_id, dst_engine_id="remote-engine", - local_block_ids=[1, 3], + local_block_ids=[[1, 3]], remote_block_ids_hint=None, layer_name="layer0", event=None, @@ -249,6 +251,7 @@ def test_finalize_single_pod_port_uses_global_rank(): def test_write_blocks_for_req_forwards_routing_to_schedule_write(): worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = False captured: dict[str, Any] = {} def _capture_schedule_write_blocks(**kwargs: Any) -> None: @@ -266,13 +269,14 @@ def test_write_blocks_for_req_forwards_routing_to_schedule_write(): assert captured["multi_pod_hosts"] == ["10.0.0.1", "10.0.0.2"] assert captured["remote_dp_size_local"] == 8 assert captured["remote_ip"] == "10.0.0.1" - assert captured["local_block_ids"] == [1] - assert captured["remote_block_ids"] == [2] + assert captured["local_block_ids"] == [[1]] + assert captured["remote_block_ids"] == [[2]] assert not hasattr(worker, "multi_pod_hosts") def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty(): worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = False captured: dict[str, Any] = {} def _capture_schedule_write_blocks(**kwargs: Any) -> None: @@ -290,8 +294,8 @@ def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty(): assert captured["multi_pod_hosts"] == ["10.0.0.9"] assert captured["remote_dp_size_local"] == 4 - assert captured["local_block_ids"] == [1] - assert captured["remote_block_ids"] == [2] + assert captured["local_block_ids"] == [[1]] + assert captured["remote_block_ids"] == [[2]] def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): @@ -373,3 +377,73 @@ def test_write_failed_is_reported_as_failed_recving_on_owning_rank(): assert results.failed_recving == {"req-1"} assert wrapper.pop_finished_write_req_ids() == set() assert list(worker._unmatched_write_failures) == ["tOtherRank"] + + +def test_hybrid_write_rejects_unequal_tp_before_scheduling_any_layer(): + from unittest.mock import Mock + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = True + worker.world_size = 2 + worker.schedule_write_blocks = Mock() + + with pytest.raises(NotImplementedError, match="heterogeneous-TP"): + worker._write_blocks_for_req("req", _req_meta(), "attention", kv_layer=None) + + worker.schedule_write_blocks.assert_not_called() + + +@pytest.mark.parametrize( + "notify,expected", + [ + # Plain WRITE peers still send one flat attention list. + ([7, 8], [[7, 8]]), + # Hybrid WRITE sends one list per transfer group, kept as-is even + # when the first group is empty (full attention hit, state only). + ([[7], [9]], [[7], [9]]), + ([[], [9]], [[], [9]]), + ], +) +def test_remote_blocks_message_normalizes_to_grouped_block_ids(notify, expected): + wrapper = _real_wrapper() + with patch.object(moriio_engine, "get_role", return_value=ROLE.PRODUCER): + wrapper._handle_remote_blocks_message( + {"transfer_id": "tA", "block_notify_list": notify, "decode_rank": 2} + ) + info = wrapper.done_remote_allocate_req_dict["tA"] + assert info.block_ids == expected + assert info.decode_dp_rank == 2 + + +def test_empty_layer_plan_posts_no_rdma_write(): + from unittest.mock import Mock + + wrapper = SimpleNamespace(write_remote_data=Mock(), write_remote_data_single=Mock()) + writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) + plan = LayerTransferPlan( + request_id="req", + transfer_id="tA", + layer_name="layer0", + sess_idx=0, + transfer_local_offsets=[], + transfer_remote_offsets=[], + transfer_sizes=[], + ) + + assert writer._do_layer_write(plan, ["session"]) == [] + wrapper.write_remote_data.assert_not_called() + wrapper.write_remote_data_single.assert_not_called() + + +@pytest.mark.parametrize("local_state", [[], [1, 2]]) +def test_hybrid_write_requires_exactly_one_running_state_slot(local_state): + worker = SimpleNamespace( + _mamba_blocks_for_layer=lambda layer_name, block_ids: block_ids[1], + ) + writer = make_moriio_writer(worker) + task = _write_task("tA", "10.0.0.1") + task.layer_name = "kda" + task.local_block_ids = [[], local_state] + + with pytest.raises(MoRIIOError, match="one running-state slot"): + writer._prepare_mamba_transfer_plans(task, RemoteAllocInfo(block_ids=[[], [9]])) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index e82d066818..04fd38be2f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -63,12 +63,13 @@ class WriteTask: request_id: ReqId transfer_id: TransferId dst_engine_id: str - local_block_ids: list[int] - remote_block_ids_hint: list[int] | None + local_block_ids: BlockIds + remote_block_ids_hint: BlockIds | None layer_name: str event: torch.cuda.Event remote_notify_port: int remote_ip: str + remote_tp_size: int | None = None multi_pod_hosts: list[str] = field(default_factory=list) remote_dp_size_local: int = 0 remote_dp_size: int = 1 @@ -94,7 +95,7 @@ class LayerTransferPlan: class RemoteAllocInfo: """Information about remote block allocation.""" - block_ids: list[int] + block_ids: BlockIds | None writes_done: int = 0 writes_expected: int | None = None decode_dp_rank: int = 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index fa75eaee3d..9760d8bef1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -338,7 +338,10 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - if self.connector_scheduler._has_mamba: + if ( + self.connector_scheduler._has_mamba + and self.connector_scheduler.mode == MoRIIOMode.READ + ): attn_block_ids, mamba_block_groups = ( self.connector_scheduler.split_block_groups(block_ids) ) @@ -420,6 +423,16 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): ) self.connector_worker.wait_for_save(self._connector_metadata) + def has_pending_push_work(self) -> bool: + return ( + self.mode == MoRIIOMode.WRITE + and self.connector_scheduler is not None + and bool(self.connector_scheduler._deferred_send_deadlines) + ) + + def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]: + return self._kv_cache_config.transfer_group_ids + def shutdown(self): if self.connector_worker is not None: self.connector_worker.shutdown() @@ -463,11 +476,10 @@ def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None: speculative_config = vllm_config.speculative_config if speculative_config is None: return - # DSpark is allowlisted because it is the only hybrid speculative method - # validated end to end with MoRIIO READ. + # Other hybrid speculative methods need their own transfer validation. if not speculative_config.use_dspark(): raise MoRIIOError( - "MoRIIO hybrid READ supports DSpark speculative decoding only, got " + "MoRIIO hybrid transfer supports DSpark speculative decoding only, got " f"method={speculative_config.method!r}" ) @@ -480,14 +492,14 @@ def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None: for spec in specs: if len(spec.shapes) != 2 or len(spec.dtypes) != 2: raise MoRIIOError( - "MoRIIO hybrid READ supports exactly two Mamba states " + "MoRIIO hybrid transfer supports exactly two Mamba states " "(conv and SSM), got " f"{len(spec.shapes)} shapes and {len(spec.dtypes)} dtypes" ) reference = specs[0] if any(spec != reference for spec in specs[1:]): raise MoRIIOError( - "MoRIIO hybrid READ requires all transferable Mamba groups " + "MoRIIO hybrid transfer requires all transferable Mamba groups " "to use the same cache spec" ) return reference @@ -515,7 +527,9 @@ class MoRIIOConnectorScheduler: self._has_mamba = bool(self._mamba_group_ids) mamba_spec = None if self._has_mamba: - if len(self._attn_group_ids) != 1: + if not self._attn_group_ids: + raise MoRIIOError("MoRIIO hybrid transfer requires an attention group") + if self.mode == MoRIIOMode.READ and len(self._attn_group_ids) != 1: raise MoRIIOError( "MoRIIO hybrid READ requires exactly one transferable " "attention group, got " @@ -529,11 +543,6 @@ class MoRIIOConnectorScheduler: ] mamba_spec = _validate_mamba_specs(mamba_specs) _validate_hybrid_speculation(vllm_config) - if self.mode != MoRIIOMode.READ: - raise MoRIIOError( - "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " - "mode only; set kv_connector_extra_config.read_mode=true" - ) self._num_ssm_scratch_blocks = ( mamba_spec.num_speculative_blocks if mamba_spec is not None else 0 ) @@ -581,11 +590,13 @@ class MoRIIOConnectorScheduler: f"cache group spec(s): {sorted(set(unsupported))}. Pass " "--disable-hybrid-kv-cache-manager to run without HMA." ) - if self.mode == MoRIIOMode.WRITE: - # TODO(simondanielsson): support HMA in WRITE mode + if self.mode == MoRIIOMode.WRITE and any( + isinstance(g.kv_cache_spec, SlidingWindowSpec) + for g in kv_cache_config.transfer_groups + ): raise NotImplementedError( - "MoRIIO WRITE mode does not support hybrid KV cache groups " - "(sliding-window attention). Use READ mode, or pass " + "MoRIIO WRITE mode does not support sliding-window " + "cache groups. Use READ mode, or pass " "--disable-hybrid-kv-cache-manager." ) @@ -784,10 +795,11 @@ class MoRIIOConnectorScheduler: if not params.get("do_remote_prefill"): return 0, False if self.mode == MoRIIOMode.WRITE: - # MoriiO in write mode, no remote prefill. Hybrid models never get - # here: register_kv_caches refuses WRITE mode for them, so there is - # no recurrent-state accounting to do. - return num_external_tokens, True + if self._has_mamba: + # P pushes h(N-1). Do not publish it under the N-token hash + # when the scheduler caches a completed asynchronous load. + num_external_tokens = max(num_external_tokens - 1, 0) + return num_external_tokens, num_external_tokens > 0 # READ mode always recomputes the last token locally on the decoder. # @@ -1178,6 +1190,16 @@ class MoRIIOConnectorScheduler: f"remote_notify_port={remote_notify_port!r})" ) + block_notify_list: list[int] | list[list[int]] + if self._has_mamba: + block_notify_list = list( + self._clamp_to_prompt_blocks( + request, blocks.get_block_ids() + ) + ) + else: + block_notify_list = blocks.get_block_ids()[0] + # Wide-EP multi-pod: a pod binds notify sockets only for # its LOCAL ranks, so the port offset must use the per-pod # local rank (% dp_local), not the global rank. Single-pod @@ -1204,7 +1226,7 @@ class MoRIIOConnectorScheduler: self.send_notify_block( req_id=request.request_id, transfer_id=request.kv_transfer_params["transfer_id"], - block_notify_list=blocks.get_block_ids()[0], + block_notify_list=block_notify_list, host=_notify_host, port=target_port, ) @@ -1227,6 +1249,26 @@ class MoRIIOConnectorScheduler: sliding-window groups are left unchanged. """ + if self._has_mamba: + # WRITE preserves transfer-group order, including a separate drafter + # attention group. READ retains its attention-first wire format. + transfer_blocks = self.kv_cache_config.select_transfer_block_ids(block_ids) + dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size + clamped_hybrid: list[list[int]] = [] + for group, blocks in zip( + self.kv_cache_config.transfer_groups, transfer_blocks, strict=True + ): + spec = group.kv_cache_spec + if isinstance(spec, MambaSpec): + clamped_hybrid.append( + clip_ssm_state_blocks(list(blocks), spec.num_speculative_blocks) + ) + else: + block_size = spec.block_size * (dcp_size if spec.dcp_sharded else 1) + clamped_hybrid.append( + list(blocks[: cdiv(req.num_prompt_tokens, block_size)]) + ) + return clamped_hybrid clamped = list(block_ids) changed = False for g, group_blocks in enumerate(block_ids): @@ -1252,6 +1294,19 @@ class MoRIIOConnectorScheduler: num_scheduled = scheduler_output.num_scheduled_tokens.get(req.request_id, 0) return req.num_computed_tokens + num_scheduled >= req.num_prompt_tokens + def _get_write_blocks( + self, req: "Request", block_ids: BlockIds, scheduler_output: SchedulerOutput + ) -> BlockIds: + if self._has_mamba: + # Align-mode state slots can move without appending a new block. + # Resolve the authoritative table at the final prefill step. + block_state = scheduler_output.kv_connector_block_state + assert block_state is not None + current_blocks = block_state.get_block_ids(req.request_id) + assert current_blocks is not None + block_ids = current_blocks + return self._clamp_to_prompt_blocks(req, block_ids) + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -1274,10 +1329,10 @@ class MoRIIOConnectorScheduler: # reserved speculative-lookahead slack), so new_block_ids is None; # still fall through to the final-chunk check below. new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i] - if new_block_ids is not None: + if new_block_ids is not None and not self._has_mamba: existing_blocks = [ - existing_blocks[g] + new_block_ids[g] - for g in range(len(new_block_ids)) + old + new + for old, new in zip(existing_blocks, new_block_ids, strict=True) ] self._reqs_need_pending_save[req_id] = (req, existing_blocks) if not self._is_final_prefill_chunk(req, scheduler_output): @@ -1287,7 +1342,9 @@ class MoRIIOConnectorScheduler: kv_params = self._req_kv_params.pop( req_id, req.kv_transfer_params or {} ) - save_block_ids = self._clamp_to_prompt_blocks(req, existing_blocks) + save_block_ids = self._get_write_blocks( + req, existing_blocks, scheduler_output + ) meta.add_new_req( request_id=req_id, local_block_ids=save_block_ids, @@ -1328,7 +1385,7 @@ class MoRIIOConnectorScheduler: if not self._is_final_prefill_chunk(req, scheduler_output): self._reqs_need_pending_save[req_id] = (req, block_ids) continue - save_block_ids = self._clamp_to_prompt_blocks(req, block_ids) + save_block_ids = self._get_write_blocks(req, block_ids, scheduler_output) meta.add_new_req( request_id=req_id, local_block_ids=save_block_ids, @@ -1425,7 +1482,7 @@ class MoRIIOConnectorScheduler: """Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later. - ``block_ids`` normally contains all cache groups. Hybrid callers from + ``block_ids`` normally contains all cache groups. Hybrid READ callers from ``request_finished_all_groups`` pass a flat attention list and the recurrent-state groups separately in ``mamba_block_groups``. """ @@ -1494,6 +1551,10 @@ class MoRIIOConnectorScheduler: list(cast(list[int], block_ids)), *[list(group) for group in mamba_block_groups], ] + elif self._has_mamba and self.mode == MoRIIOMode.WRITE: + computed_block_ids = self._clamp_to_prompt_blocks( + request, cast(BlockIds, block_ids) + ) elif self._has_mamba: raise MoRIIOError("Hybrid request completion requires all KV cache groups") else: @@ -1660,7 +1721,7 @@ class MoRIIOConnectorWorker: _, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config) self._num_mamba_transfer_groups = len(mamba_transfer_group_ids) self._mamba_payload_index_by_layer = { - layer_name: payload_index + layer_name: group_id if self.mode == MoRIIOMode.WRITE else payload_index for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1) for layer_name in kv_cache_config.transfer_groups[group_id].layer_names } @@ -1884,8 +1945,8 @@ class MoRIIOConnectorWorker: request_id: ReqId, transfer_id: TransferId, dst_engine_id: str, - local_block_ids: list[int], - remote_block_ids: list[int] | None, + local_block_ids: BlockIds, + remote_block_ids: BlockIds | None, layer_name: str, kv_layer: torch.Tensor, remote_notify_port: int, @@ -1893,6 +1954,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts: list[str], remote_dp_size_local: int, remote_dp_size: int, + remote_tp_size: int | None = None, ) -> None: """Schedule a block write operation. @@ -1909,6 +1971,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts: List of pod IPs for multi-pod Wide-EP remote_dp_size_local: Per-pod DP size for multi-pod remote_dp_size: Global DP size of the decode instance + remote_tp_size: Tensor parallel size advertised by the peer """ # synchronization to prevent dirty reads between @@ -1933,6 +1996,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts=multi_pod_hosts, remote_dp_size_local=remote_dp_size_local, remote_dp_size=remote_dp_size, + remote_tp_size=remote_tp_size, ) self._writer.schedule_write(task) @@ -2359,15 +2423,6 @@ class MoRIIOConnectorWorker: self._has_mamba = any( self._is_mamba_layer(layer_name) for layer_name in kv_caches ) - if self._has_mamba and self.mode == MoRIIOMode.WRITE: - # WRITE needs a distinct conv/ssm transfer plan. Reject before - # allocating or registering any regions rather than silently - # serving a zero recurrent state. - raise MoRIIOError( - "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " - "mode only; set kv_connector_extra_config.read_mode=true " - "for hybrid models" - ) if self._has_mamba: mamba_specs = [ cast(MambaSpec, self.layer_to_spec[layer_name]) @@ -3292,6 +3347,9 @@ class MoRIIOConnectorWorker: ) def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer): + if self._has_mamba: + # Reject before even an attention-layer write can be queued. + self._mamba_tp_ratio(int(meta.tp_size)) # Compute per-request values to pass through the task (no shared state). # This avoids race conditions when concurrent requests target different # decode pods - each WriteTask carries its own routing info. @@ -3303,16 +3361,12 @@ class MoRIIOConnectorWorker: if meta.remote_dp_size_local else int(meta.remote_dp_size) ) - # WRITE does not support HMA, so unwrap blocks into flat lists. - # remote_block_ids can itself be empty so need to be careful when unwrapping. - local_block_ids = meta.local_block_ids[0] - remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else [] self.schedule_write_blocks( request_id=req_id, transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, - local_block_ids=local_block_ids, - remote_block_ids=remote_block_ids, + local_block_ids=meta.local_block_ids, + remote_block_ids=meta.remote_block_ids, layer_name=layer_name, kv_layer=kv_layer, remote_notify_port=meta.remote_notify_port, @@ -3320,6 +3374,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts=hosts, remote_dp_size_local=dp_local, remote_dp_size=int(meta.remote_dp_size), + remote_tp_size=int(meta.tp_size), ) def merge_contiguous_blocks( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index e33a1524e1..3ffd6d10f9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -3,6 +3,7 @@ import threading import time from collections import OrderedDict, defaultdict +from math import ceil from queue import Empty, Queue from typing import TYPE_CHECKING, Any from weakref import ref as weakref_ref @@ -61,6 +62,11 @@ try: except ImportError: logger.error("MoRIIO is not available") +try: + from mori.io import StatusCode +except ImportError: + StatusCode = None + """Write task execution logic for MoRIIO connector.""" @@ -372,11 +378,15 @@ class MoRIIOWriter: task.dst_engine_id ) - # Prepare transfer plan - plan = self._prepare_transfer_plan(task, request_info, remote_moriio_meta) - - # Execute transfer - transfer_statuses = self._do_layer_write(plan, sessions) + if self.worker._is_mamba_layer(task.layer_name): + plans = self._prepare_mamba_transfer_plans(task, request_info) + else: + plans = [ + self._prepare_transfer_plan(task, request_info, remote_moriio_meta) + ] + transfer_statuses = [] + for plan in plans: + transfer_statuses.extend(self._do_layer_write(plan, sessions)) with self._write_state_lock: request_info.transfer_statuses.extend(transfer_statuses) @@ -404,10 +414,17 @@ class MoRIIOWriter: key = (task.layer_name, *_get_write_geometry_key(layer_cache)) offsets = request_info.transfer_offsets.get(key) if offsets is None: + group = self.worker.layer_to_group[task.layer_name] + assert request_info.block_ids is not None + local_blocks = task.local_block_ids[group] + remote_blocks = request_info.block_ids[group] + if self.worker._has_mamba: + # D can reserve a tail block for the token recomputed locally. + remote_blocks = remote_blocks[: len(local_blocks)] offsets = self.worker._compute_block_transfer_offsets( task.layer_name, - task.local_block_ids, - request_info.block_ids, + local_blocks, + remote_blocks, remote_moriio_meta, remote_engine_id=task.dst_engine_id, ) @@ -430,6 +447,44 @@ class MoRIIOWriter: use_batch=True, ) + def _prepare_mamba_transfer_plans( + self, task: WriteTask, request_info: RemoteAllocInfo + ) -> list[LayerTransferPlan]: + """Plan both registered regions as one scheduled layer write.""" + assert request_info.block_ids is not None + local_slots = self.worker._mamba_blocks_for_layer( + task.layer_name, task.local_block_ids + ) + remote_slots = self.worker._mamba_blocks_for_layer( + task.layer_name, request_info.block_ids + ) + if len(local_slots) != 1 or len(remote_slots) != 1: + raise MoRIIOError( + "Hybrid WRITE requires one running-state slot per Mamba group" + ) + local, remote, sizes, n_conv = self.worker._compute_mamba_transfer_offsets( + task.layer_name, + local_slots, + remote_slots, + remote_tp_size=task.remote_tp_size, + ) + conv_session, ssm_session = self.worker._region_session_indices(task.layer_name) + return [ + LayerTransferPlan( + request_id=task.request_id, + transfer_id=task.transfer_id, + layer_name=task.layer_name, + sess_idx=session, + transfer_local_offsets=local[region], + transfer_remote_offsets=remote[region], + transfer_sizes=sizes[region], + ) + for session, region in ( + (conv_session, slice(0, n_conv)), + (ssm_session, slice(n_conv, None)), + ) + ] + def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]: """Perform the actual layer write. @@ -438,6 +493,8 @@ class MoRIIOWriter: sessions: List of transfer sessions """ + if not plan.transfer_sizes: + return [] if plan.use_batch: return [ self.worker.moriio_wrapper.write_remote_data( @@ -721,7 +778,28 @@ class MoRIIOWrapper: if not transfers_to_wait: return - self._poll_transfers_until_done(transfers_to_wait) + wait_all = getattr(self.moriio_engine, "wait_all", None) + if wait_all is None or StatusCode is None: + self._poll_transfers_until_done(transfers_to_wait) + return + + # Zero means a nonblocking progress call in Mori, not an immediate + # deadline. Round up so even sub-millisecond waits remain blocking. + timeout_ms = ceil(self._transfer_timeout * 1000) + result = wait_all(transfers_to_wait, timeout_ms=timeout_ms) + if result != StatusCode.SUCCESS: + details = [ + f"{status.Message()} (code={status.Code()})" + for status in transfers_to_wait + if status.Failed() + ] + if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS): + details.append( + f"RDMA transfer timed out after {self._transfer_timeout:g}s" + ) + raise TransferError( + f"MoRIIO batch wait failed ({result}): " + "; ".join(details) + ) def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None: """Fallback for mori builds without the batched wait. @@ -858,7 +936,12 @@ class MoRIIOWrapper: ) return self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo( - block_ids=block_notify_list, decode_dp_rank=decode_dp_rank + block_ids=( + [block_notify_list] + if isinstance(block_notify_list[0], int) + else block_notify_list + ), + decode_dp_rank=decode_dp_rank, ) def _handle_write_done_message(self, data: dict):