Files
sglang/python
laoyao0822 3fc7a5c18c 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.
2026-05-12 20:31:19 +08:00
..