feat: support Kimi Linear PD disaggregation with DCP by kpham-sgl · Pull Request #32837 · sgl-project/sglang

kpham-sgl · 2026-07-30T00:43:33Z

Summary

Add PD disaggregation support for transferring dense prefill MLA rows into rank-owned decode DCP layouts.

support Mooncake and NIXL transfer backends
preserve exact final-page token counts across physical/virtual page and chunk boundaries
support heterogeneous prefill/decode attention TP and matching DCP layouts
compose pipeline-parallel layer mapping with DCP token relayout
complete zero-owned NIXL chunks with standalone notifications
add Kimi Linear PD+DCP acceptance coverage and focused control-path validation

Co-author

@YAMY1234

Supported topology

This is wider than only TP -> DCP, but it is not the full Cartesian {TP, DCP} -> {TP, DCP} matrix.

Prefill layout	Decode layout	Status
TP / DCP1	TP / DCP1	Existing PD behavior; unchanged
TP / DCP1	DCP=N	Supported by the new token relayout path
DCP=N	DCP=N	Supported when DCP size/layout match and peers connect matching DCP ranks
DCP=N	TP / DCP1	Not supported
DCP=N	DCP=M, `N != M`	Not supported

Prefill and decode attention TP sizes may differ; the implementation uses the existing heterogeneous-TP rank mapping.

Pipeline parallelism

PP is supported within the existing PD topology rules:

prefill PP=N -> decode PP=N
prefill PP=N -> decode PP=1, with decode gathering the applicable layers from all prefill stages

The DCP path resolves destination entries by model-layer ID, so PP layer partitioning is composed with token relayout for both Mooncake and NIXL.

The following are not supported:

prefill PP=1 -> decode PP=N
arbitrary unequal prefill/decode PP layouts

Runtime constraints

DCP decode requires an MLA or hybrid-MLA KV layout; the implementation is not hardcoded to Kimi, but Kimi Linear is the acceptance target.
DCP decode supports Mooncake and NIXL transfer backends only.
Prefill attention CP must be 1 and decode attention CP must be 1.
DCP decode requires chunk cache; decode radix cache and hierarchical cache are rejected.
Matching page size and KV-cache dtype are required across prefill and decode.

Implementation details

Decode registration advertises DCP size/rank and per-layer KV metadata.
Prefill detects whether the peer needs dense-to-DCP relayout or already has a matching DCP layout.
The relayout planner maps physical prefill token rows to rank-owned decode token rows while honoring exact token counts on partial final pages.
Mooncake builds contiguous transfer blocks from the relayout plan.
NIXL uses flat token-row descriptors, bypasses page-level prepared handles, and emits a standalone notification when a rank owns zero rows.
Global model-layer IDs compose PP layer ownership with DCP context ownership.

Validation scope

The registered 8-GPU acceptance test covers:

prefill: TP4 / EP4 / DCP1 / PP1
decode: TP4 / DCP4 / PP1
Kimi Linear 48B, TokenSpeed MLA, FP8 KV cache
monolithic-vs-PD token and logprob parity
physical and virtual page boundaries, chunk boundaries, mixed batches, 32K needle retrieval, CUDA-graph and eager decode batches
Mooncake when dispatched by the standard CI fixture

PP composition and the wider topology rules are implemented but are not all covered by the single registered end-to-end test.

Test plan

Pre-commit formatting and static hooks
Python syntax compilation for changed runtime modules
Focused CPU disaggregation/server-args tests on CI
Kimi Linear 8-GPU PD+DCP acceptance on CI (/rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py)

CI States

Latest PR Test (Base): ✅ Run #30590566763
Latest PR Test (Extra): ✅ Run #30590566001

gemini-code-assist · 2026-07-30T00:43:36Z

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

kpham-sgl · 2026-07-30T00:47:30Z

/rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py

github-actions · 2026-07-30T00:47:56Z

Results for /rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py:

