# NSA Prefill CP bs>1 max buffer size admission 计划

> 日期：2026-06-11
> 分支：`cjy-cp-refactor`
> 范围：在 `--enable-cp-shared-kv-prefill-bs-gt1` 下新增 batch 级峰值 buffer size gate，同时保留 `--cp-shared-kv-prefill-max-total-extend-tokens` 和 `--cp-shared-kv-prefill-max-total-cached-tokens` 的现有语义。

## 0. 目标

新增一个以 GB 为默认单位的 CP shared-KV prefill batch admission 参数：

```bash
--cp-shared-kv-prefill-max-buffer-size <gb-or-size>
```

建议字段名：

```python
cp_shared_kv_prefill_max_buffer_size: Optional[int]
```

语义：

1. 仅在 `enable_cp_shared_kv_prefill_bs_gt1 && enable_nsa_prefill_cp_shared_kv` 时生效。
2. 默认单位是 GB：裸数字 `8` 表示 `8G`，显式后缀仍支持 `8G` / `8Gi` / `8192Mi`。内部字段统一保存为 bytes。
3. 它是 batch 级 **峰值临时 buffer** admission gate，不替代 KV allocator capacity，也不替代现有 token gates。
4. 现有三个限制同时生效：
   - `cp_shared_kv_prefill_max_batch_requests`
   - `cp_shared_kv_prefill_max_total_extend_tokens`
   - `cp_shared_kv_prefill_max_total_cached_tokens`
   - 新增 `cp_shared_kv_prefill_max_buffer_size`
5. 与现有 token gates 一样，若单个 request 自身超过该 limit，允许它单独运行，避免 scheduler deadlock；如果 batch 已非空，则停止继续加入。
6. scheduler 启动完成、主要静态显存分配完成后，要按该 limit 做一次 CUDA buffer smoke allocation，提前暴露配置过大导致的 OOM。

## 1. 当前代码事实

### C1. admission 入口在 `PrefillAdder`

相关文件：

- `python/sglang/srt/managers/schedule_policy.py:378-481`
- `python/sglang/srt/managers/schedule_policy.py:554-600`
- `python/sglang/srt/managers/schedule_policy.py:638-662`
- `python/sglang/srt/managers/schedule_policy.py:916-936`

当前 CP bs>1 admission 已有：

```python
cp_shared_kv_prefill_max_batch_requests
cp_shared_kv_prefill_max_total_extend_tokens
cp_shared_kv_prefill_max_total_cached_tokens
```

其中：

- extend limit 按 `ceil_paged_tokens(extend_input_len)` 累计。
- cached limit 按 `ceil_paged_tokens(prefix_len)` 累计。
- L2/HiCache hit 会先在 `add_one_req()` 里调用 `tree_cache.init_load_back()`，更新 `prefix_indices` 和 `extend_input_len`，再进入 CP gates。

结论：新增 buffer size gate 应该放在同一个位置：`init_load_back()` 后、`_update_prefill_budget()` 前。

### C2. scheduler 已经把 CP 参数传给 `PrefillAdder`

相关文件：

- `python/sglang/srt/managers/scheduler.py:2399-2424`
- `python/sglang/srt/server_args.py:677-680`
- `python/sglang/srt/server_args.py:5885-5921`

结论：新增参数需要贯穿：

```text
ServerArgs dataclass
  -> ServerArgs._post_init validation
  -> ServerArgs.add_cli_args
  -> Scheduler.get_new_batch_prefill PrefillAdder(...)
  -> PrefillAdder constructor / gate
```

### C3. exact L1 KV capacity 仍由 allocator 管

相关文件：

- `python/sglang/srt/mem_cache/common.py:470-620`
- `python/sglang/srt/mem_cache/allocator.py:862-909`
- `python/sglang/srt/mem_cache/allocator.py:982-1045`

`alloc_paged_token_slots_extend()` 会根据 CP owner-lane page owners 分配 L1 KV pages。这个路径已经包含：

- page owner 规划；
- L1 free room eviction；
- exact capacity wait / fail-fast。

结论：`max_buffer_size` 不应该重复充当 L1 KV allocator 的 correctness gate。它主要限制 CUDA temporary / materialize / logits / descriptor 这类 batch 临时峰值。

### C4. KV pool 和 HiCache 的 per-token bytes 可从现有结构推导

