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.
This commit is contained in:
laoyao0822
2026-06-10 22:21:59 +08:00
parent 4f65d7a176
commit 9a9893e571
3 changed files with 909 additions and 0 deletions
@@ -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 延迟。