Fix CP HiCache load_cp owner-pattern mismatch (cache-hit corruption)

In CP=8 + NSA-shared-KV + HiCache disagg-prefill, cache-hit prefill produced
incoherent decode output. Cold prefill on CP was correct; pure CP without
HiCache was correct. The bug lived at the HiCache load_cp / device-alloc
interface.

Root cause: cache_controller.load_cp called the plain
mem_pool_device_allocator.alloc(logical_len), which returns logical pages
with no CP owner-pattern preservation. Cold prefill instead uses
alloc_extend_compute_owner with a zigzag owner pattern from
build_in_seq_page_compute_owners. The saved CpHiCacheNodeMetadata.owned_positions
records WHICH POSITIONS in the write-time alloc were owned by this rank. At
load time, those same positions are applied to a new alloc whose per-position
owner pattern is arbitrary -- each rank loads its host bytes into physical
slots whose corresponding logical page is owned by a DIFFERENT rank.
Attention's materialize_shared_token_kv_buffer reads from the owner's
physical slot, which was never loaded. Result: garbage.

Fix:
- CpHiCacheNodeMetadata gains two required fields: page_owners (int8 per
  logical page, identical on all CP ranks) and page_size. __post_init__
  validates; split() bisects page_owners by page index with a page-alignment
  check.
- _write_cp derives page_owners from device_indices (page-first slot of each
  page -> logical page id -> layout.owner_for_logical_pages) and stores in
  both metadata-construction sites (zero-owned and normal).
- New CPSharedPagedTokenToKVPoolAllocator.alloc_pages_with_owners() reuses
  _select_compute_owner_pages (with its tai-kernel fast path) and returns
  page-contiguous token locs whose per-page owner sequence equals the input.
- load_cp now concats page_owners across nodes_to_load and calls
  alloc_pages_with_owners. On None (lane exhausted) the caller hits the
  retry-with-eviction path; further failure returns None and degrades to
  cache miss. No silent fallback to plain alloc -- that recreated the bug.
- load_back retry path now calls _evict_for_compute_owner_lanes (module-top
  import) instead of plain evict(); this targets the deficit lane and gives
  the next alloc attempt a chance to satisfy it.
- envs import moved to module top in cache_controller.py per code-review
  feedback. Removed an over-defensive owned_check.all().item() in load_cp
  that would have re-introduced the host-sync anti-pattern 97a9f850c
  removed -- the invariant is already guaranteed by alloc_pages_with_owners.

Tests: 40 existing CpHiCacheNodeMetadata constructions migrated to pass the
new required fields. 9 new metadata tests (validators + split page-alignment).
10 new allocator tests in test_alloc_pages_with_owners.py covering input-order
preservation, lane exhaustion, release_pages fallback, debug-mode invariant.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

python/sglang/srt/mem_cache/allocator.py

@@ -763,3 +763,55 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
             assert torch.equal(selected_owners, expected_owners)
 
         return out_indices
+
+    def alloc_pages_with_owners(
+        self,
+        page_compute_owners: List[int],
+    ) -> Optional[torch.Tensor]:
+        """Allocate ``len(page_compute_owners)`` whole logical pages where
+        page ``i``'s owner equals ``page_compute_owners[i]``.
+
+        Returns a flat int64 tensor of logical token locs of length
+        ``len(owners) * page_size``, page-contiguous in input order, or
+        ``None`` when an owner lane is exhausted (caller handles
+        eviction-retry / cache-miss).
+
+        Used by CP HiCache ``load_cp`` to reproduce the write-time owner
+        pattern recorded in ``CpHiCacheNodeMetadata.page_owners``. Without
+        this, plain ``alloc()`` would return pages with an arbitrary owner
+        pattern, the saved per-position owner mask would index the wrong
+        slots, and each rank would load its host bytes into physical slots
+        whose corresponding logical page is owned by another rank → attention
+        reads garbage.
+
+        Distinguished from ``alloc_extend_compute_owner``: that variant
+        expects extend semantics (prefix_lens / seq_lens / last_loc) and
+        invokes the ``alloc_extend_naive`` Triton kernel to splice the
+        partial last page of a prefix with new pages. HiCache reload has
+        no prefix and no partial page — host pages are always page-aligned
+        — so a pure page-fresh allocation is enough.
+        """
+        if not page_compute_owners:
+            return torch.empty((0,), dtype=torch.int64, device=self.device)
+        selected = self._select_compute_owner_pages(page_compute_owners)
+        if selected is None:
+            return None
+        selected_pages, selected_free_mask, selected_release_mask = selected
+        page_size = self.page_size
+        base = selected_pages.to(torch.int64).unsqueeze(1) * page_size
+        offsets = torch.arange(
+            page_size, dtype=torch.int64, device=self.device
+        ).unsqueeze(0)
+        out_indices = (base + offsets).reshape(-1)
+        self.free_pages = self.free_pages[~selected_free_mask]
+        self.release_pages = self.release_pages[~selected_release_mask]
+        if self.debug_mode:
+            assert torch.unique(out_indices).numel() == out_indices.numel()
+            check_owners = torch.remainder(selected_pages - 1, self.cp_size)
+            expected = torch.tensor(
+                page_compute_owners,
+                dtype=check_owners.dtype,
+                device=check_owners.device,
+            )
+            assert torch.equal(check_owners, expected)
+        return out_indices