本地推理 KV cache 和 batch size 如何做取舍

本地推理服务最容易出现一种错觉：把 batch size 从 4 调到 16，吞吐上去了，但一遇到长 prompt 就显存告警，p95 延迟也突然拉长。原因不是某一个参数“太小”，而是 KV cache、批处理 token、模型权重和临时张量共用一块显存预算。

更稳妥的做法是先按 prompt 长度、生成长度和并发定义工作负载，再分别设定 KV cache 预算与 batch 上限，最后用 TTFT、TPOT、p95 和缓存驱逐记录寻找拐点。下面的数字只是压测起点，不是通用最优值。

可以先记住三个结论：短 prompt、并发稳定时，适当增大 batch 通常有利于吞吐；长上下文和突发流量时，KV cache 容量比请求数更容易成为瓶颈；当服务已经出现驱逐、重算或排队增长，继续加 batch 往往是在用长尾换吞吐。

步骤一：先固定输入长度、并发和延迟目标

不要从“显存还有多少”开始调，而要先把业务请求分成几档。例如准备短 prompt 512 tokens、中等 prompt 2048 tokens、长 prompt 8192 tokens 三档，再为每档记录最大生成长度、并发数和允许的 p95 TTFT。若只测单请求，batch size 的收益会被高估。

建议把压测输入固定为同一批问题，区分 prefill 和 decode 两段指标：TTFT 反映首 token 等待，TPOT 反映生成阶段每 token 的耗时。还要记录实际 token 数，而不是用字符数代替，因为 tokenizer 会改变显存和调度压力。

步骤二：拆开权重、KV cache 与临时张量的显存竞争

KV cache 保存已经计算过的 key/value，生成下一个 token 时可以复用它们，避免反复重算历史上下文。它会随着每个请求的上下文和生成过程增长；batch 中请求越多、序列越长，所需缓存页越多。

连续批处理还要为一次 forward 准备输入缓冲、注意力索引和 logits 相关临时空间。因此，max_batch_tokens 增大后，不只是“同时服务更多请求”，也可能减少可分配的 KV cache blocks；请求数上限过高时，词表维度相关的 logits 张量又会抬高峰值显存。

步骤三：用 KV 预算和 batch 上限给出初始配置

以 Hugging Face 的连续批处理配置为例，可以先留出一部分显存给模型运行时，再让 KV cache 使用剩余预算：

from transformers.generation import ContinuousBatchingConfig

# 这是压测起点：给 KV cache 留出余量，同时限制一次 forward 的输入规模
cb_config = ContinuousBatchingConfig(
    max_memory_percent=0.80,       # 不把可用显存全部交给缓存，给临时张量留空间
    block_size=256,                 # KV 页的 token 粒度，按框架和模型实现调整
    max_batch_tokens=4096,          # prefill 的 token 预算，不等于请求数
    max_requests_per_batch=8,       # 控制请求数与 logits 峰值，避免只盯 token 总数
    scheduler_type="fifo",         # 先采用易解释的队列策略，再根据长 prompt 调度特征调整
    safety_margin=0.15,             # 缓存接近耗尽时暂停新 prefill，优先完成活动请求
)

这里的 max_batch_tokens 与 max_requests_per_batch 要一起看：前者约束 token 总量，后者约束请求个数。短 prompt 场景可能先撞请求数上限，长 prompt 场景则可能先撞 token 或 KV 上限。若显存较小，先把 token 预算和请求数都设得保守，再逐项放大。

步骤四：用矩阵压测判断吞吐与长尾的拐点

不要只比较平均 tokens/s。固定模型、量化方式和采样参数后，按“prompt 长度 × 并发 × batch token 预算”做小矩阵，每个点至少观察 TTFT、TPOT、p95、显存峰值、KV blocks 使用率和 eviction/preemption 次数。

步骤五：把超限请求放进准入、排队与回退策略

生产环境里，参数调得再好也会遇到长 prompt 和突发并发。可以为请求设置最大输入 token、最大生成 token 和队列超时；当剩余 KV blocks 低于安全线时暂停新的 prefill，让已经开始 decode 的请求先结束。对低优先级请求可以排队或降级，对交互请求则应优先保证 TTFT 上限。

如果框架支持 paged KV cache，长度不同的请求可以共享固定页池，减少连续大块内存带来的碎片；但它不能消除总容量限制。CPU offload、prefix caching 或更小的上下文窗口都属于有代价的回退，必须单独记录带来的 PCIe 延迟、命中率或回答质量变化。

常见问题：KV cache 越大，batch size 就越应该越大吗？

不一定。KV cache 大，说明能容纳更多历史 token 或活动请求，但 batch 增大还会增加 prefill、激活和 logits 的临时开销。应以固定工作负载下的 p95 TTFT、TPOT 和驱逐次数为准；只要长尾恶化，继续扩大 batch 就不再是有效优化。

常见问题：应该先调 block size 还是先调 batch size？

通常先固定框架推荐的 block size，先调 token 预算和请求上限，因为它们直接对应业务并发和 prompt 长度。只有在请求长度差异很大、页尾浪费明显或框架文档明确指出时，才把 block size 纳入第二轮实验。

最终配置应写成一张可回放的运行清单：模型与精度、显存容量、典型 prompt/output 长度、KV 预算、batch token 上限、请求上限、队列策略、p95 目标和回退动作。这样下一次换模型或换 GPU 时，调整的是可解释的边界，而不是凭感觉改一个 batch 数字。