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.

This commit is contained in:

laoyao0822

2026-04-26 23:29:34 +08:00

parent 5af232e9de

commit d015e3fb01

5 changed files with 270 additions and 25 deletions

									
										python/sglang/srt/environ.py
									
		+1
		
												View File
												
				@@ -203,6 +203,7 @@ class Envs:

				    SGLANG_FORCE_SHUTDOWN = EnvBool(False)

				    SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)

				    SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)

				    SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)

				    SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)

				    SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)

				    SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)