Reduce CP HiCache capacity synchronization to owner-lane logic

CP shared KV and HiCache now use owner-lane metadata as the authoritative capacity view for host write admission and GPU load-back planning. This removes the debug scalar capacity env and keeps CP load-back from relying on a rank-wide scalar collective when per-owner availability is already known. The load-back planner also accounts for evicting child leaves that unlock ancestor device residency, which fixes small lane deficits despite large aggregate evictable capacity. The commit also adds gated CPU timing logs for CP shared-KV MLA/index prefetch and a CUDA microbenchmark for comparing dense all-reduce with owner-packed all-gather layouts. The timing logs are intentionally behind the existing MLA prefetch log env and should not be enabled for throughput measurements. Constraint: CP shared KV owner lanes require target/draft capacity decisions to preserve page_owners rather than total-token scalars Constraint: CUDA collective benchmarks must run on target GPU hosts, not locally Rejected: Keep SGLANG_CP_HICACHE_CAPACITY_DEBUG observer env | owner-lane admission now replaces that scalar debug path Rejected: Add a silent scalar-allreduce fallback | unexpected owner-lane mismatch should fail fast or log loudly Confidence: medium Scope-risk: moderate Directive: Do not reintroduce CP capacity collectives on the scheduler hot path without proving the owner-lane metadata is insufficient Directive: Disable SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH for end-to-end performance runs; it is diagnostic and high-volume Tested: git diff --check Tested: python -m py_compile on changed runtime/test/benchmark Python files Tested: remote pytest -q test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py (81 passed, 5 warnings) Not-tested: CUDA benchmark benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py Not-tested: full GLM5 E2E throughput after this commit
2026-05-28 08:31:49 +08:00
parent ff33446787
commit 25f2147677
8 changed files with 996 additions and 201 deletions
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -215,7 +215,6 @@ class Envs:
    SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1)
    SGLANG_CP_DRAFT_SHARED_KV = EnvBool(False)
    SGLANG_CP_DRAFT_SHARED_KV_DEBUG = EnvBool(False)
-    SGLANG_CP_HICACHE_CAPACITY_DEBUG = EnvBool(False)
    SGLANG_DISABLE_TAI_BIGRAM = EnvBool(False)
    SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
    SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
--- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
+++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py
@@ -1,6 +1,7 @@
 from __future__ import annotations

 import logging
+import time
 from dataclasses import dataclass
 from typing import Any, Optional

