diff --git a/benchmark/hicache/bench_prefill_scheduler_admission.py b/benchmark/hicache/bench_prefill_scheduler_admission.py new file mode 100755 index 000000000..1f9c545a3 --- /dev/null +++ b/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()) diff --git a/docs/advanced_features/nsa_prefill_cp_scheduler_admission_benchmark.md b/docs/advanced_features/nsa_prefill_cp_scheduler_admission_benchmark.md new file mode 100644 index 000000000..03267fd10 --- /dev/null +++ b/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 +``` + +该 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 延迟。 diff --git a/test/registered/unit/managers/test_prefill_scheduler_admission_bench.py b/test/registered/unit/managers/test_prefill_scheduler_admission_bench.py new file mode 100644 index 000000000..fb4a462fd --- /dev/null +++ b/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()