[inkcherry]

Motivation

PD decode with DCP currently rejects the Mori transfer backend because Mori does not implement the token relayout required by round-robin DCP KV shards.

Changes

• Add DCP rank and layer metadata to Mori peer registration.
• Transfer MLA KV at token granularity into each decode rank's DCP shard.
• Map PP-local KV and Mamba state tensors by global layer ID.
• Handle TP mismatch for replicated latent KV and reject unsupported head-sharded layouts.
• Skip cached-prefix early send when DCP relayout is required.
• Allow Mori in the PD decode DCP backend validation.

Validation

• Mori DCP unit tests: 6/6 passed.
• ServerArgs DCP tests: 6/6 passed.
• Kimi-K3 PP8/TP1 prefill + TP8/DCP8 decode: both servers healthy, with no KV transfer errors.
• 31K long-context retrieval and arithmetic: 8/8 passed, with no errors or truncation.
• GSM8K: 39/40 (97.5%), with no errors or truncation.
• 32K input, 1 output token, concurrency 16, 128 requests: 16,699 input tok/s and 31.11s median TTFT.
• Pre-commit and lint checks passed.

The Kimi-K3 E2E validation used an AITER DCP-enabled stack; this PR only changes the Mori transfer path.

---

CI States

Latest PR Test (Base): ✅ Run #38288271653
Latest PR Test (Extra): ❌ Run #38288270583
Latest PR Test (AMD ROCm 10): ❌ Run #38288271415

[ShangmingCai]

QQ: Does only Mori need this? If so, I think this introduces another code path divergence. Also, I think we should avoid using getattr. And the code here is confusing.

[inkcherry]

Thanks, yes, writing it that way does introduce some confusion.

#32837 — this PR supports further optimization of DCP. But it only covers Mooncake and NIXL.

• The current PR first supports the basic functionality of DCP for mori. Now I modified it here with an API supports_cached_prefix_early_send to bypass ·. No changes for nixl and mooncake.

• In subsequent separate PRs, we will try to leverage feat: support Kimi Linear PD disaggregation with DCP #32837 and the validated optimizations. If it works, we can remove this bypass.

[inkcherry]

re-verified on Kimi3 E2E.

[ShangmingCai]

No other comment, others look good. But before we merge, we need to discuss about the diff in prefill.py.

[inkcherry]

No other comment, others look good. But before we merge, we need to discuss about the diff in prefill.py.

Thank you, I have replied with the reason and revised the code. Do you think we should directly support the mori layout-related work through this PR? I am inclined to separate it.

[inkcherry]

@ShangmingCai Applied the DCP KV layout optimization to Mori without introducing backend-specific control-flow divergence now.
re-verified on Kimi3 E2E.

[ShangmingCai]

LGTM

[ShangmingCai]

/tag-and-reurn-ci

[github-actions]

⛔ /tag-and-reurn-ci isn't a recognized CI command. Did you mean /tag-and-rerun-ci?

See the command reference.

[HaiShaw]

/tag-and-rerun-ci

[HaiShaw]

@Lzy17 please have a review.

[Lzy17]

@inkcherry @HaiShaw

The port itself looks right -- every helper this calls (build_dcp_token_transfer_plan, requires_dcp_relayout, prepare_dcp_token_item_lens, resolve_dcp_dst_entry_indices) is already on main and already used by nixl and mooncake, so mori is just the third caller. pr-test-amd-mori passed on this head: 4 unit tests and 2 PD e2e tests.

One gap: nothing in CI executes the new path. test_mori_backend_basic.py is fully mocked, and test_mori_transfer_engine_e2e.py -- the one that runs on GPUs -- uses the default small model with no --dcp-size. Since requires_dcp_relayout only returns True for MLA with dcp_size 1 -> N, send_kvcache_dcp is never reached. Could you add a case with a small MLA model, prefill --dcp-size 1 and decode --dcp-size 2, on the existing linux-mi35x-gpu-8.fabric runner? If it is easier for you, we are happy to add it ourselves -- just let us know which topology you had in mind.

Two smaller things:

• _build_tp_slice_config switching to compute_head_slice_params changes the ordinary TP-mismatch path, which is unrelated to DCP. The TPMismatch e2e test does cover it, so this is not a blocker, but it would be cleaner as its own PR -- and the function-local import hides the dependency.
• test_server_args.py repoints the rejection case from mori to ascend, which is correct now that mori is allowed, but nothing asserts mori is accepted. One extra case would close that.

[inkcherry]

@inkcherry @HaiShaw

The port itself looks right -- every helper this calls (build_dcp_token_transfer_plan, requires_dcp_relayout, prepare_dcp_token_item_lens, resolve_dcp_dst_entry_indices) is already on main and already used by nixl and mooncake, so mori is just the third caller. pr-test-amd-mori passed on this head: 4 unit tests and 2 PD e2e tests.

One gap: nothing in CI executes the new path. test_mori_backend_basic.py is fully mocked, and test_mori_transfer_engine_e2e.py -- the one that runs on GPUs -- uses the default small model with no --dcp-size. Since requires_dcp_relayout only returns True for MLA with dcp_size 1 -> N, send_kvcache_dcp is never reached. Could you add a case with a small MLA model, prefill --dcp-size 1 and decode --dcp-size 2, on the existing linux-mi35x-gpu-8.fabric runner? If it is easier for you, we are happy to add it ourselves -- just let us know which topology you had in mind.

Two smaller things:

• _build_tp_slice_config switching to compute_head_slice_params changes the ordinary TP-mismatch path, which is unrelated to DCP. The TPMismatch e2e test does cover it, so this is not a blocker, but it would be cleaner as its own PR -- and the function-local import hides the dependency.
• test_server_args.py repoints the rejection case from mori to ascend, which is correct now that mori is allowed, but nothing asserts mori is accepted. One extra case would close that.

modified.

[Lzy17]

Thanks for the PR. Items:

1. Add the DCP e2e (small MLA model, prefill dcp1 / decode dcp2) and a test that mori is accepted, per my earlier comment. Neither is in the current head.
2. State in the description, or split out, that this also changes non-DCP paths (layer-id PP mapping, mamba TP-slice), which only mocked tests cover.
3. Bounds-check state_conv_shard_groups[i] / state_slice_outer_counts[i] in send_state, as is done for src_dims.
4. Surface requires_dcp_relayout errors from _add_remote_peer as a request failure. The caller only logs, so the request waits for its timeout.
5. Coordinate with #41769 (same _send_swa_dsa_state loop; dst-length lookups must use j) and #41140 (registration message needs the four new fields).

[ShangmingCai]

/rerun-failed-ci