[Feat] KVShard: PD transfer and e2e test by Shunkangz · Pull Request #38473 · sgl-project/sglang

Shunkangz · 2026-09-08T09:12:04Z

Motivation

The PR that makes the feature work end to end.

A sharded prefill rank holds only its stripe of every cached page, so the transfer changes on both ends: each prefill rank ships the pages it owns, and decode pulls from every rank of the shard group instead of one. Ownership is derived from the value (owner = logical_page % N), not the position, so the filter stays correct under any allocator placement policy — and the wire value is already the owner's local page id, so the decode side needs no change.

Validated PD-disaggregated on both arms: MLA sharded across attention-TP (GSM8K 0.82) and GQA under prefill CP sharded across attention-CP (GSM8K 0.935), both added here as registered E2E tests.

Modifications

Accuracy Tests

Speed Tests and Profiling

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 #39209811480
Latest PR Test (Extra): ✅ Run #39209811542
Latest PR Test (AMD ROCm 10): ❌ Run #39209811722

test_page_interleave_shard.py: main (sgl-project#40925) added multi-GPU gather and DSA indexer tests to the file this branch extends. Keep both: the branch's CPU tests (TestShardedPoolFootprint, TestShardSendFilter) stay with the other CPU tests before main's multi-GPU section, the docstring lists CPU items 1-7 then main's GPU paragraph, and the imports are the union of both sides. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Shunkangz · 2026-10-12T07:31:52Z

/rerun-failed-ci

Shunkangz · 2026-10-14T04:45:57Z

/rerun-failed-ci

_resolve_rank_mapping read info.enable_kv_cache_sharding before checking is_mla_backend. A non-MLA decode now stops at is_mla_backend, so it reads only the fields main reads. This fixes the hetero-TP SWA wire test, whose SimpleNamespace info has no sharding field. Behavior is unchanged for a real PrefillServerInfo, where the field defaults to False.

Shunkangz · 2026-10-15T02:08:53Z

/rerun-failed-ci

hzh0425 · 2026-10-15T18:28:18Z

/rerun-failed-ci

alphabetc1

the rest LGTM

Squashed dependency: the net change of sgl-project#38473 (open upstream) relative to main, squashed into one commit so this PR's own commits stand alone. Credit for this commit belongs to the sgl-project#38473 author; review it in sgl-project#38473. A sharded prefill rank holds only its stripe of every KV page, so each rank sends the pages it owns and decode pulls from every rank of the shard group. - disaggregation/utils.py: filter_kv_indices_for_shard_rank (logical page l lives on rank l % N as local page l // N). - disaggregation/common/conn.py: filter every send by the shard rank, skip empty non-final chunks, advertise enable_kv_cache_sharding through bootstrap; decode pulls from all prefill ranks of the shard group. - disaggregation/prefill.py: widen the bootstrap capacity bound by the shard size. - Read the sharding flag only on the MLA rank-mapping path. - Tests: send filter, empty-chunk skip, bootstrap payload, per-rank KV memory cut, and a 4-GPU PD E2E gating GSM8K accuracy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…r fit The PD bootstrap check rejected only prompts longer than N*T (N shard ranks, T tokens per rank). Sharded PrefillAdder admission needs prompt + 1 new token + one page per shard < N*T, strictly, so prompts in [N*T - N*ps - 1, N*T - 7] passed intake and bootstrap but were refused with NO_TOKEN on every pass, even on an empty pool. Under FCFS such a request blocked the prefill queue until decode's 300 s KV-wait timeout turned it into a 500. With tp=2, page 64 and 16384 tokens per rank, that is 32639..32761. The bootstrap check now rejects prompts longer than N*T - N*ps - 2 (32638 there) with a 400. Unsharded runs keep their limit. Test: a CPU unit test drives the real sharded PrefillAdder on an empty PageInterleavePoolAllocator and the real bootstrap check at the limit and one past it, for N in {2, 4}. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Sharding assigns whole logical pages to ranks; it does not split every page into per-rank stripes. Fix the comments and the E2E docstring that said otherwise. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…sion path With the radix cache disabled, PrefillAdder admits ignore_eos requests through add_one_req_ignore_eos. That path rounds the prompt up to a page and keeps 1 new token plus IGNORE_EOS_RESERVE_TOKENS spare, so on an empty pool it admits only ceil_page(L) <= N*T - N*ps - 3. The previous limit, N*T - N*ps - 2, still let [N*T - (N+1)*ps + 1, N*T - N*ps - 2] wait forever (32577..32638 for tp=2, page 64, 16384 tokens per rank). The limit is now floor_page(N*T - N*ps - 3), 32576 there. A sweep over N in {2, 4, 8} and page sizes {2, 16, 64} matches the largest prompt both real admission paths admit. The boundary test now also drives the ignore_eos path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

hzh0425 · 2026-10-19T08:50:11Z

/rerun-group disaggregation

github-actions · 2026-10-19T08:50:36Z

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

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

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

cd test/ && python3 registered/disaggregation/test_disaggregation_basic.py
cd test/ && python3 registered/disaggregation/test_disaggregation_chunked_prefill_abort.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

⛔ registered/disaggregation/test_disaggregation_xpu.py: test/registered/disaggregation/test_disaggregation_xpu.py is registered for XPU (suite stage-b-test-1-gpu-xpu), not for CUDA or CPU; rerun-test.yml has no XPU job. Rerun it with /rerun-failed-ci, or dispatch the XPU workflow manually.

Squashed dependency: the net change of sgl-project#38473 (open upstream) relative to main, squashed into one commit so this PR's own commits stand alone. Credit for this commit belongs to the sgl-project#38473 author; review it in sgl-project#38473. Each sharded prefill rank owns a subset of logical KV pages, so each rank sends the pages it owns and decode pulls from every rank of the shard group. - disaggregation/utils.py: filter_kv_indices_for_shard_rank (logical page l lives on rank l % N as local page l // N). - disaggregation/common/conn.py: filter every send by the shard rank, skip empty non-final chunks, advertise enable_kv_cache_sharding through bootstrap; decode pulls from all prefill ranks of the shard group. - disaggregation/prefill.py: widen the bootstrap capacity bound by the shard size, and reject at bootstrap any prompt that sharded prefill admission can never fit (bounded by the stricter ignore_eos path). - Read the sharding flag only on the MLA rank-mapping path. - Tests: send filter, empty-chunk skip, bootstrap payload, bootstrap limit vs sharded admission boundary, per-rank KV memory cut, and a 4-GPU PD E2E gating GSM8K accuracy. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

ShangmingCai

All changes are correct and safe, LGTM

Shunkangz requested review from ByronHsu, Duyi-Wang, Fridge003, HaiShaw, Qiaolin-Yu, ShangmingCai, Ying1123, alphabetc1, ch-wan, fzyzcjy, hanming-lu, hebiao064, hnyls2002, huangtingwei9988, hzh0425, ispobock, merrymercy, sogalin, xiezhq-hermann and yizhang2077 as code owners September 8, 2026 09:12

Shunkangz mentioned this pull request Sep 8, 2026

KV Cache Shard with Sequence Split #30501

Open

7 of 9 tasks

github-actions Bot added blackwell SM100/SM120 unified-radix-cache memory-pool labels Sep 8, 2026

huangtingwei9988 added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) labels Sep 23, 2026

