#!/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 cp_shared_kv_prefill_max_buffer_size: Optional[int] = None kv_cache_dim: int = 656 kv_dtype_bytes: int = 1 index_head_dim: int = 128 vocab_size: int = 128_000 tp_size: int = 8 enable_bs_gt1_prefetch_estimate: bool = False 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] cp_estimated_peak_buffer_bytes: int cp_buffer_breakdown: dict[str, int] 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, cfg: SchedulerBenchConfig): self.available_tokens = int(available_tokens) self.kvcache = SimpleNamespace( kv_cache_dim=cfg.kv_cache_dim, store_dtype=SimpleNamespace(itemsize=cfg.kv_dtype_bytes), index_head_dim=cfg.index_head_dim, quant_block_size=128, index_k_with_scale_buffer_dtype=SimpleNamespace(itemsize=1), ) 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 def get_kvcache(self): return self.kvcache 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.cp_shared_kv_prefill_buffer_estimator import ( CPSharedKVPrefillBufferEstimatorContext, ) 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, cp_shared_kv_prefill_max_buffer_size=cfg.cp_shared_kv_prefill_max_buffer_size, cp_shared_kv_prefill_buffer_estimator_context=( CPSharedKVPrefillBufferEstimatorContext( kvcache=allocator.get_kvcache(), model_config=SimpleNamespace(vocab_size=cfg.vocab_size), tp_size=cfg.tp_size, page_size=cfg.page_size, logprob_chunk_enabled=False, logprob_chunk_size=2048, bs_gt1_l1_prefetch_enabled=cfg.enable_bs_gt1_prefetch_estimate, ) ), ) 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, cfg) 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] peak_estimate = adder.cp_shared_kv_prefill_last_buffer_estimate peak_breakdown = ( { "layer_forward_peak_bytes": peak_estimate.layer_forward_peak_bytes, "logits_window_peak_bytes": peak_estimate.logits_window_peak_bytes, "load_back_window_peak_bytes": peak_estimate.load_back_window_peak_bytes, "materialize_peak_bytes": peak_estimate.materialize_peak_bytes, "prefetch_peak_bytes": peak_estimate.prefetch_peak_bytes, "logits_peak_bytes": peak_estimate.logits_peak_bytes, "remap_peak_bytes": peak_estimate.remap_peak_bytes, "transfer_descriptor_peak_bytes": peak_estimate.transfer_descriptor_peak_bytes, "backup_descriptor_peak_bytes": peak_estimate.backup_descriptor_peak_bytes, } if peak_estimate is not None else {} ) 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:]), cp_estimated_peak_buffer_bytes=int( adder.cp_shared_kv_prefill_estimated_peak_buffer_bytes ), cp_buffer_breakdown=peak_breakdown, 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} " f"cp_buffer_limit={trace.config.cp_shared_kv_prefill_max_buffer_size}" ) 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"cp_peak_buffer={tick.cp_estimated_peak_buffer_bytes} " 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( "--cp-max-buffer-size", type=float, default=None, help="CP shared-KV estimated temp-buffer gate in decimal GB.", ) parser.add_argument("--kv-cache-dim", type=int, default=656) parser.add_argument("--kv-dtype-bytes", type=int, default=1) parser.add_argument("--index-head-dim", type=int, default=128) parser.add_argument("--vocab-size", type=int, default=128_000) parser.add_argument("--tp-size", type=int, default=8) parser.add_argument("--enable-bs-gt1-prefetch-estimate", action="store_true") 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, cp_shared_kv_prefill_max_buffer_size=( None if args.cp_max_buffer_size is None else int(args.cp_max_buffer_size * 1e9) ), kv_cache_dim=args.kv_cache_dim, kv_dtype_bytes=args.kv_dtype_bytes, index_head_dim=args.index_head_dim, vocab_size=args.vocab_size, tp_size=args.tp_size, enable_bs_gt1_prefetch_estimate=args.enable_bs_gt1_prefetch_estimate, 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())