@@ -107,6 +108,10 @@ def _debug_owned_pages_count(
    try:
        if logical_pages.numel() == 0:
            return 0
+        if logical_pages.is_cuda:
+            # Temporary CPU timing logs must not add a CUDA synchronization via
+            # Tensor.item(); skip exact debug counts for GPU metadata.
+            return -1
        return int(layout.owned_pages_mask(logical_pages).sum().item())
    except Exception:
        logger.exception("Failed to count CP shared KV prefetch owned pages.")
@@ -124,6 +129,145 @@ def _debug_handle_keys(
    return tuple(sorted(handles.keys()))


+def _cpu_timing_start() -> float:
+    if not cp_shared_kv_mla_prefetch_log_enabled():
+        return 0.0
+    return time.perf_counter()
+
+
+def _cpu_timing_ms(start: float) -> float:
+    if start <= 0.0:
+        return -1.0
+    return (time.perf_counter() - start) * 1000.0
+
+
+@dataclass
+class _PrefetchCpuTiming:
+    start_count: int = 0
+    start_total_ms: float = 0.0
+    start_max_ms: float = 0.0
+    get_total_ms: float = 0.0
+    materialize_total_ms: float = 0.0
+    reduce_enqueue_total_ms: float = 0.0
+    consume_count: int = 0
+    consume_total_ms: float = 0.0
+    consume_wait_total_ms: float = 0.0
+    consume_suffix_total_ms: float = 0.0
+    consume_remap_total_ms: float = 0.0
+
+    def record_start(
+        self,
+        *,
+        total_ms: float,
+        get_ms: float,
+        materialize_ms: float,
+        reduce_enqueue_ms: float,
+    ) -> None:
+        if total_ms < 0.0:
+            return
+        self.start_count += 1
+        self.start_total_ms += total_ms
+        self.start_max_ms = max(self.start_max_ms, total_ms)
+        self.get_total_ms += max(get_ms, 0.0)
+        self.materialize_total_ms += max(materialize_ms, 0.0)
+        self.reduce_enqueue_total_ms += max(reduce_enqueue_ms, 0.0)
+
+    def record_consume(
+        self,
+        *,
+        total_ms: float,
+        wait_ms: float,
+        suffix_ms: float,
+        remap_ms: float,
+    ) -> None:
+        if total_ms < 0.0:
+            return
+        self.consume_count += 1
+        self.consume_total_ms += total_ms
+        self.consume_wait_total_ms += max(wait_ms, 0.0)
+        self.consume_suffix_total_ms += max(suffix_ms, 0.0)
+        self.consume_remap_total_ms += max(remap_ms, 0.0)
+
+
+def _log_prefetch_cpu_start(
+    *,
+    log_fn: Any,
+    layout: CpSharedKVLayout,
+    timing: _PrefetchCpuTiming,
+    path: str,
+    layer_id: int,
+    total_ms: float,
+    get_ms: float,
+    materialize_ms: float,
+    reduce_enqueue_ms: float,
+    prefix_pages: int,
+    total_slots: int,
+    dense_units: int,
+) -> None:
+    if not cp_shared_kv_mla_prefetch_log_enabled() or total_ms < 0.0:
+        return
+    if cp_shared_kv_mla_prefetch_should_log_layer(layer_id):
+        log_fn(
+            "cpu_timing path=%s stage=start layer=%s total_ms=%.3f "
+            "get_ms=%.3f materialize_ms=%.3f reduce_enqueue_ms=%.3f "
+            "prefix_pages=%s total_slots=%s dense_units=%s",
+            path,
+            layer_id,
+            total_ms,
+            get_ms,
+            materialize_ms,
+            reduce_enqueue_ms,
+            prefix_pages,
+            total_slots,
+            dense_units,
+        )
+    if layout.cp_rank == 0 and timing.start_count > 0 and timing.start_count % 16 == 0:
+        starts = timing.start_count
+        log_fn(
+            "cpu_summary path=%s starts=%s avg_start_ms=%.3f max_start_ms=%.3f "
+            "avg_get_ms=%.3f avg_materialize_ms=%.3f avg_reduce_enqueue_ms=%.3f",
+            path,
+            starts,
+            timing.start_total_ms / starts,
+            timing.start_max_ms,
+            timing.get_total_ms / starts,
+            timing.materialize_total_ms / starts,
+            timing.reduce_enqueue_total_ms / starts,
+        )
+
+
+def _log_prefetch_cpu_consume(
+    *,
+    log_fn: Any,
+    timing: _PrefetchCpuTiming,
+    path: str,
+    layer_id: int,
+    total_ms: float,
+    wait_ms: float,
+    suffix_ms: float,
+    remap_ms: float,
+    prefix_pages: int,
+    suffix_slots: int,
+) -> None:
+    if not cp_shared_kv_mla_prefetch_log_enabled() or total_ms < 0.0:
+        return
+    if cp_shared_kv_mla_prefetch_should_log_layer(layer_id):
+        log_fn(
+            "cpu_timing path=%s stage=consume layer=%s total_ms=%.3f "
+            "wait_event_ms=%.3f suffix_ms=%.3f remap_ms=%.3f "
+            "prefix_pages=%s suffix_slots=%s consume_count=%s",
+            path,
+            layer_id,
+            total_ms,
+            wait_ms,
+            suffix_ms,
+            remap_ms,
+            prefix_pages,
+            suffix_slots,
+            timing.consume_count,
+        )
+
+
@dataclass
 class CpSharedKVMlaPrefetchHandle:
    layer_id: int
@@ -174,6 +318,7 @@ class CpSharedKVMlaPrefetcher:
        self.handles: dict[int, CpSharedKVMlaPrefetchHandle] = {}
        self.pending_attention_handle: Optional[CpSharedKVMlaPrefetchHandle] = None
        self.disabled = False
+        self._cpu_timing = _PrefetchCpuTiming()

    @classmethod
    def maybe_create(
@@ -280,6 +425,8 @@ class CpSharedKVMlaPrefetcher:
            return None

        prefetch_stream = stream if stream is not None else torch.cuda.Stream()
+        create_cpu = _cpu_timing_start()
+        get_cpu = _cpu_timing_start()
        try:
            first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0))
            kv_cache = _prefetch_pool_get_key_buffer(
@@ -288,6 +435,8 @@ class CpSharedKVMlaPrefetcher:
                stream=prefetch_stream,
                path="mla",
            )
+            get_ms = _cpu_timing_ms(get_cpu)
+            remap_cpu = _cpu_timing_start()
            remap = get_or_build_shared_token_kv_slot_remap(
                forward_batch,
                kv_cache=kv_cache,
@@ -295,6 +444,7 @@ class CpSharedKVMlaPrefetcher:
                layout=layout,
                page_size=page_size,
            )
+            remap_ms = _cpu_timing_ms(remap_cpu)
        except Exception:
            logger.exception("Failed to initialize CP shared KV MLA prefetcher.")
            return None
@@ -316,6 +466,20 @@ class CpSharedKVMlaPrefetcher:
            remap.dense_num_pages,
            page_size,
        )
+        create_total_ms = _cpu_timing_ms(create_cpu)
+        _prefetch_log(
+            "cpu_timing path=mla stage=create cp_rank=%s cp_size=%s "
+            "total_ms=%.3f get_ms=%.3f remap_ms=%.3f prefix_pages=%s "
+            "total_slots=%s dense_pages=%s",
+            layout.cp_rank,
+            layout.cp_size,
+            create_total_ms,
+            get_ms,
+            remap_ms,
+            prefix_pages,
+            int(remap.slot_logical_pages.numel()),
+            remap.dense_num_pages,
+        )

        return cls(
            layout=layout,
@@ -386,9 +550,13 @@ class CpSharedKVMlaPrefetcher:
            )
            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_kv_cache = handle.dense_kv_cache
        suffix_slots = self.total_slots - self.prefix_pages
+        suffix_ms = 0.0

        if self.prefix_pages < self.total_slots:
            self._log_layer(
@@ -400,6 +568,7 @@ class CpSharedKVMlaPrefetcher:
                self.total_slots,
                suffix_slots,
            )
