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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user