相关文件：

- `python/sglang/srt/mem_cache/memory_pool.py:1492-1580`
- `python/sglang/srt/mem_cache/memory_pool.py:1853-1938`
- `python/sglang/srt/mem_cache/memory_pool.py:2080-2084`
- `python/sglang/srt/mem_cache/hiradix_cache.py:69-94`

NSA KV pool 中：

- MLA KV buffer 形状约为 `(size + page_size, 1, kv_cache_dim)` per layer。
- index buffer 形状约为 `(num_pages, page_size * (index_head_dim + index_head_dim / quant_block_size * 4))`。
- host HiCache 已有 `_estimate_hicache_size_per_token()`，可作为 bytes 计算参考。

结论：scheduler 可以用 `token_to_kv_pool_allocator.get_kvcache()` 推导 dtype、layer_num、kv_cache_dim、index_head_dim、active index layer count。

### C5. L1 CP shared-KV prefetch 当前仍 gate 掉 bs>1

相关文件：

- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:380-397`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:915-917`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:1191-1202`
- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py:1717-1719`

当前：

- MLA prefetch `batch_size != 1` 时 skip。
- index prefetch `batch_size != 1` 时 skip。
- MLA prefetch 若启用，会持有完整 `dense_kv_cache` handle。
- index prefetch 若启用，会持有完整 `dense_page_buffer` handle。

结论：当前 bs>1 运行时，`max_buffer_size` 不需要计入 L1 shared-KV prefetch dense buffer；但 estimator 必须预留字段，后续打开 bs>1 prefetch 时直接生效。

### C6. logits / logprob 峰值与 request 选项强相关

相关文件：

- `python/sglang/srt/layers/logits_processor.py:315-405`
- `python/sglang/srt/layers/logits_processor.py:413-545`
- `python/sglang/srt/layers/logits_processor.py:675-745`
- `python/sglang/srt/layers/logits_processor.py:849-879`

普通 prefill 不返回 input logprob 时，logits 通常只对 batch last-token rows 计算，峰值近似：

```text
batch_size * vocab_shard * logits_dtype_size
```

返回 input logprob 时，`pruned_states` 可能包含更多 extend rows；如果启用 logits chunk，峰值按 `min(pruned_rows, logprobs_chunk_size)` 估算；如果未启用 chunk，峰值按 `pruned_rows` 估算。

结论：buffer gate 至少要 conservative 地覆盖 logits/logprob 峰值，否则 `extend_tokens` 很小但 logprob 请求可能放大 CUDA temp。

## 2. 设计原则

### 2.1 估算同一时间 live 的峰值，并按 stream 并发叠加

不能按 layer 数累加 temp buffer；但也不能把所有类别简单取 `max()`。真正需要防 OOM 的是同一时间 live 的最大值。只要不同 stream 上可能同时存在，就应该在同一个并发窗口里相加。

需要区分至少三类 stream / lifetime：

```text
current/default forward stream:
  当前 layer 的 MLA/index materialize、remap、attention 使用的 dense/slot buffer

prefetch stream:
  next/next+1 layer 的 L2->L1 load、MLA/index prefix prefetch handle

backup / transfer stream:
  async backup descriptor、D2H staging、pending write ack 相关小 buffer

post-forward/logits window:
  logits/logprob/lm-head output buffer，可能与仍未完成的 prefetch/backup buffer 重叠
```

第一版估算不要写成：

```python
# 错误：会低估不同 stream 同时 live 的 buffer。
estimated_peak_buffer_bytes = max(A, B, C, D, E)
```

应该写成 stream-aware overlap windows：

```python
layer_forward_peak = (
    materialize_peak_bytes
    + remap_peak_bytes
    + prefetch_peak_bytes
    + backup_descriptor_peak_bytes
)

logits_window_peak = (
    logits_peak_bytes
    + prefetch_peak_bytes
    + backup_descriptor_peak_bytes
)

load_back_window_peak = (
    transfer_descriptor_peak_bytes
    + prefetch_peak_bytes
    + backup_descriptor_peak_bytes
)

estimated_peak_buffer_bytes = max(
    layer_forward_peak,
    logits_window_peak,
    load_back_window_peak,
)
```

