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,