Repository navigation
Conversation
Expose the flag only after scratch and scheduler contracts exist. Resolve supported configurations, reserve scratch HBM, select sharded pools and allocator, and require the rotation-aware cache together. (cherry picked from commit 40fa526)
(cherry picked from commit f025d20)
(cherry picked from commit 93c9783)
…ests Last of four PRs for logical-page KV cache sharding; stacked on "[kv-shard 3/4] Enable the feature". A sharded prefill rank holds only its stripe of every cached page, so each prefill rank ships the pages it owns and decode pulls from every rank in the shard group. UnifiedRadixCache preserves the logical Full-KV page ids and rotation invariant consumed by this tree-independent transfer path. Contents: - `disaggregation/utils.py`: add `filter_kv_indices_for_shard_rank`. Ownership comes from the value (`owner = logical_page % N`), independent of placement order. Owned entries go on the wire as `logical_page // N`, the owner's local page id. The contiguous destination slice becomes an explicit positions array for positional pairing. - `disaggregation/common/conn.py`: resolve and advertise the prefill shard identity; make decode pull from all CP ranks or register every prefill TP rank as a real sender; reject decode layouts that would require DCP relayout. - `disaggregation/nixl/conn.py`: emit a standalone notification when a rank owns no page in a non-final chunk, preventing a missing chunk id from stalling the receiver. - `disaggregation/prefill.py`: widen the bootstrap capacity bound, require a page-aligned decode-cached prefix, retain scheduler-stage metrics, and run busy invariant checks in both disaggregated prefill loops. - `disaggregation/mooncake/conn.py`: keep transfer metrics aligned with bytes actually sent when replicated state transfer is suppressed. Tests cover the send partition and wire ids, Mooncake transfer accounting, NIXL empty-chunk notification, and bootstrap flag propagation. The new 4-GPU MLA and 8-GPU GQA/CP disaggregated E2Es exercise both shard axes; the GQA case uses the canonical `--enable-prefill-cp --cp-strategy zigzag` configuration. (cherry picked from commit 9f6b072)
(cherry picked from commit ee639a5)
(cherry picked from commit f176b89)
Upstream sgl-project#41469 removed get_index_k_continuous and get_index_k_scale_continuous (and IndexKeyCache.get_k_continuous / get_k_scale_continuous) as orphaned. The sharded index cache still overrode the IndexKeyCache pair, and the indexer test called the pool pair, so CI on the main-merged branch failed with AttributeError. Drop the two overrides and read through get_index_k_scale_buffer, which still translates logical page IDs into the scratch plan and keeps the skip-topk and begin_shard_extend assertions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit e33582f)
Since sgl-project#42178, slots_per_page = page_size // index_kpool, but the sharded pool translates indexer locations in token units. With index_kpool > 1 that would address the wrong rows, so fail at construction instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> (cherry picked from commit 5a82ac9)
(cherry picked from commit 01fba66)
Preserve the final test coverage from Shunkangz/sglang PR sgl-project#40929 commits ea0efc1, 16c7162 and 156b0db while retaining current-main, PD and DSA tests. Keep current-main import boundaries in the pool implementations. Co-authored-by: shunkangz <182541032+Shunkangz@users.noreply.github.com>
Extend the existing sharded Mooncake protocol to route packed DSA indexer pages by owner. Preserve source/destination positions, reject mismatched page counts, and exclude empty indexer layers from RDMA descriptors. Count each sender only for its owned state pages. Add CPU byte-exact descriptor tests for CP4/CP8, target plus draft buffers, empty owners, prefix-hit final chunks, layout rejection and unsharded compatibility.
This was referenced Oct 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Page-sharded DSA prefill stores different indexer pages on each CP rank. The replicated-state shortcut must not suppress nonzero owners, and target/draft indexer pages must use the same source-position pairing as latent KV. Otherwise the decode worker can receive incomplete or incorrectly paired indexer state.
This increment extends Shunkang's existing PD transfer implementation:
Stack and attribution
This is a draft, dependency-stacked PR based on upstream main
492dbd0807d21e0faa72a55a2c09f541cb4696b6. It includes the prerequisites from Shunkang's #40102, #38473, #40925 and #40929, retaining their author attribution and original commit references. A small test-consolidation commit reconciles their final tests with current main.Review this increment: HZY-Wade/sglang@codex/kvshard-dsa-dependencies...codex/kvshard-dsa-pd
The functional increment is commit
a22db8cd58. The broader Files changed view also contains the unmerged upstream dependencies; those are not new implementations claimed by this PR. DSA/MTP runtime activation follows in #42972 and CPU HiCache support in #42973. The stack will be rebased as its dependencies land.Validation
git diff --checkpassed.No changes to the non-sharded state-transfer path are intended.