From dd69c5970cc4ed3b163faffdb018ecefa4eae2ec Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 06:02:13 +0800 Subject: [PATCH] Standardize CP shared-KV runtime fallback markers 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 --- ..._prefill_cp_page_aligned_cache_contract.md | 16 +++++++++ .../attention/nsa/cp_shared_kv_runtime.py | 18 ++++++++-- .../mem_cache/test_cp_shared_kv_runtime.py | 35 +++++++++++++++++++ 3 files changed, 66 insertions(+), 3 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 4bef096c8..a55717e9b 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -812,6 +812,22 @@ Implemented C11 slice: - Tightened the existing unit expectation so repeated fallback events are still emitted, but now require the production-visible fallback prefix. +Additional C11 finding: + +- Runtime fallback helpers for current reuse, TAI materialize, and TAI fused MLA + store already log at WARNING, but their messages do not use the standard + `[CP_SHARED_KV_FALLBACK]` marker. These helpers cover intended optimized hot + paths, so they should standardize the marker instead of requiring every caller + to remember the prefix. + +Implemented C11 runtime-helper slice: + +- Standardized current-reuse, TAI materialize, and TAI fused MLA store fallback + helpers to prefix warnings with `[CP_SHARED_KV_FALLBACK]` and a component + name. +- Added unit expectations for those helper prefixes and tightened the draft + partial-current-reuse fallback check to require the standard marker. + ### C12. Verification gaps before claiming completion Current state: diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 38c5fd0ab..4474e232b 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -143,7 +143,11 @@ def _log_current_reuse_fallback( if count >= limit: return _CURRENT_REUSE_FALLBACK_LOG_COUNTS[key] = count + 1 - logger.warning(message, *args) + logger.warning( + "[CP_SHARED_KV_FALLBACK][current_reuse] reason=%s " + message, + key, + *args, + ) @dataclass(frozen=True) @@ -337,7 +341,11 @@ def _log_tai_materialize_fallback( if count >= limit: return _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS[key] = count + 1 - logger.warning(message, *args) + logger.warning( + "[CP_SHARED_KV_FALLBACK][tai_materialize] reason=%s " + message, + key, + *args, + ) def _log_tai_fused_mla_store_fallback( @@ -350,7 +358,11 @@ def _log_tai_fused_mla_store_fallback( if count >= limit: return _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS[key] = count + 1 - logger.warning(message, *args) + logger.warning( + "[CP_SHARED_KV_FALLBACK][tai_fused_mla_store] reason=%s " + message, + key, + *args, + ) def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor: diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 134efadad..2dab40eba 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -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,