[ROCm][MoRIIO] Support DCP for decode instance in PD disaggregated serving with MoRI WRITE mode

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.

tests/test_config.py

@@ -436,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"),
@@ -446,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))]

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)