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 563e68b80..2c41b412c 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 @@ -211,13 +211,10 @@ def build_dense_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, t """Build a dense per-call page remap for shared-KV runtime materialization.""" dense_pages = logical_pages.clone() positive_mask = logical_pages > 0 - if not torch.any(positive_mask): - empty = logical_pages.new_empty((0,)) - return empty, dense_pages - - unique_logical_pages = torch.unique(logical_pages[positive_mask], sorted=True) + positive_pages = logical_pages[positive_mask] + unique_logical_pages = torch.unique(positive_pages, sorted=True) dense_pages[positive_mask] = remap_logical_pages_to_dense_pages( - logical_pages[positive_mask], + positive_pages, unique_logical_pages=unique_logical_pages, ) return unique_logical_pages, dense_pages @@ -229,15 +226,11 @@ def remap_logical_pages_to_dense_pages( ) -> torch.Tensor: dense_pages = logical_pages.clone() positive_mask = logical_pages > 0 - if not torch.any(positive_mask): - return dense_pages - - if unique_logical_pages.numel() == 0: - raise ValueError("unique_logical_pages is empty but logical_pages contains data") - positive_pages = logical_pages[positive_mask] insert_positions = torch.searchsorted(unique_logical_pages, positive_pages) - if insert_positions.numel() > 0: + if cp_shared_kv_debug_enabled() and insert_positions.numel() > 0: + if unique_logical_pages.numel() == 0: + raise ValueError("unique_logical_pages is empty but logical_pages contains data") if torch.any(insert_positions >= unique_logical_pages.numel()): raise ValueError("logical_pages contains entries outside unique_logical_pages") if not torch.equal(unique_logical_pages[insert_positions], positive_pages): @@ -247,6 +240,117 @@ def remap_logical_pages_to_dense_pages( return dense_pages +def build_slot_page_remap(logical_pages: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Build a fixed-shape page remap without dynamic CUDA output ops. + + The compact remap path uses boolean compaction + unique/searchsorted. Those + ops have data-dependent output shapes and can cause device-to-host + synchronization in CUDA. The slot remap instead gives every input page-table + slot a deterministic dense page id (`flat_slot + 1`) and preserves 0/-1 + sentinels. It can materialize duplicate logical pages into duplicate dense + slots, trading extra device work for avoiding CPU synchronization. + """ + + slot_logical_pages = logical_pages.reshape(-1) + dense_pages_flat = slot_logical_pages.clone() + if slot_logical_pages.numel() == 0: + return slot_logical_pages, dense_pages_flat.reshape(logical_pages.shape) + + slot_ids = torch.arange( + 1, + slot_logical_pages.numel() + 1, + device=logical_pages.device, + dtype=logical_pages.dtype, + ) + dense_pages_flat = torch.where( + slot_logical_pages > 0, + slot_ids, + dense_pages_flat, + ) + return slot_logical_pages, dense_pages_flat.reshape(logical_pages.shape) + + +def _logical_page_capacity_from_physical_page_capacity( + physical_page_capacity: int, + layout: CpSharedKVLayout, +) -> int: + # Physical page 0 is the shared dummy page. Real physical pages + # 1..N correspond to N * cp_size logical pages across the CP group. + return max(physical_page_capacity - 1, 0) * layout.cp_size + 1 + + +def build_slot_page_inverse( + slot_logical_pages: torch.Tensor, + logical_page_capacity: int, +) -> torch.Tensor: + """Build logical_page -> slot_dense_page map without unique/searchsorted.""" + + page_inverse = torch.full( + (logical_page_capacity,), + -1, + device=slot_logical_pages.device, + dtype=torch.long, + ) + if logical_page_capacity == 0: + return page_inverse + + # Page 0 is the dummy/padding page and always maps to dense page 0. + page_inverse[0] = 0 + if slot_logical_pages.numel() == 0: + return page_inverse + + flat_pages = slot_logical_pages.reshape(-1).to(torch.long) + slot_ids = torch.arange( + 1, + flat_pages.numel() + 1, + device=flat_pages.device, + dtype=torch.long, + ) + valid_pages = (flat_pages > 0) & (flat_pages < logical_page_capacity) + safe_pages = torch.where(valid_pages, flat_pages, torch.zeros_like(flat_pages)) + safe_slot_ids = torch.where(valid_pages, slot_ids, torch.zeros_like(slot_ids)) + page_inverse.scatter_(0, safe_pages, safe_slot_ids) + return page_inverse + + +def remap_logical_locs_to_slot_dense_locs( + logical_locs: torch.Tensor, + page_inverse: torch.Tensor, + page_size: int, +) -> torch.Tensor: + """Map logical token locs through a fixed-shape slot page inverse.""" + + dense_locs = torch.full_like(logical_locs, -1) + if logical_locs.numel() == 0 or page_inverse.numel() == 0: + return dense_locs + + locs_long = logical_locs.to(torch.long) + valid_locs = locs_long >= 0 + safe_locs = torch.where(valid_locs, locs_long, torch.zeros_like(locs_long)) + logical_pages = torch.div(safe_locs, page_size, rounding_mode="floor") + offsets = torch.remainder(safe_locs, page_size) + pages_in_range = logical_pages < page_inverse.numel() + safe_pages = torch.clamp(logical_pages, max=page_inverse.numel() - 1) + dense_pages = page_inverse[safe_pages] + mapped = valid_locs & pages_in_range & (dense_pages >= 0) + + if cp_shared_kv_debug_enabled() and torch.any( + valid_locs & pages_in_range & (dense_pages < 0) + ): + missing_pages = logical_pages[valid_locs & pages_in_range & (dense_pages < 0)] + raise RuntimeError( + "CP shared KV slot remap got logical locs outside remap_logical_pages. " + f"missing_page_min={int(missing_pages.min().item())} " + f"missing_page_max={int(missing_pages.max().item())} " + f"logical_locs={tensor_debug_summary(logical_locs)}" + ) + + dense_values = dense_pages.to(logical_locs.dtype) * page_size + offsets.to( + logical_locs.dtype + ) + return torch.where(mapped, dense_values, dense_locs) + + def remap_logical_locs_to_dense_locs( logical_locs: torch.Tensor, unique_logical_pages: torch.Tensor, @@ -254,9 +358,6 @@ def remap_logical_locs_to_dense_locs( ) -> torch.Tensor: dense_locs = logical_locs.clone() valid_mask = logical_locs >= 0 - if not torch.any(valid_mask): - return dense_locs - valid_locs = logical_locs[valid_mask] logical_pages = torch.div(valid_locs, page_size, rounding_mode="floor") offsets = torch.remainder(valid_locs, page_size) @@ -314,12 +415,12 @@ def build_current_loc_remap( def logical_pages_from_locs(logical_locs: torch.Tensor, page_size: int) -> torch.Tensor: logical_pages = logical_locs.clone() valid_mask = logical_locs >= 0 - if torch.any(valid_mask): - logical_pages[valid_mask] = torch.div( - logical_locs[valid_mask], - page_size, - rounding_mode="floor", - ) + valid_locs = logical_locs[valid_mask] + logical_pages[valid_mask] = torch.div( + valid_locs, + page_size, + rounding_mode="floor", + ) return logical_pages @@ -419,14 +520,13 @@ def materialize_local_token_kv_pages( page_size: int, ) -> torch.Tensor: dense_num_pages = int(unique_logical_pages.numel()) + 1 - dense_kv_cache = kv_cache.new_zeros((dense_num_pages * page_size, *kv_cache.shape[1:])) + dense_kv_cache = kv_cache.new_zeros( + (dense_num_pages * page_size, *kv_cache.shape[1:]) + ) if unique_logical_pages.numel() == 0: return dense_kv_cache owned_mask = layout.owned_pages_mask(unique_logical_pages) - if not torch.any(owned_mask): - return dense_kv_cache - owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64) owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to( torch.long @@ -439,6 +539,44 @@ def materialize_local_token_kv_pages( return dense_kv_cache +def materialize_local_token_kv_page_slots( + kv_cache: torch.Tensor, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + page_size: int, +) -> torch.Tensor: + """Materialize slot-remapped token KV with fixed-shape device ops only.""" + + dense_num_pages = int(slot_logical_pages.numel()) + 1 + dense_kv_cache = kv_cache.new_zeros( + (dense_num_pages * page_size, *kv_cache.shape[1:]) + ) + if slot_logical_pages.numel() == 0: + return dense_kv_cache + + logical_pages = slot_logical_pages.reshape(-1).to(torch.long) + owned_mask = layout.owned_pages_mask(logical_pages) + physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long) + safe_physical_pages = torch.where( + owned_mask, + physical_pages, + torch.zeros_like(physical_pages), + ) + page_offsets = torch.arange(page_size, device=kv_cache.device, dtype=torch.long) + src_tokens = (safe_physical_pages[:, None] * page_size + page_offsets).reshape(-1) + + dense_body = dense_kv_cache[page_size:].view( + dense_num_pages - 1, + page_size, + *kv_cache.shape[1:], + ) + gathered = kv_cache[src_tokens].view_as(dense_body) + owned_view = owned_mask.view(-1, *([1] * (dense_body.ndim - 1))) + zero = torch.zeros((), dtype=kv_cache.dtype, device=kv_cache.device) + dense_body.copy_(torch.where(owned_view, gathered, zero)) + return dense_kv_cache + + def token_page_copy_debug_checksum( kv_cache: torch.Tensor, dense_kv_cache: torch.Tensor, @@ -481,9 +619,6 @@ def materialize_local_paged_buffer( return dense_page_buffer owned_mask = layout.owned_pages_mask(unique_logical_pages) - if not torch.any(owned_mask): - return dense_page_buffer - owned_logical_pages = unique_logical_pages[owned_mask].to(torch.int64) owned_physical_pages = layout.logical_pages_to_physical(owned_logical_pages).to( torch.long @@ -493,6 +628,33 @@ def materialize_local_paged_buffer( return dense_page_buffer +def materialize_local_paged_buffer_page_slots( + page_buffer: torch.Tensor, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, +) -> torch.Tensor: + """Materialize slot-remapped page buffer without nonzero/boolean compaction.""" + + dense_num_pages = int(slot_logical_pages.numel()) + 1 + dense_page_buffer = page_buffer.new_zeros((dense_num_pages, *page_buffer.shape[1:])) + if slot_logical_pages.numel() == 0: + return dense_page_buffer + + logical_pages = slot_logical_pages.reshape(-1).to(torch.long) + owned_mask = layout.owned_pages_mask(logical_pages) + physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long) + safe_physical_pages = torch.where( + owned_mask, + physical_pages, + torch.zeros_like(physical_pages), + ) + gathered = page_buffer[safe_physical_pages] + owned_view = owned_mask.view(-1, *([1] * (gathered.ndim - 1))) + zero = torch.zeros((), dtype=page_buffer.dtype, device=page_buffer.device) + dense_page_buffer[1:].copy_(torch.where(owned_view, gathered, zero)) + return dense_page_buffer + + def paged_copy_debug_checksum( page_buffer: torch.Tensor, dense_page_buffer: torch.Tensor, @@ -547,6 +709,7 @@ def materialize_shared_token_kv_buffer( layout: CpSharedKVLayout, page_size: int, remap_logical_locs: torch.Tensor | None = None, + remap_logical_pages: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: _debug_assert_no_tensor_values_below( logical_locs, @@ -574,21 +737,61 @@ def materialize_shared_token_kv_buffer( physical_token_capacity=kv_cache.shape[0], ) - remap_logical_pages = logical_pages_from_locs(remap_logical_locs, page_size) - unique_logical_pages, _ = build_dense_page_remap(remap_logical_pages) - dense_locs = remap_logical_locs_to_dense_locs( - logical_locs, - unique_logical_pages=unique_logical_pages, - page_size=page_size, - ) - dense_kv_cache = materialize_local_token_kv_pages( - kv_cache=kv_cache, - unique_logical_pages=unique_logical_pages, - layout=layout, - page_size=page_size, - ) + if 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( + logical_locs, + unique_logical_pages=materialized_logical_pages, + page_size=page_size, + ) + use_slot_materialize = False + else: + _debug_assert_no_negative_tensor_values( + remap_logical_pages, + context="CP shared KV token materialize page remap", + tensor_name="remap_logical_pages", + ) + remap_logical_pages = filter_pages_mappable_to_physical_pool( + logical_pages=remap_logical_pages, + layout=layout, + physical_page_capacity=kv_cache.shape[0] // page_size, + ) + materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages) + logical_page_capacity = _logical_page_capacity_from_physical_page_capacity( + kv_cache.shape[0] // page_size, + layout, + ) + page_inverse = build_slot_page_inverse( + materialized_logical_pages, + logical_page_capacity=logical_page_capacity, + ) + dense_locs = remap_logical_locs_to_slot_dense_locs( + logical_locs, + page_inverse=page_inverse, + page_size=page_size, + ) + use_slot_materialize = True + + if use_slot_materialize: + dense_kv_cache = materialize_local_token_kv_page_slots( + kv_cache=kv_cache, + slot_logical_pages=materialized_logical_pages, + layout=layout, + page_size=page_size, + ) + else: + dense_kv_cache = materialize_local_token_kv_pages( + kv_cache=kv_cache, + unique_logical_pages=materialized_logical_pages, + layout=layout, + page_size=page_size, + ) + if cp_shared_kv_debug_enabled(): - owned_pages = unique_logical_pages[layout.owned_pages_mask(unique_logical_pages)] + owned_pages = materialized_logical_pages[ + layout.owned_pages_mask(materialized_logical_pages) + ] physical_pages = layout.logical_pages_to_physical(owned_pages) cp_shared_kv_debug_log( "materialize_token_pre", @@ -598,7 +801,7 @@ def materialize_shared_token_kv_buffer( layout.cp_rank, tensor_debug_summary(logical_locs), tensor_debug_summary(remap_logical_locs), - tensor_debug_summary(unique_logical_pages), + tensor_debug_summary(materialized_logical_pages), tensor_debug_summary(owned_pages), tensor_debug_summary(physical_pages), tensor_debug_summary(dense_locs), @@ -607,12 +810,14 @@ def materialize_shared_token_kv_buffer( token_page_copy_debug_checksum( kv_cache, dense_kv_cache, - unique_logical_pages, + materialized_logical_pages, layout, page_size, ), ) + dense_kv_cache = _all_reduce_materialized_buffer(dense_kv_cache, layout.cp_size) + if cp_shared_kv_debug_enabled(): cp_shared_kv_debug_log( "materialize_token_post", @@ -625,7 +830,6 @@ def materialize_shared_token_kv_buffer( ) return dense_kv_cache, dense_locs - def materialize_shared_paged_buffer( page_buffer: torch.Tensor, logical_pages: torch.Tensor, @@ -641,14 +845,17 @@ def materialize_shared_paged_buffer( layout=layout, physical_page_capacity=page_buffer.shape[0], ) - unique_logical_pages, dense_pages = build_dense_page_remap(logical_pages) - dense_page_buffer = materialize_local_paged_buffer( + materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages) + dense_page_buffer = materialize_local_paged_buffer_page_slots( page_buffer=page_buffer, - unique_logical_pages=unique_logical_pages, + slot_logical_pages=materialized_logical_pages, layout=layout, ) + if cp_shared_kv_debug_enabled(): - owned_pages = unique_logical_pages[layout.owned_pages_mask(unique_logical_pages)] + owned_pages = materialized_logical_pages[ + layout.owned_pages_mask(materialized_logical_pages) + ] physical_pages = layout.logical_pages_to_physical(owned_pages) cp_shared_kv_debug_log( "materialize_paged_pre", @@ -657,7 +864,7 @@ def materialize_shared_paged_buffer( "active_local_ck=%s owned_copy_ck=%s", layout.cp_rank, tensor_debug_summary(logical_pages), - tensor_debug_summary(unique_logical_pages), + tensor_debug_summary(materialized_logical_pages), tensor_debug_summary(owned_pages), tensor_debug_summary(physical_pages), tensor_debug_summary(dense_pages), @@ -666,11 +873,15 @@ def materialize_shared_paged_buffer( paged_copy_debug_checksum( page_buffer, dense_page_buffer, - unique_logical_pages, + materialized_logical_pages, layout, ), ) - dense_page_buffer = _all_reduce_materialized_buffer(dense_page_buffer, layout.cp_size) + + dense_page_buffer = _all_reduce_materialized_buffer( + dense_page_buffer, layout.cp_size + ) + if cp_shared_kv_debug_enabled(): cp_shared_kv_debug_log( "materialize_paged_post", diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index dcba05c8f..f5454d29a 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -1634,6 +1634,7 @@ class NativeSparseAttnBackend( kv_cache=kv_cache, logical_locs=page_table_1, remap_logical_locs=metadata.page_table_1, + remap_logical_pages=metadata.real_page_table, layout=forward_batch.cp_shared_kv_layout, page_size=forward_batch.token_to_kv_pool.page_size, ) 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 cb9e7a1f0..acf69c604 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 @@ -366,6 +366,35 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[8:12])) self.assertTrue(torch.equal(dense_kv[20:24], kv_cache[20:24])) + def test_materialize_token_kv_fast_path_avoids_python_tensor_predicates(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=2, cp_rank=1) + kv_cache = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1) + # Page 1 is owned by CP rank 0, so this also covers the no-local-page + # branch without using torch.any(owned_mask) in Python control flow. + logical_locs = torch.tensor([4, 5], dtype=torch.int64) + + with patch.object( + runtime, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + runtime, "_all_reduce_materialized_buffer", lambda x, _: x + ), patch.object( + runtime.torch, "any", side_effect=AssertionError("torch.any sync") + ), patch.object( + runtime.torch, "equal", side_effect=AssertionError("torch.equal sync") + ): + dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=logical_locs, + layout=layout, + page_size=4, + ) + + self.assertEqual(dense_locs.tolist(), [4, 5]) + self.assertEqual(float(dense_kv.abs().sum().item()), 0.0) + def test_materialize_token_kv_keeps_dense_shape_for_shared_remap_source(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 @@ -414,7 +443,105 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): page_buffer=page_buffer, logical_pages=logical_pages, layout=layout, - ) + ) + + def test_materialize_paged_buffer_fast_path_avoids_python_tensor_predicates(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=2, cp_rank=1) + page_buffer = torch.arange(0, 4 * 3, dtype=torch.uint8).view(4, 3) + # Page 1 is owned by CP rank 0, so this also covers the no-local-page + # branch without using torch.any(owned_mask) in Python control flow. + logical_pages = torch.tensor([1], dtype=torch.int32) + + with patch.object( + runtime, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + runtime, "_all_reduce_materialized_buffer", lambda x, _: x + ), patch.object( + runtime.torch, "any", side_effect=AssertionError("torch.any sync") + ), patch.object( + runtime.torch, "equal", side_effect=AssertionError("torch.equal sync") + ): + dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer( + page_buffer=page_buffer, + logical_pages=logical_pages, + layout=layout, + ) + + self.assertEqual(dense_pages.tolist(), [1]) + self.assertEqual(int(dense_page_buffer.sum().item()), 0) + + def test_materialize_paged_buffer_fast_path_avoids_dynamic_shape_ops(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=2, cp_rank=1) + page_buffer = torch.arange(0, 6 * 3, dtype=torch.uint8).view(6, 3) + logical_pages = torch.tensor([1, 2, 5, 6, 0], dtype=torch.int32) + + with patch.object( + runtime, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + runtime, "_all_reduce_materialized_buffer", lambda x, _: x + ), patch.object( + runtime.torch, "unique", side_effect=AssertionError("torch.unique sync") + ), patch.object( + runtime.torch, "nonzero", side_effect=AssertionError("torch.nonzero sync") + ), patch.object( + runtime.torch, + "searchsorted", + side_effect=AssertionError("torch.searchsorted sync"), + ): + dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer( + page_buffer=page_buffer, + logical_pages=logical_pages, + layout=layout, + ) + + self.assertEqual(dense_pages.tolist(), [1, 2, 3, 4, 0]) + self.assertEqual(list(dense_page_buffer.shape), [6, 3]) + self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[1])) + self.assertTrue(torch.equal(dense_page_buffer[4], page_buffer[3])) + self.assertEqual(int(dense_page_buffer[1].sum().item()), 0) + self.assertEqual(int(dense_page_buffer[3].sum().item()), 0) + + def test_materialize_token_kv_page_slot_source_avoids_dynamic_shape_ops(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=2, cp_rank=1) + kv_cache = torch.arange(0, 24, dtype=torch.float32).view(24, 1, 1) + logical_locs = torch.tensor([8, 9, 24, 25, -1], dtype=torch.int64) + remap_logical_pages = torch.tensor([1, 2, 5, 6], dtype=torch.int32) + + with patch.object( + runtime, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + runtime, "_all_reduce_materialized_buffer", lambda x, _: x + ), patch.object( + runtime.torch, "unique", side_effect=AssertionError("torch.unique sync") + ), patch.object( + runtime.torch, "nonzero", side_effect=AssertionError("torch.nonzero sync") + ), patch.object( + runtime.torch, + "searchsorted", + side_effect=AssertionError("torch.searchsorted sync"), + ): + dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer( + kv_cache=kv_cache, + logical_locs=logical_locs, + remap_logical_pages=remap_logical_pages, + layout=layout, + page_size=4, + ) + + self.assertEqual(dense_locs.tolist(), [8, 9, 16, 17, -1]) + self.assertEqual(list(dense_kv.shape), [20, 1, 1]) + self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[4:8])) + self.assertTrue(torch.equal(dense_kv[16:20], kv_cache[12:16])) + self.assertEqual(float(dense_kv[4:8].abs().sum().item()), 0.0) class TestCpSharedKVLazyDebugLogging(unittest.TestCase):