diff --git a/docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md b/docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md index e2b155974..6ed08ac1a 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md @@ -91,3 +91,36 @@ PYTHONPATH=python python3 -m pytest \ test/registered/unit/layers/test_nsa_cp_utils.py \ test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q ``` + +## 2026-05-31 update: index partial/current reuse + +The index path now follows the same target-model partial-current contract as MLA: + +```text +page-aligned cached prefix index pages ++ valid current index K/scale rows copied into their slot-dense suffix pages ++ page-tail slack left zero and invisible through valid seq/page-table lengths +``` + +Implementation notes: + +- `Indexer._can_reuse_current_index_kv()` uses `should_reuse_current_extend_kv()` instead of the old current-only gate, so target cache-hit suffixes can reuse freshly computed current index K/scale. +- Current-only batches still use the compact `current_index_kv` path and avoid page-table materialization. +- Partial cache-hit batches materialize the prefix once for the prev/next in-seq split pair, then pass the composed dense index buffer to both topk calls. +- If the async index prefetcher has a ready prefix handle, `CpSharedKVIndexPrefetcher.consume_prefix_with_current()` consumes it and only fills current rows into the suffix page slots. +- If no prefetch handle exists, `materialize_prefix_and_reuse_current_index_page_slots()` synchronously materializes only prefix slots, all-reduces that prefix range, and fills current rows locally. +- `forward_cuda()` quantizes only `valid_current_rows = extend_seq_lens_cpu[0]` current rows. This preserves current reuse when `out_cache_loc` is physically padded for page alignment. + +Debug contract correction: + +- The prefetch consume path must remap `logical_pages` through the slot `page_inverse` directly. It must not validate these logical pages against `dense_page_buffer.shape[0]` as if the dense buffer were the physical pool capacity; that debug-only check can reject valid high logical page ids that are intentionally packed into a small slot-dense buffer. + +Verification: + +```bash +# g0034 / sglang-glm5-dev-2 +cd /sgl-workspace/sglang-tai +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +# 77 passed, 5 warnings, 2 subtests passed +``` 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 977a6d78b..d5b73cf74 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 @@ -19,6 +19,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_mla_prefetch_should_log_layer, filter_locs_mappable_to_physical_pool, filter_pages_mappable_to_physical_pool, + fill_current_index_page_slots, fill_current_kv_page_slots_and_remap_locs, get_or_build_shared_paged_buffer_slot_remap, get_or_build_shared_token_kv_slot_remap, @@ -1469,6 +1470,104 @@ class CpSharedKVIndexPrefetcher: ) return dense_page_buffer, dense_pages + def consume_prefix_with_current( + self, + *, + layer_id: int, + logical_pages: torch.Tensor, + current_index_k: torch.Tensor, + current_index_scale: torch.Tensor, + current_locs: torch.Tensor, + page_size: int, + index_head_dim: int, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if self.disabled: + self._log_layer( + layer_id, + "index_consume_prefix_current_skip reason=disabled layer=%s", + layer_id, + ) + return None + + self._log_layer( + layer_id, + "index_consume_prefix_current_enter layer=%s prefix_pages=%s " + "total_slots=%s handles=%s", + layer_id, + self.prefix_pages, + self.total_slots, + _debug_handle_keys(layer_id, self.handles), + ) + handle = self.handles.get(layer_id) + if handle is None: + self._log_layer( + layer_id, + "index_consume_prefix_current_miss layer=%s", + layer_id, + ) + return None + if handle.event is None: + self.launch_pending_reduce() + handle = self.handles.get(layer_id) + if handle is None or handle.event is None: + self._log_layer( + layer_id, + "index_consume_prefix_current_miss " + "reason=prefix_reduce_not_ready layer=%s", + layer_id, + ) + return None + handle = self.handles.pop(layer_id) + if self.pending_attention_handle is handle: + self.pending_attention_handle = None + if handle.layer_id != layer_id: + self.disabled = True + self._log_layer( + layer_id, + "index_consume_prefix_current_skip reason=layer_mismatch " + "expected=%s actual=%s", + layer_id, + handle.layer_id, + ) + return None + + consume_cpu = _cpu_timing_start() + wait_cpu = _cpu_timing_start() + torch.cuda.current_stream().wait_event(handle.event) + wait_ms = _cpu_timing_ms(wait_cpu) + dense_page_buffer = handle.dense_page_buffer + + remap_cpu = _cpu_timing_start() + dense_pages = remap_logical_pages_to_slot_dense_pages( + logical_pages, + page_inverse=self.page_inverse, + ) + dense_page_buffer = fill_current_index_page_slots( + dense_page_buffer=dense_page_buffer, + current_index_k=current_index_k, + current_index_scale=current_index_scale, + current_locs=current_locs, + page_inverse=self.page_inverse, + page_size=page_size, + index_head_dim=index_head_dim, + ) + remap_ms = _cpu_timing_ms(remap_cpu) + total_ms = _cpu_timing_ms(consume_cpu) + self._log_layer( + layer_id, + "index_consume_prefix_current_hit layer=%s prefix_pages=%s " + "dense_pages=%s current_rows=%s total_ms=%.3f wait_ms=%.3f " + "remap_ms=%.3f", + layer_id, + self.prefix_pages, + int(dense_page_buffer.shape[0]), + int(current_index_k.shape[0]), + total_ms, + wait_ms, + remap_ms, + ) + return dense_page_buffer, dense_pages + def start_next_layer_prefix( self, *, 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 8a350fabe..8810a854a 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 @@ -813,6 +813,103 @@ def fill_current_kv_page_slots_and_remap_locs( return dense_kv_cache, mixed_locs, current_mask +def fill_current_index_page_slots( + *, + dense_page_buffer: torch.Tensor, + current_index_k: torch.Tensor, + current_index_scale: torch.Tensor, + current_locs: torch.Tensor, + page_inverse: torch.Tensor, + page_size: int, + index_head_dim: int, +) -> torch.Tensor: + """Fill current index K/scale rows into a slot-dense page buffer. + + The dense page buffer already has one row per logical page slot and page 0 is + the dummy page. Current rows are copied into the row/offset selected by their + logical token loc; tail-page slack remains zero and is invisible because the + index kernels still receive the valid sequence length separately. + """ + + current_locs = current_locs.reshape(-1) + current_rows = int(current_locs.numel()) + if current_rows == 0: + return dense_page_buffer + if ( + int(current_index_k.shape[0]) < current_rows + or int(current_index_scale.shape[0]) < current_rows + ): + raise ValueError( + "CP shared KV index current-slot fill got fewer current rows than " + f"locs: k_rows={int(current_index_k.shape[0])} " + f"scale_rows={int(current_index_scale.shape[0])} locs={current_rows}" + ) + + current_index_k = current_index_k[:current_rows].contiguous() + current_index_scale = current_index_scale[:current_rows].contiguous() + k_bytes = current_index_k.reshape(current_rows, -1).view(torch.uint8) + if int(k_bytes.shape[1]) != int(index_head_dim): + raise ValueError( + "CP shared KV index current-slot fill got unexpected K width: " + f"k_bytes_per_token={int(k_bytes.shape[1])} index_head_dim={index_head_dim}" + ) + scale_bytes = current_index_scale.reshape(current_rows, -1).view(torch.uint8) + scale_bytes = scale_bytes.reshape(current_rows, -1) + + page_stride = int(dense_page_buffer.shape[1]) + scale_offset = int(page_size) * int(index_head_dim) + if scale_offset + int(page_size) * int(scale_bytes.shape[1]) > page_stride: + raise ValueError( + "CP shared KV index current-slot fill got incompatible page buffer: " + f"page_stride={page_stride} scale_offset={scale_offset} " + f"scale_bytes_per_token={int(scale_bytes.shape[1])} page_size={page_size}" + ) + + current_pages = torch.div(current_locs, page_size, rounding_mode="floor") + valid_pages = (current_locs >= 0) & (current_pages >= 0) & ( + current_pages < int(page_inverse.numel()) + ) + safe_pages = torch.clamp( + current_pages, + min=0, + max=max(int(page_inverse.numel()) - 1, 0), + ) + dense_pages = page_inverse[safe_pages.to(torch.long)].to(torch.long) + valid_rows = valid_pages & (dense_pages > 0) + if not torch.any(valid_rows): + return dense_page_buffer + + valid_dense_pages = dense_pages[valid_rows] + valid_offsets = torch.remainder(current_locs[valid_rows], page_size).to(torch.long) + flat_buffer = dense_page_buffer.reshape(-1) + + k_cols = torch.arange( + int(index_head_dim), + dtype=torch.long, + device=dense_page_buffer.device, + ) + k_indices = ( + valid_dense_pages[:, None] * page_stride + + valid_offsets[:, None] * int(index_head_dim) + + k_cols[None, :] + ) + flat_buffer[k_indices.reshape(-1)] = k_bytes[valid_rows].reshape(-1) + + s_cols = torch.arange( + int(scale_bytes.shape[1]), + dtype=torch.long, + device=dense_page_buffer.device, + ) + s_indices = ( + valid_dense_pages[:, None] * page_stride + + scale_offset + + valid_offsets[:, None] * int(scale_bytes.shape[1]) + + s_cols[None, :] + ) + flat_buffer[s_indices.reshape(-1)] = scale_bytes[valid_rows].reshape(-1) + return dense_page_buffer + + def _copy_tai_dense_slot_range_body( *, tai_dense_kv_cache: torch.Tensor, @@ -2058,6 +2155,62 @@ def materialize_prefix_and_reuse_current_kv_page_slots( return mixed_kv_cache, mixed_locs +def materialize_prefix_and_reuse_current_index_page_slots( + *, + page_buffer: torch.Tensor, + current_index_k: torch.Tensor, + current_index_scale: torch.Tensor, + current_locs: torch.Tensor, + slot_remap: SharedPagedBufferSlotRemap, + layout: CpSharedKVLayout, + page_size: int, + index_head_dim: int, + prefix_pages: int, + layer_id: int | None = None, + nvtx_source: str = "index.partial_current_sync", +) -> tuple[torch.Tensor, torch.Tensor]: + """Synchronously compose prefix index materialization with current index rows.""" + + total_slots = int(slot_remap.slot_logical_pages.numel()) + 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}" + ) + + dense_page_buffer = page_buffer.new_zeros( + (slot_remap.dense_num_pages, *page_buffer.shape[1:]) + ) + materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=slot_remap.slot_logical_pages, + layout=layout, + start_slot=0, + end_slot=prefix_pages, + ) + prefix_rows = slot_range_to_page_slice(0, prefix_pages) + _all_reduce_materialized_buffer_range( + dense_page_buffer, + layout.cp_size, + prefix_rows.start, + prefix_rows.stop, + nvtx_source=nvtx_source, + nvtx_layer_id=layer_id, + nvtx_cp_rank=layout.cp_rank, + ) + dense_page_buffer = fill_current_index_page_slots( + dense_page_buffer=dense_page_buffer, + current_index_k=current_index_k, + current_index_scale=current_index_scale, + current_locs=current_locs, + page_inverse=slot_remap.page_inverse, + page_size=page_size, + index_head_dim=index_head_dim, + ) + return dense_page_buffer, slot_remap.dense_pages + + def slot_range_to_token_slice( page_size: int, start_slot: int, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index c36bea641..8386f84b8 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -17,7 +17,6 @@ from sglang.srt.layers.attention.nsa import index_buf_accessor from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, - cp_shared_kv_current_reuse_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, @@ -26,7 +25,9 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( get_or_build_shared_paged_buffer_slot_remap, is_current_only_extend_batch, log_cp_draft_shared_kv_debug, + materialize_prefix_and_reuse_current_index_page_slots, materialize_shared_paged_buffer, + should_reuse_current_extend_kv, tensor_debug_checksum, tensor_debug_summary, try_tai_prepare_cp_mqa_index, @@ -302,6 +303,7 @@ class Indexer(MultiPlatformOp): forward_batch: ForwardBatch, layer_id: int, logical_page_table: torch.Tensor, + current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> tuple[torch.Tensor, torch.Tensor]: index_buffer = forward_batch.token_to_kv_pool.get_index_k_with_scale_buffer( layer_id=layer_id @@ -314,6 +316,100 @@ class Indexer(MultiPlatformOp): index_prefetcher = getattr( forward_batch, "cp_shared_kv_index_prefetcher", None ) + if current_index_kv is not None: + page_size = int(forward_batch.token_to_kv_pool.page_size) + prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + prefix_lens = ( + [int(x) for x in prefix_lens_cpu] + if prefix_lens_cpu is not None + else None + ) + extend_lens = ( + [int(x) for x in extend_lens_cpu] + if extend_lens_cpu is not None + else None + ) + if ( + prefix_lens_cpu is None + or len(prefix_lens_cpu) != 1 + or int(prefix_lens_cpu[0]) <= 0 + or int(prefix_lens_cpu[0]) % page_size != 0 + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " + "CP shared KV index partial-current compose requires one " + "positive page-aligned prefix. " + f"cp_rank={layout.cp_rank} layer_id={layer_id} " + f"prefix_lens={prefix_lens} extend_lens={extend_lens} " + f"logical_page_table_shape={tuple(logical_page_table.shape)} " + f"page_size={page_size}" + ) + current_locs = forward_batch.out_cache_loc + if extend_lens_cpu is not None and len(extend_lens_cpu) == 1: + valid_current_rows = int(extend_lens_cpu[0]) + if ( + valid_current_rows > 0 + and valid_current_rows < int(current_locs.numel()) + and valid_current_rows <= int(current_index_kv[0].shape[0]) + and valid_current_rows <= int(current_index_kv[1].shape[0]) + ): + current_locs = current_locs[:valid_current_rows] + current_index_kv = ( + current_index_kv[0][:valid_current_rows], + current_index_kv[1][:valid_current_rows], + ) + prefix_pages = int(prefix_lens_cpu[0]) // page_size + if index_prefetcher is not None: + prefetched = index_prefetcher.consume_prefix_with_current( + layer_id=layer_id, + logical_pages=logical_page_table, + current_index_k=current_index_kv[0], + current_index_scale=current_index_kv[1], + current_locs=current_locs, + page_size=page_size, + index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, + ) + if prefetched is not None: + return prefetched + 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_prefix_and_reuse_current_index_page_slots( + page_buffer=index_buffer, + current_index_k=current_index_kv[0], + current_index_scale=current_index_kv[1], + current_locs=current_locs, + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, + prefix_pages=prefix_pages, + layer_id=layer_id, + ) + ) + if ( + cp_shared_kv_mla_prefetch_log_enabled() + and cp_shared_kv_mla_prefetch_should_log_layer(layer_id) + ): + cp_shared_kv_mla_prefetch_log( + "index_partial_current_sync_compose cp_rank=%s layer=%s " + "prefix_lens=%s extend_lens=%s prefix_pages=%s " + "current_rows=%s dense_pages=%s", + layout.cp_rank, + layer_id, + prefix_lens, + extend_lens, + prefix_pages, + int(current_index_kv[0].shape[0]), + int(materialized.shape[0]), + ) + return materialized, dense_pages + if index_prefetcher is not None: prefetched = index_prefetcher.consume( layer_id=layer_id, @@ -461,12 +557,11 @@ class Indexer(MultiPlatformOp): def _can_reuse_current_index_kv(self, forward_batch: ForwardBatch) -> bool: return ( - cp_shared_kv_current_reuse_enabled() - and forward_batch.uses_cp_shared_kv + forward_batch.uses_cp_shared_kv and self.nsa_enable_prefill_cp and forward_batch.nsa_cp_metadata is not None and is_nsa_prefill_cp_in_seq_split() - and is_current_only_extend_batch(forward_batch) + and should_reuse_current_extend_kv(forward_batch) and forward_batch.hisparse_coordinator is None and _is_cuda and not _is_fp8_fnuz @@ -1283,7 +1378,21 @@ class Indexer(MultiPlatformOp): shared_index_buffer = None shared_block_tables = None - if current_index_kv is None: + current_index_kv_for_topk = current_index_kv + if current_index_kv is not None and not is_current_only_extend_batch( + forward_batch + ): + current_index_kv_for_topk = None + shared_block_tables = metadata.get_page_table_64() + shared_index_buffer, shared_block_tables = ( + self._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id, + shared_block_tables, + current_index_kv=current_index_kv, + ) + ) + elif current_index_kv is None: shared_block_tables = metadata.get_page_table_64() shared_index_buffer, shared_block_tables = ( self._maybe_materialize_shared_index_buffer( @@ -1301,7 +1410,7 @@ class Indexer(MultiPlatformOp): metadata, kv_len_prev, actual_seq_q_prev, - current_index_kv=current_index_kv, + current_index_kv=current_index_kv_for_topk, shared_index_buffer=shared_index_buffer, shared_block_tables=shared_block_tables, actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_prev_tensor, @@ -1316,7 +1425,7 @@ class Indexer(MultiPlatformOp): metadata, kv_len_next, actual_seq_q_next, - current_index_kv=current_index_kv, + current_index_kv=current_index_kv_for_topk, shared_index_buffer=shared_index_buffer, shared_block_tables=shared_block_tables, actual_seq_q_tensor=forward_batch.nsa_cp_metadata.actual_seq_q_next_tensor, @@ -1674,9 +1783,18 @@ class Indexer(MultiPlatformOp): current_index_kv = None if self._can_reuse_current_index_kv(forward_batch): - if key.shape[0] == forward_batch.out_cache_loc.numel(): + extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) + valid_current_rows = int(forward_batch.out_cache_loc.numel()) + if extend_seq_lens_cpu is not None and len(extend_seq_lens_cpu) == 1: + valid_current_rows = min( + int(extend_seq_lens_cpu[0]), + valid_current_rows, + ) + if key.shape[0] >= valid_current_rows: current_k_fp8, current_k_scale = act_quant( - key.contiguous(), self.block_size, self.scale_fmt + key[:valid_current_rows].contiguous(), + self.block_size, + self.scale_fmt, ) current_index_kv = ( current_k_fp8.contiguous(), @@ -1690,14 +1808,25 @@ class Indexer(MultiPlatformOp): if forward_batch.cp_shared_kv_layout is not None else None, layer_id, - tensor_debug_summary(forward_batch.out_cache_loc), + tensor_debug_summary( + forward_batch.out_cache_loc[:valid_current_rows] + ), tensor_debug_checksum(current_index_kv[0]), tensor_debug_checksum(current_index_kv[1]), ) - elif cp_shared_kv_debug_enabled(): - raise RuntimeError( - "CP shared KV current index reuse shape mismatch: " - f"key_tokens={key.shape[0]} out_cache_loc={forward_batch.out_cache_loc.numel()}" + else: + _log_cp_shared_kv_index_prefetch_fallback( + "current_reuse_shape_mismatch", + "NSA index current reuse skipped because key has fewer rows " + "than valid current locs. cp_rank=%s layer=%s key_tokens=%s " + "valid_current_rows=%s out_cache_loc=%s", + forward_batch.cp_shared_kv_layout.cp_rank + if forward_batch.cp_shared_kv_layout is not None + else None, + layer_id, + int(key.shape[0]), + valid_current_rows, + int(forward_batch.out_cache_loc.numel()), ) if _is_cuda or _is_hip: 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 dfea1448c..7720389ed 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 @@ -1021,6 +1021,182 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertLess(valid_tokens, padded_pages * page_size) self.assertEqual(valid_tokens, 10) + def test_index_partial_current_sync_compose_fills_current_page_slots(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 + + page_size = 4 + index_head_dim = 4 + scale_bytes = 4 + page_bytes = page_size * index_head_dim + page_size * scale_bytes + layout = CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0) + page_buffer = torch.zeros((6, page_bytes), dtype=torch.uint8) + page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8) + page_buffer[2] = 99 + logical_pages = torch.tensor([[1, 2]], dtype=torch.int64) + slot_remap = runtime.build_shared_paged_buffer_slot_remap( + page_buffer, + logical_pages, + layout, + ) + current_k = torch.tensor( + [[1, 2, 3, 4], [5, 6, 7, 8]], + dtype=torch.uint8, + ) + current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32) + + with patch.object( + runtime, "_all_reduce_materialized_buffer_range", _identity_all_reduce + ): + dense_page_buffer, dense_pages = ( + runtime.materialize_prefix_and_reuse_current_index_page_slots( + page_buffer=page_buffer, + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=torch.tensor([8, 9], dtype=torch.int64), + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + index_head_dim=index_head_dim, + prefix_pages=1, + layer_id=2, + ) + ) + + scale_offset = page_size * index_head_dim + self.assertEqual(dense_pages.tolist(), [[1, 2]]) + self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1])) + self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0])) + self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1])) + self.assertTrue( + torch.equal( + dense_page_buffer[2, 8:scale_offset], + torch.zeros(8, dtype=torch.uint8), + ) + ) + self.assertTrue( + torch.equal( + dense_page_buffer[2, scale_offset : scale_offset + 4], + current_scale[0].view(torch.uint8), + ) + ) + self.assertTrue( + torch.equal( + dense_page_buffer[2, scale_offset + 4 : scale_offset + 8], + current_scale[1].view(torch.uint8), + ) + ) + self.assertTrue( + torch.equal( + dense_page_buffer[2, scale_offset + 8 :], + torch.zeros(8, dtype=torch.uint8), + ) + ) + + def test_index_prefetch_partial_current_compose_fills_current_page_slots(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeCurrentStream: + def __init__(self): + self.events = [] + + def wait_event(self, event): + self.events.append(event) + + page_size = 4 + index_head_dim = 4 + scale_bytes = 4 + page_bytes = page_size * index_head_dim + page_size * scale_bytes + dense_page_buffer = torch.zeros((3, page_bytes), dtype=torch.uint8) + dense_page_buffer[1] = torch.arange(10, 10 + page_bytes, dtype=torch.uint8) + fake_event = object() + prefetcher = prefetch.CpSharedKVIndexPrefetcher( + layout=CpSharedKVLayout(page_size=page_size, cp_size=1, cp_rank=0), + prefix_pages=1, + slot_logical_pages=torch.tensor([1, 20], dtype=torch.int64), + page_inverse=torch.tensor( + [-1, 1] + [-1] * 18 + [2], + dtype=torch.int64, + ), + dense_num_pages=3, + stream=object(), + ) + handle = prefetch.CpSharedKVIndexPrefetchHandle( + layer_id=1, + dense_page_buffer=dense_page_buffer, + prefix_rows=slice(1, 2), + event=fake_event, + ) + prefetcher.handles[1] = handle + prefetcher.pending_attention_handle = handle + current_k = torch.tensor( + [[11, 12, 13, 14], [15, 16, 17, 18]], + dtype=torch.uint8, + ) + current_scale = torch.tensor([[3.25], [4.5]], dtype=torch.float32) + + with envs.SGLANG_DEBUG_CP_SHARED_KV.override(True), patch.object( + prefetch.torch.cuda, + "current_stream", + return_value=FakeCurrentStream(), + ): + mixed_buffer, dense_pages = prefetcher.consume_prefix_with_current( + layer_id=1, + logical_pages=torch.tensor([[1, 20]], dtype=torch.int64), + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=torch.tensor([80, 81], dtype=torch.int64), + page_size=page_size, + index_head_dim=index_head_dim, + ) + + scale_offset = page_size * index_head_dim + self.assertEqual(dense_pages.tolist(), [[1, 2]]) + self.assertTrue(torch.equal(mixed_buffer[2, 0:4], current_k[0])) + self.assertTrue(torch.equal(mixed_buffer[2, 4:8], current_k[1])) + self.assertTrue( + torch.equal( + mixed_buffer[2, scale_offset : scale_offset + 4], + current_scale[0].view(torch.uint8), + ) + ) + self.assertEqual(prefetcher.handles, {}) + self.assertIsNone(prefetcher.pending_attention_handle) + + def test_index_current_reuse_gate_uses_partial_current_contract(self): + from pathlib import Path + + source = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py" + ).read_text() + start = source.index(" def _can_reuse_current_index_kv") + end = source.index(" @contextlib.contextmanager", start) + method_source = source[start:end] + + self.assertIn("should_reuse_current_extend_kv(forward_batch)", method_source) + self.assertNotIn("is_current_only_extend_batch(forward_batch)", method_source) + + def test_index_current_reuse_prepare_accepts_padded_out_cache_loc(self): + from pathlib import Path + + source = ( + Path(__file__).resolve().parents[4] + / "python/sglang/srt/layers/attention/nsa/nsa_indexer.py" + ).read_text() + start = source.index(" if self._can_reuse_current_index_kv") + end = source.index(" if _is_cuda or _is_hip:", start) + prepare_source = source[start:end] + + self.assertIn("valid_current_rows", prepare_source) + self.assertIn("key[:valid_current_rows]", prepare_source) + self.assertNotIn( + "key.shape[0] == forward_batch.out_cache_loc.numel()", + prepare_source, + ) + def test_materialize_local_token_kv_pages(self): from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_dense_page_remap,