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 e3cfa8da2..2f40175a1 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 @@ -443,15 +443,18 @@ class CpSharedKVMlaPrefetcher: self.owned_prefix_pages = owned_prefix_pages self.owned_total_pages = owned_total_pages self.total_slots = int(slot_logical_pages.numel()) - self.prefix_slot_spans = ( - [(0, int(prefix_pages))] if prefix_slot_spans is None and prefix_pages > 0 - else list(prefix_slot_spans or []) - ) - self.current_slot_spans = ( - [(int(prefix_pages), self.total_slots)] - if current_slot_spans is None and prefix_pages < self.total_slots - else list(current_slot_spans or []) - ) + # Spans are the canonical description, always built by + # build_batch_*_slot_spans in build(). Fail loud on a missing list rather + # than silently reconstructing from prefix_pages (an explicit [] is a valid + # no-prefix compose; only None is a caller bug). + if prefix_slot_spans is None or current_slot_spans is None: + raise ValueError( + "CP shared KV prefetcher requires explicit prefix_slot_spans and " + "current_slot_spans ([] for none); got " + f"prefix={prefix_slot_spans} current={current_slot_spans}." + ) + self.prefix_slot_spans = list(prefix_slot_spans) + self.current_slot_spans = list(current_slot_spans) self.prefix_page_count = _slot_spans_page_count(self.prefix_slot_spans) self.current_page_count = _slot_spans_page_count(self.current_slot_spans) self.stream = stream if stream is not None else torch.cuda.Stream() @@ -1305,15 +1308,18 @@ class CpSharedKVIndexPrefetcher: self.owned_prefix_pages = owned_prefix_pages self.owned_total_pages = owned_total_pages self.total_slots = int(slot_logical_pages.numel()) - self.prefix_slot_spans = ( - [(0, int(prefix_pages))] if prefix_slot_spans is None and prefix_pages > 0 - else list(prefix_slot_spans or []) - ) - self.current_slot_spans = ( - [(int(prefix_pages), self.total_slots)] - if current_slot_spans is None and prefix_pages < self.total_slots - else list(current_slot_spans or []) - ) + # Spans are the canonical description, always built by + # build_batch_*_slot_spans in build(). Fail loud on a missing list rather + # than silently reconstructing from prefix_pages (an explicit [] is a valid + # no-prefix compose; only None is a caller bug). + if prefix_slot_spans is None or current_slot_spans is None: + raise ValueError( + "CP shared KV prefetcher requires explicit prefix_slot_spans and " + "current_slot_spans ([] for none); got " + f"prefix={prefix_slot_spans} current={current_slot_spans}." + ) + self.prefix_slot_spans = list(prefix_slot_spans) + self.current_slot_spans = list(current_slot_spans) self.prefix_page_count = _slot_spans_page_count(self.prefix_slot_spans) self.current_page_count = _slot_spans_page_count(self.current_slot_spans) self.stream = stream if stream is not None else torch.cuda.Stream() 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 64f639b75..626d2d41e 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 @@ -5466,10 +5466,8 @@ def materialize_prefix_and_reuse_current_kv_page_slots( slot_remap: SharedTokenKVSlotRemap, layout: CpSharedKVLayout, page_size: int, - prefix_pages: int, loc_req_id: torch.Tensor | None = None, current_req_id: torch.Tensor | None = None, - prefix_slot_span: tuple[int, int] | None = None, prefix_slot_spans: list[tuple[int, int]] | None = None, current_slot_spans: list[tuple[int, int]] | None = None, logical_locs_row_ids: torch.Tensor | None = None, @@ -5501,35 +5499,17 @@ def materialize_prefix_and_reuse_current_kv_page_slots( ) total_slots = int(slot_remap.slot_logical_pages.numel()) - if prefix_slot_spans is not None and prefix_slot_span is not None: + # Prefix is described ONLY by the canonical per-request slot-span list, built + # once per forward by get_or_build_batch_slot_spans. [] means a genuine + # no-prefix (current-only) compose; None means the caller failed to supply + # spans -> fail loud rather than silently materialize a zero prefix. + if prefix_slot_spans is None: raise ValueError( - "Specify either prefix_slot_span or prefix_slot_spans, not both." - ) - if prefix_slot_spans is not None: - prefix_spans = _merge_slot_spans(prefix_slot_spans) - for prefix_start_slot, prefix_end_slot in prefix_spans: - if ( - prefix_start_slot < 0 - or prefix_end_slot < prefix_start_slot - or prefix_end_slot > total_slots - ): - raise ValueError( - "Invalid CP shared KV partial-current prefix slot span: " - f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} " - f"total_slots={total_slots}" - ) - elif prefix_slot_span is None: - if prefix_pages < 0 or prefix_pages > total_slots: - raise ValueError( - "Invalid CP shared KV partial-current prefix range: " - f"prefix_pages={prefix_pages} total_slots={total_slots}" - ) - prefix_spans = [(0, int(prefix_pages))] if int(prefix_pages) > 0 else [] - else: - prefix_start_slot, prefix_end_slot = ( - int(prefix_slot_span[0]), - int(prefix_slot_span[1]), + "CP shared KV partial-current compose requires explicit " + "prefix_slot_spans ([] for a no-prefix current-only compose); got None." ) + prefix_spans = _merge_slot_spans(prefix_slot_spans) + for prefix_start_slot, prefix_end_slot in prefix_spans: if ( prefix_start_slot < 0 or prefix_end_slot < prefix_start_slot @@ -5537,13 +5517,9 @@ def materialize_prefix_and_reuse_current_kv_page_slots( ): raise ValueError( "Invalid CP shared KV partial-current prefix slot span: " - f"prefix_slot_span={prefix_slot_span} total_slots={total_slots}" + f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} " + f"total_slots={total_slots}" ) - prefix_spans = ( - [(prefix_start_slot, prefix_end_slot)] - if prefix_end_slot > prefix_start_slot - else [] - ) dense_kv_cache = kv_cache.new_zeros( (slot_remap.dense_num_pages * page_size, *kv_cache.shape[1:]) @@ -5619,15 +5595,9 @@ def materialize_prefix_and_reuse_current_kv_page_slots( if layout.cp_size > 1: current_materialized_by_ipc = False if current_slot_spans is None: - if prefix_slot_span is not None or prefix_slot_spans is not None: - raise ValueError( - "CP shared KV batched current compose requires explicit " - "current_slot_spans to avoid reducing prefix slots twice." - ) - current_slot_spans = ( - [(int(prefix_pages), total_slots)] - if int(prefix_pages) < total_slots - else [] + raise ValueError( + "CP shared KV batched current compose requires explicit " + "current_slot_spans." ) merged_current_spans_for_reduce = _merge_slot_spans(current_slot_spans) if merged_current_spans_for_reduce: @@ -5747,9 +5717,7 @@ def materialize_prefix_and_reuse_current_index_page_slots( layout: CpSharedKVLayout, page_size: int, index_head_dim: int, - prefix_pages: int, current_req_id: torch.Tensor | None = None, - prefix_slot_span: tuple[int, int] | None = None, prefix_slot_spans: list[tuple[int, int]] | None = None, current_slot_spans: list[tuple[int, int]] | None = None, layer_id: int | None = None, @@ -5767,35 +5735,15 @@ def materialize_prefix_and_reuse_current_index_page_slots( ) total_slots = int(slot_remap.slot_logical_pages.numel()) - if prefix_slot_spans is not None and prefix_slot_span is not None: + # See materialize_prefix_and_reuse_current_kv_page_slots: prefix is described + # ONLY by the canonical span list. [] = genuine no-prefix; None = caller bug. + if prefix_slot_spans is None: raise ValueError( - "Specify either prefix_slot_span or prefix_slot_spans, not both." - ) - if prefix_slot_spans is not None: - prefix_spans = _merge_slot_spans(prefix_slot_spans) - for prefix_start_slot, prefix_end_slot in prefix_spans: - if ( - prefix_start_slot < 0 - or prefix_end_slot < prefix_start_slot - or prefix_end_slot > total_slots - ): - raise ValueError( - "Invalid CP shared KV index partial-current prefix slot span: " - f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} " - f"total_slots={total_slots}" - ) - elif prefix_slot_span is None: - if prefix_pages < 0 or prefix_pages > total_slots: - raise ValueError( - "Invalid CP shared KV index partial-current prefix range: " - f"prefix_pages={prefix_pages} total_slots={total_slots}" - ) - prefix_spans = [(0, int(prefix_pages))] if int(prefix_pages) > 0 else [] - else: - prefix_start_slot, prefix_end_slot = ( - int(prefix_slot_span[0]), - int(prefix_slot_span[1]), + "CP shared KV index partial-current compose requires explicit " + "prefix_slot_spans ([] for a no-prefix current-only compose); got None." ) + prefix_spans = _merge_slot_spans(prefix_slot_spans) + for prefix_start_slot, prefix_end_slot in prefix_spans: if ( prefix_start_slot < 0 or prefix_end_slot < prefix_start_slot @@ -5803,13 +5751,9 @@ def materialize_prefix_and_reuse_current_index_page_slots( ): raise ValueError( "Invalid CP shared KV index partial-current prefix slot span: " - f"prefix_slot_span={prefix_slot_span} total_slots={total_slots}" + f"prefix_slot_span={(prefix_start_slot, prefix_end_slot)} " + f"total_slots={total_slots}" ) - prefix_spans = ( - [(prefix_start_slot, prefix_end_slot)] - if prefix_end_slot > prefix_start_slot - else [] - ) dense_page_buffer = page_buffer.new_zeros( (slot_remap.dense_num_pages, *page_buffer.shape[1:]) @@ -5865,15 +5809,9 @@ def materialize_prefix_and_reuse_current_index_page_slots( if layout.cp_size > 1: current_materialized_by_ipc = False if current_slot_spans is None: - if prefix_slot_span is not None or prefix_slot_spans is not None: - raise ValueError( - "CP shared KV batched index current compose requires explicit " - "current_slot_spans to avoid reducing prefix slots twice." - ) - current_slot_spans = ( - [(int(prefix_pages), total_slots)] - if int(prefix_pages) < total_slots - else [] + raise ValueError( + "CP shared KV batched index current compose requires explicit " + "current_slot_spans." ) merged_current_spans_for_reduce = _merge_slot_spans(current_slot_spans) if merged_current_spans_for_reduce: diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 031929f50..94e5ddc3b 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -762,17 +762,13 @@ class Indexer(MultiPlatformOp): current_req_id = torch.zeros_like(current_locs, dtype=torch.long) else: current_req_id = current_req_id[: int(current_locs.shape[0])] - if len(prefix_lens_cpu) == 1: - prefix_pages = int(prefix_lens_cpu[0]) // page_size - else: - prefix_pages = 0 prefix_slot_spans, current_slot_spans = get_or_build_batch_slot_spans( forward_batch, logical_pages=logical_page_table, prefix_lens_cpu=prefix_lens_cpu, extend_lens_cpu=extend_lens_cpu, page_size=page_size, - want_prefix=len(prefix_lens_cpu) > 1, + want_prefix=True, ) if index_prefetcher is not None: prefetched = index_prefetcher.consume_prefix_with_current( @@ -819,7 +815,6 @@ class Indexer(MultiPlatformOp): layout=layout, page_size=page_size, index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, - prefix_pages=prefix_pages, current_req_id=current_req_id, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, @@ -838,7 +833,7 @@ class Indexer(MultiPlatformOp): layer_id, prefix_lens, extend_lens, - prefix_pages, + sum(int(e) - int(s) for s, e in prefix_slot_spans), prefix_slot_spans, int(current_index_kv[0].shape[0]), int(materialized.shape[0]), diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 807bbab08..93addc554 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -2085,7 +2085,7 @@ class NativeSparseAttnBackend( layout=forward_batch.cp_shared_kv_layout, page_size=page_size, ) - _, current_slot_spans = get_or_build_batch_slot_spans( + prefix_slot_spans, current_slot_spans = get_or_build_batch_slot_spans( forward_batch, logical_pages=metadata.real_page_table, prefix_lens_cpu=getattr( @@ -2095,7 +2095,11 @@ class NativeSparseAttnBackend( forward_batch, "extend_seq_lens_cpu", None ), page_size=page_size, - want_prefix=False, + # current-only batch -> no prefix; want_prefix=True still + # yields [] here but keeps want_prefix uniformly True across + # the MLA + indexer call sites so the per-forward span cache + # (keyed by want_prefix) does not thrash layer-to-layer. + want_prefix=True, ) kv_cache, page_table_1 = ( materialize_prefix_and_reuse_current_kv_page_slots( @@ -2106,7 +2110,7 @@ class NativeSparseAttnBackend( slot_remap=slot_remap, layout=forward_batch.cp_shared_kv_layout, page_size=page_size, - prefix_pages=0, + prefix_slot_spans=prefix_slot_spans, loc_req_id=cp_loc_req_id, current_req_id=cp_current_req_id_full[ : int(current_locs_for_reuse.shape[0]) @@ -2212,30 +2216,16 @@ class NativeSparseAttnBackend( f"current_locs_shape={tuple(current_locs_for_reuse.shape)} " f"page_size={page_size}" ) - if len(prefix_lens_cpu) == 1: - prefix_pages = int(prefix_lens_cpu[0]) // page_size - prefix_slot_spans, current_slot_spans = ( - get_or_build_batch_slot_spans( - forward_batch, - logical_pages=metadata.real_page_table, - prefix_lens_cpu=prefix_lens_cpu, - extend_lens_cpu=extend_lens_cpu, - page_size=page_size, - want_prefix=False, - ) - ) - else: - prefix_pages = 0 - prefix_slot_spans, current_slot_spans = ( - get_or_build_batch_slot_spans( - forward_batch, - logical_pages=metadata.real_page_table, - prefix_lens_cpu=prefix_lens_cpu, - extend_lens_cpu=extend_lens_cpu, - page_size=page_size, - want_prefix=True, - ) + prefix_slot_spans, current_slot_spans = ( + get_or_build_batch_slot_spans( + forward_batch, + logical_pages=metadata.real_page_table, + prefix_lens_cpu=prefix_lens_cpu, + extend_lens_cpu=extend_lens_cpu, + page_size=page_size, + want_prefix=True, ) + ) slot_remap = get_or_build_shared_token_kv_slot_remap( forward_batch, kv_cache=kv_cache, @@ -2252,7 +2242,6 @@ class NativeSparseAttnBackend( slot_remap=slot_remap, layout=forward_batch.cp_shared_kv_layout, page_size=page_size, - prefix_pages=prefix_pages, loc_req_id=cp_loc_req_id, current_req_id=cp_current_req_id_full[ : int(current_locs_for_reuse.shape[0]) @@ -2278,7 +2267,7 @@ class NativeSparseAttnBackend( reason, prefix_lens, extend_lens, - prefix_pages, + sum(int(e) - int(s) for s, e in prefix_slot_spans), prefix_slot_spans, int(current_kv_cache.shape[0]), int(kv_cache.shape[0]), @@ -2585,10 +2574,6 @@ class NativeSparseAttnBackend( f"extend_lens={extend_lens_cpu} " f"page_size={page_size}" ) - if len(prefix_lens_cpu) == 1: - prefix_pages = int(prefix_lens_cpu[0]) // page_size - else: - prefix_pages = 0 prefix_slot_spans, current_slot_spans = ( get_or_build_batch_slot_spans( forward_batch, @@ -2596,7 +2581,7 @@ class NativeSparseAttnBackend( prefix_lens_cpu=prefix_lens_cpu, extend_lens_cpu=extend_lens_cpu, page_size=page_size, - want_prefix=len(prefix_lens_cpu) > 1, + want_prefix=True, ) ) logical_locs_row_ids = ( @@ -2681,7 +2666,6 @@ class NativeSparseAttnBackend( slot_remap=slot_remap, layout=forward_batch.cp_shared_kv_layout, page_size=page_size, - prefix_pages=prefix_pages, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, loc_req_id=logical_locs_row_ids, diff --git a/test/manual/test_cp_shared_kv_compose_v2_8rank.py b/test/manual/test_cp_shared_kv_compose_v2_8rank.py index 47ea1cda1..8a3c406cd 100644 --- a/test/manual/test_cp_shared_kv_compose_v2_8rank.py +++ b/test/manual/test_cp_shared_kv_compose_v2_8rank.py @@ -393,7 +393,6 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None: slot_remap=kv_remap, layout=layout, page_size=page_size, - prefix_pages=0, loc_req_id=loc_req_id, current_req_id=zeros_req, layer_id=60, @@ -408,7 +407,6 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None: layout=layout, page_size=page_size, index_head_dim=128, - prefix_pages=0, current_req_id=zeros_req, layer_id=60, **spans, @@ -422,6 +420,7 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None: layout=layout, page_size=page_size, prefix_pages=prefix_pages, + **spans, slot_logical_pages=kv_remap.slot_logical_pages, page_inverse=kv_remap.page_inverse, dense_num_pages=kv_remap.dense_num_pages, @@ -452,6 +451,7 @@ def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None: idx = CpSharedKVIndexPrefetcher( layout=layout, prefix_pages=prefix_pages, + **spans, slot_logical_pages=idx_remap.slot_logical_pages, page_inverse=idx_remap.page_inverse, dense_num_pages=idx_remap.dense_num_pages, @@ -534,7 +534,6 @@ def _compose(s, layer_id: int, *, writers: list[int] | None = None): slot_remap=s["slot_remap"], layout=s["layout"], page_size=s["page_size"], - prefix_pages=0, loc_req_id=s["loc_req_id"], current_req_id=s["current_req_id"], prefix_slot_spans=s["prefix_slot_spans"], @@ -662,7 +661,6 @@ def main() -> None: layout=s["layout"], page_size=s["page_size"], index_head_dim=128, - prefix_pages=0, current_req_id=si["current_req_id"], prefix_slot_spans=s["prefix_slot_spans"], current_slot_spans=s["current_slot_spans"], diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 7301d3666..795e38d7a 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -3604,7 +3604,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): ) self.assertEqual(len(compose_calls), 1) - self.assertEqual(compose_calls[0]["prefix_pages"], 0) + self.assertEqual(compose_calls[0]["prefix_slot_spans"], []) self.assertEqual(compose_calls[0]["current_slot_spans"], [(0, 2)]) self.assertIs(compose_calls[0]["current_index_k"], current_index_kv[0]) self.assertIs(compose_calls[0]["current_index_scale"], current_index_kv[1]) 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 b49220db6..ffe1f7f03 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 @@ -218,6 +218,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), page_size=4, prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([0, 1, 2], dtype=torch.int64), page_inverse=torch.tensor([0, 1, 2], dtype=torch.int64), dense_num_pages=4, @@ -594,6 +596,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetcher = prefetch.CpSharedKVIndexPrefetcher( layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([0, 1, 2], dtype=torch.int64), page_inverse=torch.tensor([0, 1, 2], dtype=torch.int64), dense_num_pages=4, @@ -1292,7 +1296,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertIn( "materialize_prefix_and_reuse_current_kv_page_slots", branch_source ) - self.assertIn("prefix_pages=0", branch_source) + # current-only compose now passes the canonical (empty) prefix_slot_spans + # instead of the removed prefix_pages scalar fallback. + self.assertIn("prefix_slot_spans=prefix_slot_spans", branch_source) + self.assertNotIn("prefix_pages=", branch_source) self.assertNotIn("kv_cache = current_kv_cache", branch_source) def test_mla_partial_current_sync_uses_batch_prefix_slot_spans(self): @@ -2192,7 +2199,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=2, + prefix_slot_spans=[(0, 2)], ) ) @@ -2202,6 +2209,56 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(mixed_kv[12:14], current_kv)) self.assertEqual(mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1]]) + def test_materialize_prefix_requires_explicit_prefix_slot_spans(self): + # Regression (a24111a5f4): a bs=1 cache-hit passed prefix_slot_spans=[] + + # prefix_pages>0; the `is not None` guard let [] shadow the prefix_pages + # fallback so the cached prefix was never materialized (zero KV over the + # reused prefix -> garbage output). The compose now takes prefix spans as + # the sole canonical input: None must fail loud; an explicit span list must + # materialize a NON-ZERO prefix. + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + page_size = 4 + layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0) + kv_cache = torch.arange(0, 32, dtype=torch.float32).view(32, 1, 1) + logical_locs = torch.tensor([[4, 8, 20, 21, 22, 23]], dtype=torch.int64) + current_locs = torch.tensor([20, 21], dtype=torch.int64) + current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1) + slot_remap = runtime.build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=logical_locs, + remap_logical_pages=torch.tensor([[1, 2, 5]], dtype=torch.int64), + layout=layout, + page_size=page_size, + ) + + def _compose(prefix_slot_spans): + with patch.object( + runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce + ): + return runtime.materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=logical_locs, + current_kv_cache=current_kv, + current_locs=current_locs, + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + prefix_slot_spans=prefix_slot_spans, + ) + + # None -> fail loud (never silently materialize a zero prefix). + with self.assertRaises(ValueError): + _compose(None) + + # The canonical bs=1 cache-hit spans materialize the prefix as NON-ZERO + # (the rows the regression left empty). + mixed_kv, _ = _compose([(0, 2)]) + self.assertGreater(mixed_kv[4:12].abs().sum().item(), 0.0) + self.assertTrue(torch.equal(mixed_kv[4:8], kv_cache[4:8])) + self.assertTrue(torch.equal(mixed_kv[8:12], kv_cache[8:12])) + def test_batch_prefix_slot_span_covers_bounding_prefix_range(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -2956,7 +3013,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=0, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, ) @@ -3027,7 +3083,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=0, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, ) @@ -3084,7 +3139,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=layout, page_size=page_size, index_head_dim=index_head_dim, - prefix_pages=0, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, layer_id=2, @@ -3249,7 +3303,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], ) ) @@ -3303,7 +3358,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=0, + prefix_slot_spans=[], current_slot_spans=[(0, 1)], layer_id=0, ) @@ -3386,6 +3441,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), page_size=4, prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), page_inverse=page_inverse, dense_num_pages=4, @@ -3452,6 +3509,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), page_size=4, prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), dense_num_pages=4, @@ -3499,6 +3558,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), page_size=4, prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), dense_num_pages=4, @@ -3538,6 +3599,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), page_size=4, prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), dense_num_pages=4, @@ -3578,6 +3641,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetcher = prefetch.CpSharedKVIndexPrefetcher( layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), prefix_pages=2, + prefix_slot_spans=[(0, 2)], + current_slot_spans=[(2, 3)], slot_logical_pages=torch.tensor([1, 2, 5], dtype=torch.int64), page_inverse=torch.tensor([0, 1, 2, -1, -1, 3], dtype=torch.int64), dense_num_pages=4, @@ -3623,6 +3688,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetcher = prefetch.CpSharedKVIndexPrefetcher( layout=CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=0), prefix_pages=1, + prefix_slot_spans=[(0, 1)], + current_slot_spans=[(1, 3)], slot_logical_pages=torch.tensor([1, 2, 3], dtype=torch.int64), page_inverse=torch.tensor([-1, 1, 2, 3], dtype=torch.int64), dense_num_pages=padded_pages, @@ -3699,7 +3766,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=layout, page_size=page_size, index_head_dim=index_head_dim, - prefix_pages=1, + prefix_slot_spans=[(0, 1)], layer_id=2, ) ) @@ -3831,7 +3898,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=layout, page_size=page_size, index_head_dim=index_head_dim, - prefix_pages=1, + prefix_slot_spans=[(0, 1)], + current_slot_spans=[(1, 2)], layer_id=2, ) ) @@ -3889,7 +3957,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=layout, page_size=page_size, index_head_dim=index_head_dim, - prefix_pages=0, + prefix_slot_spans=[], current_slot_spans=[(0, 1)], layer_id=0, ) @@ -3968,6 +4036,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetcher = prefetch.CpSharedKVIndexPrefetcher( layout=CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0), prefix_pages=1, + prefix_slot_spans=[(0, 1)], + current_slot_spans=[(1, 2)], slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64), page_inverse=torch.tensor( [-1, 1] + [-1] * 18 + [2], @@ -4036,6 +4106,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetcher = prefetch.CpSharedKVIndexPrefetcher( layout=CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0), prefix_pages=1, + prefix_slot_spans=[(0, 1)], + current_slot_spans=[(1, 2)], slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64), page_inverse=torch.tensor( [-1, 1] + [-1] * 18 + [2], @@ -5620,7 +5692,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=0, prefix_slot_spans=prefix_slot_spans, current_slot_spans=[], layer_id=0, @@ -5997,7 +6068,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): layout=layout, page_size=page_size, index_head_dim=index_head_dim, - prefix_pages=0, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, layer_id=0, @@ -6056,7 +6126,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=1, + prefix_slot_spans=[(0, 1)], current_slot_spans=[(1, 3)], layer_id=0, ) @@ -6180,7 +6250,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): slot_remap=slot_remap, layout=layout, page_size=page_size, - prefix_pages=0, prefix_slot_spans=prefix_slot_spans, current_slot_spans=current_slot_spans, layer_id=0,