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 0e3a9c4e0..881b1b283 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 @@ -21,6 +21,7 @@ _DEBUG_LOG_COUNTS: dict[str, int] = {} _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _TAI_IPC_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {} +_TAI_INDEX_MQA_PREPARE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _CURRENT_REUSE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _SLOT_REMAP_CACHE_LOG_COUNTS: dict[str, int] = {} _TAI_IPC_PEER_PTR_CACHE: dict[tuple[object, ...], torch.Tensor] = {} @@ -240,9 +241,11 @@ def _log_slot_remap_cache_not_reused( return _SLOT_REMAP_CACHE_LOG_COUNTS[log_key] = count + 1 logger.warning( + "[CP_SHARED_KV_FALLBACK][slot_remap_cache] reason=%s " "CP shared KV %s slot remap cache not reused (%s): " "cp_rank=%s cp_size=%s batch_size=%s forward_mode=%s " "cached_key=%s new_key=%s", + reason, kind, reason, layout.cp_rank, @@ -635,6 +638,64 @@ def _log_tai_fused_mla_store_fallback( ) +def _log_tai_index_mqa_prepare_fallback( + key: str, + message: str, + *args, + limit: int = 8, +) -> None: + count = _TAI_INDEX_MQA_PREPARE_FALLBACK_LOG_COUNTS.get(key, 0) + if count >= limit: + return + _TAI_INDEX_MQA_PREPARE_FALLBACK_LOG_COUNTS[key] = count + 1 + logger.warning( + "[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare] reason=%s " + message, + key, + *args, + ) + + +def _tai_materialize_runtime_disabled_reason() -> str | None: + if not cp_shared_kv_tai_materialize_enabled(): + return "env_disabled" + if cp_shared_kv_debug_enabled(): + return "debug_enabled" + return None + + +def _log_tai_materialize_runtime_disabled( + source: str, + *, + ipc: bool = False, + limit: int = 1, +) -> None: + reason = _tai_materialize_runtime_disabled_reason() + if reason is None: + return + if ipc: + _log_tai_ipc_materialize_fallback( + reason, + "CP shared KV tai IPC materialize fast path is disabled; " + "falling back to collective materialize. source=%s " + "tai_materialize_env=%s debug=%s", + source, + cp_shared_kv_tai_materialize_enabled(), + cp_shared_kv_debug_enabled(), + limit=limit, + ) + else: + _log_tai_materialize_fallback( + reason, + "CP shared KV tai materialize fast path is disabled; " + "falling back to torch/collective materialize. source=%s " + "tai_materialize_env=%s debug=%s", + source, + cp_shared_kv_tai_materialize_enabled(), + cp_shared_kv_debug_enabled(), + limit=limit, + ) + + def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor: return tensor if tensor.is_contiguous() else tensor.contiguous() @@ -657,17 +718,63 @@ def try_tai_prepare_cp_mqa_index( """ if not cp_shared_kv_tai_index_mqa_prepare_enabled(): + _log_tai_index_mqa_prepare_fallback( + "env_disabled", + "CP shared KV TAI index MQA prepare fast path is disabled; " + "falling back to GetK/GetS plus torch range prepare. kv_len=%s " + "valid_q_count=%s page_size=%s index_head_dim=%s", + kv_len, + valid_q_count, + page_size, + index_head_dim, + limit=1, + ) return None kernel = _load_tai_index_mqa_prepare_kernel() if kernel is None: + _log_tai_index_mqa_prepare_fallback( + "kernel_missing", + "CP shared KV TAI index MQA prepare kernel is unavailable; " + "falling back to GetK/GetS plus torch range prepare. kv_len=%s " + "valid_q_count=%s page_size=%s index_head_dim=%s", + kv_len, + valid_q_count, + page_size, + index_head_dim, + limit=1, + ) return None if index_buffer.dtype != torch.uint8: + _log_tai_index_mqa_prepare_fallback( + "unsupported_dtype", + "CP shared KV TAI index MQA prepare requires uint8 page buffer; " + "falling back to GetK/GetS. dtype=%s", + index_buffer.dtype, + limit=4, + ) return None if not index_buffer.is_contiguous(): + _log_tai_index_mqa_prepare_fallback( + "non_contiguous", + "CP shared KV TAI index MQA prepare requires contiguous page buffer; " + "falling back to GetK/GetS. shape=%s stride=%s", + tuple(index_buffer.shape), + tuple(index_buffer.stride()), + limit=4, + ) return None if index_head_dim != 128 or page_size != 64: + _log_tai_index_mqa_prepare_fallback( + "unsupported_layout", + "CP shared KV TAI index MQA prepare supports page_size=64 and " + "index_head_dim=128 only; falling back to GetK/GetS. " + "page_size=%s index_head_dim=%s", + page_size, + index_head_dim, + limit=4, + ) return None try: @@ -680,7 +787,16 @@ def try_tai_prepare_cp_mqa_index( page_size=int(page_size), index_head_dim=int(index_head_dim), ) - except Exception: + except Exception as exc: + _log_tai_index_mqa_prepare_fallback( + "kernel_failed", + "CP shared KV TAI index MQA prepare failed; falling back to " + "GetK/GetS plus torch range prepare. error=%s kv_len=%s " + "valid_q_count=%s", + exc, + kv_len, + valid_q_count, + ) return None @@ -697,10 +813,26 @@ def try_tai_prepare_cp_mqa_range( """ if not cp_shared_kv_tai_index_mqa_prepare_enabled(): + _log_tai_index_mqa_prepare_fallback( + "range_env_disabled", + "CP shared KV TAI MQA range fast path is disabled; falling back to " + "torch range prepare. valid_q_count=%s ke_start=%s", + valid_q_count, + ke_start, + limit=1, + ) return None kernel = _load_tai_index_mqa_range_kernel() if kernel is None: + _log_tai_index_mqa_prepare_fallback( + "range_kernel_missing", + "CP shared KV TAI MQA range kernel is unavailable; falling back to " + "torch range prepare. valid_q_count=%s ke_start=%s", + valid_q_count, + ke_start, + limit=1, + ) return None try: @@ -709,7 +841,15 @@ def try_tai_prepare_cp_mqa_range( ke_start=int(ke_start), device=device, ) - except Exception: + except Exception as exc: + _log_tai_index_mqa_prepare_fallback( + "range_kernel_failed", + "CP shared KV TAI MQA range prepare failed; falling back to torch " + "range prepare. error=%s valid_q_count=%s ke_start=%s", + exc, + valid_q_count, + ke_start, + ) return None @@ -732,6 +872,15 @@ def try_tai_fused_mla_store( """ if not cp_shared_kv_tai_fused_mla_store_enabled(): + _log_tai_fused_mla_store_fallback( + "env_disabled", + "CP shared KV tai fused MLA store fast path is disabled; falling " + "back to torch MLA store. layer=%s cp_rank=%s tokens=%s", + getattr(layer, "layer_id", None), + layout.cp_rank, + logical_locs.numel(), + limit=1, + ) return False if cp_shared_kv_debug_enabled(): _log_tai_fused_mla_store_fallback( @@ -787,6 +936,7 @@ def _try_tai_materialize_shared_pages( layout: CpSharedKVLayout, ) -> tuple[torch.Tensor, torch.Tensor] | None: if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("paged_materialize") return None kernels = _load_tai_materialize_kernels() @@ -819,6 +969,7 @@ def _try_tai_materialize_token_kv_pages_and_locs( page_size: int, ) -> tuple[torch.Tensor, torch.Tensor] | None: if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("token_pages_and_locs") return None kernels = _load_tai_materialize_kernels() @@ -859,6 +1010,7 @@ def _try_tai_build_slot_page_inverse( logical_page_capacity: int, ) -> torch.Tensor | None: if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("slot_page_inverse") return None kernels = _load_tai_materialize_kernels() @@ -936,6 +1088,7 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs( mask_non_current_in_current_pages: bool, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("fill_current_kv_page_slots") return None kernels = _load_tai_materialize_kernels() @@ -1091,6 +1244,18 @@ def fill_current_index_page_slots( current_rows = int(current_locs.numel()) if current_rows == 0: return dense_page_buffer + if dense_page_buffer.is_cuda: + _log_tai_materialize_fallback( + "index_current_fill_torch_reference_cuda", + "CP shared KV index current-slot fill is using the torch reference " + "on CUDA; this is a fallback and should not be the steady-state hot " + "path. page_size=%s index_head_dim=%s current_rows=%s dense_pages=%s", + page_size, + index_head_dim, + current_rows, + int(dense_page_buffer.shape[0]), + limit=1, + ) if ( int(current_index_k.shape[0]) < current_rows or int(current_index_scale.shape[0]) < current_rows @@ -1194,6 +1359,7 @@ def _try_tai_materialize_token_kv_page_slots_into( end_slot: int, ) -> bool: if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("token_kv_page_slots_into") return False kernels = _load_tai_materialize_kernels() @@ -1311,6 +1477,7 @@ def _get_or_open_tai_ipc_peer_ptrs( if layout.cp_size <= 1: return None if not _tai_materialize_runtime_enabled(): + _log_tai_materialize_runtime_disabled("open_ipc_peer_ptrs", ipc=True) return None if not tensor.is_cuda or not tensor.is_contiguous(): _log_tai_ipc_materialize_fallback( @@ -1408,8 +1575,30 @@ def _try_tai_ipc_materialize_token_kv_page_slots_into( if start_slot == end_slot: return True if start_slot != 0: + _log_tai_ipc_materialize_fallback( + "start_slot_nonzero", + "CP shared KV tai IPC token materialize only supports slot ranges " + "starting at zero; falling back to local materialize plus " + "collective. cp_rank=%s cp_size=%s start_slot=%s end_slot=%s " + "page_size=%s", + layout.cp_rank, + layout.cp_size, + start_slot, + end_slot, + page_size, + ) return False if not dense_kv_cache.is_cuda or not dense_kv_cache.is_contiguous(): + _log_tai_ipc_materialize_fallback( + "dense_token_tensor_unsupported", + "CP shared KV tai IPC token materialize requires a contiguous CUDA " + "dense tensor; falling back to local materialize plus collective. " + "device=%s contiguous=%s dense_shape=%s", + dense_kv_cache.device, + dense_kv_cache.is_contiguous(), + tuple(dense_kv_cache.shape), + limit=4, + ) return False ipc_state = _get_or_open_tai_ipc_peer_ptrs(kv_cache, layout) @@ -1470,8 +1659,28 @@ def _try_tai_ipc_materialize_paged_buffer_page_slots_into( if start_slot == end_slot: return True if start_slot != 0: + _log_tai_ipc_materialize_fallback( + "paged_start_slot_nonzero", + "CP shared KV tai IPC paged materialize only supports slot ranges " + "starting at zero; falling back to local materialize plus " + "collective. cp_rank=%s cp_size=%s start_slot=%s end_slot=%s", + layout.cp_rank, + layout.cp_size, + start_slot, + end_slot, + ) return False if not dense_page_buffer.is_cuda or not dense_page_buffer.is_contiguous(): + _log_tai_ipc_materialize_fallback( + "dense_paged_tensor_unsupported", + "CP shared KV tai IPC paged materialize requires a contiguous CUDA " + "dense tensor; falling back to local materialize plus collective. " + "device=%s contiguous=%s dense_shape=%s", + dense_page_buffer.device, + dense_page_buffer.is_contiguous(), + tuple(dense_page_buffer.shape), + limit=4, + ) return False ipc_state = _get_or_open_tai_ipc_peer_ptrs(page_buffer, layout) @@ -1562,32 +1771,72 @@ def can_reuse_current_extend_kv(forward_batch) -> bool: enabled, in-seq-split mode, and tensor shape compatibility. """ - if forward_batch is None: - return False - forward_mode = getattr(forward_batch, "forward_mode", None) - if forward_mode is None or not forward_mode.is_extend_without_speculative(): - return False + return _current_extend_kv_reuse_miss_reason(forward_batch) is None - if int(getattr(forward_batch, "batch_size", 0)) != 1: - return False + +def _current_extend_kv_reuse_miss_reason(forward_batch) -> str | None: + if forward_batch is None: + return "missing_forward_batch" + forward_mode = getattr(forward_batch, "forward_mode", None) + if forward_mode is None: + return "missing_forward_mode" + if not forward_mode.is_extend_without_speculative(): + return "not_extend_without_speculative" + + batch_size = int(getattr(forward_batch, "batch_size", 0)) + if batch_size != 1: + return "batch_size_not_one" extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) out_cache_loc = getattr(forward_batch, "out_cache_loc", None) - if extend_seq_lens_cpu is None or seq_lens_cpu is None or out_cache_loc is None: - return False - if len(extend_seq_lens_cpu) != 1 or int(seq_lens_cpu.numel()) != 1: - return False + if extend_seq_lens_cpu is None: + return "missing_extend_seq_lens_cpu" + if seq_lens_cpu is None: + return "missing_seq_lens_cpu" + if out_cache_loc is None: + return "missing_out_cache_loc" + if len(extend_seq_lens_cpu) != 1: + return "extend_batch_not_one" + if int(seq_lens_cpu.numel()) != 1: + return "seq_lens_batch_not_one" extend_len = int(extend_seq_lens_cpu[0]) seq_len = int(seq_lens_cpu[0].item()) - if extend_len <= 0 or seq_len < extend_len: - return False + if extend_len <= 0: + return "non_positive_extend_len" + if seq_len < extend_len: + return "seq_len_smaller_than_extend_len" # ForwardBatch pads tensors such as out_cache_loc at the tail for CUDA graph # and CP alignment. The first extend_len rows still cover the valid current # suffix, so padded batches remain eligible for current reuse as long as the # valid suffix is fully present. - return int(out_cache_loc.numel()) >= extend_len + if int(out_cache_loc.numel()) < extend_len: + return "out_cache_loc_shorter_than_extend" + return None + + +def _log_current_extend_kv_reuse_miss(forward_batch, reason: str) -> None: + extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + out_cache_loc = getattr(forward_batch, "out_cache_loc", None) + _log_current_reuse_fallback( + reason, + "CP shared KV current/partial-current reuse fast path is not used. " + "batch_size=%s forward_mode=%s prefix_lens=%s extend_lens=%s " + "seq_lens_shape=%s out_cache_loc_shape=%s", + getattr(forward_batch, "batch_size", None), + getattr(forward_batch, "forward_mode", None), + [int(x) for x in extend_prefix_lens_cpu] + if extend_prefix_lens_cpu is not None + else None, + [int(x) for x in extend_seq_lens_cpu] + if extend_seq_lens_cpu is not None + else None, + tuple(seq_lens_cpu.shape) if seq_lens_cpu is not None else None, + tuple(out_cache_loc.shape) if out_cache_loc is not None else None, + ) def should_reuse_current_extend_kv(forward_batch) -> bool: @@ -1602,14 +1851,25 @@ def should_reuse_current_extend_kv(forward_batch) -> bool: """ if not cp_shared_kv_current_reuse_enabled(): + _log_current_reuse_fallback( + "env_disabled", + "CP shared KV current/partial-current reuse fast path is disabled " + "by SGLANG_CP_SHARED_KV_CURRENT_REUSE. batch_size=%s forward_mode=%s", + getattr(forward_batch, "batch_size", None), + getattr(forward_batch, "forward_mode", None), + limit=1, + ) return False current_only = is_current_only_extend_batch(forward_batch) if current_only: return True - partial_current = can_reuse_current_extend_kv(forward_batch) - return current_only or partial_current + reason = _current_extend_kv_reuse_miss_reason(forward_batch) + if reason is None: + return True + _log_current_extend_kv_reuse_miss(forward_batch, reason) + return False def current_extend_kv_rows_for_reuse( diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 2b50a3a50..43efdac04 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -371,6 +371,30 @@ class Indexer(MultiPlatformOp): ) if prefetched is not None: return prefetched + _log_cp_shared_kv_index_prefetch_fallback( + "current_consume_miss", + "prefetcher did not provide current index prefix buffer; " + "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), + ) + 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), + ) slot_remap = get_or_build_shared_paged_buffer_slot_remap( forward_batch, page_buffer=index_buffer, @@ -438,6 +462,15 @@ class Indexer(MultiPlatformOp): layout.cp_rank, 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 cp_shared_kv_debug_enabled(): cp_shared_kv_debug_log( diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 33e691adb..5cbb476e0 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -74,6 +74,24 @@ if TYPE_CHECKING: _is_hip = is_hip() logger = logging.getLogger(__name__) _EAGLE_ACCEPT_DRAFT_MLA_PATH_DEBUG_COUNTS: Dict[Tuple[int, int], int] = {} +_CP_SHARED_KV_MLA_PREFETCH_FALLBACK_LOG_COUNTS: Dict[str, int] = {} + + +def _log_cp_shared_kv_mla_prefetch_fallback( + reason: str, + message: str, + *args, + limit: int = 8, +) -> None: + count = _CP_SHARED_KV_MLA_PREFETCH_FALLBACK_LOG_COUNTS.get(reason, 0) + if count >= limit: + return + _CP_SHARED_KV_MLA_PREFETCH_FALLBACK_LOG_COUNTS[reason] = count + 1 + logger.warning( + "[CP_SHARED_KV_FALLBACK][mla_prefetch] reason=%s " + message, + reason, + *args, + ) if _is_hip: from sglang.srt.layers.attention.nsa.triton_kernel import get_valid_kv_indices @@ -1956,6 +1974,19 @@ class NativeSparseAttnBackend( else None ) page_size = int(forward_batch.token_to_kv_pool.page_size) + _log_cp_shared_kv_mla_prefetch_fallback( + reason, + "MLA prefetch fast path did not provide partial-current " + "prefix; using synchronous compose. cp_rank=%s " + "layer=%s prefix_lens=%s extend_lens=%s " + "current_rows=%s page_table_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + prefix_lens, + extend_lens, + int(current_kv_cache.shape[0]), + tuple(logical_page_table_1.shape), + ) if ( prefix_lens_cpu is None or len(prefix_lens_cpu) != 1 @@ -2070,6 +2101,20 @@ class NativeSparseAttnBackend( eagle_draft_mla_used_prefetch = True kv_cache, page_table_1 = prefetched_kv else: + reason = ( + "missing_prefetcher" + if mla_prefetcher is None + else "prefetch_consume_returned_none" + ) + _log_cp_shared_kv_mla_prefetch_fallback( + reason, + "MLA prefetch fast path did not provide full materialize " + "buffer; using synchronous materialize. cp_rank=%s " + "layer=%s page_table_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + tuple(page_table_1.shape), + ) slot_remap = get_or_build_shared_token_kv_slot_remap( forward_batch, kv_cache=kv_cache, 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 d6067f481..a9f91add6 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 @@ -885,7 +885,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS.clear() runtime._TAI_MATERIALIZE_FALLBACK_LOG_COUNTS.clear() + runtime._TAI_IPC_MATERIALIZE_FALLBACK_LOG_COUNTS.clear() runtime._TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS.clear() + runtime._TAI_INDEX_MQA_PREPARE_FALLBACK_LOG_COUNTS.clear() with self.assertLogs(runtime.logger.name, level="WARNING") as logs: runtime._log_current_reuse_fallback( @@ -903,6 +905,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): "store fallback %s", 3, ) + runtime._log_tai_ipc_materialize_fallback( + "ipc_reason", + "ipc fallback %s", + 4, + ) + runtime._log_tai_index_mqa_prepare_fallback( + "index_reason", + "index fallback %s", + 5, + ) self.assertIn("[CP_SHARED_KV_FALLBACK][current_reuse]", logs.output[0]) self.assertIn("current_reason", logs.output[0]) @@ -913,6 +925,115 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): 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]) + self.assertIn("[CP_SHARED_KV_FALLBACK][tai_ipc_materialize]", logs.output[3]) + self.assertIn("ipc_reason", logs.output[3]) + self.assertIn("ipc fallback 4", logs.output[3]) + self.assertIn("[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare]", logs.output[4]) + self.assertIn("index_reason", logs.output[4]) + self.assertIn("index fallback 5", logs.output[4]) + + def test_current_reuse_fast_path_miss_logs_warning(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + runtime._CURRENT_REUSE_FALLBACK_LOG_COUNTS.clear() + forward_batch = SimpleNamespace( + forward_mode=_FakeExtendForwardMode(), + batch_size=2, + extend_prefix_lens_cpu=[64, 64], + extend_seq_lens_cpu=[64, 64], + seq_lens_cpu=torch.tensor([128, 128], dtype=torch.int32), + out_cache_loc=torch.arange(128, dtype=torch.int64), + ) + + with envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.override(True): + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + self.assertFalse(runtime.should_reuse_current_extend_kv(forward_batch)) + + joined = "\n".join(logs.output) + self.assertIn("[CP_SHARED_KV_FALLBACK][current_reuse]", joined) + self.assertIn("batch_size_not_one", joined) + + def test_tai_index_mqa_prepare_fast_path_miss_logs_warning(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + runtime._TAI_INDEX_MQA_PREPARE_FALLBACK_LOG_COUNTS.clear() + + with envs.SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE.override(True): + with patch.object( + runtime, "_load_tai_index_mqa_prepare_kernel", return_value=None + ): + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + self.assertIsNone( + runtime.try_tai_prepare_cp_mqa_index( + index_buffer=torch.empty((4, 32), dtype=torch.uint8), + page_indices=torch.tensor([1], dtype=torch.int64), + kv_len=64, + valid_q_count=1, + ke_start=0, + page_size=64, + index_head_dim=128, + ) + ) + + joined = "\n".join(logs.output) + self.assertIn("[CP_SHARED_KV_FALLBACK][tai_index_mqa_prepare]", joined) + self.assertIn("kernel_missing", joined) + + def test_tai_ipc_token_materialize_fast_path_miss_logs_warning(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + runtime._TAI_IPC_MATERIALIZE_FALLBACK_LOG_COUNTS.clear() + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0) + + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + self.assertFalse( + runtime._try_tai_ipc_materialize_token_kv_page_slots_into( + kv_cache=torch.empty((16, 1), dtype=torch.float32), + dense_kv_cache=torch.empty((16, 1), dtype=torch.float32), + slot_logical_pages=torch.tensor([1, 2], dtype=torch.int64), + layout=layout, + page_size=4, + start_slot=1, + end_slot=2, + ) + ) + + joined = "\n".join(logs.output) + self.assertIn("[CP_SHARED_KV_FALLBACK][tai_ipc_materialize]", joined) + self.assertIn("start_slot_nonzero", joined) + + def test_mla_prefetch_sync_compose_paths_log_warning_in_source(self): + from pathlib import Path + + source = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa_backend.py" + ).read_text() + + self.assertIn("def _log_cp_shared_kv_mla_prefetch_fallback", source) + self.assertIn("[CP_SHARED_KV_FALLBACK][mla_prefetch]", source) + self.assertIn( + "MLA prefetch fast path did not provide partial-current", source + ) + self.assertIn( + "MLA prefetch fast path did not provide full materialize", source + ) + + def test_index_prefetch_sync_compose_paths_log_warning_in_source(self): + from pathlib import Path + + source = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py" + ).read_text() + + self.assertIn("[CP_SHARED_KV_FALLBACK][index_prefetch]", source) + self.assertIn("current_missing_prefetcher", source) + self.assertIn("current_consume_miss", source) + self.assertIn("missing_prefetcher", source) 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 ( @@ -1881,7 +2002,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=layout, page_size=4, ) - warning.assert_not_called() + self.assertFalse( + any( + "[CP_SHARED_KV_FALLBACK][slot_remap_cache]" in str(call) + for call in warning.call_args_list + ) + ) with self.assertLogs(runtime.logger.name, level="WARNING") as logs: remap_b = runtime.get_or_build_shared_token_kv_slot_remap( @@ -1993,7 +2119,12 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): logical_pages=logical_pages_a, layout=layout, ) - warning.assert_not_called() + self.assertFalse( + any( + "[CP_SHARED_KV_FALLBACK][slot_remap_cache]" in str(call) + for call in warning.call_args_list + ) + ) with self.assertLogs(runtime.logger.name, level="WARNING") as logs: remap_b = runtime.get_or_build_shared_paged_buffer_slot_remap(