3a43727216
CP shared-KV bs>1 batching was only bounded by request count, extend tokens, and cached tokens. That left temporary GPU buffers such as MLA/index materialization, remap metadata, logits windows, and transfer descriptors implicit, and raw extend-token limits could exceed the active chunked-prefill budget.\n\nThis adds an explicit max-buffer-size admission gate with a CPU-only stream-aware estimator, wires it through PrefillAdder/Scheduler, performs a startup CUDA smoke allocation when configured, and reports the estimate in the scheduler admission benchmark. When chunked prefill is active, the effective CP extend-token limit is capped by the current chunk budget so the CP path does not advertise unreachable batch capacity or lift max-prefill-tokens too far.\n\nConstraint: Admission estimation must stay CPU-only on the scheduler hot path; CUDA allocation is limited to startup smoke checking.\nConstraint: Single oversized requests must still be allowed to run alone to avoid scheduler deadlock.\nRejected: Rely only on --max-prefill-tokens | it does not reliably bound the first oversized request and does not model cache-hit/load-back pressure.\nRejected: Let CP extend limit exceed chunked-prefill size | it creates an unreachable effective capacity and misleading budget lift.\nConfidence: medium\nScope-risk: moderate\nDirective: If bs>1 L1 prefetch is enabled later, update CPSharedKVPrefillBufferEstimatorContext.bs_gt1_l1_prefetch_enabled and include the live prefetch dense buffers in overlap windows.\nTested: local py_compile for touched files\nTested: local PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py (4 passed)\nTested: remote cjy-glm5-new targeted pytest for new server_args, PrefillAdder, estimator, and benchmark cases (10 passed)\nTested: remote cjy-glm5-new PYTHONPATH=python pytest -q test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py test/registered/unit/managers/test_prefill_adder.py test/registered/unit/managers/test_prefill_scheduler_admission_bench.py (29 passed before chunk cap, then test_prefill_adder.py 21 passed after chunk cap)\nNot-tested: full server_args suite because existing TestPrepareServerArgs tries to reach HuggingFace and fails under container DNS/network\nNot-tested: GLM5 ETE smoke with --cp-shared-kv-prefill-max-buffer-size
737 lines
28 KiB
Python
Executable File
737 lines
28 KiB
Python
Executable File
#!/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())
|