+            suffix_cpu = _cpu_timing_start()
            materialize_local_token_kv_page_slots_into(
                kv_cache=kv_cache,
                dense_kv_cache=dense_kv_cache,
@@ -430,6 +599,7 @@ class CpSharedKVMlaPrefetcher:
                suffix_rows.start,
                suffix_rows.stop,
            )
+            suffix_ms = _cpu_timing_ms(suffix_cpu)

        self._log_layer(
            layer_id,
@@ -440,6 +610,7 @@ class CpSharedKVMlaPrefetcher:
            int(dense_kv_cache.shape[0]),
        )

+        remap_cpu = _cpu_timing_start()
        logical_locs = filter_locs_mappable_to_physical_pool(
            logical_locs=logical_locs,
            layout=self.layout,
@@ -450,6 +621,26 @@ class CpSharedKVMlaPrefetcher:
            page_inverse=self.page_inverse,
            page_size=self.page_size,
        )
+        remap_ms = _cpu_timing_ms(remap_cpu)
+        total_ms = _cpu_timing_ms(consume_cpu)
+        self._cpu_timing.record_consume(
+            total_ms=total_ms,
+            wait_ms=wait_ms,
+            suffix_ms=suffix_ms,
+            remap_ms=remap_ms,
+        )
+        _log_prefetch_cpu_consume(
+            log_fn=self._log,
+            timing=self._cpu_timing,
+            path="mla",
+            layer_id=layer_id,
+            total_ms=total_ms,
+            wait_ms=wait_ms,
+            suffix_ms=suffix_ms,
+            remap_ms=remap_ms,
+            prefix_pages=self.prefix_pages,
+            suffix_slots=suffix_slots,
+        )
        return dense_kv_cache, dense_locs

    def start_next_layer_prefix(
@@ -496,6 +687,8 @@ class CpSharedKVMlaPrefetcher:
            )
            return

+        start_cpu = _cpu_timing_start()
+        get_cpu = _cpu_timing_start()
        try:
            kv_cache = _prefetch_pool_get_key_buffer(
                token_to_kv_pool=token_to_kv_pool,
@@ -503,6 +696,7 @@ class CpSharedKVMlaPrefetcher:
                stream=self.stream,
                path="mla",
            )
+            get_ms = _cpu_timing_ms(get_cpu)
        except Exception:
            logger.exception(
                "Failed to get next-layer KV cache for CP shared KV MLA prefetch."
@@ -517,33 +711,36 @@ class CpSharedKVMlaPrefetcher:

        try:
            current_stream = torch.cuda.current_stream()
-            self.stream.wait_stream(current_stream)
            prefix_rows = slot_range_to_token_slice(
                self.page_size,
                0,
                self.prefix_pages,
            )
+            materialize_cpu = _cpu_timing_start()
+            dense_kv_cache = kv_cache.new_zeros(
+                (self.dense_num_pages * self.page_size, *kv_cache.shape[1:])
+            )
+            self._log_next_layer(
+                next_layer_id,
+                "start_prefix_begin next_layer=%s start_slot=0 end_slot=%s "
+                "dense_rows=%s",
+                next_layer_id,
+                self.prefix_pages,
+                int(dense_kv_cache.shape[0]),
+            )
+            materialize_local_token_kv_page_slots_into(
+                kv_cache=kv_cache,
+                dense_kv_cache=dense_kv_cache,
+                slot_logical_pages=self.slot_logical_pages,
+                layout=self.layout,
+                page_size=self.page_size,
+                start_slot=0,
+                end_slot=self.prefix_pages,
+            )
+            materialize_ms = _cpu_timing_ms(materialize_cpu)
+            reduce_cpu = _cpu_timing_start()
+            self.stream.wait_stream(current_stream)
            with torch.cuda.stream(self.stream):
-                dense_kv_cache = kv_cache.new_zeros(
-                    (self.dense_num_pages * self.page_size, *kv_cache.shape[1:])
-                )
-                self._log_next_layer(
-                    next_layer_id,
-                    "start_prefix_begin next_layer=%s start_slot=0 end_slot=%s "
-                    "dense_rows=%s",
-                    next_layer_id,
-                    self.prefix_pages,
-                    int(dense_kv_cache.shape[0]),
-                )
-                materialize_local_token_kv_page_slots_into(
-                    kv_cache=kv_cache,
-                    dense_kv_cache=dense_kv_cache,
-                    slot_logical_pages=self.slot_logical_pages,
-                    layout=self.layout,
-                    page_size=self.page_size,
-                    start_slot=0,
-                    end_slot=self.prefix_pages,
-                )
                event = _all_reduce_materialized_buffer_async(
                    dense_kv_cache[prefix_rows],
                    cp_size=self.layout.cp_size,
@@ -553,6 +750,7 @@ class CpSharedKVMlaPrefetcher:
                    nvtx_cp_rank=self.layout.cp_rank,
                    nvtx_rows=(prefix_rows.start, prefix_rows.stop),
                )
+            reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu)
            if event is None:
                self.disabled = True
                logger.warning(
@@ -571,6 +769,27 @@ class CpSharedKVMlaPrefetcher:
                    next_layer_id,
                )
                return
+            total_ms = _cpu_timing_ms(start_cpu)
+            self._cpu_timing.record_start(
+                total_ms=total_ms,
+                get_ms=get_ms,
+                materialize_ms=materialize_ms,
+                reduce_enqueue_ms=reduce_enqueue_ms,
+            )
+            _log_prefetch_cpu_start(
+                log_fn=self._log,
+                layout=self.layout,
+                timing=self._cpu_timing,
+                path="mla",
+                layer_id=next_layer_id,
+                total_ms=total_ms,
+                get_ms=get_ms,
+                materialize_ms=materialize_ms,
+                reduce_enqueue_ms=reduce_enqueue_ms,
+                prefix_pages=self.prefix_pages,
+                total_slots=self.total_slots,
+                dense_units=int(dense_kv_cache.shape[0]),
+            )
        except Exception:
            logger.exception("Failed to start CP shared KV MLA prefix prefetch.")
            self.disabled = True
@@ -717,6 +936,7 @@ class CpSharedKVIndexPrefetcher:
        self.handles: dict[int, CpSharedKVIndexPrefetchHandle] = {}
        self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None
        self.disabled = False
+        self._cpu_timing = _PrefetchCpuTiming()

    @classmethod
    def maybe_create(
@@ -864,6 +1084,8 @@ class CpSharedKVIndexPrefetcher:
            return None

        prefetch_stream = stream if stream is not None else torch.cuda.Stream()
+        create_cpu = _cpu_timing_start()
+        get_cpu = _cpu_timing_start()
        try:
            first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0))
            page_buffer = _prefetch_pool_get_index_buffer(
@@ -871,12 +1093,15 @@ class CpSharedKVIndexPrefetcher:
                layer_id=first_layer_id,
                stream=prefetch_stream,
            )
+            get_ms = _cpu_timing_ms(get_cpu)
+            remap_cpu = _cpu_timing_start()
            remap = get_or_build_shared_paged_buffer_slot_remap(
                forward_batch,
                page_buffer=page_buffer,
                logical_pages=real_page_table,
                layout=layout,
            )
+            remap_ms = _cpu_timing_ms(remap_cpu)
        except Exception as exc:
            _index_prefetch_fallback_log(
                "init_exception",
@@ -903,6 +1128,20 @@ class CpSharedKVIndexPrefetcher:
            remap.dense_num_pages,
            page_size,
        )
+        create_total_ms = _cpu_timing_ms(create_cpu)
+        _prefetch_log(
+            "cpu_timing path=index stage=create cp_rank=%s cp_size=%s "
+            "total_ms=%.3f get_ms=%.3f remap_ms=%.3f prefix_pages=%s "
+            "total_slots=%s dense_pages=%s",
+            layout.cp_rank,
+            layout.cp_size,
+            create_total_ms,
+            get_ms,
+            remap_ms,
+            prefix_pages,
+            int(remap.slot_logical_pages.numel()),
+            remap.dense_num_pages,
+        )

        return cls(
            layout=layout,
@@ -972,9 +1211,13 @@ class CpSharedKVIndexPrefetcher:
            )
            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
        suffix_slots = self.total_slots - self.prefix_pages
+        suffix_ms = 0.0

        if self.prefix_pages < self.total_slots:
            self._log_layer(
@@ -986,6 +1229,7 @@ class CpSharedKVIndexPrefetcher:
                self.total_slots,
                suffix_slots,
            )
+            suffix_cpu = _cpu_timing_start()
            materialize_local_paged_buffer_page_slots_into(
                page_buffer=page_buffer,
                dense_page_buffer=dense_page_buffer,
@@ -1014,6 +1258,7 @@ class CpSharedKVIndexPrefetcher:
                suffix_rows.start,
                suffix_rows.stop,
            )
+            suffix_ms = _cpu_timing_ms(suffix_cpu)

        self._log_layer(
            layer_id,
@@ -1024,6 +1269,7 @@ class CpSharedKVIndexPrefetcher:
            int(dense_page_buffer.shape[0]),
        )

+        remap_cpu = _cpu_timing_start()
        logical_pages = filter_pages_mappable_to_physical_pool(
            logical_pages=logical_pages,
            layout=self.layout,
@@ -1033,6 +1279,26 @@ class CpSharedKVIndexPrefetcher:
            logical_pages,
            page_inverse=self.page_inverse,
        )
+        remap_ms = _cpu_timing_ms(remap_cpu)
+        total_ms = _cpu_timing_ms(consume_cpu)
+        self._cpu_timing.record_consume(
+            total_ms=total_ms,
+            wait_ms=wait_ms,
+            suffix_ms=suffix_ms,
+            remap_ms=remap_ms,
+        )
+        _log_prefetch_cpu_consume(
+            log_fn=self._log,
+            timing=self._cpu_timing,
+            path="index",
+            layer_id=layer_id,
+            total_ms=total_ms,
+            wait_ms=wait_ms,
+            suffix_ms=suffix_ms,
+            remap_ms=remap_ms,
+            prefix_pages=self.prefix_pages,
+            suffix_slots=suffix_slots,
+        )
        return dense_page_buffer, dense_pages

    def start_next_layer_prefix(
@@ -1079,12 +1345,15 @@ class CpSharedKVIndexPrefetcher:
            )
            return

