Reduce CP shared-KV prepare overhead without diagnostic log noise

The CP shared-KV path now has a gated tai-kernel replacement for NSA index
K/scale plus MQA range preparation, and Phase8 prefetch can skip tiny prefixes
that do not cover all CP lanes. The Phase9 plan documents the next scheduler
work for overlapping CP communication with peer-request attention windows.

Temporary diagnostic logs added while validating prefetch ownership and fused
index prepare routing were removed before committing so the runtime path does
not add log-only synchronization, log counters, or shape-reporting overhead.

Constraint: Production profiling showed small per-request CPU/GPU overhead from diagnostic logging and sync-prone debug counters.
Rejected: Keep fused-index prepare fallback/used logs behind a new env var | it leaves another runtime branch and logging surface for a path that should be benchmarked with profiler evidence instead.
Rejected: Keep owned page-count prefetch logs | they require sync-prone tensor reductions and were only useful for one-off diagnosis.
Confidence: medium
Scope-risk: moderate
Directive: Reintroduce CP shared-KV diagnostics only behind explicit debug paths, and avoid .item()/shape-heavy logging in hot prefill paths.
Tested: git diff --check for staged sglang-dev changes.
Tested: AST parse for environ.py, cp_shared_kv_prefetch.py, cp_shared_kv_runtime.py, nsa_indexer.py, and test_cp_shared_kv_runtime.py.
Not-tested: Full unit test suite.
Not-tested: Multi-node GLM5 prefill/decode/router runtime after this exact commit.
This commit is contained in:
laoyao0822
2026-05-12 20:19:11 +08:00
parent c5c30a3f50
commit 96bf7a2594
8 changed files with 839 additions and 19 deletions
@@ -458,6 +458,24 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True):
self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled())
def test_mla_prefetch_min_prefix_pages_defaults_to_cp_size_and_can_override(self):
from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
cp_shared_kv_mla_prefetch_min_prefix_pages,
)
envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.clear()
self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 8)
with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(0):
self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 0)
with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(16):
self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 16)
with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(-2):
self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(4), 4)
def test_fused_mla_store_uses_tai_kernel_when_enabled(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout