[kliuae]

Overview

In WRITE mode, MoRIIOConnector copies prefill block i straight into decode block i, so when the two instances store tokens differerntly, like a decode instance with DCP and a prefill instance without DCP, ends up with its tokens on the wrong ranks. This PR makes each prefill rank write each decode rank exactly the tokens that rank holds for MLA models and at any interleave that divides the block size.

Claims

• A prefill instance without DCP can now feed a DCP decode instance for MLA models, including hybrid MLA + linear-attention models, and every token is written once per DCP group instead of once per decode rank.
• On DeepSeek-R1 (8x MI355X), prefill TP4 and decode TP4+DCP4 gives the decode side 3.9x the KV cache of a TP4 decode instance. Under a decode-heavy load it completes all 640 requests with 19% more output per GPU, where TP4 prefill and TP4 decode loses 214 of them to decode-block timeouts. Interleave 1 and 64 perform the same.
• On Kimi-Linear-48B-A3B-Instruct (MLA + KDA), prefill TP2 with decode TP2+DCP2 gives the decode instance 1.54x the KV cache of a TP2 one, completes every request with exact-length outputs matching the prefill server, and GSM8K scores are unchanged.
• Instances with the same DCP size keep main's rank-to-rank writes, in READ and WRITE mode and for hybrid models, and main's hybrid DCP unit tests pass unchanged.
• Layouts the connector can't transfer are refused at the handshake. The affected requests fail with an error and both instances keep serving, where before the request would hang or the decode engine would stop.

Validation

Unit tests

The new tests are CPU tests in a new tests/v1/kv_connector/unit/test_moriio_dcp_write.py (41 cases) and in test_moriio_hetero_tp_write.py and test_moriio_hma_scheduler.py, plus GPU test in test_moriio_connector.py. In tests/test_config.py, MoRIIO becomes a second case of main's interleave test, which is renamed to test_connector_does_not_adjust_dcp_interleave_size because it now covers two connectors. Reverting each part of the change makes its tests fail.

Behaviors	Tests	Main	With this PR
Main's pairing of prefill block j with decode block j	test_write_lands_each_token_on_its_dcp_rank, test_write_refuses_decode_allocation_short_of_its_blocks	13 failed (every case, from misplaced bytes or main's local_block_ids longer than remote_block_ids error)	13 passed
Whole-block writes only, with no token groups below the block size	test_write_lands_each_token_on_its_dcp_rank	8 failed (every case at interleave 1 and 16), 4 passed (interleave 64)	12 passed
A validator that refuses every interleave other than the block size	test_peer_dcp_layout_accepted	1 failed (decode interleave 1), 4 passed	5 passed
Main's handshake failure handling, with no write_failed and a new handshake per layer in WRITE mode and a step that raises in READ mode	test_refused_write_handshake_fails_the_request_on_every_decode_rank, test_read_refusal_fails_the_requests_on_every_tp_rank	4 failed (both WRITE cases and the 2 READ cases with a refused peer), 3 passed (READ with all peers connected, another failure, hybrid)	7 passed
Skipping a READ release that has no prefill address (build_connector_meta raises)	test_read_release_without_prefill_address_is_skipped	1 failed (no address), 1 passed (address in kv_transfer_params)	2 passed
A refusal of models with recurrent state in this layout	test_handshake_accepts_each_model_at_its_interleaves, test_hybrid_write_into_dcp_splits_mla_tokens_but_not_the_state and the validator cases	5 failed (Kimi-Linear, Kimi-K3 with its DSpark draft, the draft without MLA, both hybrid writer cases), 14 passed	19 passed
A model-level MLA check instead of a check of every DCP-sharded layer	the same tests	1 failed (a draft model without MLA is accepted), 18 passed	19 passed

The other new cases cover the API this PR adds, namely which layouts the handshake accepts and refuses (test_peer_dcp_layout_accepted, test_peer_dcp_layout_refused), which caches may be written in groups (test_only_known_token_layouts_are_written_in_runs) how a hybrid model's layers are written (test_hybrid_write_into_dcp_splits_mla_tokens_but_not_the_state), and which models the handshake accepts at each interleave, from GLM-5.x and DeepSeek-V3.2 to Kimi-Linear, Kimi-K3 with its DSpark draft and the refused non-MLA cases (test_handshake_accepts_each_model_at_its_interleaves). test_write_between_equal_dcp_sizes_is_rank_to_rank (a producer with DCP 4 writes its own blocks to the decode rank with the same index), test_connector_does_not_adjust_dcp_interleave_size[moriio] and the GPU test test_worker_accepts_decode_context_parallelism (both roles, READ and WRITE, DCP 2, interleave 1) pin main's behaviour where it should not change. At an interleave equal to the block size, the writer test also covers the whole-block path for both TP mappings.