🚀 8-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/disaggregation/test_kimi_linear_pd_dcp4.py

Transfer dense prefill MLA rows into rank-owned decode DCP layouts for Mooncake and NIXL. Support heterogeneous attention TP, matching DCP layouts, PP-aware layer mapping, exact partial-page token counts, and zero-owned transfer completion notifications. Co-authored-by: Yangmin Li <yangminl@nvidia.com>

kpham-sgl · 2026-07-30T01:27:32Z

/rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py

github-actions · 2026-07-30T01:28:00Z

Results for /rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py:

🚀 8-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/disaggregation/test_kimi_linear_pd_dcp4.py

kpham-sgl · 2026-07-30T02:14:52Z

/rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py

github-actions · 2026-07-30T02:15:15Z

Results for /rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py:

🚀 8-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/disaggregation/test_kimi_linear_pd_dcp4.py

kpham-sgl · 2026-07-30T18:36:14Z

/rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py

github-actions · 2026-07-30T18:36:40Z

Results for /rerun-test test/registered/disaggregation/test_kimi_linear_pd_dcp4.py:

🚀 8-gpu-b200 (1 test): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_kimi_linear_pd_dcp4.py

YAMY1234

LGTM

kpham-sgl · 2026-07-30T22:27:12Z

/rerun-failed-ci

ShangmingCai · 2026-07-31T08:36:33Z

+            # msg[14:16] belong to the staging field below; DCP trails it.
+            dst_dcp_size=(
+                int(msg[16].decode("ascii")) if len(msg) > 16 and msg[16] != b"" else 1
+            ),
+            dst_dcp_rank=(
+                int(msg[17].decode("ascii")) if len(msg) > 17 and msg[17] != b"" else 0
+            ),
            # Note: always put the staging field at the final
            staging=StagingRegisterInfo.from_zmq_fields(msg, 14),


Could we reorder the fields' slice to 14, 15, 16? Staging register info should be physically placed at the end since it occupies several positions.

I think I should and will refactor the metadata interpretation part, many PRs miss this note.

CC: @YAMY1234 I think putting this at the end will affect version compatibility while we append new fields, but the field number this StagingRegisterInfo contains is unclear from the name and has bad readability (that is the reason why I commented that it should be put at the last); should we unwrap it?

ShangmingCai · 2026-07-31T08:39:38Z

+                            dst_dcp_size,
+                            dst_dcp_rank,


Then adjust the code here

ShangmingCai · 2026-07-31T08:48:35Z

+    def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool:
+        if self.dcp_size == dst_dcp_size:
+            if self.dcp_rank != dst_dcp_rank:
+                raise RuntimeError(
+                    "PD peers must connect matching DCP ranks, got "
+                    f"prefill={self.dcp_rank}, decode={dst_dcp_rank}"
+                )
+            return False
+
+        if (
+            self.dcp_size == 1
+            and dst_dcp_size > 1
+            and (self.is_mla_backend or self.is_hybrid_mla_backend)
+        ):
+            return True
+
+        raise RuntimeError(
+            f"Unsupported PD DCP topology: {self.dcp_size} -> {dst_dcp_size}"
+        )


I think we should not raise RuntimeError here. If requires_dcp_relayout fails, the transfer worker or metadata listener will crash. But in our design, this Prefill instance could pair with a DCP decode instance and a regular pure decode TP instance at the same time, so it could still process the requests from the latter one. We should mark this DCP topo not-match request as failed and print some warnings/error logging, or maybe add a try block here to avoid crashing the engine.

ShangmingCai

No other comments. Looks good. Once these nits have been addressed, it could be merged after CI passes.

ShangmingCai · 2026-07-31T09:04:06Z

Or we could fix them in a following PR, considering the GPU CI has passed now.

kpham-sgl · 2026-07-31T09:14:01Z

Will address comments in a following PR

Co-authored-by: Yangmin Li <yangminl@nvidia.com>

