sglang/docs/advanced_features/nsa_prefill_cp_scheduler_admission_benchmark.md

CP shared-KV Prefill Scheduler Admission Benchmark

这个 benchmark 用来离线回答一个问题：给定一批 waiting requests，每个 request 的 L1 cache hit、L2/HiCache hit、以及实际需要 forward 的 extend 长度不同，真实 PrefillAdder 会如何组 prefill batch，最终被哪个 budget 卡住。

脚本位置：

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 是：

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 成功：

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：

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 输入：

{"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 \
  --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 对应参数是：

--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 延迟。