Reduce CP shared-KV request-boundary stalls

CP shared KV now avoids the PyTorch sort/search remap for the single-request current-only path by deriving compact rows from page-level inverse mapping. The same change keeps sort NVTX attribution gated and splits high-frequency MoE sort markers behind a separate env var so profiling does not perturb normal runs.

Decode-side disaggregation prealloc also avoids rebuilding large token index tensors and records finer allocation timing, while compute-owner allocation/free tests cover the shared-KV page-lane behavior.

Constraint: The runtime tree used for validation is the remote /sgl-workspace/sglang-tai mount, which is not itself a Git repository, so these tracked files were synchronized into the local repo before commit.

Rejected: Keep torch.sort/searchsorted for current remap | it emits ATen/CCCL radixSortKVInPlace kernels in the attention hot path.

Rejected: Enable MoE sort NVTX under the generic sort env | the MoE preprocess sort is too frequent and can make profiling look like a hang.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce token-level torch.sort/searchsorted in CP shared-KV current remap without profiling the attention hot path under Nsight.

Tested: Remote container py_compile for modified runtime files; git diff --cached --check.

Not-tested: Full multi-node GLM5 PD throughput/profile rerun after the page-inverse current remap.

This commit is contained in:

laoyao0822

2026-05-05 05:18:35 +08:00

parent a638d71d53

commit 49eaf9ffde

11 changed files with 551 additions and 92 deletions

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

				    SGLANG_FORCE_SHUTDOWN = EnvBool(False)

				    SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)

				    SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)

				    SGLANG_DEBUG_SORT_NVTX = EnvBool(False)

				    SGLANG_DEBUG_MOE_SORT_NVTX = EnvBool(False)

				    SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)

				    SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)

				    SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)