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

Standardize CP shared-KV runtime fallback markers

Browse Source 
Current reuse, TAI materialize, and TAI fused MLA store are optimized CP shared-KV paths. Their fallback helpers already emitted warnings, but the messages were not consistently grep-able with the standard CP shared-KV fallback marker.

This moves the marker into the helper layer so individual fallback call sites do not have to remember to include it, and keeps the existing per-reason rate limiting.

Constraint: Runtime fallback logs must be visible without enabling debug logging.

Rejected: Prefix only selected call sites | helper-level prefixing avoids future unmarked fallback messages.

Confidence: high

Scope-risk: narrow

Directive: New CP shared-KV runtime fallback helpers should use [CP_SHARED_KV_FALLBACK] and a component name.

Tested: Local py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py.

Tested: Local extracted assertLogs check for all three fallback helper prefixes.

Tested: Remote g0034 pytest exact runtime helper tests: 2 passed, 3 warnings.

Not-tested: Full cp_shared_kv_runtime.py suite in this slice.

Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 06:02:13 +08:00
parent e09c8256d7
commit dd69c5970c
3 changed files with 66 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
View File

@@ -477,6 +477,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
self.assertTrue(
any(
"draft_partial_current_reuse" in message
and "[CP_SHARED_KV_FALLBACK][current_reuse]" in message
for message in logs.output
)
)
@@ -489,6 +490,40 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
forward_batch.seq_lens_cpu = torch.tensor([56], dtype=torch.int32)
self.assertTrue(runtime.should_reuse_current_extend_kv(forward_batch))
def test_runtime_fallback_helpers_use_standard_warning_marker(self):
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS.clear()
runtime._TAI_MATERIALIZE_FALLBACK_LOG_COUNTS.clear()
runtime._TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS.clear()
with self.assertLogs(runtime.logger.name, level="WARNING") as logs:
runtime._log_current_reuse_fallback(
"current_reason",
"current fallback %s",
1,
)
runtime._log_tai_materialize_fallback(
"materialize_reason",
"materialize fallback %s",
2,
)
runtime._log_tai_fused_mla_store_fallback(
"store_reason",
"store fallback %s",
3,
)
self.assertIn("[CP_SHARED_KV_FALLBACK][current_reuse]", logs.output[0])
self.assertIn("current_reason", logs.output[0])
self.assertIn("current fallback 1", logs.output[0])
self.assertIn("[CP_SHARED_KV_FALLBACK][tai_materialize]", logs.output[1])
self.assertIn("materialize_reason", logs.output[1])
self.assertIn("materialize fallback 2", logs.output[1])
self.assertIn("[CP_SHARED_KV_FALLBACK][tai_fused_mla_store]", logs.output[2])
self.assertIn("store_reason", logs.output[2])
self.assertIn("store fallback 3", logs.output[2])
def test_current_loc_remap_fast_path_args_only_for_current_only_extend(self):
from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
current_loc_remap_fast_path_args,