[PD] Pack DCP1→DCP-N PD KV transfers into dest-contiguous RDMA blocks by kpham-sgl · Pull Request #35762 · sgl-project/sglang

kpham-sgl · 2026-08-20T22:32:21Z

Cyclic DCP ownership made the relayout path emit one RDMA per token. Gather owned MLA rows on prefill first so Mooncake/NIXL can send page-sized or larger blocks. DSPARK is left to a follow-up.

Motivation

DCP assigns KV ownership round-robin per token (owned_offsets = arange(rank, num_kv_tokens, dcp_size)), so on a DCP1 -> DCP-N PD transfer the source indices are strided and group_concurrent_contiguous cannot merge them: the relayout emits one RDMA descriptor per token per layer (tokens x layers per chunk). The transfer becomes descriptor-bound rather than bandwidth-bound.

Modifications

Pack cyclic DCP rows: Added a raw-pointer Triton gather and common packing path that turns strided MLA source rows into contiguous blocks, reducing transfer descriptors from one per owned token per layer to destination-contiguous runs.
Reuse a dedicated gather stream: Each worker pack buffer lazily creates and reuses its CUDA gather stream. The stream waits for default-stream KV writes, and packing completes before the transport reads the buffer.
Keep pack memory bounded and reusable: Allocate and register one buffer per transfer worker, size it from the maximum prefill chunk, and assign each DCP rank a fixed region. TP peers sharing a DCP rank reuse the same packed source.
Integrate both transports: Mooncake sends the packed blocks directly. NIXL submits immutable rank regions asynchronously and relies on the existing chunk barrier before repacking the buffer.

Accuracy Tests

test/registered/disaggregation/test_kimi_linear_pd_dcp4.py on 8xB300: 6 passed, GSM8K 0.91 (threshold 0.88). Covers token + output-logprob parity against a monolithic TP4 reference at 9 chunk/virtual-page boundary lengths, and a 32k needle-in-a-haystack at depths 0.1/0.5/0.9.

Speed Tests and Profiling

Setup. 8xB300, Kimi-Linear-48B-A3B-Instruct. main @ 20621aa14b vs this PR @ 539b78567f, on both transfer backends.

Topology: prefill TP4/EP4 dcp=1 (GPU 0-3) -> decode TP4 dcp=4 (GPU 4-7). tokenspeed_mla, fp8_e4m3 KV, --page-size 64, --chunked-prefill-size 8192, --mem-fraction-static 0.80.
Workload: OSL 32, 2 warmup requests + cache flush per point, freshly launched servers per condition. TTFT/throughput client-measured by sglang.benchmark.serving --pd-separated.
Controls: bytes moved identical per point (26.1 / 136.2 / 513.8 MB at ISL 4k / 32k / 128k); packed path verified active on the PR (4 x 31.5 MB pack buffers per prefill TP rank) and absent on main.

NIXL

ISL	conc	TTFT p50 main -> PR	out tok/s main -> PR
4096	1	115.8 -> 83.0 ms (-28.3%)	158.6 -> 187.4 (+18.2%)
4096	8	3470.2 -> 3076.3 ms (-11.3%)	69.2 -> 77.8 (+12.4%)
32768	1	511.0 -> 311.7 ms (-39.0%)	53.5 -> 79.5 (+48.8%)
32768	8	3749.3 -> 2298.4 ms (-38.7%)	64.4 -> 102.6 (+59.2%)
131072	1	2198.7 -> 1426.3 ms (-35.1%)	14.0 -> 21.0 (+50.1%)
131072	8	16300.5 -> 10280.2 ms (-36.9%)	15.0 -> 24.2 (+60.7%)

Mooncake

ISL	conc	TTFT p50 main -> PR	out tok/s main -> PR
4096	1	93.2 -> 85.5 ms (-8.4%)	172.6 -> 180.4 (+4.5%)
4096	8	3207.1 -> 3166.5 ms (-1.3%)	74.7 -> 77.2 (+3.4%)
32768	1	360.8 -> 310.6 ms (-13.9%)	71.5 -> 79.2 (+10.7%)
32768	8	2556.4 -> 2230.2 ms (-12.8%)	93.0 -> 104.7 (+12.6%)
131072	1	1575.6 -> 1429.9 ms (-9.2%)	19.2 -> 21.2 (+10.5%)
131072	8	11467.3 -> 10207.1 ms (-11.0%)	21.7 -> 24.3 (+12.1%)

