Reduce CP shared-KV request-boundary stalls

CP shared KV now avoids the PyTorch sort/search remap for the single-request current-only path by deriving compact rows from page-level inverse mapping. The same change keeps sort NVTX attribution gated and splits high-frequency MoE sort markers behind a separate env var so profiling does not perturb normal runs.

Decode-side disaggregation prealloc also avoids rebuilding large token index tensors and records finer allocation timing, while compute-owner allocation/free tests cover the shared-KV page-lane behavior.

Constraint: The runtime tree used for validation is the remote /sgl-workspace/sglang-tai mount, which is not itself a Git repository, so these tracked files were synchronized into the local repo before commit.

Rejected: Keep torch.sort/searchsorted for current remap | it emits ATen/CCCL radixSortKVInPlace kernels in the attention hot path.

Rejected: Enable MoE sort NVTX under the generic sort env | the MoE preprocess sort is too frequent and can make profiling look like a hang.

Confidence: medium

Scope-risk: moderate

Directive: Do not reintroduce token-level torch.sort/searchsorted in CP shared-KV current remap without profiling the attention hot path under Nsight.

Tested: Remote container py_compile for modified runtime files; git diff --cached --check.

Not-tested: Full multi-node GLM5 PD throughput/profile rerun after the page-inverse current remap.

test/registered/unit/mem_cache/test_cp_shared_kv_layout.py

@@ -1,4 +1,5 @@
 import unittest
+from unittest.mock import patch
 
 import numpy as np
 import torch
@@ -244,6 +245,170 @@ class TestCPSharedPagedAllocator(unittest.TestCase):
         )
         self.assertEqual(allocator.available_size(), page_size * 16)
 
+    def test_compute_owner_alloc_does_not_use_torch_isin_for_page_removal(self):
+        from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
+
+        page_size = 64
+        cp_size = 4
+        page_compute_owners = [0, 2, 0, 2, 3]
+        allocator = CPSharedPagedTokenToKVPoolAllocator(
+            logical_size=page_size * 32,
+            physical_size=page_size * 8,
+            page_size=page_size,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+            cp_size=cp_size,
+            cp_rank=0,
+        )
+        allocator.free_pages = torch.tensor(
+            [9, 1, 3, 5, 7, 4]
+            + [page for page in range(2, 33) if page not in {3, 4, 5, 7, 9}],
+            dtype=torch.int64,
+        )
+
+        with patch(
+            "torch.isin",
+            side_effect=AssertionError("compute-owner alloc should not use torch.isin"),
+        ):
+            locs = allocator.alloc_extend_compute_owner(
+                prefix_lens=torch.tensor([0], dtype=torch.int64),
+                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                seq_lens=torch.tensor(
+                    [page_size * len(page_compute_owners)], dtype=torch.int64
+                ),
+                seq_lens_cpu=torch.tensor(
+                    [page_size * len(page_compute_owners)], dtype=torch.int64
+                ),
+                last_loc=torch.tensor([-1], dtype=torch.int64),
+                extend_num_tokens=page_size * len(page_compute_owners),
+                page_compute_owners=page_compute_owners,
+            )
+
+        self.assertIsNotNone(locs)
+        logical_pages = locs.view(-1, page_size)[:, 0] // page_size
+        self.assertEqual(logical_pages.tolist(), [9, 3, 1, 7, 4])
+        self.assertEqual(
+            ((logical_pages - 1) % cp_size).tolist(),
+            page_compute_owners,
+        )
+        self.assertEqual(
+            allocator.available_size(),
+            page_size * (32 - len(page_compute_owners)),
+        )
+        for selected_page in logical_pages.tolist():
+            self.assertNotIn(selected_page, allocator.free_pages.tolist())
+
+    def test_compute_owner_alloc_can_select_release_pages_without_sort_merge(self):
+        from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
+
+        page_size = 64
+        cp_size = 4
+        allocator = CPSharedPagedTokenToKVPoolAllocator(
+            logical_size=page_size * 16,
+            physical_size=page_size * 4,
+            page_size=page_size,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=True,
+            cp_size=cp_size,
+            cp_rank=0,
+        )
+        allocator.free_pages = torch.tensor([2, 3, 4], dtype=torch.int64)
+        allocator.release_pages = torch.tensor([1, 5], dtype=torch.int64)
+
+        with patch(
+            "torch.sort",
+            side_effect=AssertionError("compute-owner alloc should not sort-merge"),
+        ):
+            locs = allocator.alloc_extend_compute_owner(
+                prefix_lens=torch.tensor([0], dtype=torch.int64),
+                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                seq_lens=torch.tensor([page_size * 2], dtype=torch.int64),
+                seq_lens_cpu=torch.tensor([page_size * 2], dtype=torch.int64),
+                last_loc=torch.tensor([-1], dtype=torch.int64),
+                extend_num_tokens=page_size * 2,
+                page_compute_owners=[0, 0],
+            )
+
+        self.assertIsNotNone(locs)
+        logical_pages = locs.view(-1, page_size)[:, 0] // page_size
+        self.assertEqual(logical_pages.tolist(), [1, 5])
+        self.assertEqual(allocator.free_pages.tolist(), [2, 3, 4])
+        self.assertEqual(allocator.release_pages.tolist(), [])
+
+    def test_compute_owner_alloc_does_not_evict_lanes_when_first_try_succeeds(self):
+        from types import SimpleNamespace
+
+        from sglang.srt.mem_cache import common
+
+        page_size = 64
+
+        class FakeAllocator:
+            def __init__(self):
+                self.page_size = page_size
+                self.cp_size = 4
+                self.calls = 0
+
+            def available_size(self):
+                return page_size * 100
+
+            def alloc_extend_compute_owner(
+                self,
+                _prefix_lens,
+                _prefix_lens_cpu,
+                _seq_lens,
+                _seq_lens_cpu,
+                _last_loc,
+                extend_num_tokens,
+                _page_compute_owners,
+            ):
+                self.calls += 1
+                return torch.arange(extend_num_tokens, dtype=torch.int64)
+
+            def alloc_extend(self, *_args, **_kwargs):
+                raise AssertionError("legacy allocation should not be used")
+
+        class FakeTreeCache:
+            def __init__(self, allocator):
+                self.token_to_kv_pool_allocator = allocator
+
+            def is_chunk_cache(self):
+                return False
+
+            def evict(self, _params):
+                raise AssertionError("tree eviction should not be used")
+
+        allocator = FakeAllocator()
+        server_args = SimpleNamespace(
+            enable_nsa_prefill_cp_shared_kv=True,
+            enable_nsa_prefill_context_parallel=True,
+            nsa_prefill_cp_mode="in-seq-split",
+        )
+
+        with (
+            patch.object(common, "get_global_server_args", return_value=server_args),
+            patch.object(
+                common,
+                "_evict_for_compute_owner_lanes",
+                side_effect=AssertionError("lane eviction should be lazy"),
+            ),
+        ):
+            out = common.alloc_paged_token_slots_extend(
+                tree_cache=FakeTreeCache(allocator),
+                prefix_lens=torch.tensor([0], dtype=torch.int64),
+                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                seq_lens=torch.tensor([page_size * 8], dtype=torch.int64),
+                seq_lens_cpu=torch.tensor([page_size * 8], dtype=torch.int64),
+                last_loc=torch.tensor([-1], dtype=torch.int64),
+                extend_num_tokens=page_size * 8,
+            )
+
+        self.assertEqual(allocator.calls, 1)
+        self.assertEqual(out.numel(), page_size * 8)
+
     def test_compute_owner_lane_eviction_recovers_exhausted_owner_lane(self):
         from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
         from sglang.srt.mem_cache.base_prefix_cache import EvictResult