+        start_cpu = _cpu_timing_start()
+        get_cpu = _cpu_timing_start()
        try:
            page_buffer = _prefetch_pool_get_index_buffer(
                token_to_kv_pool=token_to_kv_pool,
                layer_id=next_layer_id,
                stream=self.stream,
            )
+            get_ms = _cpu_timing_ms(get_cpu)
        except Exception:
            logger.exception(
                "Failed to get next-layer index buffer for CP shared KV index prefetch."
@@ -1099,28 +1368,31 @@ class CpSharedKVIndexPrefetcher:

        try:
            current_stream = torch.cuda.current_stream()
-            self.stream.wait_stream(current_stream)
            prefix_rows = slot_range_to_page_slice(0, self.prefix_pages)
+            materialize_cpu = _cpu_timing_start()
+            dense_page_buffer = page_buffer.new_zeros(
+                (self.dense_num_pages, *page_buffer.shape[1:])
+            )
+            self._log_next_layer(
+                next_layer_id,
+                "index_start_prefix_begin next_layer=%s start_slot=0 "
+                "end_slot=%s dense_pages=%s",
+                next_layer_id,
+                self.prefix_pages,
+                int(dense_page_buffer.shape[0]),
+            )
+            materialize_local_paged_buffer_page_slots_into(
+                page_buffer=page_buffer,
+                dense_page_buffer=dense_page_buffer,
+                slot_logical_pages=self.slot_logical_pages,
+                layout=self.layout,
+                start_slot=0,
+                end_slot=self.prefix_pages,
+            )
+            materialize_ms = _cpu_timing_ms(materialize_cpu)
+            reduce_cpu = _cpu_timing_start()
+            self.stream.wait_stream(current_stream)
            with torch.cuda.stream(self.stream):
-                dense_page_buffer = page_buffer.new_zeros(
-                    (self.dense_num_pages, *page_buffer.shape[1:])
-                )
-                self._log_next_layer(
-                    next_layer_id,
-                    "index_start_prefix_begin next_layer=%s start_slot=0 "
-                    "end_slot=%s dense_pages=%s",
-                    next_layer_id,
-                    self.prefix_pages,
-                    int(dense_page_buffer.shape[0]),
-                )
-                materialize_local_paged_buffer_page_slots_into(
-                    page_buffer=page_buffer,
-                    dense_page_buffer=dense_page_buffer,
-                    slot_logical_pages=self.slot_logical_pages,
-                    layout=self.layout,
-                    start_slot=0,
-                    end_slot=self.prefix_pages,
-                )
                event = _all_reduce_materialized_buffer_async(
                    dense_page_buffer[prefix_rows],
                    cp_size=self.layout.cp_size,
@@ -1130,6 +1402,7 @@ class CpSharedKVIndexPrefetcher:
                    nvtx_cp_rank=self.layout.cp_rank,
                    nvtx_rows=(prefix_rows.start, prefix_rows.stop),
                )
+            reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu)
            if event is None:
                self.disabled = True
                _index_prefetch_fallback_log(
@@ -1148,6 +1421,27 @@ class CpSharedKVIndexPrefetcher:
                    next_layer_id,
                )
                return
+            total_ms = _cpu_timing_ms(start_cpu)
+            self._cpu_timing.record_start(
+                total_ms=total_ms,
+                get_ms=get_ms,
+                materialize_ms=materialize_ms,
+                reduce_enqueue_ms=reduce_enqueue_ms,
+            )
+            _log_prefetch_cpu_start(
+                log_fn=self._log,
+                layout=self.layout,
+                timing=self._cpu_timing,
+                path="index",
+                layer_id=next_layer_id,
+                total_ms=total_ms,
+                get_ms=get_ms,
+                materialize_ms=materialize_ms,
+                reduce_enqueue_ms=reduce_enqueue_ms,
+                prefix_pages=self.prefix_pages,
+                total_slots=self.total_slots,
+                dense_units=int(dense_page_buffer.shape[0]),
+            )
        except Exception:
            logger.exception("Failed to start CP shared KV index prefix prefetch.")
            self.disabled = True
--- a/python/sglang/srt/mem_cache/hiradix_cache.py
+++ b/python/sglang/srt/mem_cache/hiradix_cache.py
@@ -303,6 +303,17 @@ class CpLoadBackEvictionPlan:
    remaining_deficit_by_owner: Tuple[int, ...]


+@dataclass(frozen=True)
+class CpWriteAdmission:
+    node_id: int
+    phase: str
+    required_by_owner: Tuple[int, ...]
+    target_available_by_owner: Tuple[int, ...]
+    draft_available_by_owner: Tuple[int, ...]
+    deficit_by_owner: Tuple[int, ...]
+    eviction_plan: CpHiCacheEvictionPlan
+
+
 class HiCachePendingBackupSplit(Exception):
    def __init__(self, node: TreeNode):
        self.node = node
@@ -523,10 +534,6 @@ class HiRadixCache(RadixCache):
        # calls are on latency-sensitive scheduler paths; keep logging coarse
        # enough to avoid per-request spam while still exposing hot collectives.
        self._cp_hicache_collective_stats = {}