End-to-end verification with MLA (DeepSeek-R1)

DeepSeek-R1 runs on one node of 8 MI355X GPUs, with the prefill instance TP4 on GPUs 0-3 and the decode instance on GPUs 4-7. The runs use MoRIIO WRITE over xGMI, the toy proxy in vLLM example, MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 on the prefill, a bf16 KV cache, block size 64, --gpu-memory-utilization 0.90 and --max-num-seqs 256. The runs include #64330, which fixes the completion-wait hang over xGMI.

Decode instance	KV cache (tokens)	Delta (%)	KV bytes written for a new 1,000-token prompt	Delta (%)	GSM8K 20-shot, flexible-extract / strict-match	Requests completed	Output tok/s per GPU	Delta (%)	TTFT p50 / p99 (s)	TPOT p50 (ms)	Preemptions
TP4, main + #64330	1,447,424		287,834,112		0.9606 / 0.8976	426 / 640	247.9		54.42 / 78.27	63.3	128
TP4+DCP4, PR, interleave 1	5,664,512	+291.4	71,958,528	-75.0	0.9636 / 0.8848	640 / 640	295.1	+19.0	1.25 / 12.11	93.1	0
TP4+DCP4, PR, interleave 64	5,664,512	+291.4	71,958,528	-75.0	0.9644 / 0.8954	640 / 640	295.6	+19.2	1.24 / 11.67	92.9	0

The benchmark is vllm bench serve with random 8K-token prompts and 4K output tokens, 640 prompts at up to 256 concurrent requests and 2 requests/s. Throughput is per GPU over all eight GPUs and counts successful requests. At 256 sequences, about 3.1 M tokens are in flight, roughly twice what a TP4 decode instance holds, so TP4 -> TP4 spends the run waiting for decode blocks. Up to 204 requests wait, and 214 of them exceed the prefill's 60 s defer_timeout and fail under the default kv_load_failure_policy. With DCP4 the decode cache never gets past half full.

DeepSeek-R1 commands

# proxy
python3 examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py --port 11500 --service-discovery-port 11503

# flags shared by every instance
FLAGS="--block-size 64 --kv-cache-dtype auto --attention-backend ROCM_AITER_MLA --gpu-memory-utilization 0.90 \
  --enable-chunked-prefill --max-num-batched-tokens 32768 --enable-prefix-caching --max-num-seqs 256 \
  --max-model-len 65536 --compilation-config {\"cudagraph_mode\":\"FULL_DECODE_ONLY\"} --reasoning-parser deepseek_r1"
export VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_USE_BREAKABLE_CUDAGRAPH=0

# prefill (both layouts)
HIP_VISIBLE_DEVICES=0,1,2,3 MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 vllm serve deepseek-ai/DeepSeek-R1 --port 11501 \
  --tensor-parallel-size 4 $FLAGS \
  --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_connector_extra_config":{"backend":"xgmi","proxy_ip":"127.0.0.1","proxy_ping_port":"11503","http_port":"11501","handshake_port":"11510","notify_port":"11518"}}'

# decode (drop the DCP and interleave flags for the TP4 rows, and only --cp-kv-cache-interleave-size for the interleave-1 row)
HIP_VISIBLE_DEVICES=4,5,6,7 vllm serve deepseek-ai/DeepSeek-R1 --port 11502 \
  --tensor-parallel-size 4 --decode-context-parallel-size 4 --cp-kv-cache-interleave-size 64 $FLAGS \
  --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"backend":"xgmi","proxy_ip":"127.0.0.1","proxy_ping_port":"11503","http_port":"11502","handshake_port":"11526","notify_port":"11534"}}'

# reference, one TP8 instance
HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 vllm serve deepseek-ai/DeepSeek-R1 --port 11501 --tensor-parallel-size 8 $FLAGS