当前 bs>1 L1 shared-KV prefetch 仍 gate 关闭时，`prefetch_peak_bytes=0`。后续打开 bs>1 prefetch 后，这个字段必须反映真实 live prefetch buffer，而不是继续当作独立窗口取 max。

### 2.2 保持 resource type 分离

内部 estimator 不应该只有一个 token scalar。至少保留这些字段，便于日志和 benchmark 判断到底被什么卡住：

```python
@dataclass(frozen=True)
class CPSharedKVPrefillBufferEstimate:
    total_peak_bytes: int
    layer_forward_peak_bytes: int
    logits_window_peak_bytes: int
    load_back_window_peak_bytes: int
    materialize_peak_bytes: int
    prefetch_peak_bytes: int
    logits_peak_bytes: int
    remap_peak_bytes: int
    transfer_descriptor_peak_bytes: int
    backup_descriptor_peak_bytes: int
    l1_load_back_bytes: int
    l1_extend_bytes: int
    host_backup_bytes: int
```

其中 gate 第一阶段只使用 `total_peak_bytes`。其它字段用于 debug、benchmark 和后续校准。

### 2.3 第一阶段宁可保守，也不能引入正式路径额外 CUDA 操作

estimator 必须是 CPU-only、无 CUDA allocation、无 collective、无 tensor dump。正式推理 hot path 不能为了 admission gate 做额外 GPU 计算。

### 2.4 single request 不 deadlock

与现有 token gates 一致：

```python
return projected > limit and len(self.can_run_list) > 0
```

单个大请求允许独占 batch。真正 capacity 仍由 allocator / CUDA OOM / fail-fast 路径处理。

## 3. 峰值估算公式

### 3.1 输入量

每个 request 在 `init_load_back()` 后已有：

```text
prefix_len = len(req.prefix_indices)
extend_len = req.extend_input_len
seq_len = len(req.fill_ids)
```

batch 级：

```text
batch_size = len(can_run_list) + 1
paged_prefix_tokens = sum(ceil_page(prefix_len_i))
paged_extend_tokens = sum(ceil_page(extend_len_i))
prefix_pages = paged_prefix_tokens / page_size
extend_pages = paged_extend_tokens / page_size
```

### 3.2 MLA materialize peak

对 CP shared-KV cache-hit/partial-current 路径，dense MLA buffer 近似：

```text
dense_mla_pages = prefix_pages + extend_pages
mla_page_bytes = page_size * kv_cache_dim * kv_dtype_size
materialize_mla_bytes = dense_mla_pages * mla_page_bytes
```

如果 current-only 且不 materialize prefix，可按 `extend_pages` 估算；第一阶段可以统一用 `prefix_pages + extend_pages` 保守估算。

### 3.3 index materialize peak

NSA index page buffer：

```text
index_page_bytes = page_size * (index_head_dim + index_head_dim / quant_block_size * 4) * uint8_size
materialize_index_bytes = dense_index_pages * index_page_bytes
```

`dense_index_pages` 第一阶段同样可用 `prefix_pages + extend_pages` 保守估算。后续 index skip 生效后，可按 active index layer 或当前 layer 是否需要 index 细化；scheduler admission 第一阶段不按 layer 区分，只估最坏 active layer。

### 3.4 remap / page_inverse peak

当前 slot remap 结构包含：

- `slot_logical_pages`
- `page_inverse`
- `dense_locs` / `dense_pages`
- sorted slot metadata

保守估算：

```text
logical_page_capacity = max_page_id_capacity_or_pages_per_request_capacity
page_inverse_bytes = batch_size * logical_page_capacity * int64_size
slot_metadata_bytes = dense_pages * int64_size * 4
remap_peak_bytes = page_inverse_bytes + slot_metadata_bytes
```

注意：这里必须避免重新引入早期 bs>1 的巨型 per-request dense inverse 设计。实现时如果无法可靠得到 `logical_page_capacity` 的紧上界，先用实际 page-table width 而不是 allocator logical capacity。

### 3.5 logits / logprob peak

普通 generation：

```text
logits_rows = batch_size
logits_peak_bytes = logits_rows * vocab_shard_size * logits_dtype_size
```

返回 input logprob：

```text
pruned_rows = sum(max(1, extend_len_i - extend_logprob_start_len_i))
logits_rows = min(pruned_rows, logprobs_chunk_size) if logprob_chunk_enabled else pruned_rows
logits_peak_bytes = logits_rows * vocab_shard_size * logits_dtype_size
```