-        self._cp_hicache_capacity_debug = (
-            envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get()
-        )
-        self._cp_hicache_capacity_debug_stats = {}
        # track per-request tokens loaded from storage (L3 hits)
        # key: request_id, value: number of tokens actually loaded from storage
        self.prefetch_loaded_tokens_by_reqid: dict[str, int] = {}
@@ -949,7 +956,7 @@ class HiRadixCache(RadixCache):
            host_hit_len=plan.host_hit_len,
        )

-    def _cp_device_leaf_is_load_back_victim(self, node: TreeNode) -> bool:
+    def _cp_device_node_is_load_back_victim_base(self, node: TreeNode) -> bool:
        if node == getattr(self, "root_node", None):
            return False
        if getattr(node, "lock_ref", 0) > 0:
@@ -969,6 +976,21 @@ class HiRadixCache(RadixCache):
            return False
        return True

+    def _cp_device_node_is_load_back_victim_after_plan(
+        self, node: TreeNode, planned_evicted_nodes: set
+    ) -> bool:
+        if not self._cp_device_node_is_load_back_victim_base(node):
+            return False
+        for child in getattr(node, "children", {}).values():
+            if child in planned_evicted_nodes:
+                continue
+            if not getattr(child, "evicted", False):
+                return False
+        return True
+
+    def _cp_device_leaf_is_load_back_victim(self, node: TreeNode) -> bool:
+        return self._cp_device_node_is_load_back_victim_after_plan(node, set())
+
    def _cp_load_back_node_owner_page_counts(
        self, node: TreeNode, cp_size: int
    ) -> Tuple[int, ...]:
@@ -987,6 +1009,32 @@ class HiRadixCache(RadixCache):
        owners = torch.remainder(logical_pages - 1, cp_size)
        return tuple(int((owners == owner).sum().item()) for owner in range(cp_size))

+    def _cp_load_back_ancestor_unlock_contribution(
+        self,
+        node: TreeNode,
+        deficits: List[int],
+        planned_evicted_nodes: set,
+        cp_size: int,
+    ) -> int:
+        ancestor = getattr(node, "parent", None)
+        while ancestor is not None and ancestor != getattr(self, "root_node", None):
+            if ancestor in planned_evicted_nodes:
+                return 0
+            if getattr(ancestor, "value", None) is None:
+                ancestor = getattr(ancestor, "parent", None)
+                continue
+            if not self._cp_device_node_is_load_back_victim_base(ancestor):
+                return 0
+            counts = self._cp_load_back_node_owner_page_counts(ancestor, cp_size)
+            contribution = sum(
+                min(int(count), int(deficit))
+                for count, deficit in zip(counts, deficits)
+            )
+            if contribution > 0:
+                return int(contribution)
+            ancestor = getattr(ancestor, "parent", None)
+        return 0
+
    def _plan_cp_load_back_owner_lane_evictions(
        self, plan: CpLoadBackPlan
    ) -> CpLoadBackEvictionPlan:
@@ -994,29 +1042,39 @@ class HiRadixCache(RadixCache):
        cp_size = len(deficits)
        planned_freed = [0 for _ in range(cp_size)]
        victims: List[TreeNode] = []
-        used_node_ids = set()
+        planned_evicted_nodes = set()
+        candidate_nodes = set(getattr(self, "evictable_leaves", set()))

        while any(v > 0 for v in deficits):
            best_node = None
            best_counts = None
            best_score = None
-            for node in list(getattr(self, "evictable_leaves", set())):
-                node_id = getattr(node, "id", None)
-                if node_id in used_node_ids:
+            for node in list(candidate_nodes):
+                if node in planned_evicted_nodes:
                    continue
-                if not self._cp_device_leaf_is_load_back_victim(node):
+                if not self._cp_device_node_is_load_back_victim_after_plan(
+                    node, planned_evicted_nodes
+                ):
                    continue
                counts = self._cp_load_back_node_owner_page_counts(node, cp_size)
                contribution = sum(
                    min(int(count), int(deficit))
                    for count, deficit in zip(counts, deficits)
                )
+                unlock_contribution = 0
                if contribution <= 0:
+                    unlock_contribution = (
+                        self._cp_load_back_ancestor_unlock_contribution(
+                            node, deficits, planned_evicted_nodes, cp_size
+                        )
+                    )
+                if contribution <= 0 and unlock_contribution <= 0:
                    continue
                score = (
                    -int(contribution),
+                    -int(unlock_contribution),
                    self.eviction_strategy.get_priority(node),
-                    int(node_id or 0),
+                    int(getattr(node, "id", 0) or 0),
                )
                if best_score is None or score < best_score:
                    best_score = score
@@ -1027,10 +1085,23 @@ class HiRadixCache(RadixCache):
                break

            victims.append(best_node)
-            used_node_ids.add(getattr(best_node, "id", None))
+            planned_evicted_nodes.add(best_node)
+            candidate_nodes.discard(best_node)
            for owner, count in enumerate(best_counts):
                planned_freed[owner] += int(count)
                deficits[owner] = max(0, deficits[owner] - int(count))
