chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Gate tiny CP shared-KV prefetch by token threshold

Browse Source 
Prefix materialize prefetch has fixed launch and coordination overhead, so tiny cache-hit prefixes should stay on the simpler synchronous path. The default threshold is now expressed as a cached import-time token threshold while the existing page override remains available for workload-specific tuning.

Both MLA and index prefix prefetchers use the same page-size-aware threshold, so CP shared-KV prefix prefetch is enabled only when the prefix covers at least one page per CP lane and at least the configured token threshold.

Constraint: The default token threshold must be read once to avoid hot-path env lookups

Constraint: Existing page threshold override must continue to disable or force the gate

Rejected: Hard-code 1024 tokens | tuning needs to be possible without code changes

Rejected: Read token threshold on every maybe_create call | unnecessary hot-path env lookup

Confidence: medium

Scope-risk: narrow

Directive: Keep MLA and index prefetch threshold logic shared; do not let their gates diverge

Tested: Not rerun per user request; prior py_compile, diff check, and isolated env-cache check had passed before commit

Not-tested: full pytest in local environment due missing optional dependencies/kernel import constraints

Not-tested: GLM5 E2E throughput after this threshold change
This commit is contained in:
laoyao0822
2026-05-28 08:57:45 +08:00
parent 25f2147677
commit 26c792939d
4 changed files with 64 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+26 -8
View File
@@ -718,23 +718,41 @@ 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):
def test_mla_prefetch_min_prefix_pages_defaults_to_1k_tokens_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,
)
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.clear()
self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 8)
self.assertEqual(runtime._MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS, 1024)
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64), 16
)
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(32, page_size=64), 32
)
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)
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64), 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)
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64),
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)
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(4, page_size=64),
16,
)
with patch.object(runtime, "_MLA_PREFETCH_DEFAULT_MIN_PREFIX_TOKENS", 2048):
self.assertEqual(
runtime.cp_shared_kv_mla_prefetch_min_prefix_pages(8, page_size=64),
32,
)
def test_fused_mla_store_uses_tai_kernel_when_enabled(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime