[AMD][PD] enable mori transfer for PD decode DCP by inkcherry · Pull Request #37358 · sgl-project/sglang

inkcherry · 2026-09-01T05:58:38Z

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

Signed-off-by: inkcherry <mingzhi.liu@amd.com> # Conflicts: # python/sglang/srt/disaggregation/mori/conn.py

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

ShangmingCai · 2026-09-13T16:56:14Z

+        requires_dcp_relayout = getattr(
+            req.disagg_kv_sender, "requires_dcp_relayout", None
+        )
+        if requires_dcp_relayout is not None and requires_dcp_relayout():
+            return
+


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.

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.

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.

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

inkcherry · 2026-09-15T03:30:51Z

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.

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

inkcherry · 2026-09-23T15:16:47Z

@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 · 2026-09-24T03:07:38Z

/tag-and-reurn-ci

github-actions · 2026-09-24T03:08:06Z

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

See the command reference.

HaiShaw · 2026-09-24T11:10:38Z

/tag-and-rerun-ci

HaiShaw · 2026-09-28T07:13:23Z

@Lzy17 please have a review.

Lzy17 · 2026-09-28T22:29:29Z

@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 · 2026-10-08T03:18:11Z

@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:

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.
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.
Bounds-check state_conv_shard_groups[i] / state_slice_outer_counts[i] in send_state, as is done for src_dims.
Surface requires_dcp_relayout errors from _add_remote_peer as a request failure. The caller only logs, so the request waits for its timeout.
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 · 2026-10-12T10:58:11Z

/rerun-failed-ci

inkcherry added 2 commits August 28, 2026 05:21

enable Mori transfer for PD decode DCP

25d2a39

merge upstream main into Mori DCP

c814307

inkcherry marked this pull request as ready for review September 1, 2026 06:05

inkcherry requested review from ByronHsu, Duyi-Wang, HaiShaw, ShangmingCai, billishyahao, hnyls2002, kkHuang-amd and sogalin as code owners September 1, 2026 06:05

Duyi-Wang added amd run-mori-pd Run mori pd related ci workflow labels Sep 3, 2026

Duyi-Wang and others added 3 commits September 3, 2026 15:53

Merge branch 'main' into feat/mori-pd-dcp

6b01eaf

Merge branch 'main' into feat/mori-pd-dcp

1373e4f

Signed-off-by: inkcherry <mingzhi.liu@amd.com> # Conflicts: # python/sglang/srt/disaggregation/mori/conn.py

drop Mori-only backend allowlist test

073d3ae

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

ShangmingCai reviewed Sep 13, 2026

View reviewed changes

inkcherry added 2 commits September 14, 2026 06:41

fix Mori DCP transfer and early-send guard

c76733c

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

Merge branch 'main' into feat/mori-pd-dcp

c6109d9

inkcherry added 2 commits September 23, 2026 06:55

enable Mori DCP cached-prefix early send

d3d00fd

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

focus Mori backend tests on DCP

57c0743

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

inkcherry requested review from Fridge003, Kangyan-Zhou, bingxche, ispobock and merrymercy as code owners September 23, 2026 07:50

inkcherry added 2 commits September 23, 2026 07:55

run Mori DCP tests in AMD CI

c115f73

Signed-off-by: inkcherry <mingzhi.liu@amd.com>

Merge branch 'main' into feat/mori-pd-dcp

71d29ef

ShangmingCai approved these changes Sep 24, 2026

View reviewed changes

Merge branch 'main' into feat/mori-pd-dcp

35fa6fc

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Sep 24, 2026

Merge branch 'main' into feat/mori-pd-dcp

5a74523

This was referenced Sep 30, 2026

[CI Monitor] Daily Report - 2026-09-30 bingxche/sglang-ci-bot#211

Open

[CI Monitor] Daily Report - 2026-10-01 bingxche/sglang-ci-bot#212

Open

merge main into mori dcp branch

4ee930c

inkcherry requested a review from Lzy17 as a code owner October 8, 2026 03:15

Lzy17 reviewed Oct 8, 2026

View reviewed changes

Lzy17 mentioned this pull request Oct 8, 2026

[PD] Consolidate sender and receiver implementations #41140

Open

3 of 4 tasks

ShangmingCai added 2 commits October 9, 2026 20:54

Merge branch 'main' into feat/mori-pd-dcp

4894ba9

Merge branch 'main' into feat/mori-pd-dcp

90ccd1f

ShangmingCai added the bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) label Oct 10, 2026

ShangmingCai mentioned this pull request Oct 11, 2026

[Roadmap] Prefill-Decode Disaggregation Roadmap (2026 Q4) #43699

Open

15 of 40 tasks

Merge branch 'main' into feat/mori-pd-dcp

cb1416b

ShangmingCai merged commit 50a1440 into sgl-project:main Oct 13, 2026
232 of 259 checks passed

Conversation

inkcherry commented Sep 1, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Changes

Validation

CI States

Uh oh!

ShangmingCai Sep 13, 2026

Choose a reason for hiding this comment

Uh oh!

inkcherry Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

inkcherry Sep 14, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai left a comment

Choose a reason for hiding this comment

Uh oh!

inkcherry commented Sep 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

inkcherry commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ShangmingCai left a comment

Choose a reason for hiding this comment

Uh oh!

ShangmingCai commented Sep 24, 2026

Uh oh!

github-actions Bot commented Sep 24, 2026

Uh oh!

HaiShaw commented Sep 24, 2026

Uh oh!

HaiShaw commented Sep 28, 2026

Uh oh!

Lzy17 commented Sep 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

inkcherry commented Oct 8, 2026

Uh oh!

Lzy17 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ShangmingCai commented Oct 12, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

inkcherry commented Sep 1, 2026 •

edited by github-actions Bot

Loading

inkcherry Sep 14, 2026 •

edited

Loading

inkcherry commented Sep 15, 2026 •

edited

Loading

inkcherry commented Sep 23, 2026 •

edited

Loading

Lzy17 commented Sep 28, 2026 •

edited

Loading

Lzy17 left a comment •

edited

Loading