+            ancestor = getattr(best_node, "parent", None)
+            while ancestor is not None and ancestor != getattr(self, "root_node", None):
+                if ancestor in planned_evicted_nodes:
+                    break
+                if self._cp_device_node_is_load_back_victim_after_plan(
+                    ancestor, planned_evicted_nodes
+                ):
+                    candidate_nodes.add(ancestor)
+                    break
+                if getattr(ancestor, "value", None) is not None:
+                    break
+                ancestor = getattr(ancestor, "parent", None)

        return CpLoadBackEvictionPlan(
            victims=tuple(victims),
@@ -1127,15 +1198,9 @@ class HiRadixCache(RadixCache):
        )
        return refreshed

-    def _cp_capacity_debug_enabled(self) -> bool:
-        return bool(
-            getattr(self, "_cp_hicache_capacity_debug", False)
-            or envs.SGLANG_CP_HICACHE_CAPACITY_DEBUG.get()
-        )
-
-    def _cp_build_write_admission_plan(
+    def _cp_build_write_admission(
        self, device_indices: torch.Tensor, *, node_id: int, phase: str
-    ) -> dict:
+    ) -> CpWriteAdmission:
        required = self._cp_required_host_tokens_by_rank(device_indices)
        snapshot = self._cp_host_capacity_snapshot()
        target_available = self._cp_host_available_tokens_by_rank(snapshot)
@@ -1147,90 +1212,15 @@ class HiRadixCache(RadixCache):
            max(0, req - avail) for req, avail in zip(required, draft_available)
        )
        deficit = tuple(max(a, b) for a, b in zip(target_deficit, draft_deficit))
-        current_compatible_need = tuple(
-            req if missing > 0 else 0 for req, missing in zip(required, deficit)
-        )
        eviction_plan = self._plan_cp_host_evictions(deficit)
