diff --git a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md index f95ed9fe8..d9c393187 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md @@ -18,6 +18,14 @@ Layer L attention 返回时不强制等待 prefetch Layer L+1 consume prefetched buffer 时 wait event,再补齐 suffix/current pages ``` +后续在 HiCache 分支补充了 batch-scoped slot remap 复用:`metadata.real_page_table` +和 index page table 在同一个 forward 内跨 layer 不变,因此只在第一次 materialize / +prefetch 时构造 `slot_logical_pages`、`page_inverse`、paged `dense_pages`,后续 layer +复用该计划;每层仍重新 materialize 本层 KV/index 数据并执行 CP all-reduce。这个优化 +只复用“逻辑页到 dense slot 的映射计划”,不复用 dense KV 内容,避免跨层 KV 数据串用。 +如果同一个 `ForwardBatch` 已经带有 remap cache 但无法复用(key mismatch 或 cache +状态不完整),运行时会打印 rate-limited warning;第一次冷启动构造 cache 不打日志。 + 保留的环境变量: ```text diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 5100f4e7b..8abd6e190 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -9,15 +9,14 @@ import torch from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( _all_reduce_materialized_buffer_async, _all_reduce_materialized_buffer_range, - build_shared_token_kv_slot_remap, - build_slot_page_inverse_optimized, - build_slot_page_remap, cp_shared_kv_debug_enabled, cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_should_log_layer, filter_locs_mappable_to_physical_pool, filter_pages_mappable_to_physical_pool, + get_or_build_shared_paged_buffer_slot_remap, + get_or_build_shared_token_kv_slot_remap, materialize_local_paged_buffer_page_slots_into, materialize_local_token_kv_page_slots_into, remap_logical_pages_to_slot_dense_pages, @@ -195,9 +194,9 @@ class CpSharedKVMlaPrefetcher: try: first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) kv_cache = token_to_kv_pool.get_key_buffer(first_layer_id) - remap = build_shared_token_kv_slot_remap( + remap = get_or_build_shared_token_kv_slot_remap( + forward_batch, kv_cache=kv_cache, - logical_locs=None, remap_logical_pages=real_page_table, layout=layout, page_size=page_size, @@ -612,18 +611,11 @@ class CpSharedKVIndexPrefetcher: page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer( layer_id=first_layer_id ) - remap_logical_pages = filter_pages_mappable_to_physical_pool( + remap = get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=page_buffer, logical_pages=real_page_table, layout=layout, - physical_page_capacity=page_buffer.shape[0], - ) - slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages) - logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * ( - layout.cp_size - ) + 1 - page_inverse = build_slot_page_inverse_optimized( - slot_logical_pages, - logical_page_capacity=logical_page_capacity, ) except Exception as exc: _index_prefetch_fallback_log( @@ -639,17 +631,17 @@ class CpSharedKVIndexPrefetcher: layout.cp_rank, layout.cp_size, prefix_pages, - int(slot_logical_pages.numel()), - int(slot_logical_pages.numel()) + 1, + int(remap.slot_logical_pages.numel()), + remap.dense_num_pages, page_size, ) return cls( layout=layout, prefix_pages=prefix_pages, - slot_logical_pages=slot_logical_pages, - page_inverse=page_inverse, - dense_num_pages=int(slot_logical_pages.numel()) + 1, + slot_logical_pages=remap.slot_logical_pages, + page_inverse=remap.page_inverse, + dense_num_pages=remap.dense_num_pages, ) def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool: 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 78bfe146c..6bd3644a4 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 @@ -3,6 +3,7 @@ from __future__ import annotations import logging from dataclasses import dataclass from functools import lru_cache +from typing import Any import torch @@ -15,6 +16,7 @@ logger = logging.getLogger(__name__) _DEBUG_LOG_COUNTS: dict[str, int] = {} _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {} _TAI_FUSED_MLA_STORE_FALLBACK_LOG_COUNTS: dict[str, int] = {} +_SLOT_REMAP_CACHE_LOG_COUNTS: dict[str, int] = {} _MLA_PREFETCH_LOG_PROBE_LAYER = 2 _SORT_NVTX_ENABLED = envs.SGLANG_DEBUG_SORT_NVTX.get() @@ -65,6 +67,102 @@ class SharedTokenKVSlotRemap: dense_num_pages: int +@dataclass(frozen=True) +class SharedPagedBufferSlotRemap: + slot_logical_pages: torch.Tensor + page_inverse: torch.Tensor + dense_pages: torch.Tensor + logical_page_capacity: int + dense_num_pages: int + + +def _tensor_identity_key(tensor: torch.Tensor) -> tuple[int, tuple[int, ...], str, str]: + return ( + int(tensor.untyped_storage().data_ptr()), + tuple(int(dim) for dim in tensor.shape), + str(tensor.dtype), + str(tensor.device), + ) + + +def _slot_remap_cache_key( + *, + logical_pages: torch.Tensor, + physical_page_capacity: int, + layout: CpSharedKVLayout, + page_size: int, +) -> tuple[object, ...]: + return ( + _tensor_identity_key(logical_pages), + int(physical_page_capacity), + int(layout.page_size), + int(layout.cp_size), + int(layout.cp_rank), + int(page_size), + ) + + +def _log_slot_remap_cache_not_reused( + *, + kind: str, + reason: str, + cached_key: tuple[object, ...] | None, + new_key: tuple[object, ...], + forward_batch: Any, + layout: CpSharedKVLayout, + limit: int = 8, +) -> None: + log_key = f"{kind}:{reason}" + count = _SLOT_REMAP_CACHE_LOG_COUNTS.get(log_key, 0) + if count >= limit: + return + _SLOT_REMAP_CACHE_LOG_COUNTS[log_key] = count + 1 + logger.warning( + "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", + kind, + reason, + layout.cp_rank, + layout.cp_size, + getattr(forward_batch, "batch_size", None), + getattr(forward_batch, "forward_mode", None), + cached_key, + new_key, + ) + + +def _maybe_log_slot_remap_cache_not_reused( + *, + kind: str, + cached_key: tuple[object, ...] | None, + cached: object | None, + new_key: tuple[object, ...], + forward_batch: Any, + layout: CpSharedKVLayout, +) -> None: + if cached is None and cached_key is None: + # First build for this forward batch. This is the expected cold path. + return + if cached is None: + reason = "missing_cached_value" + elif cached_key is None: + reason = "missing_cached_key" + elif cached_key != new_key: + reason = "key_mismatch" + else: + return + + _log_slot_remap_cache_not_reused( + kind=kind, + reason=reason, + cached_key=cached_key, + new_key=new_key, + forward_batch=forward_batch, + layout=layout, + ) + + @lru_cache(maxsize=1) def _load_tai_materialize_kernels(): try: @@ -863,6 +961,126 @@ def build_shared_token_kv_slot_remap( ) +def build_shared_paged_buffer_slot_remap( + page_buffer: torch.Tensor, + logical_pages: torch.Tensor, + layout: CpSharedKVLayout, +) -> SharedPagedBufferSlotRemap: + """Build the fixed slot-layout remap used by shared paged-buffer materialize.""" + + _debug_assert_no_negative_tensor_values( + logical_pages, + context="CP shared KV paged materialize", + tensor_name="logical_pages", + ) + logical_pages = filter_pages_mappable_to_physical_pool( + logical_pages=logical_pages, + layout=layout, + physical_page_capacity=page_buffer.shape[0], + ) + slot_logical_pages, dense_pages = build_slot_page_remap(logical_pages) + logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * ( + layout.cp_size + ) + 1 + page_inverse = build_slot_page_inverse_optimized( + slot_logical_pages, + logical_page_capacity=logical_page_capacity, + ) + return SharedPagedBufferSlotRemap( + slot_logical_pages=slot_logical_pages, + page_inverse=page_inverse, + dense_pages=dense_pages, + logical_page_capacity=logical_page_capacity, + dense_num_pages=int(slot_logical_pages.numel()) + 1, + ) + + +def get_or_build_shared_token_kv_slot_remap( + forward_batch: Any, + *, + kv_cache: torch.Tensor, + remap_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + page_size: int, +) -> SharedTokenKVSlotRemap: + """Cache token-KV materialize slot metadata on a forward batch. + + The logical page table and local physical pool capacity are batch-scoped, + not layer-scoped. The actual KV rows are still materialized per layer, but + the page inverse used to map layer-specific top-k locs into the dense view + can be reused across all layers in the same forward pass. + """ + + physical_page_capacity = kv_cache.shape[0] // page_size + key = _slot_remap_cache_key( + logical_pages=remap_logical_pages, + physical_page_capacity=physical_page_capacity, + layout=layout, + page_size=page_size, + ) + cached_key = getattr(forward_batch, "cp_shared_kv_token_slot_remap_key", None) + cached = getattr(forward_batch, "cp_shared_kv_token_slot_remap", None) + if cached is not None and cached_key == key: + return cached + _maybe_log_slot_remap_cache_not_reused( + kind="token", + cached_key=cached_key, + cached=cached, + new_key=key, + forward_batch=forward_batch, + layout=layout, + ) + + remap = build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=None, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=page_size, + ) + forward_batch.cp_shared_kv_token_slot_remap_key = key + forward_batch.cp_shared_kv_token_slot_remap = remap + return remap + + +def get_or_build_shared_paged_buffer_slot_remap( + forward_batch: Any, + *, + page_buffer: torch.Tensor, + logical_pages: torch.Tensor, + layout: CpSharedKVLayout, +) -> SharedPagedBufferSlotRemap: + """Cache paged-buffer materialize slot metadata on a forward batch.""" + + key = _slot_remap_cache_key( + logical_pages=logical_pages, + physical_page_capacity=page_buffer.shape[0], + layout=layout, + page_size=layout.page_size, + ) + cached_key = getattr(forward_batch, "cp_shared_kv_paged_slot_remap_key", None) + cached = getattr(forward_batch, "cp_shared_kv_paged_slot_remap", None) + if cached is not None and cached_key == key: + return cached + _maybe_log_slot_remap_cache_not_reused( + kind="paged", + cached_key=cached_key, + cached=cached, + new_key=key, + forward_batch=forward_batch, + layout=layout, + ) + + remap = build_shared_paged_buffer_slot_remap( + page_buffer=page_buffer, + logical_pages=logical_pages, + layout=layout, + ) + forward_batch.cp_shared_kv_paged_slot_remap_key = key + forward_batch.cp_shared_kv_paged_slot_remap = remap + return remap + + def remap_logical_locs_to_dense_locs( logical_locs: torch.Tensor, unique_logical_pages: torch.Tensor, @@ -1483,6 +1701,7 @@ def materialize_shared_token_kv_buffer( page_size: int, remap_logical_locs: torch.Tensor | None = None, remap_logical_pages: torch.Tensor | None = None, + slot_remap: SharedTokenKVSlotRemap | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: _debug_assert_no_tensor_values_below( logical_locs, @@ -1511,7 +1730,15 @@ def materialize_shared_token_kv_buffer( ) dense_kv_cache = None - if remap_logical_pages is None: + if slot_remap is not None: + materialized_logical_pages = slot_remap.slot_logical_pages + dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=slot_remap.page_inverse, + page_size=page_size, + ) + use_slot_materialize = True + elif remap_logical_pages is None: remap_pages_from_locs = logical_pages_from_locs(remap_logical_locs, page_size) materialized_logical_pages, _ = build_dense_page_remap(remap_pages_from_locs) dense_locs = remap_logical_locs_to_dense_locs( @@ -1624,6 +1851,7 @@ def materialize_shared_paged_buffer( page_buffer: torch.Tensor, logical_pages: torch.Tensor, layout: CpSharedKVLayout, + slot_remap: SharedPagedBufferSlotRemap | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: _debug_assert_no_negative_tensor_values( logical_pages, @@ -1640,16 +1868,24 @@ def materialize_shared_paged_buffer( logical_pages=logical_pages, layout=layout, ) - if tai_result is None: - materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages) + if tai_result is not None: + dense_page_buffer, dense_pages = tai_result + materialized_logical_pages = logical_pages.reshape(-1) + elif slot_remap is not None: + materialized_logical_pages = slot_remap.slot_logical_pages + dense_pages = slot_remap.dense_pages dense_page_buffer = materialize_local_paged_buffer_page_slots( page_buffer=page_buffer, slot_logical_pages=materialized_logical_pages, layout=layout, ) else: - dense_page_buffer, dense_pages = tai_result - materialized_logical_pages = logical_pages.reshape(-1) + materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages) + dense_page_buffer = materialize_local_paged_buffer_page_slots( + page_buffer=page_buffer, + slot_logical_pages=materialized_logical_pages, + layout=layout, + ) if cp_shared_kv_debug_enabled(): owned_pages = materialized_logical_pages[ diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 12b64efde..18f75eb6e 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -19,6 +19,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_debug_log, cp_shared_kv_current_reuse_enabled, filter_owned_logical_locs, + get_or_build_shared_paged_buffer_slot_remap, is_current_only_extend_batch, materialize_shared_paged_buffer, tensor_debug_checksum, @@ -344,10 +345,17 @@ class Indexer(MultiPlatformOp): layer_id, tensor_debug_summary(logical_page_table), ) + slot_remap = get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=index_buffer, + logical_pages=logical_page_table, + layout=layout, + ) materialized, dense_pages = materialize_shared_paged_buffer( page_buffer=index_buffer, logical_pages=logical_page_table, layout=layout, + slot_remap=slot_remap, ) 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 cdb909a00..e333031bf 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -22,6 +22,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, filter_owned_logical_locs, + get_or_build_shared_token_kv_slot_remap, is_current_only_extend_batch, materialize_shared_token_kv_buffer, tensor_debug_checksum, @@ -1730,6 +1731,13 @@ class NativeSparseAttnBackend( if prefetched_kv is not None: kv_cache, page_table_1 = prefetched_kv else: + slot_remap = get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=metadata.real_page_table, + layout=forward_batch.cp_shared_kv_layout, + page_size=forward_batch.token_to_kv_pool.page_size, + ) kv_cache, page_table_1 = materialize_shared_token_kv_buffer( kv_cache=kv_cache, logical_locs=page_table_1, @@ -1737,6 +1745,7 @@ class NativeSparseAttnBackend( remap_logical_pages=metadata.real_page_table, layout=forward_batch.cp_shared_kv_layout, page_size=forward_batch.token_to_kv_pool.page_size, + slot_remap=slot_remap, ) if mla_prefetcher is not None: diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 1ca1a0a95..4a553ee63 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -425,6 +425,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): cp_local_out_cache_loc: Optional[torch.Tensor] = None cp_local_physical_out_cache_loc: Optional[torch.Tensor] = None cp_shared_mla_direct_write_done: bool = False + cp_shared_kv_token_slot_remap_key: Optional[Any] = None + cp_shared_kv_token_slot_remap: Optional[Any] = None + cp_shared_kv_paged_slot_remap_key: Optional[Any] = None + cp_shared_kv_paged_slot_remap: Optional[Any] = None cp_shared_kv_mla_prefetcher: Optional[Any] = None cp_shared_kv_index_prefetcher: Optional[Any] = None 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 ad44e50df..89d3c7588 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 @@ -252,6 +252,196 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(remap.dense_num_pages, 7) self.assertEqual(remap.dense_locs.tolist(), [[4, 8, -1], [16, 0, 20]]) + def test_forward_batch_token_slot_remap_is_cached_across_layers(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 + + layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0) + forward_batch = SimpleNamespace() + kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1) + remap_logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64) + + remap_a = runtime.get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=4, + ) + remap_b = runtime.get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=4, + ) + + self.assertIs(remap_a, remap_b) + + with patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x): + dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=torch.tensor([4, 20, -1], dtype=torch.int64), + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=4, + slot_remap=remap_a, + ) + + self.assertEqual(dense_locs.tolist(), [4, 12, -1]) + self.assertEqual(list(dense_kv.shape), [16, 1, 1]) + self.assertTrue(torch.equal(dense_kv[4:8], kv_cache[4:8])) + self.assertTrue(torch.equal(dense_kv[12:16], kv_cache[20:24])) + + def test_token_slot_remap_cache_miss_logs_after_warm_cache(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._SLOT_REMAP_CACHE_LOG_COUNTS.clear() + + layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0) + forward_batch = SimpleNamespace(batch_size=1, forward_mode="extend") + kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32) + logical_pages_a = torch.tensor([[1, 2, 5]], dtype=torch.int64) + logical_pages_b = torch.tensor([[1, 2, 6]], dtype=torch.int64) + + with patch.object(runtime.logger, "warning") as warning: + remap_a = runtime.get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=logical_pages_a, + layout=layout, + page_size=4, + ) + warning.assert_not_called() + + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + remap_b = runtime.get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=logical_pages_b, + layout=layout, + page_size=4, + ) + + self.assertIsNot(remap_a, remap_b) + self.assertTrue( + any( + "token slot remap cache not reused (key_mismatch)" in message + for message in logs.output + ) + ) + + def test_token_slot_remap_incomplete_cache_state_logs(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._SLOT_REMAP_CACHE_LOG_COUNTS.clear() + + layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0) + forward_batch = SimpleNamespace(batch_size=1, forward_mode="extend") + kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32) + logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64) + forward_batch.cp_shared_kv_token_slot_remap_key = runtime._slot_remap_cache_key( + logical_pages=logical_pages, + physical_page_capacity=kv_cache.shape[0] // 4, + layout=layout, + page_size=4, + ) + forward_batch.cp_shared_kv_token_slot_remap = None + + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + runtime.get_or_build_shared_token_kv_slot_remap( + forward_batch, + kv_cache=kv_cache, + remap_logical_pages=logical_pages, + layout=layout, + page_size=4, + ) + + self.assertTrue( + any( + "token slot remap cache not reused (missing_cached_value)" + in message + for message in logs.output + ) + ) + + def test_forward_batch_paged_slot_remap_is_cached_across_layers(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 + + layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0) + forward_batch = SimpleNamespace() + page_buffer = torch.arange(0, 8 * 3, dtype=torch.uint8).view(8, 3) + logical_pages = torch.tensor([[1, 2, 5]], dtype=torch.int64) + + remap_a = runtime.get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=page_buffer, + logical_pages=logical_pages, + layout=layout, + ) + remap_b = runtime.get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=page_buffer, + logical_pages=logical_pages, + layout=layout, + ) + + self.assertIs(remap_a, remap_b) + + with patch.object(runtime, "_all_reduce_materialized_buffer", lambda x, _: x): + dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer( + page_buffer=page_buffer, + logical_pages=logical_pages, + layout=layout, + slot_remap=remap_a, + ) + + self.assertEqual(dense_pages.tolist(), [[1, 2, 3]]) + self.assertEqual(list(dense_page_buffer.shape), [4, 3]) + self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1])) + self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[2])) + self.assertTrue(torch.equal(dense_page_buffer[3], page_buffer[5])) + + def test_paged_slot_remap_cache_miss_logs_after_warm_cache(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._SLOT_REMAP_CACHE_LOG_COUNTS.clear() + + layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0) + forward_batch = SimpleNamespace(batch_size=1, forward_mode="extend") + page_buffer = torch.zeros((8, 3), dtype=torch.uint8) + logical_pages_a = torch.tensor([[1, 2, 5]], dtype=torch.int64) + logical_pages_b = torch.tensor([[1, 2, 6]], dtype=torch.int64) + + with patch.object(runtime.logger, "warning") as warning: + remap_a = runtime.get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=page_buffer, + logical_pages=logical_pages_a, + layout=layout, + ) + warning.assert_not_called() + + with self.assertLogs(runtime.logger.name, level="WARNING") as logs: + remap_b = runtime.get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=page_buffer, + logical_pages=logical_pages_b, + layout=layout, + ) + + self.assertIsNot(remap_a, remap_b) + self.assertTrue( + any( + "paged slot remap cache not reused (key_mismatch)" in message + for message in logs.output + ) + ) + def test_mla_prefetch_log_env_defaults_to_off_and_can_enable(self): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( @@ -1350,6 +1540,8 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): return self.index_buffer class FakeLayout: + page_size = 4 + cp_size = 1 cp_rank = 3 class MissingPrefetcher: @@ -1403,6 +1595,8 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): return self.index_buffer class FakeLayout: + page_size = 4 + cp_size = 1 cp_rank = 0 class MissingPrefetcher: