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

benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py

@@ -0,0 +1,444 @@
+"""Benchmark CP shared-KV prefetch collective layouts.
+
+This benchmark isolates the current CP shared-KV prefetch merge pattern:
+
+  dense_all_reduce:
+    each rank materializes its owner-lane pages into a full dense prefix
+    buffer, zeros non-owned pages, then all-reduces the full dense buffer.
+
+Against two all-gather candidates:
+
+  packed_all_gather_scatter:
+    each rank gathers only owner-lane pages, all-gathers owner-packed pages,
+    then scatters them back into the old dense slot layout.
+
+  packed_all_gather_rank_major:
+    each rank gathers only owner-lane pages and keeps the all-gather output in
+    rank-major owner-packed order.  This models the intended design where
+    dense_locs/page_inverse are remapped instead of scattering back.
+
+Run on the target CUDA machine, for example:
+
+  torchrun --nproc_per_node=8 \\
+    benchmark/hicache/bench_cp_shared_kv_prefetch_collective.py \\
+    --payload both --tokens 4096 8192 16384 32768 40320 65536
+
+Do not use this as an end-to-end throughput benchmark.  It is a microbenchmark
+for materialize + collective + optional scatter cost.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import os
+import statistics
+import time
+from dataclasses import asdict, dataclass
+from typing import Callable, Iterable
+
+import torch
+import torch.distributed as dist
+
+
+@dataclass(frozen=True)
+class CaseConfig:
+    payload: str
+    prefix_tokens: int
+    page_size: int
+    kv_dim: int
+    index_page_bytes: int
+    dtype: str
+    warmup: int
+    repeat: int
+
+
+@dataclass(frozen=True)
+class BenchResult:
+    payload: str
+    prefix_tokens: int
+    prefix_pages: int
+    path: str
+    full_mib: float
+    packed_mib_per_rank: float
+    gpu_ms_avg: float
+    gpu_ms_p50: float
+    gpu_ms_min: float
+    gpu_ms_max_rank_avg: float
+    cpu_ms_avg: float
+    cpu_ms_p50: float
+    cpu_ms_max_rank_avg: float
+
+
+def _rank_env() -> tuple[int, int, int]:
+    rank = int(os.environ.get("RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    local_rank = int(os.environ.get("LOCAL_RANK", str(rank)))
+    return rank, world_size, local_rank
+
+
+def _dtype(name: str) -> torch.dtype:
+    if name == "bf16":
+        return torch.bfloat16
+    if name == "fp16":
+        return torch.float16
+    if name == "fp32":
+        return torch.float32
+    raise ValueError(f"unsupported dtype: {name}")
+
+
+def _payload_shape_and_dtype(cfg: CaseConfig) -> tuple[tuple[int, ...], torch.dtype]:
+    if cfg.payload == "mla":
+        return (cfg.page_size, cfg.kv_dim), _dtype(cfg.dtype)
+    if cfg.payload == "index":
+        return (cfg.index_page_bytes,), torch.uint8
+    raise ValueError(f"unsupported payload: {cfg.payload}")
+
+
+def _owner_counts(prefix_pages: int, world_size: int) -> list[int]:
+    return [len(range(rank, prefix_pages, world_size)) for rank in range(world_size)]
+
+
+def _make_source_pages(
+    *,
+    max_owned_pages: int,
+    page_shape: tuple[int, ...],
+    dtype: torch.dtype,
+    device: torch.device,
+    rank: int,
+) -> torch.Tensor:
+    # Page 0 is kept as a sentinel to match CP shared-KV physical page usage.
+    shape = (max_owned_pages + 1, *page_shape)
+    if dtype is torch.uint8:
+        src = torch.empty(shape, dtype=dtype, device=device)
+        src.random_(0, 127)
+    else:
+        src = torch.empty(shape, dtype=dtype, device=device)
+        src.normal_(mean=float(rank + 1), std=0.01)
+    src[0].zero_()
+    return src
+
+
+def _copy_owned_to_dense(
+    *,
+    dense_pages: torch.Tensor,
+    src_pages: torch.Tensor,
+    owned_slots: torch.Tensor,
+    owned_count: int,
+) -> None:
+    if owned_count <= 0:
+        return
+    dense_pages[owned_slots].copy_(src_pages[1 : owned_count + 1])
+
+
+def _copy_owned_to_packed(
+    *,
+    local_packed: torch.Tensor,
+    src_pages: torch.Tensor,
+    owned_count: int,
+) -> None:
+    local_packed.zero_()
+    if owned_count > 0:
+        local_packed[:owned_count].copy_(src_pages[1 : owned_count + 1])
+
+
+def _scatter_rank_major_to_dense(
+    *,
+    dense_pages: torch.Tensor,
+    gathered: torch.Tensor,
+    prefix_pages: int,
+    owner_counts: list[int],
+    world_size: int,
+    device: torch.device,
+) -> None:
+    dense_pages.zero_()
+    for owner_rank, count in enumerate(owner_counts):
+        if count <= 0:
+            continue
+        slots = torch.arange(
+            owner_rank,
+            prefix_pages,
+            world_size,
+            device=device,
+            dtype=torch.long,
+        )
+        dense_pages[slots].copy_(gathered[owner_rank, :count])
+
+
+def _time_path(
+    fn: Callable[[], None],
+    *,
+    warmup: int,
+    repeat: int,
+    stream: torch.cuda.Stream,
+    group: dist.ProcessGroup,
+    device: torch.device,
+) -> tuple[list[float], list[float]]:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize(device)
+    dist.barrier(group=group)
+
+    gpu_ms: list[float] = []
+    cpu_ms: list[float] = []
+    for _ in range(repeat):
+        dist.barrier(group=group)
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        torch.cuda.synchronize(device)
+        t0 = time.perf_counter()
+        with torch.cuda.stream(stream):
+            start.record(stream)
+            fn()
+            end.record(stream)
+        end.synchronize()
+        t1 = time.perf_counter()
+        gpu_ms.append(float(start.elapsed_time(end)))
+        cpu_ms.append((t1 - t0) * 1000.0)
+    dist.barrier(group=group)
+    return gpu_ms, cpu_ms
+
+
+def _summarize_across_ranks(
+    *,
+    local_gpu_ms: list[float],
+    local_cpu_ms: list[float],
+    group: dist.ProcessGroup,
+    device: torch.device,
+) -> tuple[dict[str, float], dict[str, float]]:
+    local = torch.tensor(
+        [
+            statistics.mean(local_gpu_ms),
+            statistics.median(local_gpu_ms),
+            min(local_gpu_ms),
+            statistics.mean(local_cpu_ms),
+            statistics.median(local_cpu_ms),
+        ],
+        dtype=torch.float64,
+        device=device,
+    )
+    world_size = dist.get_world_size(group=group)
+    gathered = torch.empty((world_size, local.numel()), dtype=local.dtype, device=device)
+    dist.all_gather_into_tensor(gathered, local, group=group)
+    gathered_cpu = gathered.cpu()
+    gpu = {
+        "avg": float(gathered_cpu[:, 0].mean().item()),
+        "p50": float(gathered_cpu[:, 1].mean().item()),
+        "min": float(gathered_cpu[:, 2].mean().item()),
+        "max_rank_avg": float(gathered_cpu[:, 0].max().item()),
+    }
+    cpu = {
+        "avg": float(gathered_cpu[:, 3].mean().item()),
+        "p50": float(gathered_cpu[:, 4].mean().item()),
+        "max_rank_avg": float(gathered_cpu[:, 3].max().item()),
+    }
+    return gpu, cpu
+
+
+def _bench_case(
+    cfg: CaseConfig,
+    *,
+    group: dist.ProcessGroup,
+    device: torch.device,
+    rank: int,
+    world_size: int,
+) -> list[BenchResult]:
+    prefix_pages = math.ceil(cfg.prefix_tokens / cfg.page_size)
+    owner_counts = _owner_counts(prefix_pages, world_size)
+    owned_count = owner_counts[rank]
+    max_owned_pages = max(owner_counts)
+    page_shape, dtype = _payload_shape_and_dtype(cfg)
+    src_pages = _make_source_pages(
+        max_owned_pages=max_owned_pages,
+        page_shape=page_shape,
+        dtype=dtype,
+        device=device,
+        rank=rank,
+    )
+
+    dense_pages = torch.empty((prefix_pages, *page_shape), dtype=dtype, device=device)
+    local_packed = torch.empty(
+        (max_owned_pages, *page_shape), dtype=dtype, device=device
+    )
+    gathered = torch.empty(
+        (world_size, max_owned_pages, *page_shape), dtype=dtype, device=device
+    )
+    owned_slots = torch.arange(
+        rank,
+        prefix_pages,
+        world_size,
+        device=device,
+        dtype=torch.long,
+    )
+    stream = torch.cuda.Stream(device=device)
+    full_mib = dense_pages.numel() * dense_pages.element_size() / (1024 * 1024)
+    packed_mib = local_packed.numel() * local_packed.element_size() / (1024 * 1024)
+
+    def dense_all_reduce() -> None:
+        dense_pages.zero_()
+        _copy_owned_to_dense(
+            dense_pages=dense_pages,
+            src_pages=src_pages,
+            owned_slots=owned_slots,
+            owned_count=owned_count,
+        )
+        dist.all_reduce(dense_pages, op=dist.ReduceOp.SUM, group=group)
+
+    def packed_all_gather_scatter() -> None:
+        _copy_owned_to_packed(
+            local_packed=local_packed,
+            src_pages=src_pages,
+            owned_count=owned_count,
+        )
+        dist.all_gather_into_tensor(gathered, local_packed, group=group)
+        _scatter_rank_major_to_dense(
+            dense_pages=dense_pages,
+            gathered=gathered,
+            prefix_pages=prefix_pages,
+            owner_counts=owner_counts,
+            world_size=world_size,
+            device=device,
+        )
+
+    def packed_all_gather_rank_major() -> None:
+        _copy_owned_to_packed(
+            local_packed=local_packed,
+            src_pages=src_pages,
+            owned_count=owned_count,
+        )
+        dist.all_gather_into_tensor(gathered, local_packed, group=group)
+
+    paths: list[tuple[str, Callable[[], None]]] = [
+        ("dense_all_reduce", dense_all_reduce),
+        ("packed_all_gather_scatter", packed_all_gather_scatter),
+        ("packed_all_gather_rank_major", packed_all_gather_rank_major),
+    ]
+    results: list[BenchResult] = []
+    for name, fn in paths:
+        gpu_ms, cpu_ms = _time_path(
+            fn,
+            warmup=cfg.warmup,
+            repeat=cfg.repeat,
+            stream=stream,
+            group=group,
+            device=device,
+        )
+        gpu, cpu = _summarize_across_ranks(
+            local_gpu_ms=gpu_ms,
+            local_cpu_ms=cpu_ms,
+            group=group,
+            device=device,
+        )
+        results.append(
+            BenchResult(
+                payload=cfg.payload,
+                prefix_tokens=cfg.prefix_tokens,
+                prefix_pages=prefix_pages,
+                path=name,
+                full_mib=full_mib,
+                packed_mib_per_rank=packed_mib,
+                gpu_ms_avg=gpu["avg"],
+                gpu_ms_p50=gpu["p50"],
+                gpu_ms_min=gpu["min"],
+                gpu_ms_max_rank_avg=gpu["max_rank_avg"],
+                cpu_ms_avg=cpu["avg"],
+                cpu_ms_p50=cpu["p50"],
+                cpu_ms_max_rank_avg=cpu["max_rank_avg"],
+            )
+        )
+    return results
+
+
+def _print_result(result: BenchResult) -> None:
+    print(json.dumps(asdict(result), sort_keys=True), flush=True)
+    print(
+        f"{result.payload:5s} tokens={result.prefix_tokens:6d} "
+        f"pages={result.prefix_pages:5d} {result.path:28s} "
+        f"gpu_avg={result.gpu_ms_avg:8.3f}ms "
+        f"gpu_max_rank={result.gpu_ms_max_rank_avg:8.3f}ms "
+        f"cpu_avg={result.cpu_ms_avg:8.3f}ms "
+        f"full={result.full_mib:8.1f}MiB "
+        f"packed/rank={result.packed_mib_per_rank:8.1f}MiB",
+        flush=True,
+    )
+
+
+def _payloads(value: str) -> Iterable[str]:
+    if value == "both":
+        return ("mla", "index")
+    return (value,)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Benchmark CP shared-KV prefetch all-reduce vs all-gather layouts."
+    )
+    parser.add_argument("--payload", choices=("mla", "index", "both"), default="both")
+    parser.add_argument(
+        "--tokens",
+        type=int,
+        nargs="+",
+        default=[4096, 8192, 16384, 32768, 40320, 65536],
+        help="Prefix token counts to benchmark.",
+    )
+    parser.add_argument("--page-size", type=int, default=64)
+    parser.add_argument("--kv-dim", type=int, default=576)
+    parser.add_argument(
+        "--index-page-bytes",
+        type=int,
+        default=8448,
+        help="Bytes per index page. GLM5 NSA observed default is 8448.",
+    )
+    parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16")
+    parser.add_argument("--warmup", type=int, default=5)
+    parser.add_argument("--repeat", type=int, default=20)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    rank, world_size, local_rank = _rank_env()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark.")
+    torch.cuda.set_device(local_rank)
+    device = torch.device(f"cuda:{local_rank}")
+    dist.init_process_group("nccl", init_method="env://")
+    group = dist.group.WORLD
+
+    if rank == 0:
+        print(
+            "CP shared-KV prefetch collective benchmark "
+            f"world_size={world_size} tokens={args.tokens} payload={args.payload}",
+            flush=True,
+        )
+
+    try:
+        for payload in _payloads(args.payload):
+            for tokens in args.tokens:
+                cfg = CaseConfig(
+                    payload=payload,
+                    prefix_tokens=int(tokens),
+                    page_size=args.page_size,
+                    kv_dim=args.kv_dim,
+                    index_page_bytes=args.index_page_bytes,
+                    dtype=args.dtype,
+                    warmup=args.warmup,
+                    repeat=args.repeat,
+                )
+                for result in _bench_case(
+                    cfg,
+                    group=group,
+                    device=device,
+                    rank=rank,
+                    world_size=world_size,
+                ):
+                    if rank == 0:
+                        _print_result(result)
+    finally:
+        dist.barrier(group=group)
+        dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

docs/advanced_features/nsa_prefill_cp_hicache_no_collective_capacity_plan.md

@@ -185,12 +185,13 @@ waits for D2H write of layer `i`.
   - Planner skips device-valid nodes, pending backup nodes, host-protected nodes,
     pinned nodes, malformed metadata, and non-host-backed leaves.
   - Current eviction behavior is unchanged.
-- **Online debug option added:** `SGLANG_CP_HICACHE_CAPACITY_DEBUG=1`.
-  - When enabled, write admission computes the deterministic pre-reserve plan and
-    compares its eviction-needed decision with the current collective
-    `reserve_slots_max` result.
-  - Logs use `[HiCache-capacity-debug] write_admission_compare ...`.
-  - This is observer-only and does not remove or add collectives.
+- **Write admission is now owner-lane authoritative.**
+  - The old observer/debug scalar planner (`current_compatible_need` plus
+    `SGLANG_CP_HICACHE_CAPACITY_DEBUG`) has been removed from the hot path.
+  - Write admission uses the per-owner target/draft deficit vector directly,
+    evicts deterministic host victims only for lanes with deficit, and fails
+    fast if a subsequent local reservation failure contradicts the vector view.
+  - No CP capacity all-reduce is used for write admission.
 
 ### P1: Observer-only ledger
 

docs/advanced_features/nsa_prefill_cp_hicache_per_layer_backup_plan.md

@@ -226,17 +226,12 @@ all_reduce per node victim in a tight host-eviction loop
 all_reduce every scheduler tick when no completion prefix advanced
 ```
 
-Current correctness note: CP host reservation now synchronizes
-`required_host_slots` with `all_reduce(MAX)` before the host-eviction retry.
-This forces every rank into the same reserve/evict/retry branch and avoids
-collective mismatches when one rank is host-full and another rank reserves
-successfully.  It is intentionally a coarse slow-path collective, not a
-per-layer collective, but it can become a performance cost when host pressure
-is frequent because every reservation failure pays at least one rank-wide MAX
-sync and retry failures pay a second one.  Treat this as a correctness guard to
-be amortized later with batched reservation epochs, deterministic host
-watermarks, or less frequent proactive host eviction; do not move it into the
-per-layer data path.
+Current correctness note: CP host reservation no longer uses the old
+`required_host_slots` scalar `all_reduce(MAX)`.  The owner-lane vector computed
+from `CpHiCacheNodeMetadata.page_owners` is now the authoritative admission
+view for target and draft host capacity.  A reservation failure after this
+vector predicts no deficit is treated as an invariant violation rather than
+falling back to another collective.
 
 ## Per-Layer Backup Data Plane
 
@@ -466,10 +461,11 @@ paths:
 2. **Pending split behavior.**  The chosen first pass is defer/requeue the
    request that would split a node with pending backup.  Do not split the
    in-flight backup op, including for future `bs > 1`.
-3. **Ack batching threshold.**  Current `writing_check()` can all-reduce on
-   every progress poll.  The first per-layer implementation should keep one
-   final logical ack per node and check it at final visibility time, not per
-   layer; batching threshold for final commit remains a later performance pass.
+3. **Ack batching threshold.**  `writing_check()` no longer enters the TP
+   all-reduce while a write is ongoing but no final ack exists.  The remaining
+   `MIN` is the final host-visibility barrier for ack entries that do exist:
+   it should not run per layer, and removing it requires a different global
+   visibility protocol rather than a local-only commit.
 4. **Failure policy under reserve mismatch.**  Capacity pressure should skip
    backup; malformed target/draft metadata should fail fast.  If a local rank
    fails reservation after deterministic host eviction while another succeeds,

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)

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

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.

test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py

@@ -272,6 +272,47 @@ class TestCpHiCacheLoadBackOwnerLanes(CustomTestCase):
         self.assertEqual(target.value.tolist(), loaded.tolist())
         self.assertIn(target.id, cache.ongoing_load_back)
 
+    def test_load_back_evicts_blocking_leaf_to_unlock_parent_owner_lane(self):
+        allocator = _make_allocator()
+        # Only owner lane 1 is initially available.  The target needs owner lane
+        # 0.  The current leaf itself belongs to owner 1, but evicting it makes
+        # its parent evictable and the parent frees owner 0.
+        allocator.free_pages = torch.tensor([2], dtype=torch.int64)
+        cache = _make_cache(allocator)
+
+        parent = _make_node(
+            40,
+            600,
+            [0],
+            value=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
+            priority=0,
+        )
+        _attach_child(cache, cache.root_node, parent)
+        child = _make_node(
+            41,
+            700,
+            [1],
+            value=torch.tensor([8, 9, 10, 11], dtype=torch.int64),
+            priority=0,
+        )
+        _attach_child(cache, parent, child)
+        cache.evictable_leaves.add(child)
+        cache.evictable_size_ = len(parent.key) + len(child.key)
+
+        target = _make_node(42, 800, [0], value=None, priority=10)
+        _attach_child(cache, cache.root_node, target)
+
+        loaded = cache.load_back(target, mem_quota=100)
+
+        self.assertIsNotNone(loaded)
+        self.assertEqual(cache.cache_controller.load_calls, 1)
+        self.assertEqual(
+            [indices.tolist() for indices in cache.cache_controller.evicted_device_indices],
+            [[8, 9, 10, 11], [4, 5, 6, 7]],
+        )
+        self.assertEqual(target.value.tolist(), loaded.tolist())
+        self.assertIn(target.id, cache.ongoing_load_back)
+
     def test_load_back_failure_leaves_node_unassigned_and_unlocked(self):
         allocator = _make_allocator()
         allocator.free_pages = torch.tensor([1, 2], dtype=torch.int64)

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -64,6 +64,7 @@ except (ImportError, RuntimeError):
         "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
         "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
         "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+        "moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> Tensor[]",
     ):
         try:
             _sgl_kernel_lib.define(_schema)
@@ -83,6 +84,7 @@ for _schema in (
     "sgl_per_token_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s) -> ()",
     "fp8_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, Tensor? bias=None) -> Tensor",
     "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> Tensor",
+    "moe_fused_gate(Tensor input, Tensor bias, int num_expert_group, int topk_group, int topk, int num_fused_shared_experts, float routed_scaling_factor, bool apply_routed_scaling_factor_on_output) -> Tensor[]",
 ):
     try:
         _sgl_kernel_lib.define(_schema)
@@ -629,6 +631,20 @@ class FakeTokenAllocator:
     def available_size(self):
         return 0
 
+    def compute_owner_lane_stats(self, page_owners):
+        if len(page_owners) == 0:
+            return [], [], []
+        cp_size = max(int(owner) for owner in page_owners) + 1
+        required = [0] * cp_size
+        for owner in page_owners:
+            required[int(owner)] += 1
+        available = list(required)
+        deficits = [0] * cp_size
+        return required, available, deficits
+
+    def allocator_state_str(self):
+        return "FakeTokenAllocator"
+
 
 class TestHiRadixCacheCPBackup(CustomTestCase):
     def test_session_aware_cache_forwards_cp_hicache_prepare(self):
@@ -1243,7 +1259,7 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
 
         with self.assertRaisesRegex(
             RuntimeError,
-            "planner predicted no eviction need",
+            "owner-lane admission predicted no host deficit",
         ):
             cache.write_backup(node)
 
@@ -1868,6 +1884,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         cache._uses_cp_hicache = True
         cache.root_node = TreeNode()
         cache.device = "cpu"
+        cache.page_size = 1
+        cache.token_to_kv_pool_allocator = FakeTokenAllocator()
         cache.load_back_threshold = 1
         cache.evictable_size_ = 0
         cache.metrics_collector = None
@@ -1911,6 +1929,8 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
         cache._uses_cp_hicache = True
         cache.root_node = TreeNode()
         cache.root_node.value = torch.empty((0,), dtype=torch.int64)
+        cache.page_size = 1
+        cache.token_to_kv_pool_allocator = FakeTokenAllocator()
         cache.load_back_threshold = 5
         cache.evictable_size_ = 0
         cache.metrics_collector = None