-        return {
-            "node_id": node_id,
-            "phase": phase,
-            "required": required,
-            "target_available": target_available,
-            "draft_available": draft_available,
-            "deficit": deficit,
-            "current_compatible_need": current_compatible_need,
-            "eviction_plan": eviction_plan,
-        }
-
-    def _cp_debug_build_write_admission_plan(
-        self, device_indices: torch.Tensor, *, node_id: int, phase: str
-    ) -> Optional[dict]:
-        if not self._cp_capacity_debug_enabled():
-            return None
-        try:
-            return self._cp_build_write_admission_plan(
-                device_indices, node_id=node_id, phase=phase
-            )
-        except Exception as exc:
-            logger.warning(
-                "[HiCache-capacity-debug] failed to build write admission plan: node_id=%d phase=%s error=%s",
-                node_id,
-                phase,
-                exc,
-            )
-            return None
-
-    def _cp_debug_compare_write_admission(
-        self,
-        plan: Optional[dict],
-        *,
-        result,
-        planned_required_slots: int,
-    ) -> None:
-        if plan is None or not self._cp_capacity_debug_enabled():
-            return
-        stats = getattr(self, "_cp_hicache_capacity_debug_stats", None)
-        if stats is None:
-            stats = self._cp_hicache_capacity_debug_stats = {}
-        tag = "write_admission"
-        entry = stats.setdefault(tag, {"count": 0, "mismatch": 0})
-        entry["count"] += 1
-
-        planned_required_slots = max(plan["current_compatible_need"], default=0)
-        planned_needs_eviction = planned_required_slots > 0
-        reservation_failed = isinstance(result, HiCacheWriteFailure)
-        mismatch = (
-            reservation_failed
-            and not planned_needs_eviction
-            and all(v <= 0 for v in plan["eviction_plan"].remaining_deficit)
-        )
-        if mismatch:
-            entry["mismatch"] += 1
-        should_log = mismatch or entry["count"] == 1 or entry["count"] % 128 == 0
-        if not should_log:
-            return
-
-        eviction_plan = plan["eviction_plan"]
-        log_fn = logger.warning if mismatch else logger.info
-        log_fn(
-            "[HiCache-capacity-debug] write_admission_compare node_id=%d phase=%s "
-            "count=%d mismatch=%d planned_required=%d local_failed=%s "
-            "planned_current_required=%d required=%s target_avail=%s draft_avail=%s "
-            "deficit=%s planned_freed=%s remaining_deficit=%s victims=%s",
-            plan["node_id"],
-            plan["phase"],
-            entry["count"],
-            entry["mismatch"],
-            int(planned_required_slots),
-            isinstance(result, HiCacheWriteFailure),
-            int(planned_required_slots),
-            plan["required"],
-            plan["target_available"],
-            plan["draft_available"],
-            plan["deficit"],
-            eviction_plan.planned_freed,
-            eviction_plan.remaining_deficit,
-            [getattr(node, "id", None) for node in eviction_plan.victims],
+        return CpWriteAdmission(
+            node_id=node_id,
+            phase=phase,
+            required_by_owner=required,
+            target_available_by_owner=target_available,
+            draft_available_by_owner=draft_available,
+            deficit_by_owner=deficit,
+            eviction_plan=eviction_plan,
        )

    def shutdown(self):
@@ -1753,10 +1743,10 @@ class HiRadixCache(RadixCache):
            self.dec_node_lock_ref(node)
        return node

-    def _evict_cp_host_plan(
-        self, plan: dict, *, node_id: int, phase: str
+    def _evict_cp_host_for_write_admission(
+        self, admission: CpWriteAdmission, *, node_id: int, phase: str
    ) -> bool:
-        eviction_plan = plan["eviction_plan"]
+        eviction_plan = admission.eviction_plan
        if any(v > 0 for v in eviction_plan.remaining_deficit):
            logger.warning(
                "[CP_HICACHE_FALLBACK][cp_host_reservation_plan_insufficient] "
@@ -1765,10 +1755,10 @@ class HiRadixCache(RadixCache):
                "planned_freed=%s remaining_deficit=%s victims=%s",
                node_id,
                phase,
-                plan["required"],
-                plan["target_available"],
-                plan["draft_available"],
-                plan["deficit"],
+                admission.required_by_owner,
+                admission.target_available_by_owner,
+                admission.draft_available_by_owner,
+                admission.deficit_by_owner,
                eviction_plan.planned_freed,
                eviction_plan.remaining_deficit,
                [getattr(node, "id", None) for node in eviction_plan.victims],
@@ -1810,59 +1800,55 @@ class HiRadixCache(RadixCache):
    def _reserve_write_cp_indices_no_collective(
        self, device_indices: torch.Tensor, node_id: int
    ):
-        plan = self._cp_build_write_admission_plan(
+        admission = self._cp_build_write_admission(
            device_indices, node_id=node_id, phase="initial"
        )
-        planned_required_slots = max(plan["current_compatible_need"], default=0)
-        if planned_required_slots > 0:
-            if not self._evict_cp_host_plan(plan, node_id=node_id, phase="initial"):
-                return HiCacheWriteFailure(required_host_slots=planned_required_slots)
+        if any(v > 0 for v in admission.deficit_by_owner):
+            if not self._evict_cp_host_for_write_admission(
+                admission, node_id=node_id, phase="initial"
+            ):
+                return HiCacheWriteFailure(
+                    required_host_slots=max(admission.eviction_plan.remaining_deficit)
+                )

        result = self.cache_controller.reserve_write_cp(
            device_indices=device_indices,
            node_id=node_id,
        )
-        debug_plan = plan if self._cp_capacity_debug_enabled() else None
-        self._cp_debug_compare_write_admission(
-            debug_plan,
-            result=result,
-            planned_required_slots=planned_required_slots,
-        )
        if not isinstance(result, HiCacheWriteFailure):
            return result

-        retry_plan = self._cp_build_write_admission_plan(
+        retry_admission = self._cp_build_write_admission(
            device_indices, node_id=node_id, phase="retry_after_local_failure"
        )
-        retry_required_slots = max(
-            retry_plan["current_compatible_need"], default=int(result.required_host_slots)
-        )
-        if retry_required_slots <= 0:
+        if not any(v > 0 for v in retry_admission.deficit_by_owner) and not any(
+            v > 0 for v in retry_admission.eviction_plan.remaining_deficit
+        ):
            raise RuntimeError(
                "CP HiCache host reservation failed although deterministic "
-                "capacity planner predicted no eviction need: "
+                "owner-lane admission predicted no host deficit: "
                f"node_id={node_id} required_host_slots={result.required_host_slots}"
            )
-        if not self._evict_cp_host_plan(
-            retry_plan, node_id=node_id, phase="retry_after_local_failure"
+        if not self._evict_cp_host_for_write_admission(
+            retry_admission, node_id=node_id, phase="retry_after_local_failure"
        ):
-            return HiCacheWriteFailure(required_host_slots=retry_required_slots)
+            return HiCacheWriteFailure(
+                required_host_slots=max(
+                    retry_admission.eviction_plan.remaining_deficit,
+                    default=int(result.required_host_slots),
+                )
+            )

        logger.info(
            "[HiCache-write] write_backup CP retry after deterministic host eviction: "
-            "node_id=%d needed_slots=%d",
+            "node_id=%d deficit_by_owner=%s",
            node_id,
-            retry_required_slots,
+            retry_admission.deficit_by_owner,
        )
        result = self.cache_controller.reserve_write_cp(
            device_indices=device_indices,
            node_id=node_id,
        )
-        self._cp_debug_compare_write_admission(
-            retry_plan if self._cp_capacity_debug_enabled() else None,
-            result=result,
-            planned_required_slots=retry_required_slots,
-        )
        if not isinstance(result, HiCacheWriteFailure):
            return result

@@ -2258,6 +2244,16 @@ class HiRadixCache(RadixCache):
            return

        ack_queue_len = len(self.cache_controller.ack_write_queue)
+        # With per-layer CP backup, ongoing_write_through is populated before
+        # the final write ack is appended. Polling a TP all-reduce while the
+        # ack queue is empty is pure scheduler overhead: there is no candidate
+        # radix-state transition to make yet. The ack itself is appended at a
+        # deterministic layer-end / post-forward point across CP ranks; once it
+        # exists, the MIN below still gates host visibility on all ranks having
+        # completed the corresponding local transfer.
+        if ack_queue_len == 0:
+            return
+
        finish_count = 0
        for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
            if not finish_event.query():
@@ -2728,6 +2724,10 @@ class HiRadixCache(RadixCache):
            result = self.inc_lock_ref(ancester_node)
            delta = result.delta

+            if len(nodes_to_load) == 0:
+                self.dec_lock_ref(ancester_node)
+                return None
+
            # load it all or not at all.  The scalar length remains only a
            # coarse upper bound; final CP admission is the owner-lane vector
            # check below.