第一阶段如 scheduler 难以拿到 vocab shard / logits dtype，可用 conservative 默认：

```text
vocab_size from model_config.vocab_size / tp_size
logits_dtype_size = 2 unless enable_fp32_lm_head then 4
```

### 3.6 prefetch peak

当前 bs>1 L1 shared-KV prefetch disabled：

```text
prefetch_peak_bytes = 0
```

后续打开 bs>1 prefetch 后：

```text
prefetch_peak_bytes = prefetch_mla_dense_bytes + prefetch_index_dense_bytes
```

因为 prefetch handle 在 next-layer consume 前持有完整 dense buffer，必须按完整 dense pages 估算，而不是 prefix rows。

L2->L1 transfer prefetch 不一定有大 GPU staging，但必须计入：

```text
transfer_descriptor_peak_bytes = descriptor_count * descriptor_entry_size
```

如果 direct backend 只使用 CPU descriptor，GPU peak 可为 0，但 benchmark 仍应报告 CPU descriptor count/bytes。

### 3.7 backup peak

per-layer async backup 当前主要是 D2H transfer，GPU memory 大头通常不是 staging，而是 descriptor。第一阶段估算：

```text
backup_descriptor_peak_bytes = backup_page_count * descriptor_entry_size
host_backup_bytes = page_count * host_page_bytes
```

host reservation/evict 仍由 HiCache host allocator 管，`max_buffer_size` 不负责 host correctness。

## 4. 实现计划

### P0. 增加参数、GB 默认解析与校验

文件：

- `python/sglang/srt/server_args.py`
- `test/registered/unit/server_args/test_server_args.py`

步骤：

1. 在 `ServerArgs` dataclass 增加内部 bytes 字段：

```python
cp_shared_kv_prefill_max_buffer_size: Optional[int] = None
```

2. 新增 parser helper，裸数字按 decimal GB 解析，显式后缀复用现有 `human_readable_int` 语义：

```python
def human_readable_gb_size(value: str) -> int:
    value = value.strip()
    if re.fullmatch(r"\d+(?:\.\d+)?", value):
        return int(Decimal(value) * Decimal(10**9))
    return human_readable_int(value)
```

这样启动参数含义为：

```text
--cp-shared-kv-prefill-max-buffer-size 8    -> 8_000_000_000 bytes
--cp-shared-kv-prefill-max-buffer-size 8G   -> 8_000_000_000 bytes
--cp-shared-kv-prefill-max-buffer-size 8Gi  -> 8_589_934_592 bytes
```

3. 在 `_post_init` 校验：

```python
if self.cp_shared_kv_prefill_max_buffer_size is not None and self.cp_shared_kv_prefill_max_buffer_size <= 0:
    raise ValueError("cp_shared_kv_prefill_max_buffer_size must be a positive integer when specified.")
```

4. 在 CLI 增加：

```bash
--cp-shared-kv-prefill-max-buffer-size
```

help 必须明确：裸数字单位为 GB，显式后缀支持 SI/IEC。

5. 单测：

- CLI 能解析裸数字 `8` 为 `8_000_000_000`。
- CLI 能解析 `8Gi` 为 `8 * 2**30`。
- `0` 抛正数校验错误。
- 未设置时为 `None`。

### P1. 定义 CPU-only buffer estimate 数据结构

文件：

- 推荐新增：`python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py`
- 测试：`test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py`

新增：

```python
@dataclass(frozen=True)
class CPSharedKVPrefillBufferEstimate:
    total_peak_bytes: int
    layer_forward_peak_bytes: int
    logits_window_peak_bytes: int
    load_back_window_peak_bytes: int
    materialize_peak_bytes: int
    prefetch_peak_bytes: int
    logits_peak_bytes: int
    remap_peak_bytes: int
    transfer_descriptor_peak_bytes: int
    backup_descriptor_peak_bytes: int
    l1_load_back_bytes: int
    l1_extend_bytes: int
    host_backup_bytes: int
```

新增 helper：

```python
def estimate_cp_shared_kv_prefill_buffer_bytes(
    *,
    page_size: int,
    batch_size: int,
    prefix_lens: Sequence[int],
    extend_lens: Sequence[int],
    kvcache: object | None,
    model_config: object | None,
    tp_size: int,
    return_logprob_rows: int = 0,
    logprob_chunk_enabled: bool = False,
    logprob_chunk_size: int = 2048,
    bs_gt1_l1_prefetch_enabled: bool = False,
) -> CPSharedKVPrefillBufferEstimate:
```