# GSM8K 20-shot (URL = the proxy, or the TP8 instance)
lm_eval --model local-chat-completions \
  --model_args "model=deepseek-ai/DeepSeek-R1,base_url=$URL/v1/chat/completions,num_concurrent=32,max_gen_toks=16384,max_retries=3,timeout=3600,tokenizer_backend=None" \
  --tasks gsm8k --num_fewshot 20 --batch_size 1 --apply_chat_template --fewshot_as_multiturn

# benchmark
vllm bench serve --backend vllm --base-url $URL --endpoint /v1/completions --model deepseek-ai/DeepSeek-R1 \
  --dataset-name random --random-input-len 8192 --random-output-len 4096 --num-prompts 640 --max-concurrency 256 --request-rate 2 --random-range-ratio 0 --ignore-eos --seed 1029

End-to-end verification with hybrid MLA (Kimi-Linear)

Kimi-Linear mixes MLA layers with KDA (linear attention) layers that keep a recurrent state. Both instances share two MI355X GPUs at --gpu-memory-utilization 0.35 each, with MoRIIO WRITE over xGMI, the toy proxy in vLLM example, block size 2,048 on both instances (hybrid attention blocks are sized to fit the per-rank state), VLLM_SSM_CONV_STATE_LAYOUT=DS, MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 on the prefill. The runs include #64330, which fixes the completion-wait hang over xGMI.

Prefill / decode	Decode KV cache (tokens)	GSM8K 5-shot, flexible-extract / strict-match	Failed or timed-out copy waits, transfer errors
TP2 / TP2	4,830,301	0.8620 / 0.4397	0
TP2 / TP2+DCP2	7,435,995	0.8719 / 0.4458	0

With DCP2 the decode cache grows by 54 % rather than doubling, because only the MLA layers are split across the DCP ranks, while every rank keeps its whole TP share of a request's KDA state.

Kimi-Linear commands

FLAGS="--trust-remote-code --block-size 2048 --kv-cache-dtype auto --attention-backend ROCM_AITER_MLA \
  --gpu-memory-utilization 0.35 --enable-chunked-prefill --max-num-batched-tokens 32768 --enable-prefix-caching \
  --max-num-seqs 64 --max-model-len 32768 --compilation-config {\"cudagraph_mode\":\"FULL_DECODE_ONLY\"}"
export VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=0 VLLM_USE_BREAKABLE_CUDAGRAPH=0 VLLM_SSM_CONV_STATE_LAYOUT=DS

# prefill
HIP_VISIBLE_DEVICES=0,1 MORI_IO_XGMI_SCATTER_GATHER_THRESHOLD=0 vllm serve moonshotai/Kimi-Linear-48B-A3B-Instruct \
  --port 20005 --tensor-parallel-size 2 $FLAGS \
  --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_producer","kv_connector_extra_config":{"backend":"xgmi","proxy_ip":"127.0.0.1","proxy_ping_port":"36367","http_port":"20005","handshake_port":"6301","notify_port":"6105"}}'

# decode, on the same GPUs (for the TP2 -> TP2 row, drop --decode-context-parallel-size)
HIP_VISIBLE_DEVICES=0,1 vllm serve moonshotai/Kimi-Linear-48B-A3B-Instruct \
  --port 40005 --tensor-parallel-size 2 --decode-context-parallel-size 2 $FLAGS \
  --kv-transfer-config '{"kv_connector":"MoRIIOConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"backend":"xgmi","proxy_ip":"127.0.0.1","proxy_ping_port":"36367","http_port":"40005","handshake_port":"7301","notify_port":"7501"}}'

# GSM8K, through the proxy
lm_eval --model local-chat-completions \
  --model_args "model=moonshotai/Kimi-Linear-48B-A3B-Instruct,base_url=http://127.0.0.1:8000/v1/chat/completions,num_concurrent=32,max_gen_toks=16384,max_retries=1,timeout=900,tokenizer_backend=None" \
  --tasks gsm8k --num_fewshot 5 --batch_size 1 --apply_chat_template --fewshot_as_multiturn

Details

For MLA models, a TPn decode instance keeps a full copy of each request's latent KV on every TP rank. DCP splits the tokens across the ranks instead, so the same GPUs hold up to n times the tokens, and the prefill instance doesn't need DCP for that.

With this PR, a WRITE transfer into a DCP decode instance goes through four steps.

