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 ffc707934..98411989d 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 @@ -241,12 +241,21 @@ class SharedPagedBufferSlotRemap: ) -def _tensor_identity_key(tensor: torch.Tensor) -> tuple[int, tuple[int, ...], str, str]: +def _tensor_identity_key( + tensor: torch.Tensor, +) -> tuple[int, tuple[int, ...], tuple[int, ...], int, str, str, int | None]: + try: + version = int(tensor._version) + except Exception: + version = None return ( - int(tensor.untyped_storage().data_ptr()), + int(tensor.data_ptr()), tuple(int(dim) for dim in tensor.shape), + tuple(int(stride) for stride in tensor.stride()), + int(tensor.storage_offset()), str(tensor.dtype), str(tensor.device), + version, ) 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 ffe1f7f03..fa8cd4045 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 @@ -4641,6 +4641,41 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) ) + def test_token_slot_remap_cache_distinguishes_same_storage_views(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(batch_size=1, forward_mode="extend") + kv_cache = torch.zeros((64, 1, 1), dtype=torch.float32) + backing_pages = torch.tensor( + [ + [1, 2, 3], + [9, 10, 11], + ], + dtype=torch.int64, + ) + logical_pages_a = backing_pages[:1] + logical_pages_b = backing_pages[1:] + + 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, + ) + 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.assertEqual(remap_b.slot_logical_pages.tolist(), [9, 10, 11]) + 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 @@ -4715,6 +4750,39 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): 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_distinguishes_same_storage_views(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(batch_size=1, forward_mode="extend") + page_buffer = torch.zeros((16, 3), dtype=torch.uint8) + backing_pages = torch.tensor( + [ + [1, 2, 3], + [9, 10, 11], + ], + dtype=torch.int64, + ) + logical_pages_a = backing_pages[:1] + logical_pages_b = backing_pages[1:] + + remap_a = runtime.get_or_build_shared_paged_buffer_slot_remap( + forward_batch, + page_buffer=page_buffer, + logical_pages=logical_pages_a, + layout=layout, + ) + 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.assertEqual(remap_b.slot_logical_pages.tolist(), [9, 10, 11]) + 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