kpham-sgl requested review from ByronHsu, Duyi-Wang, HaiShaw, ShangmingCai, billishyahao, hnyls2002, kkHuang-amd and sogalin as code owners July 30, 2026 00:43

kpham-sgl force-pushed the feat/kimi-linear-pd-dcp-oss branch from eb6e3c9 to 1a4e550 Compare July 30, 2026 01:26

kpham-sgl force-pushed the feat/kimi-linear-pd-dcp-oss branch from 74efd78 to 1a4e550 Compare July 30, 2026 02:16

test: align Kimi Linear PD parity reference topology

560f369

sgl-project deleted a comment from github-actions Bot Jul 30, 2026

test: isolate Kimi Linear PD parity cache state

e773501

kpham-sgl assigned YAMY1234 and ShangmingCai Jul 30, 2026

YAMY1234 approved these changes Jul 30, 2026

View reviewed changes

kpham-sgl added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fastfail labels Jul 30, 2026

Merge branch 'main' into feat/kimi-linear-pd-dcp-oss

2e2f26c

ShangmingCai reviewed Jul 31, 2026

View reviewed changes

ShangmingCai approved these changes Jul 31, 2026

View reviewed changes

kpham-sgl merged commit 2573190 into sgl-project:main Jul 31, 2026
168 of 178 checks passed

0Bumblebee mentioned this pull request Aug 3, 2026

[Kimi] Support kimi-k3 #32541

Merged

kpham-sgl mentioned this pull request Aug 4, 2026

[Roadmap][DCP] Decode Context Parallelism & Helix Parallelism (2026 Q4) #29736

Open

19 of 39 tasks

saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 16, 2026

feat: support Kimi Linear PD disaggregation with DCP (sgl-project#32837)

cc73c6d

Co-authored-by: Yangmin Li <yangminl@nvidia.com>

kpham-sgl deleted the feat/kimi-linear-pd-dcp-oss branch August 24, 2026 22:25

Atituiset pushed a commit to Atituiset/sglang that referenced this pull request Sep 10, 2026

feat: support Kimi Linear PD disaggregation with DCP (sgl-project#32837)

3c33a55

Co-authored-by: Yangmin Li <yangminl@nvidia.com>

inkcherry mentioned this pull request Sep 14, 2026

[AMD][PD] enable mori transfer for PD decode DCP #37358

Merged

ShangmingCai mentioned this pull request Oct 11, 2026

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

Open

19 of 44 tasks

Conversation

kpham-sgl commented Jul 30, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Co-author

Supported topology

Pipeline parallelism

Runtime constraints

Implementation details

Validation scope

Test plan

CI States

Uh oh!

gemini-code-assist Bot commented Jul 30, 2026

Uh oh!

kpham-sgl commented Jul 30, 2026

Uh oh!

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

Uh oh!

Uh oh!

kpham-sgl commented Jul 30, 2026

Uh oh!

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

Uh oh!

Uh oh!

kpham-sgl commented Jul 30, 2026

Uh oh!

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

Uh oh!

Uh oh!

kpham-sgl commented Jul 30, 2026

Uh oh!

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

Uh oh!

Uh oh!

YAMY1234 left a comment

Choose a reason for hiding this comment

Uh oh!

kpham-sgl commented Jul 30, 2026

Uh oh!

ShangmingCai Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai Jul 31, 2026

Choose a reason for hiding this comment

Uh oh!

ShangmingCai Jul 31, 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!

ShangmingCai left a comment

Choose a reason for hiding this comment

Uh oh!

ShangmingCai commented Jul 31, 2026

Uh oh!

kpham-sgl commented Jul 31, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kpham-sgl commented Jul 30, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jul 30, 2026 •

edited

Loading

github-actions Bot commented Jul 30, 2026 •

edited

Loading

github-actions Bot commented Jul 30, 2026 •

edited

Loading

github-actions Bot commented Jul 30, 2026 •

edited

Loading

ShangmingCai Jul 31, 2026 •

edited

Loading