test/registered/unit/mem_cache/test_req_to_token_pool.py

@@ -7,7 +7,14 @@ from unittest.mock import MagicMock, patch
 
 import torch
 
-from sglang.srt.disaggregation.decode import DecodePreallocQueue, DecodeReqToTokenPool
+from sglang.srt.disaggregation.decode import (
+    DecodePreallocQueue,
+    DecodeReqToTokenPool,
+    _kv_locs_to_page_indices_cpu,
+)
+from sglang.srt.disaggregation.prefill import (
+    _kv_locs_to_page_indices_cpu as _prefill_kv_locs_to_page_indices_cpu,
+)
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -162,6 +169,58 @@ class TestDecodePreallocQueue(CustomTestCase):
         self.assertEqual(allocator_calls[1]["last_loc"].item(), -1)
         self.assertEqual(allocator_calls[1]["extend_num_tokens"], 5)
 
+    def test_kv_locs_to_page_indices_cpu_copies_only_page_starts(self):
+        kv_locs = torch.tensor(
+            [
+                128,
+                129,
+                130,
+                131,
+                448,
+                449,
+                450,
+                451,
+                640,
+            ],
+            dtype=torch.int64,
+        )
+
+        page_indices = _kv_locs_to_page_indices_cpu(kv_locs, page_size=4)
+
+        self.assertEqual(page_indices.dtype.name, "int32")
+        self.assertEqual(page_indices.tolist(), [32, 112, 160])
+
+    def test_kv_locs_to_page_indices_cpu_respects_num_tokens(self):
+        kv_locs = torch.arange(128, 128 + 16, dtype=torch.int64)
+
+        page_indices = _kv_locs_to_page_indices_cpu(
+            kv_locs, page_size=4, num_tokens=9
+        )
+
+        self.assertEqual(page_indices.dtype.name, "int32")
+        self.assertEqual(page_indices.tolist(), [32, 33, 34])
+
+    def test_prefill_kv_locs_to_page_indices_cpu_copies_only_page_starts(self):
+        kv_locs = torch.tensor(
+            [
+                256,
+                257,
+                258,
+                259,
+                512,
+                513,
+                514,
+                515,
+                768,
+            ],
+            dtype=torch.int64,
+        )
+
+        page_indices = _prefill_kv_locs_to_page_indices_cpu(kv_locs, page_size=4)
+
+        self.assertEqual(page_indices.dtype.name, "int32")
+        self.assertEqual(page_indices.tolist(), [64, 128, 192])
+
 
 if __name__ == "__main__":
     unittest.main()