第一阶段要求：

- 不 import CUDA-only package。
- 不分配 CUDA tensor。
- `kvcache is None` 时使用保守 fallback 并返回非零 remap/logits 估算。
- fp8/bf16 根据 `kvcache.store_dtype.itemsize` 计算 MLA KV bytes。
- NSA index bytes 根据 `index_head_dim` / `quant_block_size` / `index_k_with_scale_buffer_dtype.itemsize` 计算。
- `total_peak_bytes` 必须按 2.1 的 stream-aware overlap windows 计算，而不是对各字段简单取 max。

### P2. 在 `PrefillAdder` 中接入 projected peak gate

文件：

- `python/sglang/srt/managers/schedule_policy.py`
- `python/sglang/srt/managers/scheduler.py`
- `test/registered/unit/managers/test_prefill_adder.py`

新增 constructor 参数：

```python
cp_shared_kv_prefill_max_buffer_size: Optional[int] = None
cp_shared_kv_prefill_buffer_estimator_context: Optional[CPSharedKVPrefillBufferEstimatorContext] = None
```

`PrefillAdder` 内维护：

```python
self.cp_shared_kv_prefill_estimate_prefix_lens: list[int] = []
self.cp_shared_kv_prefill_estimate_extend_lens: list[int] = []
self.cp_shared_kv_prefill_estimated_peak_buffer_bytes: int = 0
```

在 `add_one_req()` 中，`init_load_back()` 后计算 candidate：

```text
candidate_prefix_lens = accepted_prefix_lens + [prefix_len]
candidate_extend_lens = accepted_extend_lens + [input_tokens]
estimate = estimate_cp_shared_kv_prefill_buffer_bytes(...)
if estimate.total_peak_bytes > max_buffer_size and can_run_list 非空:
    return AddReqResult.OTHER
```

接受 request 后，在 `_update_prefill_budget()` 附近同步更新 estimator 状态。

单测：

1. 两个小 extend request 在 token gates 未超限时，因为 projected buffer 超限而第二个返回 `OTHER`。
2. 单个 request 超过 buffer limit 仍 `CONTINUE`。
3. buffer limit 与 extend/cached token limit 同时存在时，任一超限都会停止 batch。
4. `enable_cp_shared_kv_prefill_bs_gt1=False` 时 buffer gate 不生效。

### P3. 从 scheduler 传入 estimator context

文件：

- `python/sglang/srt/managers/scheduler.py`

context 应包含：

```python
@dataclass(frozen=True)
class CPSharedKVPrefillBufferEstimatorContext:
    kvcache: object | None
    model_config: object | None
    tp_size: int
    page_size: int
    logprob_chunk_enabled: bool
    logprob_chunk_size: int
    bs_gt1_l1_prefetch_enabled: bool
```

第一阶段 `bs_gt1_l1_prefetch_enabled=False`，因为当前 prefetcher 明确拒绝 bs>1。

`kvcache` 来源：

```python
self.token_to_kv_pool_allocator.get_kvcache()
```

`model_config` 来源：

```python
self.model_config
```

`tp_size` 来源：

```python
self.tp_size
```

logprob chunk env 可从 `sglang.srt.environ.envs` 读取，或在 context 中使用与 `LogitsProcessor` 一致的默认值。

### P3.5. 启动后 CUDA buffer smoke allocation

文件：

- `python/sglang/srt/managers/scheduler.py`
- 可选新增 helper：`python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py`

目的：启动完成、主要静态显存对象已经分配后，按 `cp_shared_kv_prefill_max_buffer_size` 尝试临时分配一次，提前发现配置值超过剩余显存。

推荐接入点：`Scheduler.__init__()` 末尾、`self.is_initializing = False` 前。此时模型、KV pool、attention backend、CUDA graph、disaggregation、overlap 等主要对象已经初始化，且 scheduler 尚未 ready。

触发条件：

```python
should_smoke_check = (
    self.server_args.enable_cp_shared_kv_prefill_bs_gt1
    and self.server_args.enable_nsa_prefill_cp_shared_kv
    and self.server_args.cp_shared_kv_prefill_max_buffer_size is not None
    and self.server_args.disaggregation_mode in (None, "prefill")
)
```

实现要求：

```python
def smoke_check_cp_shared_kv_prefill_buffer_size(device: torch.device | str, size_bytes: int) -> None:
    try:
        torch.cuda.synchronize()
        probe = torch.empty(size_bytes, dtype=torch.uint8, device=device)
        # 只触碰首尾，避免 fill 整个大 buffer 带来长启动耗时；cudaMalloc 本身已能暴露大部分 OOM。
        if size_bytes > 0:
            probe[0] = 0
            probe[-1] = 0
        torch.cuda.synchronize()
    except torch.cuda.OutOfMemoryError as exc:
        raise RuntimeError(
            "[CP_SHARED_KV_FAIL_FAST][prefill_buffer_smoke] "
            f"cannot allocate cp_shared_kv_prefill_max_buffer_size={size_bytes} bytes "
            "after scheduler startup. Lower --cp-shared-kv-prefill-max-buffer-size "
            "or reduce other static GPU memory consumers."
        ) from exc
    finally:
        try:
            del probe
        except UnboundLocalError:
            pass
        torch.cuda.empty_cache()
```

注意点：

1. 这是启动期 fail-fast，不在 scheduler hot path 中运行。
2. 每个 CP/TP rank 都执行一次，验证 per-rank 剩余显存。
3. 不长期持有 probe buffer，避免实际推理容量下降。
4. 日志需要打印 rank、size GB、smoke 通过/失败。
5. 该 smoke check 只能证明启动后能分配一个同等大小的连续 CUDA block，不能替代 runtime estimator；runtime 中仍可能因为多个 stream live buffer 同时存在而 OOM，所以 admission gate 仍必须按 2.1 叠加并发窗口。

### P4. 更新 benchmark，能离线观察 buffer gate

文件：

- `benchmark/hicache/bench_prefill_scheduler_admission.py`
- `test/registered/unit/managers/test_prefill_scheduler_admission_bench.py`
- `docs/advanced_features/nsa_prefill_cp_scheduler_admission_benchmark.md`

新增 benchmark 参数：

```bash
--cp-max-buffer-size <gb-or-size>
--kv-cache-dim <int>
--kv-dtype-bytes <int>
--layer-num <int>
--index-head-dim <int>
--vocab-size <int>
--tp-size <int>
--logprob-rows-per-extend-token <float>
--enable-bs-gt1-prefetch-estimate
```

输出新增：

```json
"cp_estimated_peak_buffer_bytes": ...,
"cp_buffer_breakdown": {
  "layer_forward_peak_bytes": ...,
  "logits_window_peak_bytes": ...,
  "load_back_window_peak_bytes": ...,
  "materialize_peak_bytes": ...,
  "prefetch_peak_bytes": ...,
  "logits_peak_bytes": ...,
  "remap_peak_bytes": ...
}
```

单测：

1. `--cp-max-buffer-size 1` 按 1GB 解析，并能让 benchmark 第二个 request 被挡住。
2. text/json 输出包含 buffer bytes 与 breakdown。
3. token gates 仍保持原测试行为。

### P5. 加入限频 debug 日志

文件：

- `python/sglang/srt/managers/schedule_policy.py`

复用现有 CP bs>1 timing/debug 环境变量，不新增新 env。仅在以下情况日志：

1. buffer gate 阻止加入 request；
2. estimate 超过 slow 阈值或 debug 开启；
3. estimator fallback 到保守模型。

日志格式建议：

```text
[CP_SHARED_KV_BS_GT1_ADMISSION] stop reason=max_buffer_size projected=... limit=... batch_size=... rid=... breakdown=...
```

避免每个 request 刷屏；默认限频。

### P6. 远端验证顺序

1. 本地 CPU 单测：

```bash
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py \
  test/registered/unit/managers/test_prefill_adder.py \
  test/registered/unit/managers/test_prefill_scheduler_admission_bench.py \
  test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py
```

2. 远端容器 CPU 单测：

```bash
cd /sgl-workspace/sglang
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/server_args/test_server_args.py \
  test/registered/unit/managers/test_prefill_adder.py \
  test/registered/unit/managers/test_prefill_scheduler_admission_bench.py \
  test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py
```

