From a2ec7c68e38d867abd47afc54f8987163778b489 Mon Sep 17 00:00:00 2001 From: kliuae Date: Thu, 29 Oct 2026 16:27:32 +0800 Subject: [PATCH 1/6] support dcp for decode instance with mla Signed-off-by: kliuae Co-authored-by: Claude Opus --- tests/test_config.py | 4 + .../unit/test_moriio_dcp_write.py | 300 ++++++++++++++++++ .../unit/test_moriio_hma_scheduler.py | 2 + .../unit/test_moriio_kv_layout.py | 2 +- .../unit/test_moriio_write_done_routing.py | 3 +- vllm/config/vllm.py | 28 +- .../kv_connector/v1/moriio/moriio_common.py | 4 + .../v1/moriio/moriio_connector.py | 35 ++ .../kv_connector/v1/moriio/moriio_engine.py | 17 +- .../v1/moriio/moriio_tp_mapping.py | 56 ++++ 10 files changed, 439 insertions(+), 12 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_moriio_dcp_write.py diff --git a/tests/test_config.py b/tests/test_config.py index a53843e2a93a..b78be1c2a7e2 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -395,6 +395,10 @@ def test_per_request_spec_decode_metrics_requires_spec_decode(): kv_connector="NixlConnector", kv_role="kv_both", ), + KVTransferConfig( + kv_connector="MoRIIOConnector", + kv_role="kv_consumer", + ), KVTransferConfig( kv_connector="MultiConnector", kv_role="kv_both", diff --git a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py new file mode 100644 index 000000000000..a96e3ace9f49 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MoRIIO WRITE into a decode instance with decode context parallelism (DCP).""" + +import threading +from collections import OrderedDict +from types import SimpleNamespace + +import msgspec +import pytest +import torch + +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, + MoRIIOAgentMetadata, + RemoteAllocInfo, + WriteTask, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorWorker, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_engine import ( + MoRIIOWrapper, +) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_tp_mapping import ( + dcp_write_block_pairs, + validate_write_dcp_topology, +) +from vllm.utils.math_utils import cdiv +from vllm.v1.kv_cache_interface import MLAAttentionSpec + +from .utils import make_moriio_writer + +pytestmark = pytest.mark.cpu_test + +BLOCK_SIZE = 4 +NUM_BLOCKS = 24 +NOTIFY_PORT = 20000 +LAYERS = ("mla0", "mla1") + + +@pytest.mark.parametrize("num_tokens", [1, 65, 256, 257, 1000]) +@pytest.mark.parametrize("dcp_size", [2, 4]) +@pytest.mark.parametrize("extra_remote_blocks", [0, 1]) +def test_dcp_pairs_cover_every_block_once(num_tokens, dcp_size, extra_remote_blocks): + """Prompt lengths around the 64-token block and 256-token DCP4 boundaries.""" + block_size = 64 + num_local = cdiv(num_tokens, block_size) + num_remote = cdiv(num_tokens, block_size * dcp_size) + extra_remote_blocks + local = list(range(100, 100 + num_local)) + remote = list(range(500, 500 + num_remote)) + + written = [] + for dcp_rank in range(dcp_size): + pairs = dcp_write_block_pairs(local, remote, dcp_size, dcp_rank) + written += [(lb, dcp_rank, rb) for lb, rb in zip(*pairs)] + assert sorted(lb for lb, _, _ in written) == local + for lb, dcp_rank, rb in written: + j = local.index(lb) + assert (dcp_rank, rb) == (j % dcp_size, remote[j // dcp_size]) + + +def test_validate_write_dcp_topology(): + validate_write_dcp_topology(1, 1, 16, is_mla=False, has_mamba=True) + validate_write_dcp_topology(4, 64, 64, is_mla=True, has_mamba=False) + with pytest.raises(ValueError, match="MLA models without recurrent state"): + validate_write_dcp_topology(4, 64, 64, is_mla=False, has_mamba=False) + with pytest.raises(ValueError, match="MLA models without recurrent state"): + validate_write_dcp_topology(4, 64, 64, is_mla=True, has_mamba=True) + with pytest.raises( + ValueError, match="cp_kv_cache_interleave_size=1, block_size=64" + ): + validate_write_dcp_topology(4, 1, 64, is_mla=True, has_mamba=False) + + +def test_handshake_metadata_carries_dcp_and_defaults_for_older_peers(): + fields = dict( + engine_id="decode", + agent_metadata=b"", + kv_caches_base_addr=[], + num_blocks=NUM_BLOCKS, + block_len=0, + attn_backend_name="test", + ) + meta = MoRIIOAgentMetadata(**fields, dcp_size=4, cp_kv_cache_interleave_size=64) + decoded = msgspec.msgpack.decode( + msgspec.msgpack.encode(meta), type=MoRIIOAgentMetadata + ) + assert (decoded.dcp_size, decoded.cp_kv_cache_interleave_size) == (4, 64) + older = msgspec.msgpack.decode( + msgspec.msgpack.encode(fields), type=MoRIIOAgentMetadata + ) + assert (older.dcp_size, older.cp_kv_cache_interleave_size) == (1, 1) + + +def _spec() -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=1, head_size=8, dtype=torch.bfloat16 + ) + + +def _cache() -> torch.Tensor: + spec = _spec() + return torch.zeros( + NUM_BLOCKS, + spec.num_heads, + spec.num_states, + spec.state_content_size_bytes, + dtype=torch.uint8, + ) + + +def _worker(tp_rank: int, tp_size: int, caches: dict) -> MoRIIOConnectorWorker: + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.tp_rank = tp_rank + worker.world_size = tp_size + worker.use_mla = True + worker._has_mamba = False + worker.model_config = SimpleNamespace(get_total_num_kv_heads=lambda: 1) + worker.kv_caches = caches + worker.layer_to_spec = {name: _spec() for name in caches} + worker.layer_to_group = {name: 0 for name in caches} + worker.kv_layer_mr_offset = {} + worker._region_session_index = None + worker.layer_name_to_local_kv_cache_metadata = {name: [b""] for name in caches} + worker._record_slice_layouts(caches) + worker.shutdown = lambda: None + return worker + + +class _Session: + def __init__(self, src: torch.Tensor, dst: torch.Tensor): + self.src, self.dst = src.view(-1), dst.view(-1) + + +def _consumer_wrapper() -> MoRIIOWrapper: + wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) + wrapper.lock = threading.Lock() + wrapper.done_write_cache_req_ids = [] + wrapper.failed_write_cache_req_ids = [] + wrapper._write_done_counts = {} + wrapper._terminal_transfer_ids = OrderedDict() + return wrapper + + +def _run_write( + producer_tp: int, + consumer_tp: int, + dcp_size: int, + num_prefill_blocks: int, + decode_blocks: list[int], +): + """Run every prefill rank's writer for all layers into fake decode ranks. + + Returns the decode caches, the per-rank write_done results, the prefill + block ids, the global block contents and the bytes each prefill rank wrote. + """ + gen = torch.Generator().manual_seed(0) + content = { + name: torch.randint( + 1, 255, (num_prefill_blocks, *_cache().shape[1:]), generator=gen + ).to(torch.uint8) + for name in LAYERS + } + # Prefill block ids are not in order, so a pairing by position shows up. + prefill_blocks = [NUM_BLOCKS - 1 - j for j in range(num_prefill_blocks)] + consumer_caches = [{name: _cache() for name in LAYERS} for _ in range(consumer_tp)] + consumer_wrappers = [_consumer_wrapper() for _ in range(consumer_tp)] + metadata = msgspec.msgpack.decode( + msgspec.msgpack.encode( + MoRIIOAgentMetadata( + engine_id="decode", + agent_metadata=b"", + kv_caches_base_addr=[], + num_blocks=NUM_BLOCKS, + block_len=0, + attn_backend_name="test", + attention_block_sizes={name: BLOCK_SIZE for name in LAYERS}, + dcp_size=dcp_size, + cp_kv_cache_interleave_size=BLOCK_SIZE, + ) + ), + type=MoRIIOAgentMetadata, + ) + bytes_written = [] + done: list[set[str]] = [set() for _ in range(consumer_tp)] + for p in range(producer_tp): + caches = {name: _cache() for name in LAYERS} + for name in LAYERS: + caches[name][prefill_blocks] = content[name] + worker = _worker(p, producer_tp, caches) + mapping = worker._write_tp_mapping(consumer_tp) + + def get_built_session(engine_id, worker=worker, mapping=mapping): + if engine_id == "decode_dp0": + (target,) = mapping.targets + tp_rank = target.tp_rank + else: + tp_rank = int(engine_id.rsplit("_tp", 1)[1]) + sessions = [ + _Session(worker.kv_caches[name], consumer_caches[tp_rank][name]) + for name in worker.kv_caches + ] + return sessions, metadata + + def write_remote_data(sizes, local_offsets, remote_offsets, session): + for size, local, remote in zip(sizes, local_offsets, remote_offsets): + session.dst[remote : remote + size] = session.src[local : local + size] + return "status" + + def send_notify(transfer_id, ip, port, message_type, message_fields=None): + consumer_wrappers[port - NOTIFY_PORT]._handle_structured_message( + { + "type": message_type, + "transfer_id": transfer_id, + **(message_fields or {}), + } + ) + + worker._get_built_session = get_built_session + worker.moriio_wrapper = SimpleNamespace( + done_remote_allocate_req_dict={}, + lock=threading.Lock(), + done_req_ids=[], + write_remote_data=write_remote_data, + send_notify=send_notify, + waiting_for_transfer_complete=lambda statuses: None, + _mark_transfer_terminal_locked=lambda transfer_id: None, + _is_transfer_terminal_locked=lambda transfer_id: False, + ) + info = RemoteAllocInfo(block_ids=[list(decode_blocks)]) + worker.moriio_wrapper.done_remote_allocate_req_dict["t"] = info + writer = make_moriio_writer(worker) + for name in LAYERS: + writer._execute_write_task( + WriteTask( + request_id="req", + transfer_id="t", + dst_engine_id="decode", + local_block_ids=[prefill_blocks], + remote_block_ids_hint=None, + layer_name=name, + event=SimpleNamespace(synchronize=lambda: None), + remote_notify_port=NOTIFY_PORT, + remote_ip="decode", + remote_tp_size=consumer_tp, + tp_mapping=mapping, + ) + ) + info.writes_expected = len(LAYERS) + writer._finalize_if_complete("t", info) + # Prefill frees its blocks whether or not it had anything to write. + assert [ack.transfer_id for ack in worker.moriio_wrapper.done_req_ids] == ["t"] + bytes_written.append(info.bytes_written) + for rank_done, wrapper in zip(done, consumer_wrappers): + rank_done |= wrapper.pop_finished_write_req_ids() + return consumer_caches, done, content, bytes_written + + +@pytest.mark.parametrize("num_prefill_blocks", [1, 5, 9]) +@pytest.mark.parametrize( + ("producer_tp", "consumer_tp", "dcp_size"), + [(4, 4, 4), (1, 4, 4), (8, 4, 4), (4, 4, 2), (1, 1, 1)], +) +def test_write_lands_each_block_on_its_dcp_rank( + monkeypatch, producer_tp, consumer_tp, dcp_size, num_prefill_blocks +): + monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) + # One spare decode block, as when decode reserves lookahead slots. + num_decode_blocks = cdiv(num_prefill_blocks, dcp_size) + 1 + decode_blocks = [2 * i + 1 for i in range(num_decode_blocks)] + caches, done, content, bytes_written = _run_write( + producer_tp, consumer_tp, dcp_size, num_prefill_blocks, decode_blocks + ) + + # Every decode rank completes, including ranks that own no block of a + # short prompt. + assert done == [{"t"}] * consumer_tp + for t in range(consumer_tp): + dcp_rank = t % dcp_size + for name in LAYERS: + expected = torch.zeros_like(caches[t][name]) + for i, block in enumerate(decode_blocks): + j = i * dcp_size + dcp_rank + if j < num_prefill_blocks: + expected[block] = content[name][j] + assert torch.equal(caches[t][name], expected), (t, name) + + # Each block is written once per DCP group (consumer_tp // dcp_size + # replicas), 1 / dcp_size of what the same TP without DCP writes. + block_len = _cache()[0].numel() + total = num_prefill_blocks * block_len * len(LAYERS) + assert sum(bytes_written) == total * consumer_tp // dcp_size + + +def test_write_refuses_decode_allocation_short_of_its_blocks(monkeypatch): + monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) + with pytest.raises(ValueError, match="allocated 1"): + _run_write(4, 4, 4, num_prefill_blocks=5, decode_blocks=[3]) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 1120062e5410..44b84b4f9988 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -1008,6 +1008,8 @@ def _make_register_worker(caches, geoms, *, disable_hma: bool = True): tp_rank=0, dp_rank=0, is_producer=False, + dcp_size=1, + cp_kv_cache_interleave_size=1, vllm_config=SimpleNamespace( model_config=SimpleNamespace(hf_config=SimpleNamespace(model_type="x")), scheduler_config=SimpleNamespace( diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index b95819402bd5..9c72257a0802 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -89,7 +89,7 @@ def _worker( def _remote_meta(num_blocks: int = 16) -> SimpleNamespace: - return SimpleNamespace(num_blocks=num_blocks) + return SimpleNamespace(num_blocks=num_blocks, dcp_size=1) def _wrapper_for_messages() -> Any: diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index afd1d8ce7921..7bb4f9676971 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -295,7 +295,7 @@ def _get_built_session(engine_id: str): def _prepare_transfer_plan( task, request_info, remote_moriio_meta, remote_engine_id=None, target=None ): - return "plan" + return SimpleNamespace(transfer_sizes=[64, 32]) def _do_layer_write(plan, sessions): write_sessions.append(sessions) @@ -320,6 +320,7 @@ def _do_layer_write(plan, sessions): "status-remote-engine_dp0_tp2", "status-remote-engine_dp0_tp3", ] + assert info.bytes_written == 2 * (64 + 32) def test_fanout_write_done_notifies_every_target_after_completion(): diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 69b4d64469a2..2ba5d197fa4f 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -345,6 +345,17 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: } +def _uses_block_interleaved_pd_connector( + kv_transfer_config: KVTransferConfig | None, +) -> bool: + """NIXL and MoRIIO pair whole prefill blocks with DCP-sharded decode blocks, + which needs cp_kv_cache_interleave_size == block_size.""" + return kv_transfer_config is not None and any( + kv_transfer_config.has_connector(name) + for name in ("NixlConnector", "MoRIIOConnector") + ) + + @config(config=ConfigDict(arbitrary_types_allowed=True)) class VllmConfig: """Dataclass which contains all vllm-related configuration. This @@ -3340,7 +3351,7 @@ def _validate_v1_model_runner(self) -> None: def adjust_dcp_kv_cache_interleave_size( self, kv_cache_groups: list["KVCacheGroupSpec"] ) -> None: - """Resolve the NIXL P/D DCP interleave from the global KV cache groups. + """Resolve the NIXL/MoRIIO P/D DCP interleave from the global KV cache groups. Called once by the engine core before profiling; workers only assign the result via ``set_kv_cache_layout``. @@ -3367,7 +3378,7 @@ def adjust_dcp_kv_cache_interleave_size( "deprecated when PCP is fully supported." ) - if not self.kv_transfer_config.has_connector("NixlConnector"): + if not _uses_block_interleaved_pd_connector(self.kv_transfer_config): return if not self.parallel_config._allow_auto_resolve_cp_interleave_size: return @@ -3396,13 +3407,12 @@ def validate_block_size(self) -> None: """ block_size = self.cache_config.block_size - # Skip DCP interleave-size compatibility for NIXL P/D: the interleave - # size is finalized from the global KV cache groups. - nixl_pd_active = ( - self.kv_transfer_config is not None - and self.kv_transfer_config.has_connector("NixlConnector") - ) - if self.parallel_config.decode_context_parallel_size > 1 and not nixl_pd_active: + # Skip DCP interleave-size compatibility for NIXL/MoRIIO P/D: the + # interleave size is finalized from the global KV cache groups. + if ( + self.parallel_config.decode_context_parallel_size > 1 + and not _uses_block_interleaved_pd_connector(self.kv_transfer_config) + ): assert ( self.parallel_config.cp_kv_cache_interleave_size <= block_size and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index dadd55062397..fbabecead43b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -114,6 +114,7 @@ class RemoteAllocInfo: # producer whose copy another producer writes still has to finish). completion_targets_resolved: bool = False completion_notified: bool = False + bytes_written: int = 0 transfer_statuses: list[Any] = field(default_factory=list) transfer_offsets: dict[ TransferOffsetsKey, @@ -146,6 +147,9 @@ class MoRIIOAgentMetadata( # Tokens per block of each attention layer. Hybrid models size attention # blocks to the recurrent-state page, which depends on the TP size. attention_block_sizes: dict[str, int] = {} + # The defaults also stand for older peers that do not send these fields. + dcp_size: int = 1 + cp_kv_cache_interleave_size: int = 1 class RoleManager: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 9ae2597a9ebb..54b3cb13acc7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -89,6 +89,7 @@ MoRIIOWriteTPMapping, WriteTarget, compute_write_tp_mapping, + validate_write_dcp_topology, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( MambaConvSplitInfo, @@ -1953,6 +1954,16 @@ def __init__( self.block_window_per_layer: list[int | None] = [] self.use_mla = self.model_config.use_mla + self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size + self.cp_kv_cache_interleave_size = ( + vllm_config.parallel_config.cp_kv_cache_interleave_size + ) + if self.dcp_size > 1 and (self.is_producer or self.mode == MoRIIOMode.READ): + raise NotImplementedError( + "MoRIIO supports decode context parallelism " + f"(decode_context_parallel_size={self.dcp_size}) only on the " + "decode instance in WRITE mode" + ) # Hybrid (mamba/KDA) recurrent-state transfer. Populated in # register_kv_caches when the model has a MambaSpec kv_cache group. self._has_mamba = False @@ -2225,6 +2236,19 @@ def _moriio_handshake( metadata_bytes = received_frame[1] decoder = msgspec.msgpack.Decoder(MoRIIOAgentMetadata) metadata = decoder.decode(metadata_bytes) + if self.is_producer and self.mode == MoRIIOMode.WRITE: + try: + validate_write_dcp_topology( + metadata.dcp_size, + metadata.cp_kv_cache_interleave_size, + self.block_size, + self.use_mla, + self._has_mamba, + ) + except ValueError as e: + raise HandshakeError( + f"Cannot write into {expected_engine_id}: {e}" + ) from e got_metadata_time = time.perf_counter() logger.info( "MoRIIO handshake: get metadata took: %s", @@ -2571,6 +2595,15 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.block_shape = block_shape self.kv_element_size = kv_elem_size + if self.dcp_size > 1: + validate_write_dcp_topology( + self.dcp_size, + self.cp_kv_cache_interleave_size, + self.block_size, + self.use_mla, + self._has_mamba, + ) + self.dst_num_blocks[self.engine_id] = self.num_blocks kv_caches_base_addr = [] caches_data = [] @@ -2760,6 +2793,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): for name, layout in self.mamba_state_layouts.items() }, attention_block_sizes=self.attention_block_sizes, + dcp_size=self.dcp_size, + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, ) ready_event = threading.Event() self._moriio_handshake_listener_t = threading.Thread( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 61e2c8d38b37..6e36c4ac32a1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -40,6 +40,9 @@ pod_index, zmq_ctx, ) +from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_tp_mapping import ( + dcp_write_block_pairs, +) if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( @@ -382,6 +385,7 @@ def _execute_write_task(self, task: WriteTask) -> None: is_mamba = self.worker._is_mamba_layer(task.layer_name) transfer_statuses: list[Any] = [] + bytes_written = 0 for target in self._write_targets(task): if target is not None and not is_mamba and not target.num_heads: # Another producer rank writes this attention copy. @@ -409,8 +413,10 @@ def _execute_write_task(self, task: WriteTask) -> None: ] for plan in plans: transfer_statuses.extend(self._do_layer_write(plan, sessions)) + bytes_written += sum(plan.transfer_sizes) with self._write_state_lock: request_info.transfer_statuses.extend(transfer_statuses) + request_info.bytes_written += bytes_written # Finalize if all layers complete self._mark_write_done(task.transfer_id, request_info) @@ -452,6 +458,12 @@ def _prepare_transfer_plan( if self.worker._has_mamba: # D can reserve a tail block for the token recomputed locally. remote_blocks = remote_blocks[: len(local_blocks)] + dcp_size = remote_moriio_meta.dcp_size + if dcp_size > 1 and self.worker.layer_to_spec[task.layer_name].dcp_sharded: + tp_rank = self.worker.tp_rank if target is None else target.tp_rank + local_blocks, remote_blocks = dcp_write_block_pairs( + local_blocks, remote_blocks, dcp_size, tp_rank % dcp_size + ) if ( target is None or task.tp_mapping is None @@ -617,6 +629,7 @@ def _finalize_if_complete( transfer_statuses = list(request_info.transfer_statuses) request_info.transfer_statuses.clear() request_info.completion_notified = True + bytes_written = request_info.bytes_written # Wait for this request's transfers to complete. self.worker.moriio_wrapper.waiting_for_transfer_complete(transfer_statuses) @@ -648,9 +661,11 @@ def _finalize_if_complete( self.worker.moriio_wrapper._mark_transfer_terminal_locked(transfer_id) self._clear_transfer_state(transfer_id) logger.debug( - "Completed transfer for (request, transfer) %s, %s, notified %s", + "Completed transfer for (request, transfer) %s, %s, wrote %d bytes, " + "notified %s", request_id, transfer_id, + bytes_written, notify_targets, ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py index 9b244f68be7d..94c3bdae8c83 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py @@ -19,8 +19,15 @@ heads ``[i * n, (i + 1) * n)`` with ``n = num_kv_heads // tp_size``; with more ranks than heads every rank holds the single head ``i * num_kv_heads // tp_size``. + +A consumer with decode context parallelism (DCP) shards the tokens of its MLA +latent instead. With ``cp_kv_cache_interleave_size == block_size``, block ``j`` +of the producer, which is not DCP-sharded, lives whole on DCP rank +``j % dcp_size`` (``tp_rank % dcp_size``) in that rank's block ``j // dcp_size``, +so each block is one write (``dcp_write_block_pairs``). """ +from collections.abc import Sequence from dataclasses import dataclass from functools import cache @@ -212,3 +219,52 @@ def compute_write_tp_mapping( ) ), ) + + +def validate_write_dcp_topology( + consumer_dcp_size: int, + consumer_interleave_size: int, + block_size: int, + is_mla: bool, + has_mamba: bool, +) -> None: + if consumer_dcp_size <= 1: + return + if not is_mla or has_mamba: + raise ValueError( + "MoRIIO WRITE into a decode instance with decode context parallelism " + f"(decode_context_parallel_size={consumer_dcp_size}) supports MLA " + "models without recurrent state only" + ) + if consumer_interleave_size != block_size: + raise ValueError( + "MoRIIO WRITE into a decode instance with decode context parallelism " + f"(decode_context_parallel_size={consumer_dcp_size}) writes whole " + "blocks, so it needs cp_kv_cache_interleave_size == block_size; got " + f"cp_kv_cache_interleave_size={consumer_interleave_size}, " + f"block_size={block_size}. Leave --cp-kv-cache-interleave-size unset " + f"or set it to {block_size}" + ) + + +def dcp_write_block_pairs( + local_block_ids: Sequence[int], + remote_block_ids: Sequence[int], + dcp_size: int, + dcp_rank: int, +) -> tuple[list[int], list[int]]: + """Producer blocks that consumer DCP rank ``dcp_rank`` holds, paired with + that rank's block ids. + + A rank that holds none of them (a prompt shorter than ``dcp_size`` blocks) + gets two empty lists and must still be sent write_done. + """ + local = list(local_block_ids[dcp_rank::dcp_size]) + remote = list(remote_block_ids[: len(local)]) + if len(remote) < len(local): + raise ValueError( + f"DCP rank {dcp_rank}/{dcp_size} of the consumer needs {len(local)} " + f"blocks for {len(local_block_ids)} producer blocks, but it allocated " + f"{len(remote_block_ids)}" + ) + return local, remote From 4eda20d21a7b8c35aaca35209166b1ab930750c5 Mon Sep 17 00:00:00 2001 From: kliuae Date: Fri, 30 Oct 2026 16:39:07 +0800 Subject: [PATCH 2/6] support dcp interleave dividing block size Signed-off-by: kliuae --- docs/features/moriio_connector_usage.md | 9 + tests/test_config.py | 19 +- .../unit/test_moriio_connector.py | 27 + .../unit/test_moriio_dcp_write.py | 605 +++++++++++++----- .../unit/test_moriio_hma_scheduler.py | 28 + .../unit/test_moriio_write_done_routing.py | 4 +- tests/v1/kv_connector/unit/utils.py | 1 + vllm/config/vllm.py | 28 +- .../kv_connector/v1/moriio/moriio_common.py | 11 +- .../v1/moriio/moriio_connector.py | 180 ++++-- .../kv_connector/v1/moriio/moriio_engine.py | 96 ++- .../kv_connector/v1/moriio/moriio_layout.py | 120 ++++ .../v1/moriio/moriio_tp_mapping.py | 124 +++- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 49 ++ 14 files changed, 1039 insertions(+), 262 deletions(-) diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md index 352022c7aa6b..c591f83f4e51 100644 --- a/docs/features/moriio_connector_usage.md +++ b/docs/features/moriio_connector_usage.md @@ -162,6 +162,15 @@ The number of KV heads must be a multiple or a divisor of each TP size. Blocks a !!! note On a single node with the xGMI backend, a larger decode TP size does not work yet: MoRI cannot open a session from a prefill rank to memory on a decode rank's second or later visible GPU when the prefill rank sees fewer GPUs, and the transfer fails with an error naming the decode rank. A larger prefill TP size works on xGMI, and the RDMA backend is not affected. +### Decode context parallelism + +In WRITE mode a prefill instance without DCP can feed a decode instance with `--decode-context-parallel-size` (DCP) greater than 1, for MLA models without recurrent state. The decode instance deals the tokens of a request to the ranks of a DCP group in groups of `--cp-kv-cache-interleave-size` tokens, and each prefill rank writes a decode rank only the tokens that rank holds, so every token is written once per DCP group instead of once per decode rank. + +- `--cp-kv-cache-interleave-size` is optional on the decode instance and must divide the KV cache block size. Each group of interleaved tokens is one copy, so the block size gives the fewest copies. Speculative decoding with DCP on ROCm's MLA attention backend (`ROCM_AITER_MLA`) needs interleave 1. +- Instances with the same DCP size in WRITE or READ modes work as before with any interleave that is the same on both prefill and decode. Instances whose DCP sizes or interleaves differ in any other way refuse each other at the handshake. +- A refused handshake logs the reason and fails the requests that need it: the decode instance returns an error for each, or recomputes it with `kv_load_failure_policy` set to `recompute`. +- Both instances must run a vLLM version with this support: a prefill instance without it cannot detect a DCP decode instance and writes into the wrong blocks. + ### Transport configuration MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments. diff --git a/tests/test_config.py b/tests/test_config.py index b78be1c2a7e2..4836c74949ff 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -395,10 +395,6 @@ def test_per_request_spec_decode_metrics_requires_spec_decode(): kv_connector="NixlConnector", kv_role="kv_both", ), - KVTransferConfig( - kv_connector="MoRIIOConnector", - kv_role="kv_consumer", - ), KVTransferConfig( kv_connector="MultiConnector", kv_role="kv_both", @@ -440,7 +436,15 @@ def test_pd_dcp_interleave_size_is_adjusted_to_block_size( assert "automatically adjusted from 3 to block_size 16" in caplog.text -def test_kv_offloading_does_not_adjust_dcp_interleave_size(): +@pytest.mark.parametrize( + "kv_transfer_config", + [ + KVTransferConfig(kv_connector="OffloadingConnector", kv_role="kv_both"), + KVTransferConfig(kv_connector="MoRIIOConnector", kv_role="kv_consumer"), + ], + ids=["offloading", "moriio"], +) +def test_connector_does_not_adjust_dcp_interleave_size(kv_transfer_config): config = VllmConfig( cache_config=CacheConfig(block_size=16), device_config=DeviceConfig(device="cpu"), @@ -450,10 +454,7 @@ def test_kv_offloading_does_not_adjust_dcp_interleave_size(): cp_kv_cache_interleave_size=1, distributed_executor_backend="mp", ), - kv_transfer_config=KVTransferConfig( - kv_connector="OffloadingConnector", - kv_role="kv_both", - ), + kv_transfer_config=kv_transfer_config, ) kv_cache_groups = [SimpleNamespace(kv_cache_spec=SimpleNamespace(block_size=16))] diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 438e20bd0809..230ddaae9ecc 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -1282,3 +1282,30 @@ def test_worker_layer_to_group_routing(mock_parallel_groups): kv_cache_config=_make_hybrid_kv_cache_config(), ) assert worker.layer_to_group == {"full0": 0, "sw0": 1} + + +@pytest.mark.parametrize("read_mode", [False, True], ids=["write", "read"]) +@pytest.mark.parametrize("role", ["kv_producer", "kv_consumer"]) +def test_worker_accepts_decode_context_parallelism( + mock_parallel_groups, monkeypatch, role, read_mode +): + """Prefill and decode instances with the same DCP size start in either mode.""" + # The test model has no attention backend with DCP; an MLA model's would. + monkeypatch.setattr( + "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector." + "get_current_attn_backends", + lambda vllm_config: [SimpleNamespace(get_name=lambda: "ROCM_AITER_MLA")], + ) + vllm_config = create_vllm_config(role=role, read_mode=read_mode) + vllm_config.parallel_config.decode_context_parallel_size = 2 + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config.update( + { + "http_port": _find_free_port(), + "handshake_port": _find_free_port(), + "notify_port": _find_free_port(), + } + ) + with set_current_vllm_config(vllm_config): + worker = FakeMoRIIOConnectorWorker(vllm_config, "engine0", hand_shake_latency=0) + assert (worker.dcp_size, worker.cp_kv_cache_interleave_size) == (2, 1) diff --git a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py index a96e3ace9f49..37070dedd48e 100644 --- a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MoRIIO WRITE into a decode instance with decode context parallelism (DCP).""" +import queue import threading from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor from types import SimpleNamespace import msgspec @@ -13,9 +15,15 @@ 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, + HandshakeError, + HandshakeRefusedError, MoRIIOAgentMetadata, + MoRIIOConnectorMetadata, + MoRIIOMode, RemoteAllocInfo, + ReqMeta, WriteTask, + get_port_offset, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( MoRIIOConnectorWorker, @@ -24,84 +32,40 @@ MoRIIOWrapper, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_tp_mapping import ( - dcp_write_block_pairs, - validate_write_dcp_topology, + validate_peer_dcp_layout, ) from vllm.utils.math_utils import cdiv -from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm.v1.attention.ops.rocm_aiter_mla_sparse import indexer_k_cache_byte_ranges +from vllm.v1.kv_cache_interface import MLAAttentionSpec, SparseCacheRole from .utils import make_moriio_writer pytestmark = pytest.mark.cpu_test -BLOCK_SIZE = 4 +BLOCK_SIZE = 64 NUM_BLOCKS = 24 NOTIFY_PORT = 20000 -LAYERS = ("mla0", "mla1") -@pytest.mark.parametrize("num_tokens", [1, 65, 256, 257, 1000]) -@pytest.mark.parametrize("dcp_size", [2, 4]) -@pytest.mark.parametrize("extra_remote_blocks", [0, 1]) -def test_dcp_pairs_cover_every_block_once(num_tokens, dcp_size, extra_remote_blocks): - """Prompt lengths around the 64-token block and 256-token DCP4 boundaries.""" - block_size = 64 - num_local = cdiv(num_tokens, block_size) - num_remote = cdiv(num_tokens, block_size * dcp_size) + extra_remote_blocks - local = list(range(100, 100 + num_local)) - remote = list(range(500, 500 + num_remote)) - - written = [] - for dcp_rank in range(dcp_size): - pairs = dcp_write_block_pairs(local, remote, dcp_size, dcp_rank) - written += [(lb, dcp_rank, rb) for lb, rb in zip(*pairs)] - assert sorted(lb for lb, _, _ in written) == local - for lb, dcp_rank, rb in written: - j = local.index(lb) - assert (dcp_rank, rb) == (j % dcp_size, remote[j // dcp_size]) - - -def test_validate_write_dcp_topology(): - validate_write_dcp_topology(1, 1, 16, is_mla=False, has_mamba=True) - validate_write_dcp_topology(4, 64, 64, is_mla=True, has_mamba=False) - with pytest.raises(ValueError, match="MLA models without recurrent state"): - validate_write_dcp_topology(4, 64, 64, is_mla=False, has_mamba=False) - with pytest.raises(ValueError, match="MLA models without recurrent state"): - validate_write_dcp_topology(4, 64, 64, is_mla=True, has_mamba=True) - with pytest.raises( - ValueError, match="cp_kv_cache_interleave_size=1, block_size=64" - ): - validate_write_dcp_topology(4, 1, 64, is_mla=True, has_mamba=False) - - -def test_handshake_metadata_carries_dcp_and_defaults_for_older_peers(): - fields = dict( - engine_id="decode", - agent_metadata=b"", - kv_caches_base_addr=[], - num_blocks=NUM_BLOCKS, - block_len=0, - attn_backend_name="test", - ) - meta = MoRIIOAgentMetadata(**fields, dcp_size=4, cp_kv_cache_interleave_size=64) - decoded = msgspec.msgpack.decode( - msgspec.msgpack.encode(meta), type=MoRIIOAgentMetadata - ) - assert (decoded.dcp_size, decoded.cp_kv_cache_interleave_size) == (4, 64) - older = msgspec.msgpack.decode( - msgspec.msgpack.encode(fields), type=MoRIIOAgentMetadata +def _spec(head_size: int, **kwargs) -> MLAAttentionSpec: + return MLAAttentionSpec( + block_size=BLOCK_SIZE, + num_kv_heads=1, + head_size=head_size, + dtype=torch.uint8, + **kwargs, ) - assert (older.dcp_size, older.cp_kv_cache_interleave_size) == (1, 1) -def _spec() -> MLAAttentionSpec: - return MLAAttentionSpec( - block_size=BLOCK_SIZE, num_kv_heads=1, head_size=8, dtype=torch.bfloat16 - ) +# The fp8 caches of DeepSeek-V3.2 and GLM-5: the MLA latent (kv_lora_rank 512 + +# qk_rope_head_dim 64) and the indexer K cache (128 fp8 values, one fp32 scale). +SPECS = { + "latent": _spec(576), + "indexer": _spec(132, cache_role=SparseCacheRole.INDEXER), +} -def _cache() -> torch.Tensor: - spec = _spec() +def _cache(spec: MLAAttentionSpec) -> torch.Tensor: return torch.zeros( NUM_BLOCKS, spec.num_heads, @@ -111,16 +75,122 @@ def _cache() -> torch.Tensor: ) -def _worker(tp_rank: int, tp_size: int, caches: dict) -> MoRIIOConnectorWorker: +def _decode_position(token: int, dcp_size: int, interleave: int) -> tuple[int, int]: + """DCP rank and position of a token, as the decode slot mapping places it + (vllm/v1/worker/block_table.py).""" + rank = token // interleave % dcp_size + return rank, token // (interleave * dcp_size) * interleave + token % interleave + + +def _token_bytes(spec: MLAAttentionSpec, slot: int) -> torch.Tensor: + """Offsets of one token's bytes in a block, as the cache writers lay them + out.""" + entry = spec.state_content_size_bytes + if spec.cache_role is not SparseCacheRole.INDEXER: + return slot * entry + torch.arange(entry) + # indexer_k_quant_and_cache_triton and the fused DeepSeek-V3.2 / GLM-5 + # writer: fp8 values in tiles of 16 tokens x 16 bytes, then the scales. + head_dim = entry - 4 + dim = torch.arange(head_dim) + values = slot // 16 * 16 * head_dim + slot % 16 * 16 + dim // 16 * 256 + dim % 16 + scales = BLOCK_SIZE * head_dim + slot * 4 + torch.arange(4) + return torch.cat([values, scales]) + + +def test_indexer_byte_ranges_follow_the_tiled_layout(): + # One token: eight 16-byte pieces, 256 bytes apart, and its scale. + assert indexer_k_cache_byte_ranges(64, 132, 3, 1) == [ + *((48 + 256 * piece, 16) for piece in range(8)), + (64 * 128 + 12, 4), + ] + # Whole aligned tiles collapse to one values range and one scales range. + assert indexer_k_cache_byte_ranges(64, 132, 16, 32) == [ + (16 * 128, 32 * 128), + (64 * 128 + 16 * 4, 32 * 4), + ] + # Block size 1 is not tiled. + assert indexer_k_cache_byte_ranges(1, 132, 0, 1) == [(0, 132)] + + +def _validate(local_dcp, local_interleave, peer_dcp, peer_interleave, **kwargs): + validate_peer_dcp_layout( + local_dcp, + local_interleave, + peer_dcp, + peer_interleave, + block_size=BLOCK_SIZE, + **{"writes_into_peer": True, "is_mla": True, "has_mamba": False, **kwargs}, + ) + + +@pytest.mark.parametrize( + ("layout", "kwargs"), + [ + ((1, 1, 4, 64), {}), + ((1, 1, 4, 1), {}), + ((1, 1, 1, 64), {}), + ((8, 1, 8, 1), {"has_mamba": True}), + ((8, 1, None, None), {"writes_into_peer": False}), + ], + ids=[ + "write-into-dcp", + "write-into-dcp-unit-interleave", + "no-dcp", + "same-dcp-hybrid", + "older-peer", + ], +) +def test_peer_dcp_layout_accepted(layout, kwargs): + _validate(*layout, **kwargs) + + +@pytest.mark.parametrize( + ("layout", "kwargs", "match"), + [ + ((1, 1, 4, 48), {}, "divides the block size 64; it uses 48"), + ( + (1, 1, 4, 1), + {"whole_block_layer": "layer0"}, + "layer layer0 in whole blocks only.*needs --cp-kv-cache-interleave-size 64", + ), + ((1, 1, 4, 64), {"is_mla": False}, "MLA models without recurrent state"), + ((1, 1, 4, 1), {"has_mamba": True}, "MLA models without recurrent state"), + ((2, 1, 4, 1), {}, "between decode_context_parallel_size=2 here and 4"), + ((1, 1, 4, 64), {"writes_into_peer": False}, "=1 here and 4 on the peer"), + ((8, 64, 8, 1), {"writes_into_peer": False}, "64 here, 1 on the peer"), + ], + ids=[ + "write-into-dcp-interleave-not-dividing-block", + "write-into-dcp-whole-block-layer", + "write-into-dcp-gqa", + "write-into-dcp-hybrid", + "dcp2-into-dcp4", + "read-no-dcp-from-dcp", + "same-dcp-interleave", + ], +) +def test_peer_dcp_layout_refused(layout, kwargs, match): + with pytest.raises(ValueError, match=match): + _validate(*layout, **kwargs) + + +def _worker( + tp_rank: int, tp_size: int, specs: dict[str, MLAAttentionSpec], dcp_size: int = 1 +) -> MoRIIOConnectorWorker: + caches = {name: _cache(spec) for name, spec in specs.items()} worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) worker.tp_rank = tp_rank worker.world_size = tp_size + worker.dcp_size = dcp_size + worker.block_size = BLOCK_SIZE worker.use_mla = True worker._has_mamba = False - worker.model_config = SimpleNamespace(get_total_num_kv_heads=lambda: 1) + worker.model_config = SimpleNamespace( + get_total_num_kv_heads=lambda: 1, get_head_size=lambda: 576 + ) worker.kv_caches = caches - worker.layer_to_spec = {name: _spec() for name in caches} - worker.layer_to_group = {name: 0 for name in caches} + worker.layer_to_spec = dict(specs) + worker.layer_to_group = dict.fromkeys(caches, 0) worker.kv_layer_mr_offset = {} worker._region_session_index = None worker.layer_name_to_local_kv_cache_metadata = {name: [b""] for name in caches} @@ -129,14 +199,30 @@ def _worker(tp_rank: int, tp_size: int, caches: dict) -> MoRIIOConnectorWorker: return worker -class _Session: - def __init__(self, src: torch.Tensor, dst: torch.Tensor): - self.src, self.dst = src.view(-1), dst.view(-1) +@pytest.mark.parametrize( + ("specs", "whole_block_layer"), + [ + (SPECS, None), + ({**SPECS, "indexer": _spec(132)}, "indexer"), + ({"latent": _spec(576, tokens_per_state=4)}, "latent"), + ], + ids=["latent-and-indexer", "indexer-without-indexer-role", "compressed-states"], +) +def test_only_known_token_layouts_are_written_in_runs(specs, whole_block_layer): + """Runs of tokens are written only for an MLA latent of the model's latent + width and for the indexer K cache; any other layer, such as DeepSeek-V4's + indexer cache with the default cache role, needs whole blocks.""" + worker = _worker(0, 1, specs) + for interleave in (1, 16): + assert worker._whole_block_dcp_layer(interleave) == whole_block_layer + assert worker._whole_block_dcp_layer(BLOCK_SIZE) is None -def _consumer_wrapper() -> MoRIIOWrapper: +def _wrapper() -> MoRIIOWrapper: wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) wrapper.lock = threading.Lock() + wrapper.done_req_ids = [] + wrapper.done_remote_allocate_req_dict = {} wrapper.done_write_cache_req_ids = [] wrapper.failed_write_cache_req_ids = [] wrapper._write_done_counts = {} @@ -144,52 +230,70 @@ def _consumer_wrapper() -> MoRIIOWrapper: return wrapper +def _decode_metadata(dcp_size: int, interleave: int) -> MoRIIOAgentMetadata: + metadata = MoRIIOAgentMetadata( + engine_id="decode", + agent_metadata=b"", + kv_caches_base_addr=[], + num_blocks=NUM_BLOCKS, + block_len=0, + attn_backend_name="test", + attention_block_sizes=dict.fromkeys(SPECS, BLOCK_SIZE), + dcp_size=dcp_size, + cp_kv_cache_interleave_size=interleave, + ) + return msgspec.msgpack.decode( + msgspec.msgpack.encode(metadata), type=MoRIIOAgentMetadata + ) + + +def _write_task(name: str, prefill_blocks: list[int], mapping) -> WriteTask: + return WriteTask( + request_id="req", + transfer_id="t", + dst_engine_id="decode", + local_block_ids=[prefill_blocks], + remote_block_ids_hint=None, + layer_name=name, + event=SimpleNamespace(synchronize=lambda: None), + remote_notify_port=NOTIFY_PORT, + remote_ip="decode", + remote_tp_size=mapping.consumer_tp_size, + tp_mapping=mapping, + ) + + def _run_write( producer_tp: int, consumer_tp: int, dcp_size: int, num_prefill_blocks: int, decode_blocks: list[int], + interleave: int, + producer_dcp_size: int = 1, ): - """Run every prefill rank's writer for all layers into fake decode ranks. - - Returns the decode caches, the per-rank write_done results, the prefill - block ids, the global block contents and the bytes each prefill rank wrote. - """ + """Run every prefill rank's writer for all layers into fake decode ranks; + return their caches and write_done results, the prefill block contents and + the bytes each prefill rank wrote.""" gen = torch.Generator().manual_seed(0) content = { name: torch.randint( - 1, 255, (num_prefill_blocks, *_cache().shape[1:]), generator=gen + 1, 255, (num_prefill_blocks, *_cache(spec).shape[1:]), generator=gen ).to(torch.uint8) - for name in LAYERS + for name, spec in SPECS.items() } # Prefill block ids are not in order, so a pairing by position shows up. prefill_blocks = [NUM_BLOCKS - 1 - j for j in range(num_prefill_blocks)] - consumer_caches = [{name: _cache() for name in LAYERS} for _ in range(consumer_tp)] - consumer_wrappers = [_consumer_wrapper() for _ in range(consumer_tp)] - metadata = msgspec.msgpack.decode( - msgspec.msgpack.encode( - MoRIIOAgentMetadata( - engine_id="decode", - agent_metadata=b"", - kv_caches_base_addr=[], - num_blocks=NUM_BLOCKS, - block_len=0, - attn_backend_name="test", - attention_block_sizes={name: BLOCK_SIZE for name in LAYERS}, - dcp_size=dcp_size, - cp_kv_cache_interleave_size=BLOCK_SIZE, - ) - ), - type=MoRIIOAgentMetadata, - ) + consumer_caches = [ + {name: _cache(spec) for name, spec in SPECS.items()} for _ in range(consumer_tp) + ] + consumer_wrappers = [_wrapper() for _ in range(consumer_tp)] + metadata = _decode_metadata(dcp_size, interleave) bytes_written = [] - done: list[set[str]] = [set() for _ in range(consumer_tp)] for p in range(producer_tp): - caches = {name: _cache() for name in LAYERS} - for name in LAYERS: - caches[name][prefill_blocks] = content[name] - worker = _worker(p, producer_tp, caches) + worker = _worker(p, producer_tp, SPECS, producer_dcp_size) + for name, cache in worker.kv_caches.items(): + cache[prefill_blocks] = content[name] mapping = worker._write_tp_mapping(consumer_tp) def get_built_session(engine_id, worker=worker, mapping=mapping): @@ -199,7 +303,10 @@ def get_built_session(engine_id, worker=worker, mapping=mapping): else: tp_rank = int(engine_id.rsplit("_tp", 1)[1]) sessions = [ - _Session(worker.kv_caches[name], consumer_caches[tp_rank][name]) + SimpleNamespace( + src=worker.kv_caches[name].view(-1), + dst=consumer_caches[tp_rank][name].view(-1), + ) for name in worker.kv_caches ] return sessions, metadata @@ -210,17 +317,15 @@ def write_remote_data(sizes, local_offsets, remote_offsets, session): return "status" def send_notify(transfer_id, ip, port, message_type, message_fields=None): + message = {"type": message_type, "transfer_id": transfer_id} consumer_wrappers[port - NOTIFY_PORT]._handle_structured_message( - { - "type": message_type, - "transfer_id": transfer_id, - **(message_fields or {}), - } + {**message, **(message_fields or {})} ) worker._get_built_session = get_built_session + info = RemoteAllocInfo(block_ids=[list(decode_blocks)]) worker.moriio_wrapper = SimpleNamespace( - done_remote_allocate_req_dict={}, + done_remote_allocate_req_dict={"t": info}, lock=threading.Lock(), done_req_ids=[], write_remote_data=write_remote_data, @@ -229,72 +334,278 @@ def send_notify(transfer_id, ip, port, message_type, message_fields=None): _mark_transfer_terminal_locked=lambda transfer_id: None, _is_transfer_terminal_locked=lambda transfer_id: False, ) - info = RemoteAllocInfo(block_ids=[list(decode_blocks)]) - worker.moriio_wrapper.done_remote_allocate_req_dict["t"] = info writer = make_moriio_writer(worker) - for name in LAYERS: - writer._execute_write_task( - WriteTask( - request_id="req", - transfer_id="t", - dst_engine_id="decode", - local_block_ids=[prefill_blocks], - remote_block_ids_hint=None, - layer_name=name, - event=SimpleNamespace(synchronize=lambda: None), - remote_notify_port=NOTIFY_PORT, - remote_ip="decode", - remote_tp_size=consumer_tp, - tp_mapping=mapping, - ) - ) - info.writes_expected = len(LAYERS) + for name in worker.kv_caches: + writer._execute_write_task(_write_task(name, prefill_blocks, mapping)) + info.writes_expected = len(worker.kv_caches) writer._finalize_if_complete("t", info) # Prefill frees its blocks whether or not it had anything to write. assert [ack.transfer_id for ack in worker.moriio_wrapper.done_req_ids] == ["t"] bytes_written.append(info.bytes_written) - for rank_done, wrapper in zip(done, consumer_wrappers): - rank_done |= wrapper.pop_finished_write_req_ids() + done = [wrapper.pop_finished_write_req_ids() for wrapper in consumer_wrappers] return consumer_caches, done, content, bytes_written -@pytest.mark.parametrize("num_prefill_blocks", [1, 5, 9]) +def _expected_decode_cache( + spec: MLAAttentionSpec, + content: torch.Tensor, + decode_blocks: list[int], + dcp_size: int, + dcp_rank: int, + interleave: int, +) -> torch.Tensor: + """A decode DCP rank's cache once every token it holds has arrived.""" + expected = torch.zeros_like(_cache(spec)).view(NUM_BLOCKS, -1) + source = content.view(len(content), -1) + for token in range(len(content) * BLOCK_SIZE): + rank, position = _decode_position(token, dcp_size, interleave) + if rank == dcp_rank: + expected[ + decode_blocks[position // BLOCK_SIZE], + _token_bytes(spec, position % BLOCK_SIZE), + ] = source[token // BLOCK_SIZE, _token_bytes(spec, token % BLOCK_SIZE)] + return expected.view_as(_cache(spec)) + + +@pytest.mark.parametrize("interleave", [1, 16, BLOCK_SIZE]) +@pytest.mark.parametrize("num_prefill_blocks", [1, 5]) @pytest.mark.parametrize( - ("producer_tp", "consumer_tp", "dcp_size"), - [(4, 4, 4), (1, 4, 4), (8, 4, 4), (4, 4, 2), (1, 1, 1)], + ("producer_tp", "consumer_tp", "dcp_size"), [(1, 4, 4), (8, 4, 2)] ) -def test_write_lands_each_block_on_its_dcp_rank( - monkeypatch, producer_tp, consumer_tp, dcp_size, num_prefill_blocks +def test_write_lands_each_token_on_its_dcp_rank( + monkeypatch, producer_tp, consumer_tp, dcp_size, num_prefill_blocks, interleave ): + """A sparse MLA model writes its latent and indexer caches token by token + at any interleave, from one prefill rank into several decode ranks and + from several prefill ranks into one.""" monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) # One spare decode block, as when decode reserves lookahead slots. num_decode_blocks = cdiv(num_prefill_blocks, dcp_size) + 1 decode_blocks = [2 * i + 1 for i in range(num_decode_blocks)] caches, done, content, bytes_written = _run_write( - producer_tp, consumer_tp, dcp_size, num_prefill_blocks, decode_blocks + producer_tp, + consumer_tp, + dcp_size, + num_prefill_blocks, + decode_blocks, + interleave, ) - # Every decode rank completes, including ranks that own no block of a + # Every decode rank completes, including ranks that hold no token of a # short prompt. assert done == [{"t"}] * consumer_tp for t in range(consumer_tp): - dcp_rank = t % dcp_size - for name in LAYERS: - expected = torch.zeros_like(caches[t][name]) - for i, block in enumerate(decode_blocks): - j = i * dcp_size + dcp_rank - if j < num_prefill_blocks: - expected[block] = content[name][j] + for name, spec in SPECS.items(): + expected = _expected_decode_cache( + spec, content[name], decode_blocks, dcp_size, t % dcp_size, interleave + ) assert torch.equal(caches[t][name], expected), (t, name) - - # Each block is written once per DCP group (consumer_tp // dcp_size + # Each token is written once per DCP group (consumer_tp // dcp_size # replicas), 1 / dcp_size of what the same TP without DCP writes. - block_len = _cache()[0].numel() - total = num_prefill_blocks * block_len * len(LAYERS) + block_len = sum(_cache(spec)[0].numel() for spec in SPECS.values()) + total = num_prefill_blocks * block_len assert sum(bytes_written) == total * consumer_tp // dcp_size def test_write_refuses_decode_allocation_short_of_its_blocks(monkeypatch): monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) with pytest.raises(ValueError, match="allocated 1"): - _run_write(4, 4, 4, num_prefill_blocks=5, decode_blocks=[3]) + _run_write(1, 4, 4, num_prefill_blocks=5, decode_blocks=[3], interleave=1) + + +def test_write_between_equal_dcp_sizes_is_rank_to_rank(monkeypatch): + """Prefill and decode with the same DCP size hold the same shard on the same + rank, so each rank writes its own blocks unchanged to its peer.""" + monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) + decode_blocks = [2 * i + 1 for i in range(5)] + caches, done, content, _ = _run_write( + 4, 4, 4, 5, decode_blocks, interleave=1, producer_dcp_size=4 + ) + + assert done == [{"t"}] * 4 + for t in range(4): + for name in SPECS: + expected = torch.zeros_like(caches[t][name]) + expected[decode_blocks] = content[name] + assert torch.equal(caches[t][name], expected), (t, name) + + +def _handshake_worker(is_producer: bool, handshake) -> MoRIIOConnectorWorker: + """Rank 0 of a prefill instance with TP 1 in WRITE mode or of a decode + instance with TP 2 in READ mode; ``handshake`` stands in for the peer.""" + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker.is_producer = is_producer + worker.mode = MoRIIOMode.WRITE if is_producer else MoRIIOMode.READ + worker.tp_rank = 0 + worker.world_size = 1 if is_producer else 2 + worker.use_mla = True + worker._has_mamba = False + worker.layer_to_spec = {} + worker.moriio_config = SimpleNamespace(backend="xgmi", transfer_timeout=30.0) + worker._moriio_handshake = handshake + worker._handshake_lock = threading.RLock() + worker._handshake_initiation_executor = ThreadPoolExecutor(1) + worker._handshake_futures = {} + worker._eager_handshaked_engines = set() + worker._remote_agents = {} + worker.remote_moriio_metadata = {} + worker.write_ready_flags = {} + worker.load_ready_flag = {} + worker._ready_requests = queue.Queue() + worker._invalid_block_ids = set() + worker._reqs_to_send = {} + worker.tp_group = SimpleNamespace(cpu_group=None) + worker.shutdown = lambda: None + return worker + + +def _req_meta(remote_host: str, tp_size: int) -> ReqMeta: + return ReqMeta( + transfer_id="t", + local_block_ids=[[5, 6]], + remote_block_ids=[[7, 8]], + remote_host=remote_host, + remote_port=0, + remote_handshake_port=0, + remote_notify_port=NOTIFY_PORT, + remote_engine_id=f"{remote_host}:0", + tp_size=tp_size, + remote_dp_size=1, + ) + + +def _refuse(host, port, tp_size, engine_id, dp_rank, tp_rank): + raise HandshakeRefusedError(f"Cannot transfer KV with {engine_id}") + + +@pytest.mark.parametrize("targets_connected", [False, True]) +def test_refused_write_handshake_fails_the_request_on_every_decode_rank( + targets_connected, +): + """A refused handshake is tried once per request, not once per layer. The + request's first write task frees the prefill blocks and tells each decode + rank it would have written that the transfer failed, unless those ranks + are connected through another handshake.""" + handshakes = [] + + def refuse(*args): + handshakes.append(args[3]) + _refuse(*args) + + worker = _handshake_worker(True, refuse) + wrapper = _wrapper() + sent = [] + wrapper.send_notify = lambda transfer_id, ip, port, message_type: sent.append( + (transfer_id, port, message_type) + ) + worker.moriio_wrapper = wrapper + writer = worker._writer = make_moriio_writer(worker) + tasks = [] + + def write_blocks(req_id, meta, layer_name, kv_layer): + task = _write_task(layer_name, [5, 6], worker._write_tp_mapping(meta.tp_size)) + task.dst_engine_id = meta.remote_engine_id + if writer.schedule_write(task): + tasks.append(task) + + worker._write_blocks_for_req = write_blocks + metadata = MoRIIOConnectorMetadata() + metadata.reqs_to_save["req"] = _req_meta("decode", tp_size=4) + + for layer in ("mla0", "mla1", "mla2"): + worker.save_kv_layer(metadata, layer, kv_layer=None, attn_metadata=None) + worker._handshake_initiation_executor.shutdown() + + # One handshake per decode rank of the MLA fan-out, for the first layer. + assert sorted(handshakes) == [f"decode:0_dp0_tp{t}" for t in range(4)] + assert {task.layer_name for task in tasks} >= {"mla1", "mla2"} + + wrapper.done_remote_allocate_req_dict["t"] = RemoteAllocInfo( + block_ids=[[7, 8]], decode_dp_rank=0 + ) + written = [] + if targets_connected: + worker.remote_moriio_metadata = dict.fromkeys(handshakes) + worker._get_built_session = lambda engine_id: ([engine_id], None) + writer._prepare_transfer_plan = lambda *args, **kwargs: SimpleNamespace( + transfer_sizes=[64] + ) + + def do_layer_write(plan, sessions): + written.append(sessions) + return [] + + writer._do_layer_write = do_layer_write + writer._execute_write_task(tasks[0]) + + if targets_connected: + assert written == [[f"decode:0_dp0_tp{t}"] for t in range(4)] + assert sent == [] and wrapper.done_req_ids == [] + return + assert [ack.transfer_id for ack in wrapper.done_req_ids] == ["t"] + assert sent == [ + ("t", NOTIFY_PORT + get_port_offset(0, t), "write_failed") for t in range(4) + ] + # The request's other write tasks are dropped. + assert writer._is_transfer_terminal("t") + assert not writer.is_unconnected("t") + + +VOTES = {"connected": 2, "refused": 1, "failed": 0} + + +@pytest.mark.parametrize( + ("here", "peer", "has_mamba", "outcome"), + [ + ("connected", "connected", False, "read"), + ("refused", "connected", False, "fail"), + ("connected", "refused", False, "fail"), + ("failed", "connected", False, "raise"), + ("refused", "refused", True, "raise"), + ], +) +def test_read_refusal_fails_the_requests_on_every_tp_rank( + monkeypatch, here, peer, has_mamba, outcome +): + """A decode instance in READ mode fails the requests from a prefill + instance whose layout any TP rank refused, the same way on every rank, + and keeps running; any other handshake failure still fails the step.""" + + def handshake(*args): + if here == "refused": + _refuse(*args) + if here == "failed": + raise HandshakeError("connection reset") + return {"agent"} + + def all_reduce(vote, group, op): + vote[0] = min(int(vote[0]), VOTES[peer]) + + monkeypatch.setattr(torch.distributed, "all_reduce", all_reduce) + worker = _handshake_worker(False, handshake) + worker._has_mamba = has_mamba + reads = [] + worker._read_blocks_for_req = lambda req_id, meta: reads.append( + (req_id, meta.local_block_ids) + ) + worker._await_reads_issued_this_step = lambda: None + metadata = MoRIIOConnectorMetadata() + metadata.reqs_to_recv["req"] = _req_meta("prefill", tp_size=2) + + if outcome == "raise": + with pytest.raises(HandshakeError, match="failing this step fast"): + worker.start_load_kv(metadata) + else: + worker.start_load_kv(metadata) + worker._handshake_initiation_executor.shutdown() + + # A refused peer gets a read into no blocks, which only releases its blocks. + expected_reads = {"read": [("req", [[5, 6]])], "fail": [("req", [[]])]} + assert reads == expected_reads.get(outcome, []) + assert worker.get_block_ids_with_load_errors() == ( + {5, 6} if outcome == "fail" else set() + ) + # A refused peer is handshaked again for the next requests. + assert worker._eager_handshaked_engines == ( + {"prefill:0"} if outcome == "read" else set() + ) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 44b84b4f9988..da7ebb61a058 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -757,6 +757,34 @@ def _make_read_request(remote_block_ids): ) +@pytest.mark.parametrize("peer_known", [True, False]) +def test_read_release_without_prefill_address_is_skipped(peer_known): + """After a failed KV load the frontend finishes a cleanup copy of the + request under an id that embeds no prefill address. Its release cannot be + routed, so it is skipped instead of failing the next connector metadata.""" + sched = _make_read_scheduler() + sched._has_mamba = False + sched.blocks_per_sw = [0] + params = { + "do_remote_prefill": True, + "transfer_id": "tx", + "remote_engine_id": "prefill", + "remote_block_ids": [[10, 11]], + } + if peer_known: + params.update( + remote_host="127.0.0.1", remote_handshake_port=6000, remote_notify_port=6001 + ) + request = SimpleNamespace( + request_id="ad780c39f0ad1090", kv_transfer_params=params, status=None + ) + + assert sched.request_finished(request, []) == (False, None) + meta = sched.build_connector_meta(SchedulerOutput.make_empty()) + + assert list(meta.reqs_to_recv) == (["ad780c39f0ad1090"] if peer_known else []) + + @pytest.mark.parametrize("external_tokens", [0, 256]) @pytest.mark.parametrize("has_mamba", [False, True]) @pytest.mark.parametrize("zero_ids", [None, [], [101], [101, 102, 202, 999]]) diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index 7bb4f9676971..58f0c025a0f1 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -295,7 +295,7 @@ def _get_built_session(engine_id: str): def _prepare_transfer_plan( task, request_info, remote_moriio_meta, remote_engine_id=None, target=None ): - return SimpleNamespace(transfer_sizes=[64, 32]) + return SimpleNamespace(transfer_sizes=[]) def _do_layer_write(plan, sessions): write_sessions.append(sessions) @@ -320,7 +320,6 @@ def _do_layer_write(plan, sessions): "status-remote-engine_dp0_tp2", "status-remote-engine_dp0_tp3", ] - assert info.bytes_written == 2 * (64 + 32) def test_fanout_write_done_notifies_every_target_after_completion(): @@ -463,6 +462,7 @@ def _fanout_save_worker(remote_agents: set[str]) -> MoRIIOConnectorWorker: worker._write_blocks_for_req = ( # type: ignore[method-assign] lambda req_id, meta, layer_name, kv_layer: worker.writes.append(req_id) ) + worker._writer = make_moriio_writer(worker) return worker diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 903ef4e0e327..0d7336b73db5 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -681,6 +681,7 @@ def make_moriio_writer(fake_worker: Any) -> Any: writer._scheduled_writes = defaultdict(int) writer._scheduled_layers = defaultdict(set) writer._sealed_writes = {} + writer._unconnected = set() writer._deferred_tasks = [] writer._defer_timeout = 60.0 writer.ensure_worker_started = lambda: None diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 2ba5d197fa4f..69b4d64469a2 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -345,17 +345,6 @@ def enable_qk_norm_rope_kvcache(cfg: "VllmConfig") -> bool: } -def _uses_block_interleaved_pd_connector( - kv_transfer_config: KVTransferConfig | None, -) -> bool: - """NIXL and MoRIIO pair whole prefill blocks with DCP-sharded decode blocks, - which needs cp_kv_cache_interleave_size == block_size.""" - return kv_transfer_config is not None and any( - kv_transfer_config.has_connector(name) - for name in ("NixlConnector", "MoRIIOConnector") - ) - - @config(config=ConfigDict(arbitrary_types_allowed=True)) class VllmConfig: """Dataclass which contains all vllm-related configuration. This @@ -3351,7 +3340,7 @@ def _validate_v1_model_runner(self) -> None: def adjust_dcp_kv_cache_interleave_size( self, kv_cache_groups: list["KVCacheGroupSpec"] ) -> None: - """Resolve the NIXL/MoRIIO P/D DCP interleave from the global KV cache groups. + """Resolve the NIXL P/D DCP interleave from the global KV cache groups. Called once by the engine core before profiling; workers only assign the result via ``set_kv_cache_layout``. @@ -3378,7 +3367,7 @@ def adjust_dcp_kv_cache_interleave_size( "deprecated when PCP is fully supported." ) - if not _uses_block_interleaved_pd_connector(self.kv_transfer_config): + if not self.kv_transfer_config.has_connector("NixlConnector"): return if not self.parallel_config._allow_auto_resolve_cp_interleave_size: return @@ -3407,12 +3396,13 @@ def validate_block_size(self) -> None: """ block_size = self.cache_config.block_size - # Skip DCP interleave-size compatibility for NIXL/MoRIIO P/D: the - # interleave size is finalized from the global KV cache groups. - if ( - self.parallel_config.decode_context_parallel_size > 1 - and not _uses_block_interleaved_pd_connector(self.kv_transfer_config) - ): + # Skip DCP interleave-size compatibility for NIXL P/D: the interleave + # size is finalized from the global KV cache groups. + nixl_pd_active = ( + self.kv_transfer_config is not None + and self.kv_transfer_config.has_connector("NixlConnector") + ) + if self.parallel_config.decode_context_parallel_size > 1 and not nixl_pd_active: assert ( self.parallel_config.cp_kv_cache_interleave_size <= block_size and block_size % self.parallel_config.cp_kv_cache_interleave_size == 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index fbabecead43b..ac6c5b3c9ba3 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -147,9 +147,9 @@ class MoRIIOAgentMetadata( # Tokens per block of each attention layer. Hybrid models size attention # blocks to the recurrent-state page, which depends on the TP size. attention_block_sizes: dict[str, int] = {} - # The defaults also stand for older peers that do not send these fields. - dcp_size: int = 1 - cp_kv_cache_interleave_size: int = 1 + # None: a peer that does not advertise its DCP layout. + dcp_size: int | None = None + cp_kv_cache_interleave_size: int | None = None class RoleManager: @@ -220,6 +220,11 @@ class HandshakeError(MoRIIOError): pass +class HandshakeRefusedError(HandshakeError): + """The peer's KV cache layout cannot be transferred with; a retry fails + the same way.""" + + class TransferError(MoRIIOError): """Exception raised when transfer fails.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index 54b3cb13acc7..ab4906a885d4 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import contextlib +import dataclasses import logging import math import queue @@ -34,6 +35,7 @@ ROLE, EngineId, HandshakeError, + HandshakeRefusedError, MoRIIOAgentMetadata, MoRIIOConfig, MoRIIOConnectorMetadata, @@ -77,19 +79,22 @@ compute_head_slice_offsets, compute_mamba_conv_split_count, compute_mamba_slice_offsets, + compute_token_run_offsets, get_layer_head_layout, get_layer_transfer_geometry, get_mamba_state_layout, + get_token_run_table, is_mla_cache_layer, iter_layer_registration_regions, kda_conv_ssm, validate_transfer_block_ids, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_tp_mapping import ( + DCPWriteRuns, MoRIIOWriteTPMapping, WriteTarget, compute_write_tp_mapping, - validate_write_dcp_topology, + validate_peer_dcp_layout, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( MambaConvSplitInfo, @@ -1565,6 +1570,24 @@ def request_finished( # embedded MoRIIO ZMQ addresses. if self.mode == MoRIIOMode.WRITE: self._release_write_prefill_blocks(request.request_id, params) + elif get_peer_zmq_from_request_id( + request.request_id, is_producer=False + ) is None and not all( + params.get(key) + for key in ( + "remote_host", + "remote_handshake_port", + "remote_notify_port", + ) + ): + # E.g. the frontend's cleanup copy of a request whose KV load + # failed, under an id without the prefill address. The prefill + # instance frees the blocks after its defer_timeout. + logger.warning( + "Cannot release the prefill blocks of request %s: no prefill " + "address in its request id or kv_transfer_params", + request.request_id, + ) else: # Carry the actually-allocated blocks (incl. the per-request KDA # mamba slot) instead of []: a do_remote_prefill request finishing @@ -1958,12 +1981,6 @@ def __init__( self.cp_kv_cache_interleave_size = ( vllm_config.parallel_config.cp_kv_cache_interleave_size ) - if self.dcp_size > 1 and (self.is_producer or self.mode == MoRIIOMode.READ): - raise NotImplementedError( - "MoRIIO supports decode context parallelism " - f"(decode_context_parallel_size={self.dcp_size}) only on the " - "decode instance in WRITE mode" - ) # Hybrid (mamba/KDA) recurrent-state transfer. Populated in # register_kv_caches when the model has a MambaSpec kv_cache group. self._has_mamba = False @@ -2236,19 +2253,24 @@ def _moriio_handshake( metadata_bytes = received_frame[1] decoder = msgspec.msgpack.Decoder(MoRIIOAgentMetadata) metadata = decoder.decode(metadata_bytes) - if self.is_producer and self.mode == MoRIIOMode.WRITE: - try: - validate_write_dcp_topology( - metadata.dcp_size, - metadata.cp_kv_cache_interleave_size, - self.block_size, - self.use_mla, - self._has_mamba, - ) - except ValueError as e: - raise HandshakeError( - f"Cannot write into {expected_engine_id}: {e}" - ) from e + try: + validate_peer_dcp_layout( + self.dcp_size, + self.cp_kv_cache_interleave_size, + metadata.dcp_size, + metadata.cp_kv_cache_interleave_size, + writes_into_peer=self.is_producer and self.mode == MoRIIOMode.WRITE, + block_size=self.block_size, + is_mla=self.use_mla, + has_mamba=self._has_mamba, + whole_block_layer=self._whole_block_dcp_layer( + metadata.cp_kv_cache_interleave_size + ), + ) + except ValueError as e: + raise HandshakeRefusedError( + f"Cannot transfer KV with {expected_engine_id}: {e}" + ) from e got_metadata_time = time.perf_counter() logger.info( "MoRIIO handshake: get metadata took: %s", @@ -2323,8 +2345,12 @@ def _background_moriio_handshake( pod_hosts = list(meta.multi_pod_hosts) if meta.multi_pod_hosts else [host] remote_dp_size_local = int(meta.remote_dp_size_local) or remote_dp_size - def request_ready(_f: Future[Any], entry=(req_id, meta)): - logger.info("MoRIIO handshake done for request %s", req_id) + def request_ready(f: Future[Any], entry=(req_id, meta)): + if f.exception() is None: + logger.info("MoRIIO handshake done for request %s", req_id) + elif self.is_producer and self.mode == MoRIIOMode.WRITE: + logger.error("MoRIIO handshake failed for request %s", req_id) + self._writer.mark_unconnected(entry[1].transfer_id) self._ready_requests.put(entry) self.load_ready_flag[remote_engine_id] = True self.write_ready_flags[remote_engine_id] = True @@ -2595,15 +2621,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.block_shape = block_shape self.kv_element_size = kv_elem_size - if self.dcp_size > 1: - validate_write_dcp_topology( - self.dcp_size, - self.cp_kv_cache_interleave_size, - self.block_size, - self.use_mla, - self._has_mamba, - ) - self.dst_num_blocks[self.engine_id] = self.num_blocks kv_caches_base_addr = [] caches_data = [] @@ -3113,7 +3130,11 @@ def save_kv_layer( if is_fanout else not write_targets or dp0_remote_engine_id in self._remote_agents ) - if not handshake_ready: + # A request whose handshake failed is not retried; its writes + # reach the decode ranks that are connected or fail the request. + if not handshake_ready and not self._writer.is_unconnected( + meta.transfer_id + ): # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: if remote_engine_id not in self._remote_agents: @@ -3160,7 +3181,9 @@ def get_engine_name_with_dp_tp(self, engine_name, dp_rank, tp_rank): # rank and spread reads across them; other configs keep the DP-only key. return f"{engine_name}_dp{dp_rank}_tp{tp_rank}" - def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> None: + def _eager_handshake_all_dp_ranks( + self, metadata: MoRIIOConnectorMetadata + ) -> set[str]: """Handshake EVERY remote prefill DP rank BEFORE the decode forward pass, identically across all local TP workers. @@ -3188,12 +3211,16 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No Every worker reaches the all-reduce(MIN) vote; if ANY worker failed, ALL raise the same error AFTER the collective, so the step fails fast and uniformly in ~seconds instead of one rank hanging the forward for 600s. + When every failure is a refusal of the peer's layout, every worker + returns the peer instead and its requests fail; hybrid models still + raise, since their load errors cannot be reported. """ import torch.distributed as dist # Distinct remote engines referenced this step, in metadata (== # scheduler) order so every TP worker iterates engines identically. engines: dict[str, ReqMeta] = {} + refused_engines: set[str] = set() for _req_id, meta in metadata.reqs_to_recv.items(): remote_engine_id = ( str(meta.remote_host) + ":" + str(meta.remote_handshake_port) @@ -3263,13 +3290,15 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No # Join outside the lock. Bounded handshake errors are recorded here # and reported after the all-reduce. all_ok = True + refused = True results: dict[str, set[str]] = {} for eid, fut in futures: try: results[eid] = fut.result() - except Exception: + except Exception as e: logger.exception("Eager MoRIIO handshake failed for %s", eid) all_ok = False + refused = refused and isinstance(e, HandshakeRefusedError) with self._handshake_lock: for eid, agents in results.items(): @@ -3287,9 +3316,16 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No # CPU all-reduce = TP-uniform success vote AND lockstep barrier: it # blocks until every TP worker arrives, gives them the same verdict, # and stays off the model compute stream. - vote = torch.tensor([1 if all_ok else 0], device="cpu", dtype=torch.int32) + # 2: connected, 1: only refusals of the peer's layout, 0: failed. + vote = torch.tensor( + [2 if all_ok else int(refused)], device="cpu", dtype=torch.int32 + ) dist.all_reduce(vote, group=self.tp_group.cpu_group, op=dist.ReduceOp.MIN) - if int(vote.item()) == 0: + verdict = int(vote.item()) + if verdict == 1 and not self._has_mamba: + refused_engines.add(remote_engine_id) + continue + if verdict < 2: raise HandshakeError( f"Eager MoRIIO handshake failed for {remote_engine_id} on " "at least one TP rank; failing this step fast to avoid a " @@ -3297,6 +3333,7 @@ def _eager_handshake_all_dp_ranks(self, metadata: MoRIIOConnectorMetadata) -> No ) self._eager_handshaked_engines.add(remote_engine_id) + return refused_engines def start_load_kv(self, metadata: MoRIIOConnectorMetadata): """Start loading by triggering non-blocking moriio_xfer. @@ -3322,7 +3359,7 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): # read enters the forward pass. A lazy per-rank handshake on the read # path lets TP workers diverge into a forward collective while a peer is # still blocked handshaking -> NCCL hang (see below). - self._eager_handshake_all_dp_ranks(metadata) + refused_engines = self._eager_handshake_all_dp_ranks(metadata) wait_handshake_readd_req = False remote_engine_id = None @@ -3332,6 +3369,19 @@ def start_load_kv(self, metadata: MoRIIOConnectorMetadata): str(meta.remote_host) + ":" + str(meta.remote_handshake_port) ) meta.remote_engine_id = remote_engine_id + if remote_engine_id in refused_engines: + # Scheduler recovery fails or recomputes the request; reading + # into no blocks only releases the prefill instance's blocks. + self._invalid_block_ids.update( + block_id for group in meta.local_block_ids for block_id in group + ) + self._read_blocks_for_req( + req_id, + dataclasses.replace( + meta, local_block_ids=[[] for _ in meta.local_block_ids] + ), + ) + continue # The eager handshake above already covered every referenced engine # (and keys the mirror per (dp, tp), which the DP-only dp0 probe below # would miss). Only fall back to the lazy background handshake for an @@ -3626,6 +3676,62 @@ def _compute_block_transfer_offsets( layer_name, (local, remote, sizes), remote_moriio_meta, remote_engine_id ) + def _compute_dcp_run_offsets( + self, + layer_name: str, + runs: DCPWriteRuns, + run_tokens: int, + remote_moriio_meta: MoRIIOAgentMetadata, + remote_engine_id: EngineId | None = None, + ) -> tuple[list[int], list[int], list[int]]: + """Offsets writing runs of ``run_tokens`` tokens into the blocks of + one consumer DCP rank.""" + offsets = compute_token_run_offsets( + layer_name, + self.kv_caches[layer_name], + self.layer_to_spec, + runs, + run_tokens, + remote_moriio_meta.num_blocks, + self._mla_latent_size, + merge_fn=lambda local, remote, sizes: self.merge_contiguous_blocks( + local, remote, sizes, assume_sorted=False + ), + ) + return self._apply_layer_mr_offsets( + layer_name, offsets, remote_moriio_meta, remote_engine_id + ) + + def _whole_block_dcp_layer(self, interleave_size: int | None) -> str | None: + """A DCP-sharded layer that cannot be written in runs of + ``interleave_size`` tokens, only in whole blocks.""" + if ( + not interleave_size + or interleave_size >= self.block_size + or self.block_size % interleave_size + ): + return None + return next( + ( + layer_name + for layer_name, kv_cache in self.kv_caches.items() + if self.layer_to_spec[layer_name].dcp_sharded + and get_token_run_table( + layer_name, + kv_cache, + self.layer_to_spec, + interleave_size, + self._mla_latent_size, + ) + is None + ), + None, + ) + + @property + def _mla_latent_size(self) -> int | None: + return self.model_config.get_head_size() if self.use_mla else None + def _apply_layer_mr_offsets( self, layer_name: str, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index 6e36c4ac32a1..ba8a6a875be7 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -41,7 +41,7 @@ zmq_ctx, ) from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_tp_mapping import ( - dcp_write_block_pairs, + dcp_write_runs, ) if TYPE_CHECKING: @@ -112,6 +112,8 @@ def __init__(self, worker: "MoRIIOConnectorWorker"): self._scheduled_writes: dict[TransferId, int] = defaultdict(int) self._scheduled_layers: dict[TransferId, set[str]] = defaultdict(set) self._sealed_writes: dict[TransferId, int] = {} + # Transfers whose handshake failed. + self._unconnected: set[TransferId] = set() self._defer_timeout = worker.moriio_config.defer_timeout @property @@ -167,6 +169,16 @@ def is_scheduled(self, transfer_id: TransferId, layer_name: str) -> bool: with self._write_state_lock: return layer_name in self._scheduled_layers.get(transfer_id, set()) + def mark_unconnected(self, transfer_id: TransferId) -> None: + """Record that a transfer's handshake failed. Its writes still go to + decode ranks that are connected; otherwise the transfer fails.""" + with self._write_state_lock: + self._unconnected.add(transfer_id) + + def is_unconnected(self, transfer_id: TransferId) -> bool: + with self._write_state_lock: + return transfer_id in self._unconnected + def seal_pending_transfers(self) -> None: """Seal expected WRITE counts after the model forward has run. @@ -267,8 +279,12 @@ def _notify_write_failed(self, task: WriteTask, age: float) -> None: task.request_id, age, ) - # No remote allocation arrived, so the decode DP rank that owns the - # request is unknown. Tell every rank; the others drop it. + self._send_write_failed(task) + + def _send_write_failed(self, task: WriteTask) -> None: + # Only the remote allocation names the decode DP rank that owns the + # request, and it may not have arrived. Tell every rank; the others + # drop it. for decode_dp_rank in range(max(task.remote_dp_size, 1)): for remote_ip, remote_port in self._resolve_notify_endpoints( task, decode_dp_rank @@ -282,7 +298,7 @@ def _notify_write_failed(self, task: WriteTask, age: float) -> None: ) except Exception: logger.exception( - "Failed to notify decode DP rank %d that transfer %s timed out", + "Failed to notify decode DP rank %d that transfer %s failed", decode_dp_rank, task.transfer_id, ) @@ -292,6 +308,7 @@ def _clear_transfer_state(self, transfer_id: TransferId) -> None: self._scheduled_writes.pop(transfer_id, None) self._scheduled_layers.pop(transfer_id, None) self._sealed_writes.pop(transfer_id, None) + self._unconnected.discard(transfer_id) def _is_transfer_terminal(self, transfer_id: TransferId) -> bool: wrapper = self.worker.moriio_wrapper @@ -376,6 +393,31 @@ def _execute_write_task(self, task: WriteTask) -> None: ) return + is_mamba = self.worker._is_mamba_layer(task.layer_name) + targets = [ + target + for target in self._write_targets(task) + # Another producer rank writes the attention copies of the rest. + if target is None or is_mamba or target.num_heads + ] + engine_ids = [ + self._target_engine_id(task, request_info.decode_dp_rank, target) + for target in targets + ] + if self.is_unconnected(task.transfer_id) and not all( + engine_id in self.worker.remote_moriio_metadata for engine_id in engine_ids + ): + # Tell the decode ranks, so they do not wait for KV that cannot + # arrive. + logger.error( + "Failing the KV transfer of request %s: the handshake with its " + "decode instance failed", + task.request_id, + ) + self._mark_request_done(task.transfer_id) + self._send_write_failed(task) + return + # Wait for CUDA event # The attention computation of the current layer cannot # overlap with the kv transfer task, @@ -383,16 +425,9 @@ def _execute_write_task(self, task: WriteTask) -> None: # This event is used to synchronize the kv transfer and computation tasks. task.event.synchronize() - is_mamba = self.worker._is_mamba_layer(task.layer_name) transfer_statuses: list[Any] = [] bytes_written = 0 - for target in self._write_targets(task): - if target is not None and not is_mamba and not target.num_heads: - # Another producer rank writes this attention copy. - continue - target_engine_id = self._target_engine_id( - task, request_info.decode_dp_rank, target - ) + for target, target_engine_id in zip(targets, engine_ids): # Get or create sessions sessions, remote_moriio_meta = self.worker._get_built_session( target_engine_id @@ -458,13 +493,40 @@ def _prepare_transfer_plan( if self.worker._has_mamba: # D can reserve a tail block for the token recomputed locally. remote_blocks = remote_blocks[: len(local_blocks)] - dcp_size = remote_moriio_meta.dcp_size - if dcp_size > 1 and self.worker.layer_to_spec[task.layer_name].dcp_sharded: + dcp_size = remote_moriio_meta.dcp_size or 1 + runs = None + if ( + dcp_size > 1 + and self.worker.dcp_size == 1 + and self.worker.layer_to_spec[task.layer_name].dcp_sharded + ): + block_size = self.worker.layer_to_spec[task.layer_name].block_size + run_tokens = ( + remote_moriio_meta.cp_kv_cache_interleave_size or block_size + ) tp_rank = self.worker.tp_rank if target is None else target.tp_rank - local_blocks, remote_blocks = dcp_write_block_pairs( - local_blocks, remote_blocks, dcp_size, tp_rank % dcp_size + runs = dcp_write_runs( + local_blocks, + remote_blocks, + block_size, + dcp_size, + run_tokens, + tp_rank % dcp_size, ) - if ( + if run_tokens == block_size: + # Each run is a whole block. + local_blocks = runs.local_blocks.tolist() + remote_blocks = runs.remote_blocks.tolist() + runs = None + if runs is not None: + offsets = self.worker._compute_dcp_run_offsets( + task.layer_name, + runs, + run_tokens, + remote_moriio_meta, + remote_engine_id=remote_engine_id, + ) + elif ( target is None or task.tp_mapping is None or ( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index 6a0e51023ab1..19f95f3a667b 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -2,8 +2,10 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable, Mapping +from functools import cache from typing import NamedTuple, cast +import numpy as np import torch from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( @@ -16,6 +18,7 @@ MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, + SparseCacheRole, UniformTypeKVCacheSpecs, ) @@ -491,6 +494,123 @@ def compute_block_transfer_offsets( return merge_fn(offset_local, offset_remote, sizes) +@cache +def _token_run_table( + is_indexer: bool, block_size: int, entry_bytes: int, run_tokens: int +) -> tuple[np.ndarray, np.ndarray] | None: + if is_indexer: + from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( + indexer_k_cache_byte_ranges, + ) + + runs = [ + indexer_k_cache_byte_ranges(block_size, entry_bytes, slot, run_tokens) + for slot in range(0, block_size, run_tokens) + ] + else: + # MLA latents are stored token by token. + runs = [ + [(slot * entry_bytes, run_tokens * entry_bytes)] + for slot in range(0, block_size, run_tokens) + ] + sizes = [size for _, size in runs[0]] + if any([size for _, size in run] != sizes for run in runs): + return None + offsets = [[offset for offset, _ in run] for run in runs] + return np.array(offsets, dtype=np.int64), np.array(sizes, dtype=np.int64) + + +def get_token_run_table( + layer_name: str, + kv_cache: torch.Tensor, + layer_to_spec: Mapping[str, KVCacheSpec], + run_tokens: int, + latent_size: int | None, +) -> tuple[np.ndarray, np.ndarray] | None: + """Byte ranges of a run of ``run_tokens`` tokens in a block of this layer. + + Returns ``(offsets, sizes)``: ``offsets[i]`` are the ranges, relative to + the block, of the run that starts at slot ``i * run_tokens``, and every run + has ranges of the same ``sizes``. Only an MLA latent of ``latent_size`` + elements per token and the indexer K cache qualify; any other layer + returns None and can only be copied in whole blocks. + """ + spec = layer_to_spec[layer_name] + if ( + type(spec) is not MLAAttentionSpec + or spec.tokens_per_state != 1 + or spec.num_heads != 1 + or spec.state_content_bytes is not None + ): + return None + geometry = cast( + LayerTransferGeometry, + get_layer_transfer_geometry(layer_name, kv_cache, layer_to_spec), + ) + if geometry.block_len != geometry.block_size * geometry.slot_size_bytes: + return None + is_indexer = spec.cache_role is SparseCacheRole.INDEXER + # The indexer cache holds 128 fp8 values and one fp32 scale per token. + if is_indexer and geometry.slot_size_bytes != 132: + return None + if not is_indexer and spec.head_size != latent_size: + return None + return _token_run_table( + is_indexer, geometry.block_size, geometry.slot_size_bytes, run_tokens + ) + + +def compute_token_run_offsets( + layer_name: str, + kv_cache: torch.Tensor, + layer_to_spec: Mapping[str, KVCacheSpec], + runs: tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray], + run_tokens: int, + remote_num_blocks: int, + latent_size: int | None, + merge_fn: Callable[ + [list[int], list[int], list[int]], tuple[list[int], list[int], list[int]] + ] = merge_contiguous_offsets, +) -> tuple[list[int], list[int], list[int]]: + """Offsets copying runs of ``run_tokens`` tokens between blocks of the + same layout. + + ``runs`` holds the local block id and slot and the remote block id and + slot of each run's first token. + """ + table = get_token_run_table( + layer_name, kv_cache, layer_to_spec, run_tokens, latent_size + ) + if table is None: + raise ValueError( + f"MoRIIO cannot copy layer {layer_name} in runs of {run_tokens} tokens" + ) + offsets, sizes = table + geometry = cast( + LayerTransferGeometry, + get_layer_transfer_geometry( + layer_name, kv_cache, layer_to_spec, remote_num_blocks + ), + ) + local_blocks, local_slots, remote_blocks, remote_slots = runs + for side, blocks, num_blocks in ( + ("local", local_blocks, geometry.num_blocks), + ("remote", remote_blocks, remote_num_blocks), + ): + if len(blocks) and not 0 <= blocks.min() <= blocks.max() < num_blocks: + raise ValueError(f"{side} block ids are outside [0, {num_blocks})") + block_stride = geometry.block_stride * kv_cache.element_size() + local = (local_blocks * block_stride)[:, None] + offsets[local_slots // run_tokens] + remote = (remote_blocks * block_stride)[:, None] + offsets[ + remote_slots // run_tokens + ] + return merge_fn( + local.ravel().tolist(), + remote.ravel().tolist(), + np.tile(sizes, len(local_blocks)).tolist(), + ) + + class MambaOffsetTemplate(NamedTuple): """Slot-independent conv+ssm offset decomposition for a KDA layer. diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py index 94c3bdae8c83..c2042b351e37 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py @@ -21,15 +21,24 @@ tp_size``. A consumer with decode context parallelism (DCP) shards the tokens of its MLA -latent instead. With ``cp_kv_cache_interleave_size == block_size``, block ``j`` -of the producer, which is not DCP-sharded, lives whole on DCP rank -``j % dcp_size`` (``tp_rank % dcp_size``) in that rank's block ``j // dcp_size``, -so each block is one write (``dcp_write_block_pairs``). +caches instead. A producer with the same DCP size holds the same shard on the +same rank. A producer without DCP writes each consumer DCP rank +(``tp_rank % dcp_size``) the tokens it holds: with ``I = +cp_kv_cache_interleave_size``, token ``t`` lives on DCP rank ``(t // I) % +dcp_size`` at position ``p = (t // (I * dcp_size)) * I + t % I`` of that rank's +share, i.e. in its block ``p // block_size``. ``I`` divides the block size, so +each group of ``I`` tokens is contiguous on both sides and is one write +(``dcp_write_runs``); with ``I == block_size`` a group is a whole block. """ from collections.abc import Sequence from dataclasses import dataclass from functools import cache +from typing import NamedTuple + +import numpy as np + +from vllm.utils.math_utils import cdiv @dataclass(frozen=True) @@ -221,50 +230,109 @@ def compute_write_tp_mapping( ) -def validate_write_dcp_topology( - consumer_dcp_size: int, - consumer_interleave_size: int, +def validate_peer_dcp_layout( + local_dcp_size: int, + local_interleave_size: int, + peer_dcp_size: int | None, + peer_interleave_size: int | None, + *, + writes_into_peer: bool, block_size: int, is_mla: bool, has_mamba: bool, + whole_block_layer: str | None = None, ) -> None: - if consumer_dcp_size <= 1: + """Refuse a peer whose DCP layout this worker cannot transfer KV with. + + Equal DCP sizes transfer rank to rank. A WRITE producer without DCP can + also write into a DCP consumer, one write per group of interleaved tokens; + ``whole_block_layer`` names a layer whose blocks can only be written whole, + which needs the consumer's interleave to be the block size. A peer that + does not advertise its layout predates these checks and is accepted as + before. + """ + if peer_dcp_size is None: + return + if peer_dcp_size == local_dcp_size: + if ( + local_dcp_size > 1 + and peer_interleave_size is not None + and peer_interleave_size != local_interleave_size + ): + raise ValueError( + "MoRIIO needs the same --cp-kv-cache-interleave-size on both " + f"instances with decode_context_parallel_size={local_dcp_size}: " + f"{local_interleave_size} here, {peer_interleave_size} on the peer" + ) return + if not (writes_into_peer and local_dcp_size == 1): + raise ValueError( + "MoRIIO cannot transfer KV between decode_context_parallel_size=" + f"{local_dcp_size} here and {peer_dcp_size} on the peer. Use the same " + "size on both instances, or a prefill instance without decode " + "context parallelism in WRITE mode" + ) if not is_mla or has_mamba: raise ValueError( - "MoRIIO WRITE into a decode instance with decode context parallelism " - f"(decode_context_parallel_size={consumer_dcp_size}) supports MLA " - "models without recurrent state only" + "MoRIIO WRITE from a prefill instance without decode context " + "parallelism into a decode instance with " + f"decode_context_parallel_size={peer_dcp_size} supports MLA models " + "without recurrent state only" ) - if consumer_interleave_size != block_size: + if peer_interleave_size is None or block_size % peer_interleave_size: raise ValueError( "MoRIIO WRITE into a decode instance with decode context parallelism " - f"(decode_context_parallel_size={consumer_dcp_size}) writes whole " - "blocks, so it needs cp_kv_cache_interleave_size == block_size; got " - f"cp_kv_cache_interleave_size={consumer_interleave_size}, " - f"block_size={block_size}. Leave --cp-kv-cache-interleave-size unset " - f"or set it to {block_size}" + "needs a --cp-kv-cache-interleave-size that divides the block size " + f"{block_size}; it uses {peer_interleave_size}" + ) + if whole_block_layer is not None and peer_interleave_size < block_size: + raise ValueError( + f"MoRIIO writes layer {whole_block_layer} in whole blocks only, so " + "the decode instance with decode_context_parallel_size=" + f"{peer_dcp_size} needs --cp-kv-cache-interleave-size {block_size} " + f"(the block size); it uses {peer_interleave_size}" ) -def dcp_write_block_pairs( +class DCPWriteRuns(NamedTuple): + """First token of each run: producer block id and slot, consumer block id + and slot.""" + + local_blocks: np.ndarray + local_slots: np.ndarray + remote_blocks: np.ndarray + remote_slots: np.ndarray + + +def dcp_write_runs( local_block_ids: Sequence[int], remote_block_ids: Sequence[int], + block_size: int, dcp_size: int, + interleave_size: int, dcp_rank: int, -) -> tuple[list[int], list[int]]: - """Producer blocks that consumer DCP rank ``dcp_rank`` holds, paired with - that rank's block ids. +) -> DCPWriteRuns: + """Runs of ``interleave_size`` producer tokens that consumer DCP rank + ``dcp_rank`` holds, in the rank's token order. - A rank that holds none of them (a prompt shorter than ``dcp_size`` blocks) - gets two empty lists and must still be sent write_done. + A rank that holds none of them (a prompt shorter than ``dcp_size`` runs) + gets empty runs and must still be sent write_done. """ - local = list(local_block_ids[dcp_rank::dcp_size]) - remote = list(remote_block_ids[: len(local)]) - if len(remote) < len(local): + runs = np.arange( + dcp_rank, len(local_block_ids) * block_size // interleave_size, dcp_size + ) + tokens = runs * interleave_size + positions = runs // dcp_size * interleave_size + needed = cdiv(len(runs) * interleave_size, block_size) + if len(remote_block_ids) < needed: raise ValueError( - f"DCP rank {dcp_rank}/{dcp_size} of the consumer needs {len(local)} " + f"DCP rank {dcp_rank}/{dcp_size} of the consumer needs {needed} " f"blocks for {len(local_block_ids)} producer blocks, but it allocated " f"{len(remote_block_ids)}" ) - return local, remote + return DCPWriteRuns( + np.asarray(local_block_ids, dtype=np.int64)[tokens // block_size], + tokens % block_size, + np.asarray(remote_block_ids, dtype=np.int64)[positions // block_size], + positions % block_size, + ) diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 0db22e03c0bf..40ec25e2e55b 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -215,6 +215,55 @@ def indexer_k_quant_and_cache_triton( ) +def indexer_k_cache_byte_ranges( + block_size: int, + entry_bytes: int, + first_slot: int, + num_slots: int, + block_tile_size: int = 16, + head_tile_size: int = 16, +) -> list[tuple[int, int]]: + """(offset in block, size) of the bytes that tokens ``[first_slot, + first_slot + num_slots)`` occupy in one block of the indexer K cache + written by ``indexer_k_quant_and_cache_triton``. + + A block holds the fp8 K values of all its tokens, then one fp32 scale per + token. Above block size 1 the values are shuffled in tiles of + ``block_tile_size`` tokens: a token's ``head_tile_size``-byte pieces lie + ``block_tile_size * head_tile_size`` bytes apart. + """ + head_dim = entry_bytes - 4 + end = first_slot + num_slots + ranges = [] + if block_size == 1: + ranges.append((first_slot * head_dim, num_slots * head_dim)) + else: + tile_bytes = block_tile_size * head_dim + for tile in range( + first_slot // block_tile_size, (end - 1) // block_tile_size + 1 + ): + tile_start = tile * block_tile_size + lo = max(first_slot, tile_start) - tile_start + hi = min(end, tile_start + block_tile_size) - tile_start + ranges += [ + ( + tile * tile_bytes + + piece * block_tile_size * head_tile_size + + lo * head_tile_size, + (hi - lo) * head_tile_size, + ) + for piece in range(head_dim // head_tile_size) + ] + ranges.append((block_size * head_dim + first_slot * 4, num_slots * 4)) + merged: list[tuple[int, int]] = [] + for offset, size in ranges: + if merged and sum(merged[-1]) == offset: + merged[-1] = (merged[-1][0], merged[-1][1] + size) + else: + merged.append((offset, size)) + return merged + + @triton.jit def _cp_gather_indexer_quant_cache_kernel( kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B] From 82a6fc9d44dc9b0bb10c7ad4710c75dd55f89b95 Mon Sep 17 00:00:00 2001 From: kliuae Date: Fri, 30 Oct 2026 16:51:08 +0800 Subject: [PATCH 3/6] write the sparse mla indexer cache in whole blocks only Signed-off-by: kliuae Co-authored-by: Claude Opus --- .../unit/test_moriio_dcp_write.py | 48 +++++------------- .../kv_connector/v1/moriio/moriio_layout.py | 37 ++++---------- .../v1/attention/ops/rocm_aiter_mla_sparse.py | 49 ------------------- 3 files changed, 23 insertions(+), 111 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py index 37070dedd48e..e29da895672e 100644 --- a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -35,7 +35,6 @@ validate_peer_dcp_layout, ) from vllm.utils.math_utils import cdiv -from vllm.v1.attention.ops.rocm_aiter_mla_sparse import indexer_k_cache_byte_ranges from vllm.v1.kv_cache_interface import MLAAttentionSpec, SparseCacheRole from .utils import make_moriio_writer @@ -57,11 +56,10 @@ def _spec(head_size: int, **kwargs) -> MLAAttentionSpec: ) -# The fp8 caches of DeepSeek-V3.2 and GLM-5: the MLA latent (kv_lora_rank 512 + -# qk_rope_head_dim 64) and the indexer K cache (128 fp8 values, one fp32 scale). +# The fp8 MLA latent of DeepSeek-V3.2 and GLM-5 (kv_lora_rank 512 + +# qk_rope_head_dim 64). SPECS = { "latent": _spec(576), - "indexer": _spec(132, cache_role=SparseCacheRole.INDEXER), } @@ -86,30 +84,7 @@ def _token_bytes(spec: MLAAttentionSpec, slot: int) -> torch.Tensor: """Offsets of one token's bytes in a block, as the cache writers lay them out.""" entry = spec.state_content_size_bytes - if spec.cache_role is not SparseCacheRole.INDEXER: - return slot * entry + torch.arange(entry) - # indexer_k_quant_and_cache_triton and the fused DeepSeek-V3.2 / GLM-5 - # writer: fp8 values in tiles of 16 tokens x 16 bytes, then the scales. - head_dim = entry - 4 - dim = torch.arange(head_dim) - values = slot // 16 * 16 * head_dim + slot % 16 * 16 + dim // 16 * 256 + dim % 16 - scales = BLOCK_SIZE * head_dim + slot * 4 + torch.arange(4) - return torch.cat([values, scales]) - - -def test_indexer_byte_ranges_follow_the_tiled_layout(): - # One token: eight 16-byte pieces, 256 bytes apart, and its scale. - assert indexer_k_cache_byte_ranges(64, 132, 3, 1) == [ - *((48 + 256 * piece, 16) for piece in range(8)), - (64 * 128 + 12, 4), - ] - # Whole aligned tiles collapse to one values range and one scales range. - assert indexer_k_cache_byte_ranges(64, 132, 16, 32) == [ - (16 * 128, 32 * 128), - (64 * 128 + 16 * 4, 32 * 4), - ] - # Block size 1 is not tiled. - assert indexer_k_cache_byte_ranges(1, 132, 0, 1) == [(0, 132)] + return slot * entry + torch.arange(entry) def _validate(local_dcp, local_interleave, peer_dcp, peer_interleave, **kwargs): @@ -203,15 +178,18 @@ def _worker( ("specs", "whole_block_layer"), [ (SPECS, None), - ({**SPECS, "indexer": _spec(132)}, "indexer"), + ( + {**SPECS, "indexer": _spec(132, cache_role=SparseCacheRole.INDEXER)}, + "indexer", + ), ({"latent": _spec(576, tokens_per_state=4)}, "latent"), ], - ids=["latent-and-indexer", "indexer-without-indexer-role", "compressed-states"], + ids=["latent", "indexer", "compressed-states"], ) def test_only_known_token_layouts_are_written_in_runs(specs, whole_block_layer): """Runs of tokens are written only for an MLA latent of the model's latent - width and for the indexer K cache; any other layer, such as DeepSeek-V4's - indexer cache with the default cache role, needs whole blocks.""" + width; any other layer, such as the indexer K cache of DeepSeek-V3.2 and + GLM-5, needs whole blocks, i.e. an interleave equal to the block size.""" worker = _worker(0, 1, specs) for interleave in (1, 16): assert worker._whole_block_dcp_layer(interleave) == whole_block_layer @@ -375,9 +353,9 @@ def _expected_decode_cache( def test_write_lands_each_token_on_its_dcp_rank( monkeypatch, producer_tp, consumer_tp, dcp_size, num_prefill_blocks, interleave ): - """A sparse MLA model writes its latent and indexer caches token by token - at any interleave, from one prefill rank into several decode ranks and - from several prefill ranks into one.""" + """An MLA model writes its latent cache token by token at any interleave, + from one prefill rank into several decode ranks and from several prefill + ranks into one.""" monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) # One spare decode block, as when decode reserves lookahead slots. num_decode_blocks = cdiv(num_prefill_blocks, dcp_size) + 1 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py index 19f95f3a667b..f9bcf915c6c5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py @@ -18,7 +18,6 @@ MambaSpec, MLAAttentionSpec, SlidingWindowMLASpec, - SparseCacheRole, UniformTypeKVCacheSpecs, ) @@ -496,23 +495,13 @@ def compute_block_transfer_offsets( @cache def _token_run_table( - is_indexer: bool, block_size: int, entry_bytes: int, run_tokens: int + block_size: int, entry_bytes: int, run_tokens: int ) -> tuple[np.ndarray, np.ndarray] | None: - if is_indexer: - from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( - indexer_k_cache_byte_ranges, - ) - - runs = [ - indexer_k_cache_byte_ranges(block_size, entry_bytes, slot, run_tokens) - for slot in range(0, block_size, run_tokens) - ] - else: - # MLA latents are stored token by token. - runs = [ - [(slot * entry_bytes, run_tokens * entry_bytes)] - for slot in range(0, block_size, run_tokens) - ] + # MLA latents are stored token by token. + runs = [ + [(slot * entry_bytes, run_tokens * entry_bytes)] + for slot in range(0, block_size, run_tokens) + ] sizes = [size for _, size in runs[0]] if any([size for _, size in run] != sizes for run in runs): return None @@ -532,8 +521,8 @@ def get_token_run_table( Returns ``(offsets, sizes)``: ``offsets[i]`` are the ranges, relative to the block, of the run that starts at slot ``i * run_tokens``, and every run has ranges of the same ``sizes``. Only an MLA latent of ``latent_size`` - elements per token and the indexer K cache qualify; any other layer - returns None and can only be copied in whole blocks. + elements per token qualifies; any other layer returns None and can only be + copied in whole blocks. """ spec = layer_to_spec[layer_name] if ( @@ -549,15 +538,9 @@ def get_token_run_table( ) if geometry.block_len != geometry.block_size * geometry.slot_size_bytes: return None - is_indexer = spec.cache_role is SparseCacheRole.INDEXER - # The indexer cache holds 128 fp8 values and one fp32 scale per token. - if is_indexer and geometry.slot_size_bytes != 132: - return None - if not is_indexer and spec.head_size != latent_size: + if spec.head_size != latent_size: return None - return _token_run_table( - is_indexer, geometry.block_size, geometry.slot_size_bytes, run_tokens - ) + return _token_run_table(geometry.block_size, geometry.slot_size_bytes, run_tokens) def compute_token_run_offsets( diff --git a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py index 40ec25e2e55b..0db22e03c0bf 100644 --- a/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py +++ b/vllm/v1/attention/ops/rocm_aiter_mla_sparse.py @@ -215,55 +215,6 @@ def indexer_k_quant_and_cache_triton( ) -def indexer_k_cache_byte_ranges( - block_size: int, - entry_bytes: int, - first_slot: int, - num_slots: int, - block_tile_size: int = 16, - head_tile_size: int = 16, -) -> list[tuple[int, int]]: - """(offset in block, size) of the bytes that tokens ``[first_slot, - first_slot + num_slots)`` occupy in one block of the indexer K cache - written by ``indexer_k_quant_and_cache_triton``. - - A block holds the fp8 K values of all its tokens, then one fp32 scale per - token. Above block size 1 the values are shuffled in tiles of - ``block_tile_size`` tokens: a token's ``head_tile_size``-byte pieces lie - ``block_tile_size * head_tile_size`` bytes apart. - """ - head_dim = entry_bytes - 4 - end = first_slot + num_slots - ranges = [] - if block_size == 1: - ranges.append((first_slot * head_dim, num_slots * head_dim)) - else: - tile_bytes = block_tile_size * head_dim - for tile in range( - first_slot // block_tile_size, (end - 1) // block_tile_size + 1 - ): - tile_start = tile * block_tile_size - lo = max(first_slot, tile_start) - tile_start - hi = min(end, tile_start + block_tile_size) - tile_start - ranges += [ - ( - tile * tile_bytes - + piece * block_tile_size * head_tile_size - + lo * head_tile_size, - (hi - lo) * head_tile_size, - ) - for piece in range(head_dim // head_tile_size) - ] - ranges.append((block_size * head_dim + first_slot * 4, num_slots * 4)) - merged: list[tuple[int, int]] = [] - for offset, size in ranges: - if merged and sum(merged[-1]) == offset: - merged[-1] = (merged[-1][0], merged[-1][1] + size) - else: - merged.append((offset, size)) - return merged - - @triton.jit def _cp_gather_indexer_quant_cache_kernel( kv_cache_ptr, # [n_blks,blk_size//tile_blk,head_dim//16B,tile_blk,16B] From 76da4fb72442a5af8c0ec3d36b5e1287a3be9af2 Mon Sep 17 00:00:00 2001 From: kliuae Date: Fri, 30 Oct 2026 17:26:33 +0800 Subject: [PATCH 4/6] support hybrid mla Signed-off-by: kliuae Co-authored-by: Claude Opus --- docs/features/moriio_connector_usage.md | 2 +- .../unit/test_moriio_dcp_write.py | 10 +-- .../unit/test_moriio_hetero_tp_write.py | 80 ++++++++++++++++++- .../v1/moriio/moriio_connector.py | 1 - .../v1/moriio/moriio_tp_mapping.py | 6 +- 5 files changed, 86 insertions(+), 13 deletions(-) diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md index c591f83f4e51..25d73abaedd0 100644 --- a/docs/features/moriio_connector_usage.md +++ b/docs/features/moriio_connector_usage.md @@ -164,7 +164,7 @@ The number of KV heads must be a multiple or a divisor of each TP size. Blocks a ### Decode context parallelism -In WRITE mode a prefill instance without DCP can feed a decode instance with `--decode-context-parallel-size` (DCP) greater than 1, for MLA models without recurrent state. The decode instance deals the tokens of a request to the ranks of a DCP group in groups of `--cp-kv-cache-interleave-size` tokens, and each prefill rank writes a decode rank only the tokens that rank holds, so every token is written once per DCP group instead of once per decode rank. +In WRITE mode a prefill instance without DCP can feed a decode instance with `--decode-context-parallel-size` (DCP) greater than 1, for MLA models, including hybrid models whose other layers keep a recurrent state. The decode instance deals the tokens of a request to the ranks of a DCP group in groups of `--cp-kv-cache-interleave-size` tokens, and each prefill rank writes a decode rank only the tokens that rank holds, so every token is written once per DCP group instead of once per decode rank. A recurrent state is not split by token: each decode rank receives the state of its TP heads, as without DCP. - `--cp-kv-cache-interleave-size` is optional on the decode instance and must divide the KV cache block size. Each group of interleaved tokens is one copy, so the block size gives the fewest copies. Speculative decoding with DCP on ROCm's MLA attention backend (`ROCM_AITER_MLA`) needs interleave 1. - Instances with the same DCP size in WRITE or READ modes work as before with any interleave that is the same on both prefill and decode. Instances whose DCP sizes or interleaves differ in any other way refuse each other at the handshake. diff --git a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py index e29da895672e..e7ec0af97526 100644 --- a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -94,7 +94,7 @@ def _validate(local_dcp, local_interleave, peer_dcp, peer_interleave, **kwargs): peer_dcp, peer_interleave, block_size=BLOCK_SIZE, - **{"writes_into_peer": True, "is_mla": True, "has_mamba": False, **kwargs}, + **{"writes_into_peer": True, "is_mla": True, **kwargs}, ) @@ -104,14 +104,14 @@ def _validate(local_dcp, local_interleave, peer_dcp, peer_interleave, **kwargs): ((1, 1, 4, 64), {}), ((1, 1, 4, 1), {}), ((1, 1, 1, 64), {}), - ((8, 1, 8, 1), {"has_mamba": True}), + ((8, 1, 8, 1), {"is_mla": False}), ((8, 1, None, None), {"writes_into_peer": False}), ], ids=[ "write-into-dcp", "write-into-dcp-unit-interleave", "no-dcp", - "same-dcp-hybrid", + "same-dcp-gqa", "older-peer", ], ) @@ -128,8 +128,7 @@ def test_peer_dcp_layout_accepted(layout, kwargs): {"whole_block_layer": "layer0"}, "layer layer0 in whole blocks only.*needs --cp-kv-cache-interleave-size 64", ), - ((1, 1, 4, 64), {"is_mla": False}, "MLA models without recurrent state"), - ((1, 1, 4, 1), {"has_mamba": True}, "MLA models without recurrent state"), + ((1, 1, 4, 64), {"is_mla": False}, "supports MLA models only"), ((2, 1, 4, 1), {}, "between decode_context_parallel_size=2 here and 4"), ((1, 1, 4, 64), {"writes_into_peer": False}, "=1 here and 4 on the peer"), ((8, 64, 8, 1), {"writes_into_peer": False}, "64 here, 1 on the peer"), @@ -138,7 +137,6 @@ def test_peer_dcp_layout_accepted(layout, kwargs): "write-into-dcp-interleave-not-dividing-block", "write-into-dcp-whole-block-layer", "write-into-dcp-gqa", - "write-into-dcp-hybrid", "dcp2-into-dcp4", "read-no-dcp-from-dcp", "same-dcp-interleave", diff --git a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py index a1aeedf2fff5..4a0aa3d262c7 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py @@ -253,7 +253,7 @@ def __init__(self, src: torch.Tensor, dst: torch.Tensor): self.src, self.dst = src, dst -def _connect(producer: Rank, consumers: list[Rank], dp_engine: str) -> None: +def _connect(producer: Rank, consumers: list[Rank], dp_engine: str, **dcp) -> None: """Route the producer's sessions and notifications to the decode ranks.""" consumer = consumers[0].worker metadata = MoRIIOAgentMetadata( @@ -271,6 +271,7 @@ def _connect(producer: Rank, consumers: list[Rank], dp_engine: str) -> None: for name, layout in consumer.mamba_state_layouts.items() }, attention_block_sizes=consumer.attention_block_sizes, + **dcp, ) # The handshake carries the metadata as msgpack. metadata = msgspec.msgpack.decode( @@ -443,6 +444,83 @@ def test_write_lands_every_shard_byte_exact( assert torch.equal(received, expected), (layer.name, consumer.tp_rank) +@pytest.mark.parametrize( + ("producer_tp", "interleave"), [(1, 1), (4, BLOCK_SIZE)], ids=["fan-out", "fan-in"] +) +def test_hybrid_write_into_dcp_splits_mla_tokens_but_not_the_state( + monkeypatch, producer_tp, interleave +): + """Prefill without DCP -> decode TP2 x DCP2 of an MLA + KDA model: each + decode rank receives the MLA tokens of its DCP rank and the whole recurrent + state of its TP heads, and completes once each of its writers is done.""" + monkeypatch.setattr( + ssm_conv_transfer_utils, "is_conv_state_dim_first", lambda: True + ) + monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) + heads, layers = MODELS["mla_hybrid"] + state = _global_state(layers, heads) + groups = _groups("mla_hybrid", True) + consumers = [ + _build_rank("mla_hybrid", d, 2, KVCacheLayout.LBNHC, True, state, False) + for d in range(2) + ] + for consumer in consumers: + consumer.wrapper = _consumer_wrapper() + + done_after = {} + for p in range(producer_tp): + producer = _build_rank( + "mla_hybrid", p, producer_tp, KVCacheLayout.LBNHC, True, state, True + ) + producer.worker.dcp_size = 1 + producer.worker.model_config.get_head_size = lambda: 8 + _connect( + producer, + consumers, + "decode", + dcp_size=2, + cp_kv_cache_interleave_size=interleave, + ) + info = RemoteAllocInfo(block_ids=[DECODE_BLOCKS[g] for g in groups]) + producer.wrapper.done_remote_allocate_req_dict["t"] = info + writer = make_moriio_writer(producer.worker) + for layer in layers: + writer._execute_write_task( + WriteTask( + request_id="req", + transfer_id="t", + dst_engine_id="decode", + local_block_ids=[PREFILL_BLOCKS[g] for g in groups], + remote_block_ids_hint=None, + layer_name=layer.name, + event=SimpleNamespace(synchronize=lambda: None), + remote_notify_port=NOTIFY_PORT, + remote_ip="decode", + remote_tp_size=2, + tp_mapping=producer.worker._write_tp_mapping(2), + ) + ) + info.writes_expected = len(layers) + writer._finalize_if_complete("t", info) + for d, consumer in enumerate(consumers): + if consumer.wrapper.pop_finished_write_req_ids(): + done_after[d] = p + + # With fan-in, the second writer of a decode rank writes only state. + assert done_after == ({0: 0, 1: 0} if producer_tp == 1 else {0: 1, 1: 3}) + for d, consumer in enumerate(consumers): + kda = consumer.caches["kda0"].view(NUM_BLOCKS, -1) + assert torch.equal(kda, _expected(layers[1], "mamba", state, heads, d, 2, kda)) + mla = torch.zeros_like(consumer.caches["mla0"]) + for token in range(len(PREFILL_BLOCKS["full"]) * BLOCK_SIZE): + if token // interleave % 2 == d: + pos = token // (2 * interleave) * interleave + token % interleave + src = PREFILL_BLOCKS["full"][token // BLOCK_SIZE] + dst = DECODE_BLOCKS["full"][pos // BLOCK_SIZE] + mla[dst, :, pos % BLOCK_SIZE] = state["mla0"][src, :, token % BLOCK_SIZE] + assert torch.equal(consumer.caches["mla0"], mla), d + + def test_write_refuses_peer_with_different_block_size(monkeypatch): """Hybrid MLA models get TP-dependent attention block sizes; copying blocks between them would misplace tokens, so the write must refuse.""" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index ab4906a885d4..dc26cce524a2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -2262,7 +2262,6 @@ def _moriio_handshake( writes_into_peer=self.is_producer and self.mode == MoRIIOMode.WRITE, block_size=self.block_size, is_mla=self.use_mla, - has_mamba=self._has_mamba, whole_block_layer=self._whole_block_dcp_layer( metadata.cp_kv_cache_interleave_size ), diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py index c2042b351e37..5ef8d4a006dc 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py @@ -239,7 +239,6 @@ def validate_peer_dcp_layout( writes_into_peer: bool, block_size: int, is_mla: bool, - has_mamba: bool, whole_block_layer: str | None = None, ) -> None: """Refuse a peer whose DCP layout this worker cannot transfer KV with. @@ -272,12 +271,11 @@ def validate_peer_dcp_layout( "size on both instances, or a prefill instance without decode " "context parallelism in WRITE mode" ) - if not is_mla or has_mamba: + if not is_mla: raise ValueError( "MoRIIO WRITE from a prefill instance without decode context " "parallelism into a decode instance with " - f"decode_context_parallel_size={peer_dcp_size} supports MLA models " - "without recurrent state only" + f"decode_context_parallel_size={peer_dcp_size} supports MLA models only" ) if peer_interleave_size is None or block_size % peer_interleave_size: raise ValueError( From e50e61bc7272635a8a78dba24dc86699bc8a1091 Mon Sep 17 00:00:00 2001 From: kliuae Date: Fri, 30 Oct 2026 17:54:28 +0800 Subject: [PATCH 5/6] refuse non-mla layers Signed-off-by: kliuae Co-authored-by: Claude Opus --- .../unit/test_moriio_dcp_write.py | 82 ++++++++++++++++++- .../unit/test_moriio_hetero_tp_write.py | 37 +++++---- .../v1/moriio/moriio_connector.py | 15 +++- .../v1/moriio/moriio_tp_mapping.py | 10 ++- 4 files changed, 121 insertions(+), 23 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py index e7ec0af97526..32b035f0bfeb 100644 --- a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -35,7 +35,12 @@ validate_peer_dcp_layout, ) from vllm.utils.math_utils import cdiv -from vllm.v1.kv_cache_interface import MLAAttentionSpec, SparseCacheRole +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + MLAAttentionSpec, + SparseCacheRole, +) from .utils import make_moriio_writer @@ -94,7 +99,7 @@ def _validate(local_dcp, local_interleave, peer_dcp, peer_interleave, **kwargs): peer_dcp, peer_interleave, block_size=BLOCK_SIZE, - **{"writes_into_peer": True, "is_mla": True, **kwargs}, + **{"writes_into_peer": True, **kwargs}, ) @@ -104,7 +109,7 @@ def _validate(local_dcp, local_interleave, peer_dcp, peer_interleave, **kwargs): ((1, 1, 4, 64), {}), ((1, 1, 4, 1), {}), ((1, 1, 1, 64), {}), - ((8, 1, 8, 1), {"is_mla": False}), + ((8, 1, 8, 1), {"non_mla_layer": "attn0"}), ((8, 1, None, None), {"writes_into_peer": False}), ], ids=[ @@ -128,7 +133,7 @@ def test_peer_dcp_layout_accepted(layout, kwargs): {"whole_block_layer": "layer0"}, "layer layer0 in whole blocks only.*needs --cp-kv-cache-interleave-size 64", ), - ((1, 1, 4, 64), {"is_mla": False}, "supports MLA models only"), + ((1, 1, 4, 64), {"non_mla_layer": "attn0"}, "layer attn0 is not MLA"), ((2, 1, 4, 1), {}, "between decode_context_parallel_size=2 here and 4"), ((1, 1, 4, 64), {"writes_into_peer": False}, "=1 here and 4 on the peer"), ((8, 64, 8, 1), {"writes_into_peer": False}, "64 here, 1 on the peer"), @@ -194,6 +199,75 @@ def test_only_known_token_layouts_are_written_in_runs(specs, whole_block_layer): assert worker._whole_block_dcp_layer(BLOCK_SIZE) is None +KDA = MambaSpec( + block_size=BLOCK_SIZE, shapes=((4, 3), (2, 8, 8)), dtypes=(torch.uint8,) * 2 +) +GQA = FullAttentionSpec( + block_size=BLOCK_SIZE, num_kv_heads=8, head_size=64, dtype=torch.uint8 +) + + +@pytest.mark.parametrize( + ("mla_specs", "other_specs", "refused"), + [ + ( + {**SPECS, "indexer": _spec(132, cache_role=SparseCacheRole.INDEXER)}, + {}, + {1: "layer indexer in whole blocks only.*interleave-size 64"}, + ), + ( + {"latent": _spec(576, tokens_per_state=4)}, + {}, + {1: "layer latent in whole blocks only.*interleave-size 64"}, + ), + (SPECS, {"kda": KDA}, {}), + ( + {**SPECS, "draft": _spec(576, non_causal_multi_token_decode=True)}, + {"kda": KDA}, + {}, + ), + ( + SPECS, + {"kda": KDA, "draft": GQA}, + dict.fromkeys((1, BLOCK_SIZE), "layer draft is not MLA"), + ), + ({}, {"attn": GQA}, dict.fromkeys((1, BLOCK_SIZE), "layer attn is not MLA")), + ], + ids=[ + "glm5-deepseek-v32", + "compressed-states", + "kimi-linear", + "kimi-k3-dspark", + "non-mla-draft", + "gqa", + ], +) +def test_handshake_accepts_each_model_at_its_interleaves( + mla_specs, other_specs, refused +): + """A prefill instance without DCP writes a DCP decode instance's MLA + layers. DSpark and DFlash run a draft without MLA without DCP, so the + decode instance keeps its cache whole, which the prefill instance cannot + see: such a draft is refused.""" + worker = _worker(0, 1, mla_specs) + worker.use_mla = bool(mla_specs) + for name, spec in other_specs.items(): + worker.kv_caches[name] = None + worker.layer_to_spec[name] = spec + # A layer whose KV is not transferred is never written. + worker.layer_to_spec["not_transferred"] = GQA + for interleave in (1, BLOCK_SIZE): + kwargs = { + "non_mla_layer": worker._non_mla_dcp_layer(), + "whole_block_layer": worker._whole_block_dcp_layer(interleave), + } + if interleave in refused: + with pytest.raises(ValueError, match=refused[interleave]): + _validate(1, 1, 4, interleave, **kwargs) + else: + _validate(1, 1, 4, interleave, **kwargs) + + def _wrapper() -> MoRIIOWrapper: wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) wrapper.lock = threading.Lock() diff --git a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py index 4a0aa3d262c7..8cecf1a213fc 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py @@ -450,18 +450,23 @@ def test_write_lands_every_shard_byte_exact( def test_hybrid_write_into_dcp_splits_mla_tokens_but_not_the_state( monkeypatch, producer_tp, interleave ): - """Prefill without DCP -> decode TP2 x DCP2 of an MLA + KDA model: each - decode rank receives the MLA tokens of its DCP rank and the whole recurrent + """Prefill without DCP -> decode TP2 x DCP2 of an MLA + KDA model with an + MLA DSpark draft in its own group, as Kimi-K3: each decode rank receives + the MLA tokens of its DCP rank, target and draft, and the whole recurrent state of its TP heads, and completes once each of its writers is done.""" monkeypatch.setattr( ssm_conv_transfer_utils, "is_conv_state_dim_first", lambda: True ) monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) - heads, layers = MODELS["mla_hybrid"] + draft = Layer("draft0", "mla", "draft") + monkeypatch.setitem(MODELS, "k3", (1, [*MODELS["mla_hybrid"][1], draft])) + monkeypatch.setitem(PREFILL_BLOCKS, "draft", [0, 5]) + monkeypatch.setitem(DECODE_BLOCKS, "draft", [4, 1]) + heads, layers = MODELS["k3"] state = _global_state(layers, heads) - groups = _groups("mla_hybrid", True) + groups = _groups("k3", True) consumers = [ - _build_rank("mla_hybrid", d, 2, KVCacheLayout.LBNHC, True, state, False) + _build_rank("k3", d, 2, KVCacheLayout.LBNHC, True, state, False) for d in range(2) ] for consumer in consumers: @@ -470,10 +475,11 @@ def test_hybrid_write_into_dcp_splits_mla_tokens_but_not_the_state( done_after = {} for p in range(producer_tp): producer = _build_rank( - "mla_hybrid", p, producer_tp, KVCacheLayout.LBNHC, True, state, True + "k3", p, producer_tp, KVCacheLayout.LBNHC, True, state, True ) producer.worker.dcp_size = 1 producer.worker.model_config.get_head_size = lambda: 8 + assert producer.worker._non_mla_dcp_layer() is None _connect( producer, consumers, @@ -511,14 +517,17 @@ def test_hybrid_write_into_dcp_splits_mla_tokens_but_not_the_state( for d, consumer in enumerate(consumers): kda = consumer.caches["kda0"].view(NUM_BLOCKS, -1) assert torch.equal(kda, _expected(layers[1], "mamba", state, heads, d, 2, kda)) - mla = torch.zeros_like(consumer.caches["mla0"]) - for token in range(len(PREFILL_BLOCKS["full"]) * BLOCK_SIZE): - if token // interleave % 2 == d: - pos = token // (2 * interleave) * interleave + token % interleave - src = PREFILL_BLOCKS["full"][token // BLOCK_SIZE] - dst = DECODE_BLOCKS["full"][pos // BLOCK_SIZE] - mla[dst, :, pos % BLOCK_SIZE] = state["mla0"][src, :, token % BLOCK_SIZE] - assert torch.equal(consumer.caches["mla0"], mla), d + for name, group in (("mla0", "full"), ("draft0", "draft")): + mla = torch.zeros_like(consumer.caches[name]) + for token in range(len(PREFILL_BLOCKS[group]) * BLOCK_SIZE): + if token // interleave % 2 == d: + pos = token // (2 * interleave) * interleave + token % interleave + src = PREFILL_BLOCKS[group][token // BLOCK_SIZE] + dst = DECODE_BLOCKS[group][pos // BLOCK_SIZE] + mla[dst, :, pos % BLOCK_SIZE] = state[name][ + src, :, token % BLOCK_SIZE + ] + assert torch.equal(consumer.caches[name], mla), (name, d) def test_write_refuses_peer_with_different_block_size(monkeypatch): diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index dc26cce524a2..fee308b5042e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -2261,7 +2261,7 @@ def _moriio_handshake( metadata.cp_kv_cache_interleave_size, writes_into_peer=self.is_producer and self.mode == MoRIIOMode.WRITE, block_size=self.block_size, - is_mla=self.use_mla, + non_mla_layer=self._non_mla_dcp_layer(), whole_block_layer=self._whole_block_dcp_layer( metadata.cp_kv_cache_interleave_size ), @@ -3727,6 +3727,19 @@ def _whole_block_dcp_layer(self, interleave_size: int | None) -> str | None: None, ) + def _non_mla_dcp_layer(self) -> str | None: + """A DCP-sharded layer that is not MLA, e.g. of a draft model that the + decode instance may run without DCP and so keep unsharded.""" + return next( + ( + layer_name + for layer_name in self.kv_caches + if self.layer_to_spec[layer_name].dcp_sharded + and not self._is_mla_cache_layer(layer_name) + ), + None, + ) + @property def _mla_latent_size(self) -> int | None: return self.model_config.get_head_size() if self.use_mla else None diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py index 5ef8d4a006dc..aea2efebc092 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_tp_mapping.py @@ -238,13 +238,14 @@ def validate_peer_dcp_layout( *, writes_into_peer: bool, block_size: int, - is_mla: bool, + non_mla_layer: str | None = None, whole_block_layer: str | None = None, ) -> None: """Refuse a peer whose DCP layout this worker cannot transfer KV with. Equal DCP sizes transfer rank to rank. A WRITE producer without DCP can - also write into a DCP consumer, one write per group of interleaved tokens; + also write into a DCP consumer, one write per group of interleaved tokens, + for MLA layers only (``non_mla_layer`` names another DCP-sharded layer); ``whole_block_layer`` names a layer whose blocks can only be written whole, which needs the consumer's interleave to be the block size. A peer that does not advertise its layout predates these checks and is accepted as @@ -271,11 +272,12 @@ def validate_peer_dcp_layout( "size on both instances, or a prefill instance without decode " "context parallelism in WRITE mode" ) - if not is_mla: + if non_mla_layer is not None: raise ValueError( "MoRIIO WRITE from a prefill instance without decode context " "parallelism into a decode instance with " - f"decode_context_parallel_size={peer_dcp_size} supports MLA models only" + f"decode_context_parallel_size={peer_dcp_size} supports MLA layers " + f"only; layer {non_mla_layer} is not MLA" ) if peer_interleave_size is None or block_size % peer_interleave_size: raise ValueError( From c365ec6c6b9a767ec4b6cb8b53c964f57bbb4521 Mon Sep 17 00:00:00 2001 From: kliuae Date: Mon, 9 Nov 2026 12:10:05 +0000 Subject: [PATCH 6/6] wait for every writer for transfer reporting and drop handshake for terminal transfer Signed-off-by: kliuae --- .../unit/test_moriio_dcp_write.py | 49 ++++++++++++-- .../unit/test_moriio_tp_mapping.py | 20 ++++-- .../unit/test_moriio_write_done_routing.py | 9 +-- .../v1/moriio/moriio_connector.py | 4 +- .../kv_connector/v1/moriio/moriio_engine.py | 67 ++++++++++++------- 5 files changed, 110 insertions(+), 39 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py index 32b035f0bfeb..88782fd79eb1 100644 --- a/tests/v1/kv_connector/unit/test_moriio_dcp_write.py +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -545,9 +545,11 @@ def refuse(*args): worker = _handshake_worker(True, refuse) wrapper = _wrapper() sent = [] - wrapper.send_notify = lambda transfer_id, ip, port, message_type: sent.append( - (transfer_id, port, message_type) - ) + + def send_notify(transfer_id, ip, port, message_type, message_fields=None): + sent.append((transfer_id, port, message_type, message_fields)) + + wrapper.send_notify = send_notify worker.moriio_wrapper = wrapper writer = worker._writer = make_moriio_writer(worker) tasks = [] @@ -564,7 +566,8 @@ def write_blocks(req_id, meta, layer_name, kv_layer): for layer in ("mla0", "mla1", "mla2"): worker.save_kv_layer(metadata, layer, kv_layer=None, attn_metadata=None) - worker._handshake_initiation_executor.shutdown() + # One worker, so this returns after the handshake callbacks. + worker._handshake_initiation_executor.submit(lambda: None).result() # One handshake per decode rank of the MLA fan-out, for the first layer. assert sorted(handshakes) == [f"decode:0_dp0_tp{t}" for t in range(4)] @@ -589,17 +592,53 @@ def do_layer_write(plan, sessions): writer._execute_write_task(tasks[0]) if targets_connected: + worker._handshake_initiation_executor.shutdown() assert written == [[f"decode:0_dp0_tp{t}"] for t in range(4)] assert sent == [] and wrapper.done_req_ids == [] return assert [ack.transfer_id for ack in wrapper.done_req_ids] == ["t"] assert sent == [ - ("t", NOTIFY_PORT + get_port_offset(0, t), "write_failed") for t in range(4) + ("t", NOTIFY_PORT + get_port_offset(0, t), "write_failed", None) + for t in range(4) ] # The request's other write tasks are dropped. assert writer._is_transfer_terminal("t") assert not writer.is_unconnected("t") + worker.save_kv_layer(metadata, "mla3", kv_layer=None, attn_metadata=None) + worker._handshake_initiation_executor.shutdown() + assert len(handshakes) == 4 + writer.mark_unconnected("t") + assert not writer.is_unconnected("t") + + +def test_unconnected_fan_in_writer_tells_the_owning_rank_how_many_write(): + worker = _handshake_worker(True, _refuse) + worker.tp_rank, worker.world_size, worker._has_mamba = 1, 2, True + worker._is_mamba_layer = lambda layer_name: True + wrapper = _wrapper() + sent = [] + + def send_notify(transfer_id, ip, port, message_type, message_fields=None): + sent.append((port, message_type, message_fields)) + + wrapper.send_notify = send_notify + worker.moriio_wrapper = wrapper + writer = make_moriio_writer(worker) + task = _write_task("kda0", [5, 6], worker._write_tp_mapping(1)) + task.remote_dp_size = task.remote_dp_size_local = 2 + wrapper.done_remote_allocate_req_dict["t"] = RemoteAllocInfo( + block_ids=[[7, 8]], decode_dp_rank=1 + ) + writer.mark_unconnected("t") + + writer._execute_write_task(task) + + assert sent == [ + (NOTIFY_PORT + get_port_offset(1, 0), "write_failed", {"num_writers": 2}) + ] + assert [ack.transfer_id for ack in wrapper.done_req_ids] == ["t"] + VOTES = {"connected": 2, "refused": 1, "failed": 0} diff --git a/tests/v1/kv_connector/unit/test_moriio_tp_mapping.py b/tests/v1/kv_connector/unit/test_moriio_tp_mapping.py index 5cddac938eac..4e02b8b51cca 100644 --- a/tests/v1/kv_connector/unit/test_moriio_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_moriio_tp_mapping.py @@ -367,12 +367,24 @@ def test_consumer_completes_after_every_writer_reported(monkeypatch): assert wrapper.pop_finished_write_req_ids() == {"u"} -def test_consumer_write_failure_drops_partial_write_count(monkeypatch): +@pytest.mark.parametrize("first", ["write_done", "write_failed"]) +def test_consumer_write_failure_waits_for_every_writer(monkeypatch, first): monkeypatch.setattr(moriio_engine, "get_role", lambda: ROLE.CONSUMER) wrapper = _consumer_wrapper() + second = "write_failed" if first == "write_done" else "write_done" + wrapper._handle_structured_message( - {"type": "write_done", "transfer_id": "t", "num_writers": 2} + {"type": first, "transfer_id": "t", "num_writers": 2} + ) + assert wrapper.pop_failed_write_req_ids() == set() + assert wrapper.pop_finished_write_req_ids() == set() + + wrapper._handle_structured_message( + {"type": second, "transfer_id": "t", "num_writers": 2} ) - wrapper._handle_structured_message({"type": "write_failed", "transfer_id": "t"}) - assert wrapper._write_done_counts == {} assert wrapper.pop_failed_write_req_ids() == {"t"} + assert wrapper.pop_finished_write_req_ids() == set() + assert wrapper._write_done_counts == {} + + wrapper._handle_structured_message({"type": "write_failed", "transfer_id": "u"}) + assert wrapper.pop_failed_write_req_ids() == {"u"} diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index 58f0c025a0f1..6f3181485356 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -462,6 +462,7 @@ def _fanout_save_worker(remote_agents: set[str]) -> MoRIIOConnectorWorker: worker._write_blocks_for_req = ( # type: ignore[method-assign] lambda req_id, meta, layer_name, kv_layer: worker.writes.append(req_id) ) + worker.moriio_wrapper = _real_wrapper() worker._writer = make_moriio_writer(worker) return worker @@ -543,8 +544,8 @@ def test_release_reaches_every_prefill_rank_and_no_other(params, expected_tp_ran def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): wrapper = _real_wrapper() sent: list[tuple[str, int, str]] = [] - wrapper.send_notify = lambda tid, ip, port, message_type: sent.append( - (ip, port, message_type) + wrapper.send_notify = lambda tid, ip, port, message_type, message_fields=None: ( + sent.append((ip, port, message_type)) ) writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper, tp_rank=0)) task = _write_task( @@ -579,8 +580,8 @@ def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): def test_deferred_write_timeout_notifies_every_fanout_target(): wrapper = _real_wrapper() sent: list[tuple[str, int, str]] = [] - wrapper.send_notify = lambda tid, ip, port, message_type: sent.append( - (ip, port, message_type) + wrapper.send_notify = lambda tid, ip, port, message_type, message_fields=None: ( + sent.append((ip, port, message_type)) ) writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper, tp_rank=0)) task = _write_task( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index fee308b5042e..7f24ac648ad9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -3131,9 +3131,7 @@ def save_kv_layer( ) # A request whose handshake failed is not retried; its writes # reach the decode ranks that are connected or fail the request. - if not handshake_ready and not self._writer.is_unconnected( - meta.transfer_id - ): + if not handshake_ready and self._writer.may_handshake(meta.transfer_id): # Initiate handshake with remote engine to exchange metadata. with self._handshake_lock: if remote_engine_id not in self._remote_agents: diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index ba8a6a875be7..0105ff2d3284 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -173,12 +173,20 @@ def mark_unconnected(self, transfer_id: TransferId) -> None: """Record that a transfer's handshake failed. Its writes still go to decode ranks that are connected; otherwise the transfer fails.""" with self._write_state_lock: - self._unconnected.add(transfer_id) + # A handshake callback may run after the transfer has ended. + if not self._is_transfer_terminal(transfer_id): + self._unconnected.add(transfer_id) def is_unconnected(self, transfer_id: TransferId) -> bool: with self._write_state_lock: return transfer_id in self._unconnected + def may_handshake(self, transfer_id: TransferId) -> bool: + with self._write_state_lock: + return transfer_id not in self._unconnected and not ( + self._is_transfer_terminal(transfer_id) + ) + def seal_pending_transfers(self) -> None: """Seal expected WRITE counts after the model forward has run. @@ -281,25 +289,32 @@ def _notify_write_failed(self, task: WriteTask, age: float) -> None: ) self._send_write_failed(task) - def _send_write_failed(self, task: WriteTask) -> None: - # Only the remote allocation names the decode DP rank that owns the - # request, and it may not have arrived. Tell every rank; the others - # drop it. - for decode_dp_rank in range(max(task.remote_dp_size, 1)): - for remote_ip, remote_port in self._resolve_notify_endpoints( - task, decode_dp_rank - ): + def _send_write_failed( + self, task: WriteTask, decode_dp_rank: int | None = None + ) -> None: + # Decode ranks that do not own the request drop the message. + decode_dp_ranks = ( + range(max(task.remote_dp_size, 1)) + if decode_dp_rank is None + else (decode_dp_rank,) + ) + for dp_rank in decode_dp_ranks: + endpoints = self._resolve_notify_endpoints(task, dp_rank) + for (remote_ip, remote_port), num_writers in endpoints.items(): try: self.worker.moriio_wrapper.send_notify( task.transfer_id, remote_ip, remote_port, message_type="write_failed", + message_fields=( + {"num_writers": num_writers} if num_writers != 1 else None + ), ) except Exception: logger.exception( "Failed to notify decode DP rank %d that transfer %s failed", - decode_dp_rank, + dp_rank, task.transfer_id, ) @@ -415,7 +430,7 @@ def _execute_write_task(self, task: WriteTask) -> None: task.request_id, ) self._mark_request_done(task.transfer_id) - self._send_write_failed(task) + self._send_write_failed(task, request_info.decode_dp_rank) return # Wait for CUDA event @@ -827,8 +842,8 @@ def __init__( self.done_remote_allocate_req_dict: dict[TransferId, RemoteAllocInfo] = {} self.done_write_cache_req_ids: list[str] = [] self.failed_write_cache_req_ids: list[str] = [] - # write_done messages received per transfer whose writers are not all in. - self._write_done_counts: dict[TransferId, int] = {} + # (writers reported, any failed) per transfer still waiting for writers. + self._write_done_counts: dict[TransferId, tuple[int, bool]] = {} self._terminal_transfer_ids: OrderedDict[TransferId, None] = OrderedDict() self._transfer_timeout = transfer_timeout self.notify_thread: threading.Thread | None = None @@ -1158,22 +1173,28 @@ def _handle_write_done_message(self, data: dict): assert get_role() != ROLE.PRODUCER, ( "Only decode can get WRITE completion messages" ) - transfer_id = data["transfer_id"] - num_writers = int(data.get("num_writers", 1)) - with self.lock: - received = self._write_done_counts.get(transfer_id, 0) + 1 - if received < num_writers: - self._write_done_counts[transfer_id] = received - return - self._write_done_counts.pop(transfer_id, None) - self.done_write_cache_req_ids.append(transfer_id) + self._record_writer_report(data, failed=False) def _handle_write_failed_message(self, data: dict): assert get_role() != ROLE.PRODUCER, "Only decode can get WRITE failure messages" + self._record_writer_report(data, failed=True) + + def _record_writer_report(self, data: dict, failed: bool) -> None: transfer_id = data["transfer_id"] + num_writers = int(data.get("num_writers", 1)) with self.lock: + received, any_failed = self._write_done_counts.get(transfer_id, (0, False)) + received += 1 + any_failed = any_failed or failed + # Another writer may still be writing into the blocks after a failure. + if received < num_writers: + self._write_done_counts[transfer_id] = (received, any_failed) + return self._write_done_counts.pop(transfer_id, None) - self.failed_write_cache_req_ids.append(transfer_id) + if any_failed: + self.failed_write_cache_req_ids.append(transfer_id) + else: + self.done_write_cache_req_ids.append(transfer_id) def _handle_release_message(self, data: dict): assert get_role() == ROLE.PRODUCER, (