Gate CP shared-KV prefill batching behind explicit limits
The scheduler can now admit multi-request NSA in-seq CP shared-KV prefill batches only when the shared-KV bs>1 flag is explicitly enabled. The gate is still disabled by default and is scoped to CP shared-KV so ordinary CP is not widened accidentally. Batch admission is bounded by optional request-count and page-aligned extend-token limits while real memory capacity remains allocator-owned. This keeps bf16 and fp8 on the same scheduler path because dtype differences are already reflected in KV pool token/page capacity. Constraint: bs>1 runtime paths remain guarded by existing CP shared-KV fail-fast checks. Constraint: Scheduler must not duplicate bf16/fp8 byte-level capacity estimation. Rejected: Open the old CP gate unconditionally | ordinary CP would inherit an unverified shared-KV-specific batching path. Rejected: Treat the extend-token cap as a hard per-request limit | a single large request could deadlock the scheduler. Confidence: medium Scope-risk: moderate Directive: Keep CP shared-KV batching gated until ETE validates EAGLE accept length, output length, and HiCache load/backup behavior under real traffic. Tested: local py_compile for server_args, schedule_policy, scheduler, prefill_adder tests, and server_args tests. Tested: remote g0034 py_compile for the same files. Tested: remote g0034 pytest target set: 5 passed for parser, parameter validation, default single-request CP gate, enabled bs>1 gate, and page-aligned extend cap. Tested: remote g0034 pytest test_prefill_adder.py => 13 passed. Not-tested: full server_args test file has an unrelated HuggingFace DNS/config-download failure in TestPrepareServerArgs.test_prepare_server_args. Not-tested: ETE production traffic with --enable-cp-shared-kv-prefill-bs-gt1.
This commit is contained in:
@@ -903,3 +903,35 @@ PYTHONPATH=python python -m pytest -q \
|
||||
test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py
|
||||
=> 10 passed, 3 warnings
|
||||
```
|
||||
|
||||
## 18. 2026-06-04 scheduler admission gate 首版打开
|
||||
|
||||
目标:先允许 NSA in-seq CP shared-KV 在 scheduler 层组成真实 multi-request prefill batch,同时用显式参数把首版风险收窄。
|
||||
|
||||
### 参数
|
||||
|
||||
1. `--enable-cp-shared-kv-prefill-bs-gt1`
|
||||
- 默认关闭;关闭时 CP prefill 仍保持旧的 `batch_size == 1` admission gate。
|
||||
- 打开后,且 `--enable-nsa-prefill-cp-shared-kv` 同时开启时,scheduler 可以把多个 waiting requests 放进同一个 prefill batch;普通 CP 不因该参数打开 bs>1。
|
||||
2. `--cp-shared-kv-prefill-max-batch-requests`
|
||||
- 可选正整数;限制同一个 CP shared-KV prefill batch 的 request 数。
|
||||
- 不设置时只受通用 `--prefill-max-requests` 约束。
|
||||
3. `--cp-shared-kv-prefill-max-total-extend-tokens`
|
||||
- 可选正整数;限制同一个 CP shared-KV prefill batch 的 **page-aligned extend tokens** 总量。
|
||||
- 单个 request 自身超过该值时仍允许独占 batch,避免 scheduler deadlock;真正是否能运行仍由 KV allocator 容量判断。
|
||||
|
||||
### 容量语义
|
||||
|
||||
scheduler admission 不做 bf16/fp8 字节级估算。原因是实际 KV pool 容量已经在初始化时按 dtype、page size、model shape、mem fraction 转换为 token/page capacity。admission 只消费现有 `PrefillAdder` 的 token/page budget:
|
||||
|
||||
- `rem_input_tokens` / `rem_total_tokens` / `cur_rem_tokens` 继续作为通用容量预算;
|
||||
- `alloc_for_extend()` 仍是最终容量仲裁点;
|
||||
- CP owner-lane 不足仍通过 `KVCapacityWaitError` deferred,不在 scheduler 新增 collective 或重复推导 dtype bytes。
|
||||
|
||||
这样 bf16 与 fp8 都走同一个 admission 逻辑,差异由底层 allocator 的实际可用 token capacity 体现。
|
||||
|
||||
### 风险与后续
|
||||
|
||||
1. `cp_shared_kv_prefill_max_total_extend_tokens` 首版使用 page-aligned extend 累计,和 allocator 的最小 page 单位一致;日志/指标里如果要展示用户侧 token,需要另加 valid-token 统计。
|
||||
2. 该 gate 只控制 scheduler 组 batch,不解决所有 bs>1 runtime correctness;打开 ETE 时仍应保留现有 fail-fast,特别是 draft/EAGLE、logprob、hidden capture、compute padding 路径。
|
||||
3. 如果 ETE 中 `KVCapacityWaitError` 频繁出现,下一步应把 owner-lane free pages/evictable pages 的 batch precheck 前移到 scheduler,而不是引入 all-reduce。
|
||||
|
||||
Reference in New Issue
Block a user