1. The decode instance deals each request's tokens to its DCP ranks in groups of --cp-kv-cache-interleave-size (I) tokens. With DCP N and block size B, one scheduler block covers N x B tokens, and each rank keeps its share in its own physical block under the same block id.
2. The decode scheduler sends that block list to every prefill rank, as on main.
3. Each prefill rank, which holds the prompt unsharded, writes each decode rank only the tokens that rank holds, one write per group of I tokens.
4. Each decode rank finishes the transfer once its writers have sent write_done.

On main the prefill rank pairs its block i with decode block i.

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_layout.py's compute_block_transfer_offsets
if len(local_block_ids) > len(remote_block_ids):
    raise ValueError("local_block_ids longer than remote_block_ids: ...")

So a prompt longer than one decode block fails outright, and a shorter one puts its tokens on the wrong DCP ranks. At MoRIIO's default interleave of 1, a single prefill block is spread over all the decode ranks.

This PR makes the following changes:

• It maps tokens instead of blocks for a producer without DCP (dcp_write_runs). Token t goes to DCP rank (t // I) % N at position p = (t // (I * N)) * I + t % I, which is block ids[p // B], slot p % B, the placement vllm/v1/worker/block_table.py uses. Since I divides B, each group of I tokens is one write, and at I = B the writes are the existing whole-block writes. MLA latents are written this way. Other DCP-sharded MLA caches, such as the sparse MLA indexer cache of DeepSeek-V3.2, GLM-5.2 and GLM-5.3, are written in whole blocks and need I = B. Recurrent-state layers of hybrid models (Mamba, KDA) are not split by DCP and keep main's TP-slice writes.
• It advertises dcp_size and cp_kv_cache_interleave_size in the handshake, and the instance that opens the connection (the prefill instance in WRITE mode, the decode instance in READ mode) refuses layouts it can't transfer. Those are a DCP-sharded layer that isn't MLA (a GQA model, or a draft model without MLA, which the decode instance runs without DCP), an interleave that doesn't divide the block size or is below it for a whole-block layer, equal DCP sizes with different interleaves, and any other DCP size mismatch. A peer that doesn't advertise its layout is accepted as before.
• It fails only the affected requests after a refusal. In WRITE mode the prefill instance frees the request's blocks and sends write_failed once the decode allocation arrives, and the decode instance fails the request per kv_load_failure_policy. In READ mode the TP ranks agree on the refusal (HandshakeRefusedError), report the requests' blocks as invalid and release the prefill's blocks with an empty read, so the engine keeps running. A later release of such a request that carries no prefill address is skipped with a warning instead of raising in build_connector_meta.

A producer with the same DCP size as the decode instance keeps the original rank-to-rank writes in READ and WRITE mode. MoRIIO also keeps its default interleave of 1, which speculative decoding with DCP on ROCM_AITER_MLA needs, while the block size gives the fewest writes.

Notes:

• Over xGMI, main's completion wait can sleep out the full transfer_timeout, or fail with copies still in flight and send no write_failed. [Bugfix][ROCm][MoRIIO] Wait for MoRIIO WRITE mode completions in short chunks #64330 fixes the wait and the end-to-end runs were performed with the fix.
• Every DCP-sharded layer must be MLA, so GQA models and draft models without MLA are refused.

AI assistant was used for this PR.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

Documentation preview: https://vllm--64722.org.readthedocs.build/en/64722/

[tjtanaa]

This can report the transfer finished while another producer is still writing its share. For a hybrid TP2 prefill → TP1 decode transfer, the decode rank has two writers. If rank 0 has posted writes and rank 1’s handshake fails, rank 1 sends write_failed here; the consumer immediately reports both failed_recving and finished_recving, clears its writer count, and allows the scheduler to free/reuse the destination blocks. Rank 0’s writes may still be in flight. This path also handles transient handshake errors, not just uniform layout refusals.

Please retain the failure state but wait for every expected writer to become terminal before reporting receiving finished. A focused fan-in regression with one pending writer and one failed handshake would cover this; the new refusal test currently uses one producer and therefore cannot expose it.

[tjtanaa]

Please also check terminal transfer state before starting a handshake here. The writer can process a refused request’s first task between layer hooks: _mark_request_done() marks it terminal and clears _unconnected, so the next hook enters this branch and starts another handshake. Its callback then re-adds the terminal transfer to _unconnected, but schedule_write() drops terminal tasks and nothing clears that entry afterward.

A CPU probe of this ordering produces two handshakes and a retained marker. Please prevent terminal callbacks from resurrecting the state, and extend the existing test to execute the first task before invoking the next layer hook.