Keep CP current IPC staging proportional to touched pages

Cache-hit bs>1 current reuse can create very large dense attention buffers
while touching only a small set of current pages. The previous SGLang runtime
asked tai-kernel for a staging buffer sized like the full dense tensor, which
caused CUDA OOM before the current IPC fast path could run.

Switch token and index current IPC helpers to descriptor-compact staging: publish
the dense destination pages into compact staging slots and materialize peers
from compact source page ids back to the original dense destination pages.
Document the failure mode and the compact-staging contract so the dense-sized
contract is not reintroduced.

Constraint: CUDA + SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 must fail fast instead of silently falling back to current-slot all_reduce
Rejected: Let staging allocation failure fall back to all_reduce | hides the bug and restores the expensive collective path
Rejected: Size staging by the full dense tensor | reproduces the 965MB staging OOM on long-prefix cache-hit batches
Confidence: high
Scope-risk: moderate
Directive: Current IPC helper source ids are compact staging ids; destination ids remain dense slot pages
Tested: Remote cjy-glm5-new PYTHONPATH=python:/mnt/beegfs/cjy/tai-kernel/python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 144 passed, 2 subtests passed
Tested: Local py_compile cp_shared_kv_runtime.py
Not-tested: Full ETE service restart with production traffic after this commit
(cherry picked from commit 906ecbe5d4f08b73242e98e2b628e26516d5b04a)

python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py

@@ -460,6 +460,7 @@ def _load_tai_ipc_kernels():
             "materialize_cuda_ipc_peer_pages_slot_indices",
             "materialize_cuda_ipc_peer_pages_slot_indices_wait_ready",
             "publish_cuda_ipc_slot_pages_and_mark_ready",
+            "publish_cuda_ipc_slot_pages_compact_and_mark_ready",
         )
         missing = [name for name in required if not hasattr(tai_ipc, name)]
         if missing:
@@ -2534,15 +2535,17 @@ def _build_current_staging_ipc_descriptors(
     layout: CpSharedKVLayout,
     spans: list[tuple[int, int]],
     device: torch.device,
+    slot_indices: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     flat_slot_logical_pages = _contiguous_for_tai(
         slot_logical_pages.reshape(-1).to(device=device)
     )
-    slot_indices = _slot_spans_to_cuda_slot_indices(
-        spans,
-        total_slots=int(flat_slot_logical_pages.numel()),
-        device=device,
-    )
+    if slot_indices is None:
+        slot_indices = _slot_spans_to_cuda_slot_indices(
+            spans,
+            total_slots=int(flat_slot_logical_pages.numel()),
+            device=device,
+        )
     if slot_indices.numel() == 0:
         empty = torch.empty((0,), dtype=torch.long, device=device)
         return empty, empty, empty
@@ -2559,6 +2562,21 @@ def _build_current_staging_ipc_descriptors(
     return owner_ranks.contiguous(), src_page_indices, dense_page_indices
 
 
+def _build_compact_current_staging_src_page_indices(
+    src_page_indices: torch.Tensor,
+) -> torch.Tensor:
+    compact_src = torch.arange(
+        int(src_page_indices.numel()),
+        dtype=torch.long,
+        device=src_page_indices.device,
+    )
+    return torch.where(
+        src_page_indices >= 0,
+        compact_src,
+        torch.full_like(compact_src, -1),
+    ).contiguous()
+
+
 def _try_tai_ipc_materialize_token_kv_page_slot_spans_into(
     *,
     kv_cache: torch.Tensor,
@@ -2646,9 +2664,15 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into(
     if not spans:
         return True
     page_nbytes = _token_kv_page_nbytes(dense_kv_cache, page_size)
-    required_nbytes = int(dense_kv_cache.shape[0]) * _page_nbytes_from_page_tensor(
-        dense_kv_cache
+    slot_indices = _slot_spans_to_cuda_slot_indices(
+        spans,
+        total_slots=int(slot_logical_pages.reshape(-1).numel()),
+        device=dense_kv_cache.device,
     )
+    if slot_indices.numel() == 0:
+        return True
+    dense_page_indices = (slot_indices + 1).to(torch.long).contiguous()
+    required_nbytes = int(dense_page_indices.numel()) * int(page_nbytes)
     staging_state = _get_or_create_tai_ipc_current_staging(
         kind="token",
         dense_tensor=dense_kv_cache,
@@ -2665,13 +2689,15 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into(
                 layout=layout,
                 spans=spans,
                 device=dense_kv_cache.device,
+                slot_indices=slot_indices,
             )
         )
-        if dense_page_indices.numel() == 0:
-            return True
+        compact_src_page_indices = _build_compact_current_staging_src_page_indices(
+            src_page_indices
+        )
         state.ready_seq += 1
         ready_seq = int(state.ready_seq)
-        kernels.publish_cuda_ipc_slot_pages_and_mark_ready(
+        kernels.publish_cuda_ipc_slot_pages_compact_and_mark_ready(
             dense_kv_cache,
             state.staging,
             dense_page_indices,
@@ -2684,7 +2710,7 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into(
             state.ready_peer_ptrs,
             dense_kv_cache,
             owner_ranks,
-            src_page_indices,
+            compact_src_page_indices,
             dense_page_indices,
             ready_seq=ready_seq,
             page_nbytes=page_nbytes,
@@ -2866,7 +2892,15 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into(
     if not spans:
         return True
     page_nbytes = _page_nbytes_from_page_tensor(dense_page_buffer)
-    required_nbytes = int(dense_page_buffer.shape[0]) * page_nbytes
+    slot_indices = _slot_spans_to_cuda_slot_indices(
+        spans,
+        total_slots=int(slot_logical_pages.reshape(-1).numel()),
+        device=dense_page_buffer.device,
+    )
+    if slot_indices.numel() == 0:
+        return True
+    dense_page_indices = (slot_indices + 1).to(torch.long).contiguous()
+    required_nbytes = int(dense_page_indices.numel()) * int(page_nbytes)
     staging_state = _get_or_create_tai_ipc_current_staging(
         kind="paged",
         dense_tensor=dense_page_buffer,
@@ -2883,13 +2917,15 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into(
                 layout=layout,
                 spans=spans,
                 device=dense_page_buffer.device,
+                slot_indices=slot_indices,
             )
         )
-        if dense_page_indices.numel() == 0:
-            return True
+        compact_src_page_indices = _build_compact_current_staging_src_page_indices(
+            src_page_indices
+        )
         state.ready_seq += 1
         ready_seq = int(state.ready_seq)
-        kernels.publish_cuda_ipc_slot_pages_and_mark_ready(
+        kernels.publish_cuda_ipc_slot_pages_compact_and_mark_ready(
             dense_page_buffer,
             state.staging,
             dense_page_indices,
@@ -2902,7 +2938,7 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into(
             state.ready_peer_ptrs,
             dense_page_buffer,
             owner_ranks,
-            src_page_indices,
+            compact_src_page_indices,
             dense_page_indices,
             ready_seq=ready_seq,
             page_nbytes=page_nbytes,