3. 远端短 ETE：

- 使用原启动参数；
- 追加 `--cp-shared-kv-prefill-max-buffer-size 8` 起步；裸数字表示 8GB。若要二进制 GiB，显式写 `8Gi`。
- 保留：
  - `--cp-shared-kv-prefill-max-total-extend-tokens 65536`
  - `--cp-shared-kv-prefill-max-total-cached-tokens` 如已配置继续保留；
- 用 50/200 条 GSM8K 或 replay 小样本确认 correctness 不掉点；
- 再跑 replay 观察 batch size、吞吐、OOM/卡死情况。

## 5. 风险与边界

### R1. 估算不应替代 allocator correctness

L1 device KV pages、owner-lane free room、HiCache host reservation 仍由现有 allocator / radix / HiCache 路径保证。buffer gate 只减少 batch 过大导致的临时 CUDA OOM 和 CPU descriptor pressure。

### R2. bs>1 prefetch 未来打开时需要更新 context 和并发窗口

当前代码明确 gate 掉 bs>1 prefetch，因此第一版 `prefetch_peak_bytes=0` 是符合当前运行时的。后续打开 bs>1 prefetch 必须把 `bs_gt1_l1_prefetch_enabled` 改成真实状态，并将 prefetch MLA/index dense handle 加到可能重叠的 layer/logits/load-back windows 中，不能把它作为独立 max window。

### R3. logits 估算只能 conservative

scheduler admission 阶段拿不到所有 logits processor runtime 状态。第一版按 model_config、request logprob flags、env chunk size 做 conservative estimate。若发现过度保守，再用实际 runtime telemetry 校准。

### R4. 不要引入 CUDA 操作到 admission

所有 estimator 逻辑必须 CPU-only。不能在 scheduler scan waiting queue 时创建 CUDA tensor、调用 collective 或触发 TAI kernel。

### R5. 参数默认不改变现有行为

`cp_shared_kv_prefill_max_buffer_size=None` 时完全保持当前行为。只有用户显式设置时才新增 gate。

## 6. 推荐初始参数

对当前 GLM5 FP8 + CP shared-KV + bs>1，建议从保守值开始：

```bash
--cp-shared-kv-prefill-max-buffer-size 8
```

裸数字单位是 GB。若需要二进制 GiB，可显式写：

```bash
--cp-shared-kv-prefill-max-buffer-size 8Gi
```

如果 replay 中 batch size 被频繁 buffer gate 卡住且无 OOM，可逐步调到：

```bash
--cp-shared-kv-prefill-max-buffer-size 12
--cp-shared-kv-prefill-max-buffer-size 16
```

不要直接去掉 extend/cached token gates。三个 gate 分别约束不同风险：

```text
max_total_extend_tokens -> compute / logits / fresh KV growth
max_total_cached_tokens -> cache-hit materialize / load-back / descriptor pressure
max_buffer_size         -> dtype/model/config-aware CUDA temp 峰值
```

## 7. 当前实现状态（2026-06-11）

已完成：

1. `--cp-shared-kv-prefill-max-buffer-size` 已接入 `ServerArgs`，内部保存 bytes；裸数字按 decimal GB 解析，显式 `G/Gi/Mi` 后缀复用现有 size parser。
2. 新增 CPU-only estimator：`python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py`。
   - 按 page 粒度估算 prefix/extend materialize。
   - 同时估算 MLA KV、NSA index、remap、logits、load-back/backup descriptor。
   - `total_peak_bytes` 使用 stream-aware window：layer-forward / logits / load-back 三个并发窗口取最大；窗口内会叠加可重叠的 prefetch/backup descriptor。
3. `PrefillAdder` 已接入 projected buffer gate。
   - 仍保留 request-count / total-extend / total-cached 三个 gate。
   - `max_buffer_size=None` 时不维护 prefix/extend estimator list，避免默认 hot path 额外 CPU 开销。
   - 单个超限 request 仍允许独占 batch，避免 scheduler deadlock。
   - buffer gate 命中时复用 `SGLANG_CP_SHARED_KV_BS_GT1_DEBUG` 做限频 debug 日志，不新增 env。
4. `Scheduler` 已接入 estimator context。
   - 只有 `cp_shared_kv_prefill_max_buffer_size` 设置时才构建 context。
   - 当前 bs>1 L1 prefetch 仍关闭，因此 context 中 `bs_gt1_l1_prefetch_enabled=False`。
