Reduce CP shared KV materialize and direct-write overhead

Shared KV now relies on page-aligned CP metadata and compute-owner page allocation so persistent MLA KV and NSA index shards can be written by the rank that computed them. The compatibility read path keeps the dense full-view contract for existing topk and attention kernels, but removes duplicated prev/next index materialize, adds optional tai materialize integration, and tightens tests/docs around the fallback boundaries.

Constraint: Decode remains non-CP while prefill CP owns the shared-KV changes

Constraint: Existing attention/topk kernels still expect dense full-view KV/index inputs

Rejected: Change attention kernels to read owner-sharded KV directly | larger semantic change reserved for later phases

Rejected: Merge index K/scale storage with MLA KV storage | would couple topk and attention cache lifecycles before materialize overhead is isolated

Confidence: medium

Scope-risk: broad

Directive: Do not remove fallback logging or debug-gated assertions without reproducing long-context chunked/radix-hit paths

Tested: git diff --check --cached

Not-tested: Local pytest/runtime server verification not run in this commit step per current workflow constraints

This commit is contained in:

laoyao0822

2026-05-02 07:07:28 +08:00

parent 2317952a01

commit 5769b63082

17 changed files with 2524 additions and 96 deletions

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

				    SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)

				    SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)

				    SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)

				    SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)

				    SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)

				    SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)

				    SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)