Model CP scheduler admission with cache-hit pressure

Add an offline benchmark that reuses PrefillAdder to model how L1 cached tokens, L2 HiCache hits, and extend tokens shape CP shared-KV prefill batch admission. The tool makes scheduler stop reasons and fake L2 load-back capacity pressure observable without starting a model.

Constraint: The benchmark must stay CPU/offline and avoid depending on CUDA execution or live services.

Constraint: L2 cached tokens are modeled as host_hit_length, so successful load-back both increases prefix_len and consumes fake L1 capacity.

Rejected: Build an ETE benchmark first | too slow for isolating scheduler admission behavior.

Rejected: Reimplement scheduler logic from scratch | would drift from PrefillAdder semantics.

Confidence: high

Scope-risk: narrow

Directive: Treat duration_us as Python admission overhead only; it is not an ETE latency metric.

Tested: Remote pytest test/registered/unit/managers/test_prefill_scheduler_admission_bench.py: 4 passed as part of 6 targeted tests.

Tested: Remote synthetic benchmark run with --cp-max-total-cached-tokens showed second 4096-token cached request stopped with OTHER.

Not-tested: Real traffic trace import from production logs.

benchmark/hicache/bench_prefill_scheduler_admission.py

@@ -0,0 +1,657 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+"""Offline benchmark/model for SGLang prefill scheduler admission.
+
+This script answers a narrow question: given an ordered waiting queue with
+per-request L1 cache hit, L2/HiCache hit, and compute extend lengths, what would
+PrefillAdder admit into a prefill batch and which budget stops the scan?
+
+It intentionally reuses the production PrefillAdder admission logic instead of
+reimplementing the scheduler.  CUDA/model execution is not required.  L2 load
+back is modeled by a fake tree cache that consumes fake L1/device allocator
+capacity and records each load-back event.
+
+Examples:
+
+  PYTHONPATH=python python benchmark/hicache/bench_prefill_scheduler_admission.py \
+    --synthetic-grid --l1-cached-tokens 0,4096 --l2-cached-tokens 0,4096 \
+    --extend-tokens 256,1024,4096 --available-tokens 200000 \
+    --max-prefill-tokens 16384 --cp-max-total-extend-tokens 65536 \
+    --output text
+
+  cat requests.jsonl
+  {"rid":"r0","l1_cached_tokens":40320,"l2_cached_tokens":0,"extend_tokens":128}
+  {"rid":"r1","l1_cached_tokens":0,"l2_cached_tokens":32768,"extend_tokens":512}
+
+  PYTHONPATH=python python benchmark/hicache/bench_prefill_scheduler_admission.py \
+    --requests-jsonl requests.jsonl --output json
+"""
+
+import argparse
+import json
+import math
+import sys
+import time
+import types
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Iterable, Optional
+
+import torch
+
+
+_SGL_KERNEL_LIBRARIES = []
+
+
+def _install_sgl_kernel_stubs() -> None:
+    """Install minimal sgl_kernel stubs for CPU-only scheduler imports."""
+
+    if "sgl_kernel" not in sys.modules:
+        sys.modules["sgl_kernel"] = types.ModuleType("sgl_kernel")
+    sys.modules["sgl_kernel"].__file__ = "sgl_kernel_stub.py"
+    sys.modules["sgl_kernel"].__path__ = []
+
+    if not hasattr(sys.modules["sgl_kernel"], "__getattr__"):
+
+        def _sgl_kernel_getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *args, **kwargs: None
+            setattr(sys.modules["sgl_kernel"], name, fn)
+            return fn
+
+        sys.modules["sgl_kernel"].__getattr__ = _sgl_kernel_getattr
+
+    if "sgl_kernel.kvcacheio" not in sys.modules:
+        sys.modules["sgl_kernel.kvcacheio"] = types.ModuleType("sgl_kernel.kvcacheio")
+
+    for name in (
+        "sgl_per_token_group_quant_8bit",
+        "sgl_per_token_group_quant_fp8",
+        "sgl_per_token_quant_fp8",
+        "fp8_blockwise_scaled_mm",
+        "fp8_scaled_mm",
+        "silu_and_mul",
+    ):
+        if not hasattr(sys.modules["sgl_kernel"], name):
+            setattr(sys.modules["sgl_kernel"], name, lambda *args, **kwargs: None)
+
+    if "sgl_kernel.quantization" not in sys.modules:
+        quantization_stub = types.ModuleType("sgl_kernel.quantization")
+        for name in (
+            "ggml_dequantize",
+            "ggml_moe_a8",
+            "ggml_moe_a8_vec",
+            "ggml_moe_get_block_size",
+            "ggml_mul_mat_a8",
+            "ggml_mul_mat_vec_a8",
+        ):
+            setattr(quantization_stub, name, lambda *args, **kwargs: None)
+        sys.modules["sgl_kernel.quantization"] = quantization_stub
+
+    sgl_kernel_lib = torch.library.Library("sgl_kernel", "FRAGMENT")
+    _SGL_KERNEL_LIBRARIES.append(sgl_kernel_lib)
+    for schema in (
+        "sgl_per_token_group_quant_8bit(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+        "sgl_per_token_group_quant_fp8(Tensor input, Tensor(a!) output_q, Tensor(b!) output_s, int group_size, float eps, float fp8_min, float fp8_max, bool scale_ue8m0) -> ()",
+        "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",
+    ):
+        try:
+            sgl_kernel_lib.define(schema)
+        except RuntimeError as exc:
+            if "already" not in str(exc).lower() and "duplicate" not in str(exc).lower():
+                raise
+
+
+@dataclass(frozen=True)
+class RequestSpec:
+    rid: str
+    l1_cached_tokens: int
+    l2_cached_tokens: int
+    extend_tokens: int
+    max_new_tokens: int = 1
+    output_tokens: int = 0
+
+    def __post_init__(self) -> None:
+        for field_name in (
+            "l1_cached_tokens",
+            "l2_cached_tokens",
+            "extend_tokens",
+            "max_new_tokens",
+            "output_tokens",
+        ):
+            value = getattr(self, field_name)
+            if value < 0:
+                raise ValueError(f"{field_name} must be non-negative, got {value}")
+
+
+@dataclass(frozen=True)
+class SchedulerBenchConfig:
+    page_size: int = 64
+    available_tokens: int = 1_000_000
+    evictable_tokens: int = 0
+    max_prefill_tokens: int = 16_384
+    chunked_prefill_size: Optional[int] = None
+    mixed_with_decode_tokens: int = 0
+    new_token_ratio: float = 1.0
+    enable_cp_context: bool = True
+    enable_cp_shared_kv_prefill_bs_gt1: bool = True
+    cp_shared_kv_prefill_max_batch_requests: Optional[int] = None
+    cp_shared_kv_prefill_max_total_extend_tokens: Optional[int] = None
+    cp_shared_kv_prefill_max_total_cached_tokens: Optional[int] = None
+    max_ticks: int = 1
+    consume_l2_load_back_capacity: bool = True
+
+    def __post_init__(self) -> None:
+        if self.page_size <= 0:
+            raise ValueError(f"page_size must be positive, got {self.page_size}")
+        if self.available_tokens < 0:
+            raise ValueError("available_tokens must be non-negative")
+        if self.evictable_tokens < 0:
+            raise ValueError("evictable_tokens must be non-negative")
+        if self.max_prefill_tokens < 0:
+            raise ValueError("max_prefill_tokens must be non-negative")
+        if self.max_ticks <= 0:
+            raise ValueError("max_ticks must be positive")
+
+
+@dataclass(frozen=True)
+class LoadBackEvent:
+    rid: str
+    requested_tokens: int
+    paged_tokens: int
+    loaded_tokens: int
+    mem_quota: Optional[int]
+    available_before: int
+    available_after: int
+    skipped_reason: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class AcceptedRequest:
+    rid: str
+    l1_cached_tokens: int
+    l2_cached_tokens: int
+    loaded_l2_tokens: int
+    compute_extend_tokens: int
+    initial_extend_tokens: int
+    effective_extend_tokens: int
+    max_new_tokens: int
+
+
+@dataclass(frozen=True)
+class TickResult:
+    tick: int
+    accepted: list[AcceptedRequest]
+    stopped_on_rid: Optional[str]
+    stopped_result: Optional[str]
+    rem_input_tokens_after_tick: int
+    rem_total_tokens_after_tick: float
+    cur_rem_tokens_after_tick: float
+    cp_total_extend_tokens: int
+    cp_total_cached_tokens: int
+    log_hit_tokens: int
+    log_input_tokens: int
+    allocator_available_after_tick: int
+    load_back_events: list[LoadBackEvent]
+    duration_us: float
+
+
+@dataclass(frozen=True)
+class TraceResult:
+    config: SchedulerBenchConfig
+    request_count: int
+    ticks: list[TickResult]
+    remaining_rids: list[str]
+    blocked: bool
+
+
+class FakeTokenAllocator:
+    def __init__(self, available_tokens: int):
+        self.available_tokens = int(available_tokens)
+
+    def available_size(self) -> int:
+        return self.available_tokens
+
+    def full_available_size(self) -> int:
+        return self.available_tokens
+
+    def swa_available_size(self) -> int:
+        return self.available_tokens
+
+    def consume(self, tokens: int) -> bool:
+        if tokens < 0:
+            raise ValueError(f"tokens must be non-negative, got {tokens}")
+        if tokens > self.available_tokens:
+            return False
+        self.available_tokens -= tokens
+        return True
+
+
+class FakeTreeCache:
+    def __init__(
+        self,
+        *,
+        allocator: FakeTokenAllocator,
+        page_size: int,
+        evictable_tokens: int,
+        consume_l2_load_back_capacity: bool,
+    ):
+        self.allocator = allocator
+        self.page_size = int(page_size)
+        self._evictable_tokens = int(evictable_tokens)
+        self.consume_l2_load_back_capacity = bool(consume_l2_load_back_capacity)
+        self.disable = False
+        self.load_back_events: list[LoadBackEvent] = []
+
+    def supports_mamba(self) -> bool:
+        return False
+
+    def supports_swa(self) -> bool:
+        return False
+
+    def is_tree_cache(self) -> bool:
+        return True
+
+    def full_evictable_size(self) -> int:
+        return self._evictable_tokens
+
+    def swa_evictable_size(self) -> int:
+        return self._evictable_tokens
+
+    def evictable_size(self) -> int:
+        return self._evictable_tokens
+
+    def inc_lock_ref(self, _node):
+        from sglang.srt.mem_cache.base_prefix_cache import IncLockRefResult
+
+        return IncLockRefResult()
+
+    def dec_lock_ref(self, _node, *_args, **_kwargs):
+        from sglang.srt.mem_cache.base_prefix_cache import DecLockRefResult
+
+        return DecLockRefResult()
+
+    def init_load_back(self, params):
+        rid = getattr(params.last_host_node, "rid", "unknown")
+        requested = int(params.host_hit_length)
+        if requested <= 0:
+            return torch.empty((0,), dtype=torch.int64), params.last_host_node
+
+        paged = _ceil_to_page(requested, self.page_size)
+        before = self.allocator.available_size()
+        skipped_reason: Optional[str] = None
+        loaded = requested
+
+        if params.mem_quota is not None and paged > int(params.mem_quota):
+            skipped_reason = "over_mem_quota"
+            loaded = 0
+        elif self.consume_l2_load_back_capacity and not self.allocator.consume(paged):
+            skipped_reason = "allocator_capacity"
+            loaded = 0
+
+        after = self.allocator.available_size()
+        self.load_back_events.append(
+            LoadBackEvent(
+                rid=rid,
+                requested_tokens=requested,
+                paged_tokens=paged,
+                loaded_tokens=loaded,
+                mem_quota=params.mem_quota,
+                available_before=before,
+                available_after=after,
+                skipped_reason=skipped_reason,
+            )
+        )
+
+        if loaded <= 0:
+            return torch.empty((0,), dtype=torch.int64), params.last_host_node
+        return torch.arange(loaded, dtype=torch.int64), params.last_host_node
+
+
+class FakeRunningBatch:
+    reqs: list = []
+    batch_is_full: bool = False
+
+    def release_req(self, _req):
+        return None
+
+    def filter_batch(self, *_args, **_kwargs):
+        return None
+
+    def is_empty(self) -> bool:
+        return not self.reqs
+
+    def batch_size(self) -> int:
+        return len(self.reqs)
+
+
+class _FakeReq:
+    def __init__(self, spec: RequestSpec):
+        self.rid = spec.rid
+        self.priority = 0
+        self.output_ids = [0] * spec.output_tokens
+        self.sampling_params = SimpleNamespace(
+            max_new_tokens=spec.max_new_tokens,
+            ignore_eos=False,
+        )
+        self.time_stats = SimpleNamespace(wait_queue_entry_time=0.0)
+        self.host_hit_length = spec.l2_cached_tokens
+        self.prefix_indices = torch.arange(spec.l1_cached_tokens, dtype=torch.int64)
+        self.fill_ids = list(
+            range(spec.l1_cached_tokens + spec.l2_cached_tokens + spec.extend_tokens)
+        )
+        self.extend_input_len = spec.l2_cached_tokens + spec.extend_tokens
+        self.extend_logprob_start_len = 0
+        self.last_node = SimpleNamespace(rid=spec.rid)
+        self.last_host_node = SimpleNamespace(rid=spec.rid)
+        self.cache_protected_len = 0
+
+    def set_extend_input_len(self, value: int) -> None:
+        self.extend_input_len = int(value)
+
+    def finished(self) -> bool:
+        return False
+
+
+def _ceil_to_page(tokens: int, page_size: int) -> int:
+    if tokens <= 0:
+        return 0
+    return int(math.ceil(tokens / float(page_size)) * page_size)
+
+
+def _configure_scheduler_globals(enable_cp_context: bool) -> None:
+    _install_sgl_kernel_stubs()
+    from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+
+    if enable_cp_context:
+        set_global_server_args_for_scheduler(
+            ServerArgs(
+                model_path="dummy",
+                enable_nsa_prefill_context_parallel=True,
+                nsa_prefill_cp_mode="in-seq-split",
+            )
+        )
+    else:
+        set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+
+def _make_prefill_adder(cfg: SchedulerBenchConfig, tree_cache: FakeTreeCache, allocator: FakeTokenAllocator):
+    _install_sgl_kernel_stubs()
+    from sglang.srt.managers.schedule_policy import PrefillAdder
+
+    return PrefillAdder(
+        page_size=cfg.page_size,
+        tree_cache=tree_cache,
+        token_to_kv_pool_allocator=allocator,
+        running_batch=FakeRunningBatch(),
+        new_token_ratio=cfg.new_token_ratio,
+        rem_input_tokens=cfg.max_prefill_tokens,
+        rem_chunk_tokens=cfg.chunked_prefill_size,
+        mixed_with_decode_tokens=cfg.mixed_with_decode_tokens,
+        priority_scheduling_preemption_threshold=0,
+        enable_cp_shared_kv_prefill_bs_gt1=cfg.enable_cp_shared_kv_prefill_bs_gt1,
+        cp_shared_kv_prefill_max_batch_requests=cfg.cp_shared_kv_prefill_max_batch_requests,
+        cp_shared_kv_prefill_max_total_extend_tokens=cfg.cp_shared_kv_prefill_max_total_extend_tokens,
+        cp_shared_kv_prefill_max_total_cached_tokens=cfg.cp_shared_kv_prefill_max_total_cached_tokens,
+    )
+
+
+def _result_name(result) -> str:
+    return getattr(result, "name", str(result))
+
+
+def _accepted_request(spec: RequestSpec, req: _FakeReq, loaded_l2_tokens: int) -> AcceptedRequest:
+    return AcceptedRequest(
+        rid=spec.rid,
+        l1_cached_tokens=spec.l1_cached_tokens,
+        l2_cached_tokens=spec.l2_cached_tokens,
+        loaded_l2_tokens=loaded_l2_tokens,
+        compute_extend_tokens=spec.extend_tokens,
+        initial_extend_tokens=spec.l2_cached_tokens + spec.extend_tokens,
+        effective_extend_tokens=int(req.extend_input_len),
+        max_new_tokens=spec.max_new_tokens,
+    )
+
+
+def run_scheduler_admission_trace(
+    requests: list[RequestSpec], cfg: SchedulerBenchConfig
+) -> TraceResult:
+    _configure_scheduler_globals(cfg.enable_cp_context)
+    from sglang.srt.managers.schedule_policy import AddReqResult
+
+    pending = list(requests)
+    ticks: list[TickResult] = []
+    allocator = FakeTokenAllocator(cfg.available_tokens)
+    tree_cache = FakeTreeCache(
+        allocator=allocator,
+        page_size=cfg.page_size,
+        evictable_tokens=cfg.evictable_tokens,
+        consume_l2_load_back_capacity=cfg.consume_l2_load_back_capacity,
+    )
+
+    for tick_idx in range(cfg.max_ticks):
+        if not pending:
+            break
+
+        adder = _make_prefill_adder(cfg, tree_cache, allocator)
+        req_by_obj: dict[object, tuple[RequestSpec, _FakeReq]] = {}
+        stopped_on_rid: Optional[str] = None
+        stopped_result: Optional[str] = None
+        load_event_start = len(tree_cache.load_back_events)
+        start = time.perf_counter()
+
+        for spec in pending:
+            req = _FakeReq(spec)
+            before_events = len(tree_cache.load_back_events)
+            result = adder.add_one_req(
+                req,
+                has_chunked_req=False,
+                truncation_align_size=None,
+            )
+            after_events = len(tree_cache.load_back_events)
+
+            if req in adder.can_run_list:
+                req_by_obj[req] = (spec, req)
+
+            if result != AddReqResult.CONTINUE:
+                stopped_on_rid = spec.rid
+                stopped_result = _result_name(result)
+                # If the stopping request was still accepted, keep it in the
+                # batch just like the real scheduler does before breaking.
+                if req in adder.can_run_list:
+                    req_by_obj[req] = (spec, req)
+                break
+
+            # Keep loop variables observable under debugger without changing
+            # behavior; this also makes the loadback event span explicit.
+            _ = before_events, after_events
+
+        duration_us = (time.perf_counter() - start) * 1_000_000.0
+        loaded_by_rid: dict[str, int] = {}
+        for event in tree_cache.load_back_events[load_event_start:]:
+            loaded_by_rid[event.rid] = loaded_by_rid.get(event.rid, 0) + event.loaded_tokens
+
+        accepted = [
+            _accepted_request(spec, req, loaded_by_rid.get(spec.rid, 0))
+            for req in adder.can_run_list
+            for spec, req in [req_by_obj[req]]
+        ]
+        accepted_rids = {req.rid for req in accepted}
+        pending = [spec for spec in pending if spec.rid not in accepted_rids]
+
+        ticks.append(
+            TickResult(
+                tick=tick_idx,
+                accepted=accepted,
+                stopped_on_rid=stopped_on_rid,
+                stopped_result=stopped_result,
+                rem_input_tokens_after_tick=int(adder.rem_input_tokens),
+                rem_total_tokens_after_tick=float(adder.rem_total_tokens),
+                cur_rem_tokens_after_tick=float(adder.cur_rem_tokens),
+                cp_total_extend_tokens=int(adder.cp_shared_kv_prefill_total_extend_tokens),
+                cp_total_cached_tokens=int(adder.cp_shared_kv_prefill_total_cached_tokens),
+                log_hit_tokens=int(adder.log_hit_tokens),
+                log_input_tokens=int(adder.log_input_tokens),
+                allocator_available_after_tick=allocator.available_size(),
+                load_back_events=list(tree_cache.load_back_events[load_event_start:]),
+                duration_us=duration_us,
+            )
+        )
+
+        if not accepted:
+            break
+
+    return TraceResult(
+        config=cfg,
+        request_count=len(requests),
+        ticks=ticks,
+        remaining_rids=[spec.rid for spec in pending],
+        blocked=bool(pending),
+    )
+
+
+def _parse_int_list(value: str | Iterable[int]) -> list[int]:
+    if isinstance(value, str):
+        return [int(item.strip()) for item in value.split(",") if item.strip()]
+    return [int(item) for item in value]
+
+
+def _load_requests_jsonl(path: Path) -> list[RequestSpec]:
+    requests: list[RequestSpec] = []
+    with path.open("r", encoding="utf-8") as f:
+        for line_no, line in enumerate(f, start=1):
+            line = line.strip()
+            if not line:
+                continue
+            data = json.loads(line)
+            try:
+                requests.append(RequestSpec(**data))
+            except TypeError as exc:
+                raise ValueError(f"invalid request at {path}:{line_no}: {data}") from exc
+    return requests
+
+
+def _build_synthetic_requests(args) -> list[RequestSpec]:
+    requests: list[RequestSpec] = []
+    rid = 0
+    for l1 in _parse_int_list(args.l1_cached_tokens):
+        for l2 in _parse_int_list(args.l2_cached_tokens):
+            for extend in _parse_int_list(args.extend_tokens):
+                for _ in range(args.repeat_per_case):
+                    requests.append(
+                        RequestSpec(
+                            rid=f"r{rid}_l1{l1}_l2{l2}_e{extend}",
+                            l1_cached_tokens=l1,
+                            l2_cached_tokens=l2,
+                            extend_tokens=extend,
+                            max_new_tokens=args.max_new_tokens,
+                        )
+                    )
+                    rid += 1
+    return requests
+
+
+def _trace_to_dict(trace: TraceResult) -> dict:
+    return asdict(trace)
+
+
+def _print_text(trace: TraceResult) -> None:
+    print(
+        "config "
+        f"page_size={trace.config.page_size} available={trace.config.available_tokens} "
+        f"evictable={trace.config.evictable_tokens} max_prefill={trace.config.max_prefill_tokens} "
+        f"cp_extend_limit={trace.config.cp_shared_kv_prefill_max_total_extend_tokens} "
+        f"cp_cached_limit={trace.config.cp_shared_kv_prefill_max_total_cached_tokens}"
+    )
+    for tick in trace.ticks:
+        accepted = ",".join(
+            f"{req.rid}(l1={req.l1_cached_tokens},l2={req.l2_cached_tokens},"
+            f"loaded={req.loaded_l2_tokens},extend={req.effective_extend_tokens})"
+            for req in tick.accepted
+        )
+        print(
+            f"tick={tick.tick} bs={len(tick.accepted)} accepted=[{accepted}] "
+            f"stop={tick.stopped_on_rid}:{tick.stopped_result} "
+            f"cp_extend={tick.cp_total_extend_tokens} cp_cached={tick.cp_total_cached_tokens} "
+            f"log_hit={tick.log_hit_tokens} "
+            f"log_input={tick.log_input_tokens} rem_input={tick.rem_input_tokens_after_tick} "
+            f"rem_total={tick.rem_total_tokens_after_tick:.1f} "
+            f"cur_rem={tick.cur_rem_tokens_after_tick:.1f} "
+            f"allocator_available={tick.allocator_available_after_tick} "
+            f"duration_us={tick.duration_us:.1f}"
+        )
+        for event in tick.load_back_events:
+            print(
+                f"  load_back rid={event.rid} requested={event.requested_tokens} "
+                f"paged={event.paged_tokens} loaded={event.loaded_tokens} "
+                f"quota={event.mem_quota} avail={event.available_before}->{event.available_after} "
+                f"skip={event.skipped_reason}"
+            )
+    if trace.remaining_rids:
+        print(f"remaining={','.join(trace.remaining_rids)} blocked={trace.blocked}")
+
+
+def build_arg_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description=__doc__)
+    source = parser.add_mutually_exclusive_group(required=True)
+    source.add_argument("--requests-jsonl", type=Path)
+    source.add_argument("--synthetic-grid", action="store_true")
+    parser.add_argument("--l1-cached-tokens", default="0,4096,32768")
+    parser.add_argument("--l2-cached-tokens", default="0,4096,32768")
+    parser.add_argument("--extend-tokens", default="128,512,2048,8192")
+    parser.add_argument("--repeat-per-case", type=int, default=1)
+    parser.add_argument("--max-new-tokens", type=int, default=1)
+    parser.add_argument("--page-size", type=int, default=64)
+    parser.add_argument("--available-tokens", type=int, default=1_000_000)
+    parser.add_argument("--evictable-tokens", type=int, default=0)
+    parser.add_argument("--max-prefill-tokens", type=int, default=16_384)
+    parser.add_argument("--chunked-prefill-size", type=int, default=None)
+    parser.add_argument("--mixed-with-decode-tokens", type=int, default=0)
+    parser.add_argument("--max-ticks", type=int, default=1)
+    parser.add_argument("--disable-cp-context", action="store_true")
+    parser.add_argument("--disable-cp-bs-gt1", action="store_true")
+    parser.add_argument("--cp-max-batch-requests", type=int, default=8)
+    parser.add_argument("--cp-max-total-extend-tokens", type=int, default=65_536)
+    parser.add_argument("--cp-max-total-cached-tokens", type=int, default=None)
+    parser.add_argument("--no-consume-l2-load-back-capacity", action="store_true")
+    parser.add_argument("--output", choices=("text", "json"), default="text")
+    return parser
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    args = build_arg_parser().parse_args(argv)
+    if args.requests_jsonl is not None:
+        requests = _load_requests_jsonl(args.requests_jsonl)
+    else:
+        requests = _build_synthetic_requests(args)
+
+    cfg = SchedulerBenchConfig(
+        page_size=args.page_size,
+        available_tokens=args.available_tokens,
+        evictable_tokens=args.evictable_tokens,
+        max_prefill_tokens=args.max_prefill_tokens,
+        chunked_prefill_size=args.chunked_prefill_size,
+        mixed_with_decode_tokens=args.mixed_with_decode_tokens,
+        enable_cp_context=not args.disable_cp_context,
+        enable_cp_shared_kv_prefill_bs_gt1=not args.disable_cp_bs_gt1,
+        cp_shared_kv_prefill_max_batch_requests=args.cp_max_batch_requests,
+        cp_shared_kv_prefill_max_total_extend_tokens=args.cp_max_total_extend_tokens,
+        cp_shared_kv_prefill_max_total_cached_tokens=args.cp_max_total_cached_tokens,
+        max_ticks=args.max_ticks,
+        consume_l2_load_back_capacity=not args.no_consume_l2_load_back_capacity,
+    )
+    trace = run_scheduler_admission_trace(requests, cfg)
+    if args.output == "json":
+        print(json.dumps(_trace_to_dict(trace), indent=2, sort_keys=True))
+    else:
+        _print_text(trace)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