5. 启动期 CUDA smoke allocation 已接入。
   - 触发条件：CP shared-KV + bs>1 + max-buffer-size 设置 + prefill/null disagg mode。
   - 每个 rank 在 scheduler ready 前临时分配一次同等大小 `uint8` CUDA tensor，触碰首尾后释放，用于提前暴露配置过大。
6. `benchmark/hicache/bench_prefill_scheduler_admission.py` 已增加 buffer limit、KV/index/logits 参数和 breakdown 输出。

验证：

```bash
# 本地 py_compile
python -m py_compile \
  python/sglang/srt/server_args.py \
  python/sglang/srt/managers/cp_shared_kv_prefill_buffer_estimator.py \
  python/sglang/srt/managers/schedule_policy.py \
  python/sglang/srt/managers/scheduler.py \
  benchmark/hicache/bench_prefill_scheduler_admission.py \
  test/registered/unit/server_args/test_server_args.py \
  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

# 本地可运行的 estimator 单测
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py
# 结果：4 passed

# 远端 cjy-glm5-new /sgl-workspace/sglang-tai targeted 单测
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_cp_shared_kv_prefill_buffer_estimator.py \
  test/registered/unit/server_args/test_server_args.py::test_cp_shared_kv_prefill_max_buffer_size_defaults_to_gb \
  test/registered/unit/server_args/test_server_args.py::test_cp_shared_kv_prefill_max_buffer_size_accepts_iec_suffix \
  test/registered/unit/server_args/test_server_args.py::test_cp_shared_kv_prefill_max_buffer_size_rejects_non_positive_value \
  test/registered/unit/managers/test_prefill_adder.py::TestPrefillAdder::test_cp_prefill_buffer_limit_stops_second_request_without_token_gate \
  test/registered/unit/managers/test_prefill_adder.py::TestPrefillAdder::test_cp_prefill_buffer_limit_allows_single_oversized_request \
  test/registered/unit/managers/test_prefill_scheduler_admission_bench.py::TestPrefillSchedulerAdmissionBench::test_max_buffer_size_is_observable_in_scheduler_trace
# 结果：10 passed

# 远端 managers 相关全文件
PYTHONPATH=python python -m 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
```

未完全验证：

- `test/registered/unit/server_args/test_server_args.py` 全文件在远端有 1 个既有网络依赖失败：`TestPrepareServerArgs.test_prepare_server_args` 需要从 HuggingFace 拉 `Qwen/Qwen2.5-1.5B-Instruct/config.json`，当前容器 DNS/HF 网络失败；本次新增的 3 个 server_args targeted tests 已通过。
- 尚未跑真实 ETE 启动验证 `--cp-shared-kv-prefill-max-buffer-size` smoke allocation 对当前 GLM5 配置的实际可用值。

## 8. Chunked prefill 下的 effective extend limit（2026-06-11）

新增修正：当 CP shared-KV bs>1 与 chunked prefill 同时开启时，`PrefillAdder` 中实际使用的 extend grouping limit 取：

```python
effective_extend_limit = min(
    cp_shared_kv_prefill_max_total_extend_tokens,
    current_rem_chunk_tokens,
)
```

原因：chunked prefill 已经限制一个 batch 当前最多消耗的 extend token。如果 CP 专用 extend limit 大于当前 chunk budget，较大的值不可达；继续用 raw CP limit 还会把 generic `max_prefill_tokens` lift 到过大的数值，造成 admission debug/benchmark 表达的 batch 容量与真实 chunk 容量不一致。

实现位置：

- `python/sglang/srt/managers/schedule_policy.py` 的 `PrefillAdder.__init__()`。

验证：

```bash
# 先确认 red test：旧实现下 256 != 128，失败。
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_prefill_adder.py::TestPrefillAdder::test_cp_prefill_total_extend_limit_is_capped_by_chunked_prefill_size

# 修复后远端 cjy-glm5-new：
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_prefill_adder.py::TestPrefillAdder::test_cp_prefill_total_extend_limit_is_capped_by_chunked_prefill_size
# 结果：1 passed

PYTHONPATH=python python -m pytest -q \
  test/registered/unit/managers/test_prefill_adder.py
# 结果：21 passed
```