Fridge003 mentioned this pull request Sep 28, 2026

[Roadmap] Prefill Context Parallelism (2026 Q4) #21788

Open

15 of 21 tasks

This was referenced Oct 7, 2026

[KV Sharding][1/3] Transfer DSA target and draft indexer state over PD #42969

Closed

[KV Sharding][2/3] Enable GLM DSA and MTP prefill runtime #42972

Closed

[Feat] KVShard: Integrate GLM DSA/MTP with PD and CPU HiCache #42973

Merged

Shunkangz added 2 commits October 14, 2026 15:34

Merge branch 'main' into kvshard/4-pd-transfer-e2e

871f670

Merge branch 'main' into kvshard/4-pd-transfer-e2e

8edc149

hzh0425 self-assigned this Oct 15, 2026

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

Merge branch 'main' into kvshard/4-pd-transfer-e2e

fe2cfa8

hzh0425 approved these changes Oct 15, 2026

View reviewed changes

alphabetc1 reviewed Oct 16, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/prefill.py Resolved

Comment thread python/sglang/srt/disaggregation/common/conn.py Resolved

alphabetc1 approved these changes Oct 16, 2026

View reviewed changes

AMD-yanfeiwang mentioned this pull request Oct 16, 2026

[Feat] [AMD] kv-shard: Support KV cache sharding with aiter and MoRI #44660

Open

4 tasks done

Merge branch 'main' into kvshard/4-pd-transfer-e2e

58ad21d

This comment has been minimized.

Sign in to view

fivechenxi mentioned this pull request Oct 19, 2026

GLM-5.3 6000D：7K 多前缀稳态优化补丁与复现记录 fivechenxi/sglang#28

Open

5 of 6 tasks

Shunkangz changed the title ~~[kv-shard 4/4] PD transfer and e2e test~~ [Feat] KVShard: PD transfer and e2e test Oct 19, 2026

Shunkangz and others added 3 commits October 18, 2026 23:16

[kv-shard 4/4] Say each rank owns a subset of logical KV pages

cbdcfcc

Sharding assigns whole logical pages to ranks; it does not split every page into per-rank stripes. Fix the comments and the E2E docstring that said otherwise. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

hzh0425 approved these changes Oct 19, 2026

View reviewed changes

hzh0425 mentioned this pull request Oct 19, 2026

Agentic Cache Scheduling System Roadmap (Q4) #43098

Open

ShangmingCai approved these changes Oct 20, 2026

View reviewed changes

hzh0425 merged commit eb4b167 into sgl-project:main Oct 20, 2026
295 of 312 checks passed

Navigation Menu

Repository navigation

[Feat] KVShard: PD transfer and e2e test - #38473

[Feat] KVShard: PD transfer and e2e test#38473
hzh0425 merged 11 commits into
sgl-project:mainfrom
Shunkangz:kvshard/4-pd-transfer-e2e

Shunkangz commented Sep 8, 2026 •

edited by github-actions Bot

Loading

Uh oh!

Shunkangz commented Oct 12, 2026

Uh oh!

Shunkangz commented Oct 14, 2026

Uh oh!

Shunkangz commented Oct 15, 2026

Uh oh!

hzh0425 commented Oct 15, 2026

Uh oh!

Uh oh!

Uh oh!

alphabetc1 left a comment

Uh oh!

This comment has been minimized.

hzh0425 commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 2026 •

edited

Loading

Uh oh!

ShangmingCai left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

Shunkangz commented Sep 8, 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

Checklist

Review and Merge Process

CI States

Uh oh!

Shunkangz commented Oct 12, 2026

Uh oh!

Shunkangz commented Oct 14, 2026

Uh oh!

Shunkangz commented Oct 15, 2026

Uh oh!

hzh0425 commented Oct 15, 2026

Uh oh!

Uh oh!

Uh oh!

alphabetc1 left a comment

Choose a reason for hiding this comment

Uh oh!

This comment has been minimized.

hzh0425 commented Oct 19, 2026

Uh oh!

github-actions Bot commented Oct 19, 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!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Shunkangz commented Sep 8, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Oct 19, 2026 •

edited

Loading