docs/advanced_features/nsa_prefill_cp_scheduler_admission_benchmark.md

@@ -0,0 +1,111 @@
+# CP shared-KV Prefill Scheduler Admission Benchmark
+
+这个 benchmark 用来离线回答一个问题：给定一批 waiting requests，每个 request 的 L1 cache hit、L2/HiCache hit、以及实际需要 forward 的 extend 长度不同，真实 `PrefillAdder` 会如何组 prefill batch，最终被哪个 budget 卡住。
+
+脚本位置：
+
+```bash
+benchmark/hicache/bench_prefill_scheduler_admission.py
+```
+
+## 建模语义
+
+每条输入 request 使用三个 token 维度：
+
+- `l1_cached_tokens`：已经在 L1/device radix cache 命中的 token，映射到 `prefix_indices` 长度。
+- `l2_cached_tokens`：在 HiCache/L2 命中的 token，映射到 `req.host_hit_length`。
+- `extend_tokens`：L1/L2 cache 都不能覆盖，需要当前 prefill forward 计算的 token。
+
+因此 benchmark 构造的初始 scheduler request 是：
+
+```text
+fill_ids length      = l1_cached_tokens + l2_cached_tokens + extend_tokens
+prefix_indices len   = l1_cached_tokens
+req.extend_input_len = l2_cached_tokens + extend_tokens
+req.host_hit_length  = l2_cached_tokens
+```
+
+进入 `PrefillAdder.add_one_req()` 后，如果 L2 load-back 成功：
+
+```text
+prefix_indices += loaded_l2_tokens
+effective extend_input_len = extend_tokens
+```
+
+所以 L2 hit 的双重影响是：
+
+1. 减少当前 forward 需要计算的 token。
+2. 需要 load-back 到 L1/device cache，仍会消耗 L1 allocator capacity。
+
+benchmark 用 fake tree-cache 显式记录 load-back event，并默认按 page 对齐消耗 fake allocator available tokens。
+
+## 使用示例
+
+Synthetic grid：
+
+```bash
+PYTHONPATH=python python benchmark/hicache/bench_prefill_scheduler_admission.py \
+  --synthetic-grid \
+  --l1-cached-tokens 0,4096 \
+  --l2-cached-tokens 0,4096 \
+  --extend-tokens 128,2048 \
+  --available-tokens 20000 \
+  --max-prefill-tokens 16384 \
+  --cp-max-total-extend-tokens 65536 \
+  --cp-max-batch-requests 8 \
+  --output text
+```
+
+JSONL 输入：
+
+```jsonl
+{"rid":"r0","l1_cached_tokens":40320,"l2_cached_tokens":0,"extend_tokens":128}
+{"rid":"r1","l1_cached_tokens":0,"l2_cached_tokens":32768,"extend_tokens":512}
+```
+
+```bash
+PYTHONPATH=python python benchmark/hicache/bench_prefill_scheduler_admission.py \
+  --requests-jsonl requests.jsonl \
+  --available-tokens 200000 \
+  --cp-max-total-extend-tokens 65536 \
+  --cp-max-total-cached-tokens 131072 \
+  --output json
+```
+
+## 输出字段重点
+
+每个 tick 输出：
+
+- `accepted`：本 tick 被 `PrefillAdder` 接收入 batch 的 request。
+- `stopped_on_rid` / `stopped_result`：scan waiting queue 时第一个挡住的 request 和原因。
+- `log_hit_tokens`：PrefillAdder 统计的 prefix hit token，包含 L1 + 成功 load-back 的 L2。
+- `log_input_tokens`：PrefillAdder 统计的需要 forward 的 paged input token。
+- `cp_total_extend_tokens`：CP bs>1 total extend budget 的累计值。
+- `cp_total_cached_tokens`：CP bs>1 total cached/hit budget 的累计值；对应 `--cp-shared-kv-prefill-max-total-cached-tokens` 的 admission 视角。
+- `allocator_available_after_tick`：fake L1 allocator 在 L2 load-back 后的剩余容量。
+- `load_back_events`：每次 L2 load-back 的 requested/paged/loaded/quota/available 变化。
+
+
+## Cached-token batch limit
+
+`--cp-shared-kv-prefill-max-total-cached-tokens` 用来限制单个 CP shared-KV bs>1 prefill batch 中累计 cached/hit tokens，避免大量高 cache-hit 请求虽然 `extend_tokens` 很小，但 prefix materialize、index/top-k、L2 load-back、descriptor 构造等 cached-token 相关工作过重。
+
+benchmark 对应参数是：
+
+```bash
+--cp-max-total-cached-tokens <tokens>
+```
+
+该 limit 的语义与 total extend limit 一致：
+
+- 只在 CP shared-KV bs>1 admission 中生效。
+- 按 page 对齐累计 accepted request 的 cached tokens。
+- 如果加入新 request 会超过 limit 且当前 batch 已非空，则停止组 batch。
+- 单个 cached token 超过 limit 的 request 仍允许单独运行，避免 scheduler deadlock。
+
+## 当前边界
+
+- 复用真实 `PrefillAdder` admission 逻辑。
+- 不启动真实模型，不测 CUDA kernel，不测真实 attention/transfer。
+- fake allocator 只模拟 scheduler admission 期间的 L2 load-back capacity 消耗；真实 `prepare_for_extend()` 和执行后的 release 行为不在本 benchmark 内。
+- `duration_us` 只是 Python admission 路径耗时，不能代表 ETE 延迟。