Summary. Backend-dependent: NIXL -28% to -39% TTFT and up to 1.6x output throughput at ISL >= 32k; mooncake -9% to -14% and ~1.1x. Mooncake's baseline is already faster at every point, and packing converges the two to within ~3% -- per-token descriptors penalise NIXL far more, and this PR removes that penalty. ISL 4096 / conc 8 barely moves on either backend (~3.1-3.5 s in all conditions), so something other than descriptor count dominates there. Cost: 4 x 31.5 MB pack buffers per prefill TP rank, leaving the KV pool essentially unchanged.

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review and Merge Process

Ping Merge Oncalls to start the process. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ✅ Run #33031619005
Latest PR Test (Extra): ✅ Run #33064293710
Latest PR Test (AMD ROCm 7.2): ❌ Run #33031619022

Cyclic DCP ownership made the relayout path emit one RDMA per token. Gather owned MLA rows on prefill first so Mooncake/NIXL can send page-sized or larger blocks. DSPARK is left to a follow-up. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep the feature description on the DCP page; the long-form design note stays local, not in the published docs. Co-authored-by: Cursor <cursoragent@cursor.com>

The packed-transfer design note is local-only; dcp.mdx does not need a docs update in this PR. Co-authored-by: Cursor <cursoragent@cursor.com>

Fold the duplicated Mooncake/NIXL ensure helpers into CommonKVManager and drop the unused lock. Each backend only registers the buffer with its transfer engine. Co-authored-by: Cursor <cursoragent@cursor.com>

Attach kv_data_tensors next to kv_data_ptrs when the MLA pool lines up, and skip the pointer reverse-lookup setter. Co-authored-by: Cursor <cursoragent@cursor.com>

Centralize lazy gather-stream ownership in StagingBuffer so DCP packing follows the existing staging synchronization path without redundant GPU IDs. Co-authored-by: Cursor <cursoragent@cursor.com>

Remove the rollout toggle so eligible DCP transfers consistently use the packed path while retaining the buffer-sizing override. Co-authored-by: Cursor <cursoragent@cursor.com>

Drop SGLANG_DISAGG_DCP_PACK_MAX_TOKENS; the pack buffer already cannot exceed a prefill chunk, so the env override was unused. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep KVArgs backend-neutral by moving DCP packing into a fused Triton KV-cache kernel that consumes registered addresses directly, while protecting NIXL's asynchronous source-buffer lifetime. Co-authored-by: Cursor <cursoragent@cursor.com>

Name the KV-cache kernel after its PD-to-DCP gather role so it is distinct from transfer orchestration. Co-authored-by: Cursor <cursoragent@cursor.com>

Reuse backend staging-memory hooks for DCP pack buffers and document the expected buffer footprint. Co-authored-by: Cursor <cursoragent@cursor.com>

Protect the descriptor-collapse benefit and clarify invalid KV item-length diagnostics. Co-authored-by: Cursor <cursoragent@cursor.com>

`_register_staging_memory` dropped its `gpu_id` parameter in favour of `self.kv_args.gpu_id`, but the test still passed it positionally and asserted the explicitly-passed value, failing with a TypeError. Drop the third argument at both call sites and expect the manager's own gpu_id (1), which also pins the new behaviour. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

kpham-sgl · 2026-08-27T01:58:59Z

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

github-actions · 2026-08-27T01:59:21Z

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-08-27T02:20:23Z

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

github-actions · 2026-08-27T02:20:51Z

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

ShangmingCai

Looks good to me, very clean. cc: @YAMY1234 for staging_buffer change review.

ShangmingCai · 2026-08-27T10:21:01Z

/rerun-group disaggregation

ShangmingCai · 2026-08-27T10:21:07Z

/tag-and-rerun-ci

github-actions · 2026-08-27T10:21:23Z

Results for /rerun-group disaggregation:

🚀 4-gpu-gb300 (1 test): ✅ View workflow run

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

