Files
sglang/python
laoyao0822 d015e3fb01 Reduce CP shared KV current-chunk materialization
Phase 2 shared KV keeps persistent KV/index sharded across CP ranks but
uses a full-view compatibility layer before NSA topk and MLA attention.
For current-only prefill chunks, the current KV/index tensors have already
been CP all-gathered and reranged before being written to the sharded
persistent pool. This change adds a guarded current-reuse path that remaps
logical current locs to compact tensor rows and skips the shared-KV
materialize path for current-only MLA and NSA indexer reads.

Constraint: Existing NSA/MLA kernels still consume full-view/compact page tables; history and mixed current/history batches must keep the Phase 2 fallback.
Rejected: Make all history attention shard-aware in this patch | that requires global topk merge and distributed sparse attention and belongs to a later phase.
Confidence: medium
Scope-risk: moderate
Reversibility: clean
Directive: Do not remove the Phase 2 fallback until mixed/history shared-KV paths have correctness and performance coverage.
Tested: python -m py_compile on modified Python files
Tested: git diff --check on staged modified files
Not-tested: local pytest collection is blocked by missing pybase64 in this environment.
Not-tested: full long-context chunked prefill/decode performance in this commit step.
2026-04-26 23:29:34 +08:00
..