diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md index d806e04cde8c..0117737dd1e4 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, 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. +- 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 64808b61078e..d33b19aa77ce 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -516,7 +516,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"), @@ -526,10 +534,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 71259674cb77..53ba86ea0021 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -1276,3 +1276,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 new file mode 100644 index 000000000000..74aa691d6a56 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_moriio_dcp_write.py @@ -0,0 +1,700 @@ +# 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 queue +import threading +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +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, + HandshakeError, + HandshakeRefusedError, + MoRIIOAgentMetadata, + MoRIIOConnectorMetadata, + MoRIIOMode, + RemoteAllocInfo, + ReqMeta, + WriteTask, + get_port_offset, +) +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 ( + validate_peer_dcp_layout, +) +from vllm.utils.math_utils import cdiv +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + MambaSpec, + MLAAttentionSpec, + SparseCacheRole, +) + +from .utils import make_moriio_writer + +pytestmark = pytest.mark.cpu_test + +BLOCK_SIZE = 64 +NUM_BLOCKS = 24 +NOTIFY_PORT = 20000 + + +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, + ) + + +# The fp8 MLA latent of DeepSeek-V3.2 and GLM-5 (kv_lora_rank 512 + +# qk_rope_head_dim 64). +SPECS = { + "latent": _spec(576), +} + + +def _cache(spec: MLAAttentionSpec) -> torch.Tensor: + return torch.zeros( + NUM_BLOCKS, + spec.num_heads, + spec.num_states, + spec.state_content_size_bytes, + dtype=torch.uint8, + ) + + +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 + return slot * entry + torch.arange(entry) + + +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, **kwargs}, + ) + + +@pytest.mark.parametrize( + ("layout", "kwargs"), + [ + ((1, 1, 4, 64), {}), + ((1, 1, 4, 1), {}), + ((1, 1, 1, 64), {}), + ((8, 1, 8, 1), {"non_mla_layer": "attn0"}), + ((8, 1, None, None), {"writes_into_peer": False}), + ], + ids=[ + "write-into-dcp", + "write-into-dcp-unit-interleave", + "no-dcp", + "same-dcp-gqa", + "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), {"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"), + ], + ids=[ + "write-into-dcp-interleave-not-dividing-block", + "write-into-dcp-whole-block-layer", + "write-into-dcp-gqa", + "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, get_head_size=lambda: 576 + ) + worker.kv_caches = 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} + worker._record_slice_layouts(caches) + worker.shutdown = lambda: None + return worker + + +@pytest.mark.parametrize( + ("specs", "whole_block_layer"), + [ + (SPECS, None), + ( + {**SPECS, "indexer": _spec(132, cache_role=SparseCacheRole.INDEXER)}, + "indexer", + ), + ({"latent": _spec(576, tokens_per_state=4)}, "latent"), + ], + 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; 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 + 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() + 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 = {} + wrapper._terminal_transfer_ids = OrderedDict() + 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; + 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(spec).shape[1:]), generator=gen + ).to(torch.uint8) + 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(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 = [] + for p in range(producer_tp): + 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): + if engine_id == "decode_dp0": + (target,) = mapping.targets + tp_rank = target.tp_rank + else: + tp_rank = int(engine_id.rsplit("_tp", 1)[1]) + sessions = [ + 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 + + 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): + message = {"type": message_type, "transfer_id": transfer_id} + consumer_wrappers[port - NOTIFY_PORT]._handle_structured_message( + {**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={"t": info}, + 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, + ) + writer = make_moriio_writer(worker) + 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) + done = [wrapper.pop_finished_write_req_ids() for wrapper in consumer_wrappers] + return consumer_caches, done, content, bytes_written + + +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"), [(1, 4, 4), (8, 4, 2)] +) +def test_write_lands_each_token_on_its_dcp_rank( + monkeypatch, producer_tp, consumer_tp, dcp_size, num_prefill_blocks, interleave +): + """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 + 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, + interleave, + ) + + # 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): + 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 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 = 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(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 = [] + + 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 = [] + + 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) + # 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)] + 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: + 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", 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} + + +@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_hetero_tp_write.py b/tests/v1/kv_connector/unit/test_moriio_hetero_tp_write.py index ff40382fb870..dd6bef372394 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,92 @@ 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 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) + 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("k3", True) + consumers = [ + _build_rank("k3", 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( + "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, + "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)) + 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): """Hybrid MLA models get TP-dependent attention block sizes; copying blocks between them would misplace tokens, so the write must refuse.""" 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 d506c74f6969..698cb26bd5d4 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -1013,6 +1013,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]]) @@ -1264,6 +1292,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 84aa615a41e9..0abb4f8e87d5 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_tp_mapping.py b/tests/v1/kv_connector/unit/test_moriio_tp_mapping.py index 31f5f1666fe8..efdd50dee5b6 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 be19e589e213..8bcf4f4af1fe 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=[]) def _do_layer_write(plan, sessions): write_sessions.append(sessions) @@ -462,6 +462,8 @@ 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 @@ -542,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( @@ -578,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/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index 2aa21b0e5098..10366838d547 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -713,6 +713,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/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index 2668c1370b6b..938802bbfe87 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] = {} + # 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: @@ -216,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 fd5bcf04a2c3..9e90b75de127 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,18 +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_peer_dcp_layout, ) from vllm.distributed.kv_transfer.kv_connector.v1.ssm_conv_transfer_utils import ( MambaConvSplitInfo, @@ -1614,6 +1620,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 @@ -2009,6 +2033,10 @@ 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 + ) # Hybrid (mamba/KDA) recurrent-state transfer. Populated in # register_kv_caches when the model has a MambaSpec kv_cache group. self._has_mamba = False @@ -2281,6 +2309,23 @@ def _moriio_handshake( metadata_bytes = received_frame[1] decoder = msgspec.msgpack.Decoder(MoRIIOAgentMetadata) metadata = decoder.decode(metadata_bytes) + 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, + non_mla_layer=self._non_mla_dcp_layer(), + 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", @@ -2355,8 +2400,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 @@ -2816,6 +2865,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( @@ -3134,7 +3185,9 @@ 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 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: @@ -3181,7 +3234,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. @@ -3209,12 +3264,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) @@ -3284,13 +3343,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(): @@ -3308,9 +3369,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 " @@ -3318,6 +3386,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. @@ -3343,7 +3412,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 @@ -3353,6 +3422,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 @@ -3647,6 +3729,75 @@ 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, + ) + + 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 + 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 d43d6b2c5199..94f61d20bfc7 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_runs, +) if TYPE_CHECKING: from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( @@ -110,6 +113,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 @@ -165,6 +170,24 @@ 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: + # 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. @@ -265,23 +288,34 @@ 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. - 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 - ): + self._send_write_failed(task) + + 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 timed out", - decode_dp_rank, + "Failed to notify decode DP rank %d that transfer %s failed", + dp_rank, task.transfer_id, ) @@ -290,6 +324,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 @@ -374,6 +409,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, request_info.decode_dp_rank) + return + # Wait for CUDA event # The attention computation of the current layer cannot # overlap with the kv transfer task, @@ -381,15 +441,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] = [] - 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 - ) + bytes_written = 0 + 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 @@ -410,8 +464,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) @@ -453,7 +509,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 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 + runs = dcp_write_runs( + local_blocks, + remote_blocks, + block_size, + dcp_size, + run_tokens, + tp_rank % dcp_size, + ) + 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 ( @@ -618,6 +707,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) @@ -649,9 +739,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, ) @@ -751,8 +843,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 @@ -1099,22 +1191,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, ( 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 b26e82ef8433..c2f92693574c 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 ( @@ -491,6 +493,107 @@ def compute_block_transfer_offsets( return merge_fn(offset_local, offset_remote, sizes) +@cache +def _token_run_table( + block_size: int, entry_bytes: int, run_tokens: int +) -> tuple[np.ndarray, np.ndarray] | None: + # 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 qualifies; 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 + if spec.head_size != latent_size: + return None + return _token_run_table(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 6faae891e421..15de7f6586ca 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,10 +19,26 @@ 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 +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) @@ -212,3 +228,111 @@ def compute_write_tp_mapping( ) ), ) + + +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, + 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, + 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 + 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 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 layers " + f"only; layer {non_mla_layer} is not MLA" + ) + if peer_interleave_size is None or block_size % peer_interleave_size: + raise ValueError( + "MoRIIO WRITE into a decode instance with decode context parallelism " + "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}" + ) + + +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, +) -> 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` runs) + gets empty runs and must still be sent write_done. + """ + 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 {needed} " + f"blocks for {len(local_block_ids)} producer blocks, but it allocated " + f"{len(remote_block_ids)}" + ) + 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, + )