🚀 2-gpu-h100 (5 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_basic.py
cd test/ && python3 registered/disaggregation/test_disaggregation_decode_offload.py
cd test/ && python3 registered/disaggregation/test_disaggregation_optimistic_prefill.py
cd test/ && python3 registered/disaggregation/test_disaggregation_rust_server.py
cd test/ && python3 registered/disaggregation/test_disaggregation_unified_memory.py

🚀 8-gpu-h20 (5 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_decode_radix_cache.py
cd test/ && python3 registered/disaggregation/test_disaggregation_different_tp.py
cd test/ && python3 registered/disaggregation/test_disaggregation_dp_attention.py
cd test/ && python3 registered/disaggregation/test_disaggregation_nixl.py
cd test/ && python3 registered/disaggregation/test_disaggregation_pp.py

🚀 8-gpu-h200 (4 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_decode_radix_cache_swa.py
cd test/ && python3 registered/disaggregation/test_disaggregation_dsv4.py
cd test/ && python3 registered/disaggregation/test_disaggregation_hisparse.py
cd test/ && python3 registered/disaggregation/test_disaggregation_hybrid_attention.py

🚀 4-gpu-b200 (2 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_dwdp_gpt_oss.py
cd test/ && python3 registered/disaggregation/test_disaggregation_inkling_mxfp8.py

🚀 4-gpu-h100 (2 tests): ✅ View workflow run

cd test/ && python3 registered/disaggregation/test_disaggregation_kimi_linear.py
cd test/ && python3 registered/disaggregation/test_epd_disaggregation.py

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

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

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

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

YAMY1234

Staging buffer change LGTM!

kpham-sgl · 2026-08-29T04:41:27Z

Passed all based CI and PD disaggs test

Resolved two conflicts introduced by upstream: * mori/conn.py: sgl-project#29133 landed an equivalent Mori ABORT handler (_TAG_ABORT / _handle_abort_message) that holds transfer_lock across the status read-modify-write and returns early on an already-Failed room. Kept upstream's handler and dropped this branch's _handle_abort_notification, which duplicated it without the lock. sgl-project#29133 also covers the Mori _submit_kv_transfer cleared-room guard this branch carried. * nixl/conn.py: kept the cleared-room guard's room_transfer_infos local alongside the packed_source_by_dcp_rank map added by sgl-project#35762; the two changes are adjacent but independent.

…ontract `send_kvcache_dcp` now returns one handle per transfer part (target + draft), and the worker extends its handle list with the result. The test merged from main (sgl-project#35762) still mocked a single string handle, so `extend` split it into characters and the chunk barrier polled 41 fake handles instead of 3. The fake `kv_args` also lacked `num_draft_entries`, which the pack path now reads. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

github-actions Bot added the documentation Improvements or additions to documentation label Aug 20, 2026

kpham-sgl force-pushed the dcp-pd-packed-transfer branch from 2e593f4 to fcdc512 Compare August 24, 2026 22:45

github-actions Bot added the jit-kernel label Aug 25, 2026

kpham-sgl force-pushed the dcp-pd-packed-transfer branch 3 times, most recently from 8fc1cc1 to f77adfc Compare August 25, 2026 21:52

kpham-sgl marked this pull request as ready for review August 27, 2026 01:39

kpham-sgl requested review from BBuf, ByronHsu, DarkSharpness, Duyi-Wang, HaiShaw, HydraQYH, ShangmingCai, celve, hnyls2002, sogalin and yuan-luo as code owners August 27, 2026 01:39

kpham-sgl changed the title ~~[WIP][DCP] Pack DCP1→DCP-N PD KV transfers into dest-contiguous RDMA blocks~~ [PD] Pack DCP1→DCP-N PD KV transfers into dest-contiguous RDMA blocks Aug 27, 2026

kpham-sgl and others added 11 commits August 26, 2026 18:51

[DCP] Drop packed-transfer design page from the docs site

232c8f2

Keep the feature description on the DCP page; the long-form design note stays local, not in the published docs. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Leave the published DCP page unchanged

f7e8d8a

The packed-transfer design note is local-only; dcp.mdx does not need a docs update in this PR. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Init pack buffers once on the common manager at peer register

974cc30

Fold the duplicated Mooncake/NIXL ensure helpers into CommonKVManager and drop the unused lock. Each backend only registers the buffer with its transfer engine. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Drive pack gather from kv_args tensors and drop leftover comments

939d151

Attach kv_data_tensors next to kv_data_ptrs when the MLA pool lines up, and skip the pointer reverse-lookup setter. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Reuse staging gather stream for packed transfers

ec5c355

Centralize lazy gather-stream ownership in StagingBuffer so DCP packing follows the existing staging synchronization path without redundant GPU IDs. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Enable packed transfers by default

f7f4637

Remove the rollout toggle so eligible DCP transfers consistently use the packed path while retaining the buffer-sizing override. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Size pack buffers from chunked_prefill_size

509b247

Drop SGLANG_DISAGG_DCP_PACK_MAX_TOKENS; the pack buffer already cannot exceed a prefill chunk, so the env override was unused. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Clarify packed gather kernel naming

3870b57

Name the KV-cache kernel after its PD-to-DCP gather role so it is distinct from transfer orchestration. Co-authored-by: Cursor <cursoragent@cursor.com>

[DCP] Unify staging memory registration

a6f6bed

Reuse backend staging-memory hooks for DCP pack buffers and document the expected buffer footprint. Co-authored-by: Cursor <cursoragent@cursor.com>

kpham-sgl and others added 2 commits August 26, 2026 18:51

[DCP] Restore packed grouping coverage

609158d

Protect the descriptor-collapse benefit and clarify invalid KV item-length diagnostics. Co-authored-by: Cursor <cursoragent@cursor.com>

kpham-sgl force-pushed the dcp-pd-packed-transfer branch from 6d29344 to c18c9d3 Compare August 27, 2026 01:54

sgl-project deleted a comment from github-actions Bot Aug 27, 2026

ShangmingCai approved these changes Aug 27, 2026

View reviewed changes

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

ShangmingCai added the bypass-fastfail label Aug 27, 2026

foraxe mentioned this pull request Aug 28, 2026

[PD] Add GPUNETIO direct peer rows for DCP #36774

Closed

3 of 6 tasks

YAMY1234 approved these changes Aug 29, 2026

View reviewed changes

kpham-sgl merged commit 3760296 into sgl-project:main Aug 29, 2026
333 of 386 checks passed

kpham-sgl mentioned this pull request Sep 3, 2026

[PD] Transfer the DCP-replicated DSPARK draft KV in DCP1->DCP-N relayouts #37709

Merged

AMD-yanfeiwang mentioned this pull request Sep 7, 2026

[PD][MORI] Enable packed DCP1-to-DCP-N KV transfers #38317

Closed

AMD-yanfeiwang mentioned this pull request Sep 12, 2026

[PD][MORI] Transfer DCP-replicated DSPARK draft KV #39197

Closed

kpham-sgl mentioned this pull request Sep 14, 2026

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

Open

19 of 39 tasks

kalenforn mentioned this pull request Sep 16, 2026

[DSA][DCP] Decode context parallel for the DSA attention family (GLM-5.3 / DeepSeek-V3.2) on Hopper #39330

Open

hrpang0203 mentioned this pull request Sep 16, 2026

[PD][GLM-5.x] Improve PP prefill to DCP decode KV transfer #39768

Open

kpham-sgl mentioned this pull request Sep 19, 2026

[PD] Bound cached-prefix DCP transfers by pack capacity #40376

Merged

4 tasks done

ShangmingCai mentioned this pull request Oct 11, 2026

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

Open

19 of 44 tasks

Navigation Menu

Repository navigation

[PD] Pack DCP1→DCP-N PD KV transfers into dest-contiguous RDMA blocks - #35762

[PD] Pack DCP1→DCP-N PD KV transfers into dest-contiguous RDMA blocks#35762
kpham-sgl merged 16 commits into
sgl-project:mainfrom
kpham-sgl:dcp-pd-packed-transfer

kpham-sgl commented Aug 20, 2026 •

edited by github-actions Bot

Loading

Uh oh!

kpham-sgl commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026 •

edited

Loading

Uh oh!

kpham-sgl commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026 •

edited

Loading

Uh oh!

ShangmingCai left a comment

Uh oh!

ShangmingCai commented Aug 27, 2026

Uh oh!

ShangmingCai commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 2026 •

edited

Loading

Uh oh!

YAMY1234 left a comment

Uh oh!

kpham-sgl commented Aug 29, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

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

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

NIXL

Mooncake

Checklist

Review and Merge Process

CI States

Uh oh!

kpham-sgl commented Aug 27, 2026

Uh oh!

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

Uh oh!

Uh oh!

kpham-sgl commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 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 Aug 27, 2026

Uh oh!

ShangmingCai commented Aug 27, 2026

Uh oh!

github-actions Bot commented Aug 27, 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 Aug 29, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kpham-sgl commented Aug 20, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Aug 27, 2026 •

edited

Loading

github-actions Bot commented Aug 27, 2026 •

edited

Loading

github-actions Bot commented Aug 27, 2026 •

edited

Loading