[Feature][MoRIIO] Support hybrid WRITE transfers with DSpark (#60544)

Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: OpenAI Codex <noreply@openai.com>
Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com>
Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

.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

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.

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

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

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."""

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"
+    )

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)

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(
+
+@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(

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"]

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 == []

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]]))

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

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(

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
 
+        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):