Fix num_tokens_after_batch for different termination #49285

PR vllm-project/vllm#49285 ——《[KV Offload] Fix num_tokens_after_batch for different termination types》。 +17/-18，单文件改动 scheduler.py。

1. 两层架构

┌─────────────────────────────────────────────────────────┐
│  Layer 1: 语义层  num_tokens_after_batch →  "想 store 多少" │
│  基于终止类型决定 token 边界                                │
└─────────────────────────────────────────────────────────┘
                          ↓
                  _calc_num_offloadable_tokens
                          ↓
┌─────────────────────────────────────────────────────────┐
│  Layer 2: 物理层  storable_chunks 的 clamp → "能 store 多少" │
│  基于实际 block_ids 防止越界                                │
└─────────────────────────────────────────────────────────┘

语义层定意图，物理层兜底
只靠物理层能跑（不 crash），但语义错会产生”想 store 不存在的 KV”
这个 PR 修的是语义层

2. 两个方向的语义错

正向（abort）：num_computed_tokens < num_tokens，用 num_tokens 想 store 不存在的数据 反向（EOS）：async scheduling 下 num_computed_tokens 可能漏算最后几个 confirmed tokens

3. 最终修复

if req.status is RequestStatus.FINISHED_ABORTED:
    num_tokens_after_batch = req.num_computed_tokens   # 只要实际 KV
elif req.is_finished():
    num_tokens_after_batch = req.num_tokens            # 其他都用 num_tokens

只特判 ABORTED：只有它”已计算 < 提示”，其他终止都跑完了，clamp 兜底。

4. 边界情况与处理原理

场景	语义层	物理层（clamp）	结果
Queued abort（从未调度）	`num_computed_tokens = 0`	`block_ids = []` → `num_chunks = 0`	无 store job ✓
Preempt 后 abort（waiting）	`num_computed_tokens = 0`	`block_ids = []` → `num_chunks = 0`	无 store job ✓
Running abort	`num_computed_tokens`（部分 KV）	clamp 截断到实际 blocks	只 store 有 KV 的 chunks ✓
EOS	`num_tokens`（含最后 block）	clamp 截断	完整 store ✓
Length cap / Repetition	`num_tokens`	clamp 截断	完整 store ✓
Error（mid-execution）	`num_tokens`（含未确认的 KV）	clamp 截断到实际 blocks	多余的 chunks 被截掉 ✓

所有场景下 clamp 都能防止 crash——block_ids 决定上限，语义层只影响”想 store 多少”。

5. 关键原理：为什么两层缺一不可

只有物理层：能跑，但 queued abort 时仍会”想 store prompt 长度的 KV”，浪费 _calc_num_offloadable_tokens 的计算，且意图不清晰
只有语义层：EOS 用 num_computed_tokens 会少 store 最后 chunk；abort 用 num_tokens 会触发 assert（assert len(offload_keys) == len(offload_block_ids)）
两层协同：语义层决定意图边界，物理层保证不越界

6. Reviewer 的设计哲学

orozery 把改动压到 4 行，背后原则：

YAGNI：LENGTH/REPETITION 罕见，用 num_computed_tokens 也能跑（clamp 兜底）
正交分解：语义层只管意图，物理层只管能力，互不污染
最小变更：preempt 清 offload_keys、entry guard 等”看似相关”的改动都 revert

7. 小结

根因：语义层一刀切 num_tokens，与 abort 语义不匹配
修复：ABORTED 用 num_computed_tokens，其他用 num_tokens
防御：物理层 clamp 处理所有 block_ids 边界
架构启示：意图与能力分层，修改意图不破坏能力，增强能力不污染意图