From eabe0bbfad459241f46c605d919f2ef6fe34b0f4 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Tue, 9 Jun 2026 00:19:15 +0800 Subject: [PATCH] Preserve CP slot-remap reference paths after SYH pick The SYH per-request page_inverse migration fixes bs>1 cache contamination, but the conflict resolution made several helper/reference paths require explicit request ids and broke existing unit coverage. This keeps production bs>1 callers on explicit req-id routing while allowing bs=1/reference helpers to infer or default request ids without reintroducing the hot-path batch-global inverse. Constraint: Runtime bs>1 materialize paths must route through per-request page_inverse rows to avoid cross-request KV aliasing. Constraint: The remote test container may still have single-row tai-kernel materialize helpers while production bs>1 requires the new req-id ABI. Rejected: Revert to batch-global page_inverse | it is the GSM8K cache-hit corruption root cause. Rejected: Update tests only | the helper API is still useful for bs=1/reference callers and py-level regression coverage. Confidence: medium Scope-risk: moderate Directive: Do not remove explicit loc_req_id/current_req_id from production bs>1 call sites; default inference is for legacy/reference use only. Tested: Local py_compile for cp_shared_kv_runtime.py and nsa_indexer.py. Tested: Remote container py_compile for cp_shared_kv_runtime.py. Tested: Remote container pytest: test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py => 263 passed, 5 warnings, 2 subtests passed. Not-tested: Full ETE GSM8K/cache-hit run after SYH pick. --- .../attention/nsa/cp_shared_kv_runtime.py | 313 ++++++++++++++---- .../srt/layers/attention/nsa/nsa_indexer.py | 5 +- 2 files changed, 249 insertions(+), 69 deletions(-) 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 0c7fcde1a..1bebb91bb 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 @@ -531,17 +531,29 @@ def _tai_current_slot_fill_sparse_pages_self_test( (int(current_locs.numel()),), device=device, dtype=torch.long ) - mixed_kv, mixed_locs, current_mask = fill_kernel( - dense_kv, - materialized_locs, - current_kv, - logical_locs, - current_locs, - page_inverse, - current_req_id, - page_size=page_size, - mask_non_current_in_current_pages=True, - ) + try: + mixed_kv, mixed_locs, current_mask = fill_kernel( + dense_kv, + materialized_locs, + current_kv, + logical_locs, + current_locs, + page_inverse, + current_req_id, + page_size=page_size, + mask_non_current_in_current_pages=True, + ) + except TypeError: + mixed_kv, mixed_locs, current_mask = fill_kernel( + dense_kv, + materialized_locs, + current_kv, + logical_locs, + current_locs, + page_inverse, + page_size=page_size, + mask_non_current_in_current_pages=True, + ) expected_locs = torch.tensor( [[4, 5, 8, 9, 12, 13, -1, -1]], device=device, @@ -1345,17 +1357,34 @@ def _try_tai_materialize_token_kv_pages_and_locs( try: tai_slot_logical_pages = _contiguous_for_tai(slot_logical_pages.reshape(-1)) - page_inverse = kernels.build_slot_page_inverse( - tai_slot_logical_pages, - logical_page_capacity, - int(max(int(batch_rows), 1)), - ) - dense_locs = kernels.remap_logical_locs_to_slot_dense_locs( - _contiguous_for_tai(logical_locs), - page_inverse, - _row_req_id_for_tai(loc_req_id, logical_locs), - page_size=page_size, - ) + try: + page_inverse = kernels.build_slot_page_inverse( + tai_slot_logical_pages, + logical_page_capacity, + int(max(int(batch_rows), 1)), + ) + except TypeError: + if int(batch_rows) != 1: + raise + page_inverse = kernels.build_slot_page_inverse( + tai_slot_logical_pages, + logical_page_capacity, + ) + try: + dense_locs = kernels.remap_logical_locs_to_slot_dense_locs( + _contiguous_for_tai(logical_locs), + page_inverse, + _row_req_id_for_tai(loc_req_id, logical_locs), + page_size=page_size, + ) + except TypeError: + if int(batch_rows) != 1: + raise + dense_locs = kernels.remap_logical_locs_to_slot_dense_locs( + _contiguous_for_tai(logical_locs), + page_inverse, + page_size=page_size, + ) dense_kv_cache = kernels.materialize_shared_token_kv_pages( kv_cache, tai_slot_logical_pages, @@ -1388,11 +1417,20 @@ def _try_tai_build_slot_page_inverse( return None try: - return kernels.build_slot_page_inverse( - _contiguous_for_tai(slot_logical_pages.reshape(-1)), - logical_page_capacity, - int(max(int(batch_rows), 1)), - ) + tai_slot_logical_pages = _contiguous_for_tai(slot_logical_pages.reshape(-1)) + try: + return kernels.build_slot_page_inverse( + tai_slot_logical_pages, + logical_page_capacity, + int(max(int(batch_rows), 1)), + ) + except TypeError: + if int(batch_rows) != 1: + raise + return kernels.build_slot_page_inverse( + tai_slot_logical_pages, + logical_page_capacity, + ) except Exception as exc: _log_tai_materialize_fallback( "page_inverse_failed", @@ -1515,10 +1553,76 @@ def build_page_table_row_req_id(logical_pages: torch.Tensor) -> torch.Tensor: return torch.zeros_like(logical_pages, dtype=torch.long) +def _slot_page_inverse_2d(page_inverse: torch.Tensor) -> torch.Tensor: + if page_inverse.dim() == 1: + return page_inverse.unsqueeze(0) + return page_inverse + + +def _infer_req_id_for_locs_from_slot_remap( + logical_locs: torch.Tensor, + slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap, + page_size: int, + *, + strict: bool = False, +) -> torch.Tensor: + """Infer row ids for legacy/reference callers that omit explicit req ids. + + Runtime bs>1 paths pass request ids explicitly. This helper keeps older + tests and bs=1 helper use working without reintroducing the batch-global + page inverse: for multi-row slot remaps it resolves each loc by membership + in exactly one row of the request-scoped slot map. + """ + + if logical_locs.numel() == 0: + return torch.zeros_like(logical_locs, dtype=torch.long) + page_inverse = _slot_page_inverse_2d(slot_remap.page_inverse) + rows = int(page_inverse.shape[0]) + if rows <= 1: + return torch.zeros_like(logical_locs, dtype=torch.long) + + sorted_pages = getattr(slot_remap, "slot_sorted_logical_pages_by_row", None) + if sorted_pages is None or sorted_pages.dim() != 2: + raise ValueError( + "CP shared KV per-request slot remap requires explicit req ids when " + "row-scoped slot pages are unavailable." + ) + + locs_long = logical_locs.to(torch.long) + valid_locs = locs_long >= 0 + pages = torch.div( + torch.where(valid_locs, locs_long, torch.zeros_like(locs_long)), + page_size, + rounding_mode="floor", + ) + req_id = torch.zeros_like(locs_long, dtype=torch.long) + match_count = torch.zeros_like(locs_long, dtype=torch.long) + for row in range(rows): + row_pages = sorted_pages[row] + row_pages = row_pages[row_pages >= 0] + if row_pages.numel() == 0: + continue + row_match = valid_locs & torch.isin(pages, row_pages) + req_id = torch.where(row_match, torch.full_like(req_id, row), req_id) + match_count = match_count + row_match.to(torch.long) + + ambiguous = valid_locs & (match_count != 1) + if torch.any(ambiguous): + if strict: + bad_pages = pages[ambiguous] + raise ValueError( + "CP shared KV could not infer unique per-request ids for slot remap. " + f"bad_page_min={int(bad_pages.min().item())} " + f"bad_page_max={int(bad_pages.max().item())}" + ) + req_id = torch.where(ambiguous, torch.zeros_like(req_id), req_id) + return req_id + + def build_slot_page_inverse_optimized( slot_logical_pages: torch.Tensor, logical_page_capacity: int, - batch_rows: int, + batch_rows: int = 1, ) -> torch.Tensor: tai_result = _try_tai_build_slot_page_inverse( slot_logical_pages, @@ -1538,8 +1642,10 @@ def remap_logical_locs_to_slot_dense_locs_optimized( logical_locs: torch.Tensor, page_inverse: torch.Tensor, page_size: int, - loc_req_id: torch.Tensor, + loc_req_id: torch.Tensor | None = None, ) -> torch.Tensor: + if loc_req_id is None: + loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long) if _tai_materialize_runtime_enabled(): kernels = _load_tai_materialize_kernels() if kernels is not None: @@ -1550,6 +1656,14 @@ def remap_logical_locs_to_slot_dense_locs_optimized( _row_req_id_for_tai(loc_req_id, logical_locs), page_size=page_size, ) + except TypeError: + if int(_slot_page_inverse_2d(page_inverse).shape[0]) != 1: + raise + return kernels.remap_logical_locs_to_slot_dense_locs( + _contiguous_for_tai(logical_locs), + _contiguous_for_tai(page_inverse), + page_size=page_size, + ) except Exception as exc: _log_tai_materialize_fallback( "loc_remap_failed", @@ -1574,8 +1688,8 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs( current_locs: torch.Tensor, page_inverse: torch.Tensor, page_size: int, - current_req_id: torch.Tensor, - mask_non_current_in_current_pages: bool, + current_req_id: torch.Tensor | None = None, + mask_non_current_in_current_pages: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: if not _tai_materialize_runtime_enabled(): _log_tai_materialize_runtime_disabled("fill_current_kv_page_slots") @@ -1602,6 +1716,9 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs( ) return None + if current_req_id is None: + current_req_id = torch.zeros_like(current_locs.reshape(-1), dtype=torch.long) + if dense_kv_cache.is_cuda and not _tai_current_slot_fill_supports_sparse_pages( dense_kv_cache.device ): @@ -1619,17 +1736,31 @@ def _try_tai_fill_current_kv_page_slots_and_remap_locs( return None try: - return fill_kernel( - _contiguous_for_tai(dense_kv_cache), - _contiguous_for_tai(materialized_dense_locs), - _contiguous_for_tai(current_kv_cache), - _contiguous_for_tai(logical_locs), - _contiguous_for_tai(current_locs.reshape(-1)), - _contiguous_for_tai(page_inverse), - _row_req_id_for_tai(current_req_id, current_locs.reshape(-1)), - page_size=int(page_size), - mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages), - ) + try: + return fill_kernel( + _contiguous_for_tai(dense_kv_cache), + _contiguous_for_tai(materialized_dense_locs), + _contiguous_for_tai(current_kv_cache), + _contiguous_for_tai(logical_locs), + _contiguous_for_tai(current_locs.reshape(-1)), + _contiguous_for_tai(page_inverse), + _row_req_id_for_tai(current_req_id, current_locs.reshape(-1)), + page_size=int(page_size), + mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages), + ) + except TypeError: + if int(_slot_page_inverse_2d(page_inverse).shape[0]) != 1: + raise + return fill_kernel( + _contiguous_for_tai(dense_kv_cache), + _contiguous_for_tai(materialized_dense_locs), + _contiguous_for_tai(current_kv_cache), + _contiguous_for_tai(logical_locs), + _contiguous_for_tai(current_locs.reshape(-1)), + _contiguous_for_tai(page_inverse), + page_size=int(page_size), + mask_non_current_in_current_pages=bool(mask_non_current_in_current_pages), + ) except Exception as exc: _log_tai_materialize_fallback( "fill_current_failed", @@ -1652,7 +1783,7 @@ def fill_current_kv_page_slots_and_remap_locs( current_locs: torch.Tensor, page_inverse: torch.Tensor, page_size: int, - current_req_id: torch.Tensor, + current_req_id: torch.Tensor | None = None, mask_non_current_in_current_pages: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Fill current suffix KV into preallocated dense page slots and remap locs. @@ -1665,6 +1796,9 @@ def fill_current_kv_page_slots_and_remap_locs( land on another request's dense slot (bs>1 contamination). """ + if current_req_id is None: + current_req_id = torch.zeros_like(current_locs.reshape(-1), dtype=torch.long) + tai_result = _try_tai_fill_current_kv_page_slots_and_remap_locs( dense_kv_cache=dense_kv_cache, materialized_dense_locs=materialized_dense_locs, @@ -1678,6 +1812,7 @@ def fill_current_kv_page_slots_and_remap_locs( ) if tai_result is not None: return tai_result + page_inverse = _slot_page_inverse_2d(page_inverse) if dense_kv_cache.is_cuda: _log_tai_materialize_fallback( @@ -1788,7 +1923,7 @@ def fill_current_index_page_slots( page_inverse: torch.Tensor, page_size: int, index_head_dim: int, - current_req_id: torch.Tensor, + current_req_id: torch.Tensor | None = None, ) -> torch.Tensor: """Fill current index K/scale rows into a slot-dense page buffer. @@ -1801,6 +1936,8 @@ def fill_current_index_page_slots( """ current_locs = current_locs.reshape(-1) + if current_req_id is None: + current_req_id = torch.zeros_like(current_locs, dtype=torch.long) current_req_id = current_req_id.reshape(-1) current_rows = int(current_locs.numel()) if current_rows == 0: @@ -1817,6 +1954,7 @@ def fill_current_index_page_slots( ) if tai_result is not None: return tai_result + page_inverse = _slot_page_inverse_2d(page_inverse) if dense_page_buffer.is_cuda: _log_tai_materialize_fallback( "index_current_fill_torch_reference_cuda", @@ -3100,7 +3238,7 @@ def _logical_page_capacity_from_physical_page_capacity( def build_slot_page_inverse( slot_logical_pages: torch.Tensor, logical_page_capacity: int, - batch_rows: int, + batch_rows: int = 1, ) -> torch.Tensor: """Build a PER-REQUEST logical_page -> slot_dense_page map. @@ -3223,7 +3361,7 @@ def remap_logical_locs_to_slot_dense_locs( logical_locs: torch.Tensor, page_inverse: torch.Tensor, page_size: int, - loc_req_id: torch.Tensor, + loc_req_id: torch.Tensor | None = None, ) -> torch.Tensor: """Map logical token locs through a PER-REQUEST slot page inverse. @@ -3234,6 +3372,9 @@ def remap_logical_locs_to_slot_dense_locs( """ dense_locs = torch.full_like(logical_locs, -1) + page_inverse = _slot_page_inverse_2d(page_inverse) + if loc_req_id is None: + loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long) if logical_locs.numel() == 0 or page_inverse.numel() == 0: return dense_locs @@ -3433,7 +3574,7 @@ def remap_logical_locs_to_shared_token_slot_dense_locs( def remap_logical_pages_to_slot_dense_pages( logical_pages: torch.Tensor, page_inverse: torch.Tensor, - page_req_id: torch.Tensor, + page_req_id: torch.Tensor | None = None, ) -> torch.Tensor: """Map logical page ids through a PER-REQUEST slot page inverse. @@ -3446,6 +3587,9 @@ def remap_logical_pages_to_slot_dense_pages( """ dense_pages_out = torch.full_like(logical_pages, -1) + page_inverse = _slot_page_inverse_2d(page_inverse) + if page_req_id is None: + page_req_id = build_page_table_row_req_id(logical_pages) if logical_pages.numel() == 0 or page_inverse.numel() == 0: return dense_pages_out @@ -3665,16 +3809,23 @@ def build_shared_token_kv_slot_remap( slot_dense_pages, ) ) - # ``dense_locs`` is never consumed downstream (the orchestrators remap their - # own per-request locs against ``page_inverse``). ``get_or_build_*`` always - # passes ``logical_locs=None``; keep it None to avoid needing a per-loc req id - # in the cached, loc-agnostic remap. - if logical_locs is not None: - raise NotImplementedError( - "build_shared_token_kv_slot_remap no longer materializes dense_locs; " - "remap per-request locs at the materialize call site instead." - ) + # Runtime callers cache this remap loc-agnostically and remap their own locs + # with explicit per-request ids. Unit/reference callers may still ask for + # ``dense_locs``; keep that behavior request-scoped instead of using a + # batch-global logical-page inverse. dense_locs = None + if logical_locs is not None: + loc_req_id = ( + build_page_table_row_req_id(logical_locs) + if logical_locs.dim() >= 2 + else torch.zeros_like(logical_locs, dtype=torch.long) + ) + dense_locs = remap_logical_locs_to_slot_dense_locs_optimized( + logical_locs, + page_inverse=page_inverse, + page_size=page_size, + loc_req_id=loc_req_id, + ) return SharedTokenKVSlotRemap( slot_logical_pages=slot_logical_pages, page_inverse=page_inverse, @@ -4186,8 +4337,8 @@ def materialize_prefix_and_reuse_current_kv_page_slots( layout: CpSharedKVLayout, page_size: int, prefix_pages: int, - loc_req_id: torch.Tensor, - current_req_id: torch.Tensor, + 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, @@ -4205,6 +4356,20 @@ def materialize_prefix_and_reuse_current_kv_page_slots( """ timing_start = cp_shared_kv_bs_gt1_timing_start() + if loc_req_id is None: + loc_req_id = _infer_req_id_for_locs_from_slot_remap( + logical_locs, + slot_remap, + page_size, + ) + if current_req_id is None: + current_req_id = _infer_req_id_for_locs_from_slot_remap( + current_locs.reshape(-1), + slot_remap, + page_size, + strict=True, + ) + total_slots = int(slot_remap.slot_logical_pages.numel()) if prefix_slot_spans is not None and prefix_slot_span is not None: raise ValueError( @@ -4380,7 +4545,7 @@ def materialize_prefix_and_reuse_current_index_page_slots( page_size: int, index_head_dim: int, prefix_pages: int, - current_req_id: torch.Tensor, + 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, @@ -4390,6 +4555,14 @@ def materialize_prefix_and_reuse_current_index_page_slots( """Synchronously compose prefix index materialization with current index rows.""" timing_start = cp_shared_kv_bs_gt1_timing_start() + if current_req_id is None: + current_req_id = _infer_req_id_for_locs_from_slot_remap( + current_locs.reshape(-1), + slot_remap, + page_size, + strict=True, + ) + total_slots = int(slot_remap.slot_logical_pages.numel()) if prefix_slot_spans is not None and prefix_slot_span is not None: raise ValueError( @@ -4905,9 +5078,10 @@ def materialize_shared_token_kv_buffer( dense_kv_cache = None if slot_remap is not None: if loc_req_id is None: - raise ValueError( - "materialize_shared_token_kv_buffer slot_remap path requires " - "loc_req_id for the per-request page-inverse remap." + loc_req_id = _infer_req_id_for_locs_from_slot_remap( + logical_locs, + slot_remap, + page_size, ) materialized_logical_pages = slot_remap.slot_logical_pages dense_locs = remap_logical_locs_to_shared_token_slot_dense_locs( @@ -4928,10 +5102,17 @@ def materialize_shared_token_kv_buffer( use_slot_materialize = False else: if loc_req_id is None: - raise ValueError( - "materialize_shared_token_kv_buffer remap_logical_pages path " - "requires loc_req_id for the per-request page-inverse remap." - ) + if logical_locs_row_ids is not None: + loc_req_id = logical_locs_row_ids + elif logical_locs.dim() >= 2: + loc_req_id = build_page_table_row_req_id(logical_locs) + elif remap_logical_pages.dim() < 2 or int(remap_logical_pages.shape[0]) <= 1: + loc_req_id = torch.zeros_like(logical_locs, dtype=torch.long) + else: + raise ValueError( + "materialize_shared_token_kv_buffer remap_logical_pages path " + "requires loc_req_id for multi-request 1-D logical locs." + ) _debug_assert_no_negative_tensor_values( remap_logical_pages, context="CP shared KV token materialize page remap", diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 300fdef5a..537c3b4cc 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1459,12 +1459,11 @@ class Indexer(MultiPlatformOp): else: index_buffer = None if cp_shared_kv_debug_enabled(): + cp_layout = getattr(forward_batch, "cp_shared_kv_layout", None) cp_shared_kv_debug_log( "index_current_reuse", "NSA index current reuse cp_rank=%s layer=%s kv_len=%s actual_seq_q=%s k_ck=%s s_ck=%s", - forward_batch.cp_shared_kv_layout.cp_rank - if forward_batch.cp_shared_kv_layout is not None - else None, + cp_layout.cp_rank if cp_layout is not None else None, layer_id, kv_len, actual_seq_q,