test/registered/unit/managers/test_prefill_scheduler_admission_bench.py

@@ -0,0 +1,141 @@
+import importlib.util
+import sys
+import unittest
+from pathlib import Path
+
+_REPO_ROOT = Path(__file__).resolve().parents[4]
+sys.path.insert(0, str(_REPO_ROOT / "python"))
+
+
+_BENCH_PATH = (
+    _REPO_ROOT
+    / "benchmark"
+    / "hicache"
+    / "bench_prefill_scheduler_admission.py"
+)
+
+
+def _load_bench_module():
+    spec = importlib.util.spec_from_file_location(
+        "bench_prefill_scheduler_admission", _BENCH_PATH
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+class TestPrefillSchedulerAdmissionBench(unittest.TestCase):
+    def test_l1_l2_and_extend_tokens_are_reported_with_real_prefill_semantics(self):
+        bench = _load_bench_module()
+        cfg = bench.SchedulerBenchConfig(
+            page_size=64,
+            available_tokens=10000,
+            evictable_tokens=0,
+            max_prefill_tokens=1024,
+            enable_cp_shared_kv_prefill_bs_gt1=True,
+            cp_shared_kv_prefill_max_batch_requests=8,
+            cp_shared_kv_prefill_max_total_extend_tokens=4096,
+        )
+        trace = bench.run_scheduler_admission_trace(
+            [
+                bench.RequestSpec(
+                    rid="hit-l1-l2",
+                    l1_cached_tokens=128,
+                    l2_cached_tokens=128,
+                    extend_tokens=64,
+                    max_new_tokens=1,
+                )
+            ],
+            cfg,
+        )
+
+        self.assertEqual(len(trace.ticks), 1)
+        tick = trace.ticks[0]
+        self.assertEqual([req.rid for req in tick.accepted], ["hit-l1-l2"])
+        accepted = tick.accepted[0]
+        self.assertEqual(accepted.initial_extend_tokens, 192)
+        self.assertEqual(accepted.effective_extend_tokens, 64)
+        self.assertEqual(accepted.l1_cached_tokens, 128)
+        self.assertEqual(accepted.l2_cached_tokens, 128)
+        self.assertEqual(accepted.loaded_l2_tokens, 128)
+        self.assertEqual(tick.log_hit_tokens, 256)
+        self.assertEqual(tick.log_input_tokens, 64)
+
+    def test_cp_total_extend_limit_controls_batching_not_generic_max_prefill_tokens(self):
+        bench = _load_bench_module()
+        cfg = bench.SchedulerBenchConfig(
+            page_size=64,
+            available_tokens=10000,
+            evictable_tokens=0,
+            max_prefill_tokens=192,
+            enable_cp_shared_kv_prefill_bs_gt1=True,
+            cp_shared_kv_prefill_max_batch_requests=8,
+            cp_shared_kv_prefill_max_total_extend_tokens=256,
+        )
+        trace = bench.run_scheduler_admission_trace(
+            [
+                bench.RequestSpec("a", 0, 0, 128, max_new_tokens=1),
+                bench.RequestSpec("b", 0, 0, 128, max_new_tokens=1),
+            ],
+            cfg,
+        )
+
+        self.assertEqual(len(trace.ticks), 1)
+        self.assertEqual([req.rid for req in trace.ticks[0].accepted], ["a", "b"])
+        self.assertEqual(trace.ticks[0].cp_total_extend_tokens, 256)
+
+    def test_l2_load_back_consumes_l1_capacity_and_can_stop_later_requests(self):
+        bench = _load_bench_module()
+        cfg = bench.SchedulerBenchConfig(
+            page_size=64,
+            available_tokens=320,
+            evictable_tokens=0,
+            max_prefill_tokens=4096,
+            enable_cp_shared_kv_prefill_bs_gt1=True,
+            cp_shared_kv_prefill_max_batch_requests=8,
+            cp_shared_kv_prefill_max_total_extend_tokens=4096,
+        )
+        trace = bench.run_scheduler_admission_trace(
+            [
+                bench.RequestSpec("l2-heavy", 0, 128, 64, max_new_tokens=1),
+                bench.RequestSpec("next", 0, 0, 64, max_new_tokens=1),
+            ],
+            cfg,
+        )
+
+        self.assertEqual([req.rid for req in trace.ticks[0].accepted], ["l2-heavy"])
+        self.assertEqual(trace.ticks[0].stopped_on_rid, "next")
+        self.assertEqual(trace.ticks[0].stopped_result, "NO_TOKEN")
+        self.assertEqual(trace.ticks[0].allocator_available_after_tick, 192)
+        self.assertEqual(trace.ticks[0].load_back_events[0].loaded_tokens, 128)
+
+    def test_total_cached_limit_is_observable_in_scheduler_trace(self):
+        bench = _load_bench_module()
+        cfg = bench.SchedulerBenchConfig(
+            page_size=64,
+            available_tokens=10000,
+            evictable_tokens=0,
+            max_prefill_tokens=4096,
+            enable_cp_shared_kv_prefill_bs_gt1=True,
+            cp_shared_kv_prefill_max_batch_requests=8,
+            cp_shared_kv_prefill_max_total_extend_tokens=4096,
+            cp_shared_kv_prefill_max_total_cached_tokens=4096,
+        )
+        trace = bench.run_scheduler_admission_trace(
+            [
+                bench.RequestSpec("a", 4096, 0, 64, max_new_tokens=1),
+                bench.RequestSpec("b", 4096, 0, 64, max_new_tokens=1),
+            ],
+            cfg,
+        )
+
+        self.assertEqual([req.rid for req in trace.ticks[0].accepted], ["a"])
+        self.assertEqual(trace.ticks[0].stopped_on_rid, "b")
+        self.assertEqual(trace.ticks[0].stopped_result, "OTHER")
+        self.assertEqual(trace.ticks[0].cp_total_cached_tokens, 4096)
+
+
+if __name__ == "__main__":
+    unittest.main()