Move CP shared-KV prefetch reduce into attention overlap

Phase8 prefetch was starting local materialization and the CP all-reduce from the same late hook, after rank-skewed MQA/topk/materialize work had already separated CP ranks. This commit separates local prefix materialization from reduce launch, starts next-layer prefetch earlier in MLA prepare, and explicitly launches pending MLA/index reduces at attention boundaries so the collectives can overlap the intended window instead of drifting behind skewed ranks.

Constraint: Prefetch remains optional and gated by the existing CP shared-KV prefetch environment controls.

Constraint: Profiling required per-source materialize all-reduce ranges, so NVTX labels are gated behind SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX.

Rejected: Launch all-reduce immediately in start_next_layer_prefix | this reproduced the old skew-sensitive timing and could put prefetch behind current-layer work.

Rejected: Remove later hooks entirely | they are still needed as fallbacks when the early MLA prepare hook is bypassed.

Confidence: medium

Scope-risk: moderate

Directive: Preserve the consume-time fallback launch; otherwise missed launch paths silently lose correctness or overlap.

Tested: Remote g0034 docker py_compile for changed SGLang files; python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -q (50 passed).

Not-tested: Full GLM5 multi-node prefill/decode profile after this exact commit.

This commit is contained in:

laoyao0822

2026-05-12 20:31:19 +08:00

parent aa27a444f6

commit 3fc7a5c18c

7 changed files with 850 additions and 107 deletions

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

				    SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False)

				    SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)

				    SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)

				    SGLANG_CP_SHARED_KV_MATERIALIZE_NVTX = EnvBool(False)

				    SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)

				    SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)

				    SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)