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:
@@ -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 延迟。
|
||||
Reference in New Issue
Block a user