From 937f89ef8988f569e65075b1f37993c9f0305173 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Tue, 2 Jun 2026 23:52:22 +0800 Subject: [PATCH] Avoid index prefetch fallback noise when disabled The missing index prefetcher state is expected when MLA/index prefetch is disabled. Logging it as an index-prefetch fallback made the normal synchronous materialize path look like a failed fast path and created noisy reports. This keeps consume-miss warnings for enabled prefetchers, but suppresses missing-prefetcher warnings when the prefetch feature gate is off. Constraint: Index prefetch is controlled by the same MLA prefetch enable gate in the current pipeline. Rejected: Suppress all index prefetch fallback warnings | consume misses still indicate an enabled prefetch pipeline failed to serve a requested buffer. Confidence: high Scope-risk: narrow Directive: Missing prefetcher is not a fallback when SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0; do not re-add warning noise without a separate enable gate. Tested: Local py_compile for nsa_indexer.py and test_cp_shared_kv_runtime.py. Tested: Remote g0034 targeted regression set passed 4 tests: disabled sync materialize no warning, consume miss warning, first-layer no-warning, enabled create-skip warning. Not-tested: Full CUDA ETE run with logs. --- ..._prefill_cp_page_aligned_cache_contract.md | 27 ++++++++++ .../srt/layers/attention/nsa/nsa_indexer.py | 49 ++++++++++++------- .../mem_cache/test_cp_shared_kv_runtime.py | 44 +++++++++++++++++ 3 files changed, 101 insertions(+), 19 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 e03d11f15..d7bacded2 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 @@ -5462,3 +5462,30 @@ Remaining risk: requires higher-level eviction/allocation policy to preserve larger free runs; the allocator now avoids the worst CPU scan cost but cannot create physical contiguity that the free set does not contain. + +### C127 — 2026-06-02 Disabled index prefetch must not emit fallback warnings + +Finding: + +- `CpSharedKVIndexPrefetcher.maybe_create()` correctly returns `None` without a + warning when `SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0`. +- The later synchronous materialization path in `nsa_indexer.py` still logged + `[CP_SHARED_KV_FALLBACK][index_prefetch] reason=missing_prefetcher` whenever + `forward_batch.cp_shared_kv_index_prefetcher` was `None`. +- With MLA/index prefetch disabled, `None` is the expected state, so this warning + made normal sync materialize look like a failed fast path and spammed logs. + +Correction: + +- Gate only the `missing_prefetcher` / `current_missing_prefetcher` warnings on + `cp_shared_kv_mla_prefetch_enabled()`. +- Keep `consume_miss` warnings unchanged: if a prefetcher exists but cannot serve + the requested layer/current buffer, that is still a real enabled-prefetch miss. + +Validation: + +- Remote RED/GREEN test added: + `test_index_sync_materialize_does_not_warn_when_prefetch_disabled`. +- Remote targeted regression set passed (`4 passed`): disabled sync materialize, + consume miss warning, first-layer no-warning, and create-skip warning when + enabled. diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 43efdac04..aabd191d2 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -17,6 +17,7 @@ from sglang.srt.layers.attention.nsa import index_buf_accessor from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, + cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, @@ -131,6 +132,14 @@ def _log_cp_shared_kv_index_prefetch_fallback( ) +def _should_log_missing_index_prefetcher() -> bool: + # When MLA/index prefetch is disabled, a missing index prefetcher is the + # expected sync-materialize path, not a fast-path fallback. Still log consume + # misses when a prefetcher exists because those indicate an enabled prefetch + # pipeline failed to provide the requested layer/current buffer. + return cp_shared_kv_mla_prefetch_enabled() + + class BaseIndexerMetadata(ABC): @abstractmethod def get_seqlens_int32(self) -> torch.Tensor: @@ -384,17 +393,18 @@ class Indexer(MultiPlatformOp): tuple(logical_page_table.shape), ) else: - _log_cp_shared_kv_index_prefetch_fallback( - "current_missing_prefetcher", - "index prefetcher is unavailable; falling back to sync " - "partial-current index compose. layer=%s cp_rank=%s " - "prefix_lens=%s extend_lens=%s logical_page_table_shape=%s", - layer_id, - layout.cp_rank, - prefix_lens, - extend_lens, - tuple(logical_page_table.shape), - ) + if _should_log_missing_index_prefetcher(): + _log_cp_shared_kv_index_prefetch_fallback( + "current_missing_prefetcher", + "index prefetcher is unavailable; falling back to sync " + "partial-current index compose. layer=%s cp_rank=%s " + "prefix_lens=%s extend_lens=%s logical_page_table_shape=%s", + layer_id, + layout.cp_rank, + prefix_lens, + extend_lens, + tuple(logical_page_table.shape), + ) slot_remap = get_or_build_shared_paged_buffer_slot_remap( forward_batch, page_buffer=index_buffer, @@ -463,14 +473,15 @@ class Indexer(MultiPlatformOp): tuple(logical_page_table.shape), ) else: - _log_cp_shared_kv_index_prefetch_fallback( - "missing_prefetcher", - "index prefetcher is unavailable; falling back to sync paged " - "materialize. layer=%s cp_rank=%s logical_page_table_shape=%s", - layer_id, - layout.cp_rank, - tuple(logical_page_table.shape), - ) + if _should_log_missing_index_prefetcher(): + _log_cp_shared_kv_index_prefetch_fallback( + "missing_prefetcher", + "index prefetcher is unavailable; falling back to sync paged " + "materialize. layer=%s cp_rank=%s logical_page_table_shape=%s", + layer_id, + layout.cp_rank, + tuple(logical_page_table.shape), + ) if cp_shared_kv_debug_enabled(): cp_shared_kv_debug_log( 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 a9f91add6..d14b47a61 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 @@ -3658,6 +3658,50 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): ) self.assertIn("cuda_unavailable_or_stream_capturing", logger.call_args.args[1]) + def test_index_sync_materialize_does_not_warn_when_prefetch_disabled(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import nsa_indexer + + class FakePool: + page_size = 4 + + def get_index_k_with_scale_buffer(self, layer_id): + return torch.arange(0, 12, dtype=torch.float32).view(4, 3) + + class FakeLayout: + page_size = 4 + cp_size = 1 + cp_rank = 0 + + fallback_buffer = torch.full((3, 3), 9.0) + fallback_pages = torch.tensor([[1, 2]], dtype=torch.int32) + forward_batch = SimpleNamespace( + token_to_kv_pool=FakePool(), + uses_cp_shared_kv=True, + cp_shared_kv_layout=FakeLayout(), + cp_shared_kv_index_prefetcher=None, + ) + logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) + indexer = object.__new__(nsa_indexer.Indexer) + + with envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.override(False), patch.object( + nsa_indexer, + "materialize_shared_paged_buffer", + return_value=(fallback_buffer, fallback_pages), + ), patch( + "sglang.srt.layers.attention.nsa.nsa_indexer.logger", + create=True, + ) as logger: + dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id=7, + logical_page_table=logical_pages, + ) + + self.assertIs(dense_buffer, fallback_buffer) + self.assertIs(dense_pages, fallback_pages) + logger.warning.assert_not_called() + if __name__ == "__main__": unittest.main()