自托管LLM必读：如何计算KV Cache预算并实现请求准入控制

自托管LLM必读：如何计算KV Cache预算并实现请求准入控制 | 前端进阶之旅

Full translation

完整中文译文

为什么"能跑"的服务器在高并发下还是会崩？

加载权重只能证明模型能装进显存，但这和服务器能不能对外服务是两回事。推理时真正决定存活的是 KV 缓存：它把请求上下文中已生成的 token 对应的 key/value 张量缓存起来，这样模型每一步解码就不用重新计算。单个请求的 KV 缓存大小约为：

kv_bytes_per_request =
    2 × num_layers × num_kv_heads × head_dim × seq_len × bytes_per_element

前面的 2 是 key 和 value。注意这里用的是 num_kv_heads，而不是总的 attention head 数量——现代模型使用分组查询注意力（GQA），相比 naive 公式可以把体积缩小 4–8 倍。真正要命的是 seq_len，它还要乘以并发请求数。权重是一次性的固定成本；KV 缓存是每个请求都要付的成本，每个并发用户、每增加一个 token 的上下文都要再付一次。

这就是为什么故障模式是非线性的。八个短对话可能跑得稳稳当当；但一个 6,000 token 的文档乘以八的并发，可能瞬间就把同样的预算打爆——因为这八个请求同时背负着巨大的 KV 占用。

结论：权重只能告诉你模型能不能加载；KV 缓存告诉你能同时持有多少并发 token——而这个数字才决定你能不能撑住。

如何计算 KV 缓存预算？

从权重和运行时开销之后剩下的显存出发，再除以每个 token 的成本。下面是一个在租卡之前就可以跑的 calculator：

def kv_bytes_per_token(num_layers, num_kv_heads, head_dim, dtype_bytes=2):
    # 2 = keys + values. dtype_bytes: 2 for FP16/BF16 KV, 1 for FP8/INT8 KV.
    return 2 * num_layers * num_kv_heads * head_dim * dtype_bytes

def token_budget(total_gib, weight_gib, overhead_gib, per_token_bytes):
    free_bytes = (total_gib - weight_gib - overhead_gib) * (1024 ** 3)
    if free_bytes <= 0:
        raise ValueError("No memory left for KV cache after weights + overhead")
    return int(free_bytes // per_token_bytes)

# Llama-3-8B-class: 32 layers, 8 KV heads (GQA), head_dim 128, BF16 KV
per_tok = kv_bytes_per_token(num_layers=32, num_kv_heads=8, head_dim=128)
budget = token_budget(total_gib=24, weight_gib=16, overhead_gib=2, per_token_bytes=per_tok)

print(f"{per_tok/1024:.1f} KiB per token")      # ~128 KiB/token
print(f"~{budget:,} total KV tokens in flight") # the number that matters

输出——你能同时持有的 KV token 总数——才是你的真实容量。如果算出来大约是 48,000 个 token，那就是六个 8k 上下文的请求，或者四十八个 1k 的请求，或者一个 48k 上下文的大请求把所有余量吃光。并发数和上下文长度在同一个池子里互相挤占。

结论：你的容量不是"N 个请求"，而是一个固定的 KV token 在飞数量，每个请求按其上下文长度从这个池子里支取。

在哪里拒绝：API 边缘层还是调度器？

两层都要，各司其职。评论者的问题——是在边缘层截断上下文，还是让调度器在缓存预算耗尽时再拒绝——答案是"都要"：边缘层做便宜、可预知的拒绝；调度器做实时背压，而这是边缘层看不到的。

在边缘层设置硬性上下文上限是单位价值最高的控制，因为无界上下文的请求是唯一能把健康服务器变成 OOM 的元凶。在它接触到 GPU 之前就用明确的 400 拒绝它。然后用一个并发门控来粗粒度地卸掉负载——返回 429，这样调度器永远不会被要求做不可能的事。最后，调度器（vLLM 的 PagedAttention 分配器，或 TGI 的 max-concurrent-requests 限制）来处理细粒度的、每一毫秒都在变化的真实情况。截至 2026 年中，vLLM 在 KV block 不足时甚至会抢占并重新计算一个正在运行的请求——这是正确的，但代价是延迟，所以你需要上游的门控来吸收大部分压力。

结论：在边缘层截断上下文以保证正确性，用并发门控卸负载，调度器保有最后一毫秒的真相——不要让某一层单独干三件事。

一个最小的准入门控长什么样？

下面是一个边缘门控，它在代理到推理后端之前先强制执行上下文上限和 token 在飞预算。这里故意写得简单——用信号量控制 token 计数——因为目标是快速、便宜地失败：

import asyncio
from fastapi import FastAPI, HTTPException, Request

app = FastAPI()

MAX_CONTEXT = 8192          # hard per-request cap: input + generation
KV_TOKEN_BUDGET = 48_000    # from token_budget() above, with headroom
_in_flight = 0
_lock = asyncio.Lock()

async def admit(cost_tokens: int):
    global _in_flight
    async with _lock:
        if _in_flight + cost_tokens > KV_TOKEN_BUDGET:
            raise HTTPException(429, "KV budget exhausted",
                                headers={"Retry-After": "2"})
        _in_flight += cost_tokens
    return cost_tokens

async def release(cost_tokens: int):
    global _in_flight
    async with _lock:
        _in_flight -= cost_tokens

@app.post("/generate")
async def generate(req: Request):
    body = await req.json()
    prompt_tokens = body["prompt_tokens"]          # count upstream, not len(text)
    max_new = min(body.get("max_new_tokens", 512), MAX_CONTEXT - prompt_tokens)
    cost = prompt_tokens + max_new
    if cost > MAX_CONTEXT:
        raise HTTPException(400, f"context {cost} exceeds cap {MAX_CONTEXT}")

    await admit(cost)
    try:
        return await call_backend(body)   # your vLLM/TGI call
    finally:
        await release(cost)

这里有两个容易踩的坑。第一，用模型的真 tokenizer 在上游计数 prompt_tokens——len(text) 不是 token 数量，会放行 GPU 实际装不下的请求。第二，为 max_new_tokens 预留，而不是只算 prompt：生成过程每步都会让 KV 缓存增长一个 token，所以短 prompt 加大生成预算仍然是一个昂贵的请求。按你允许的最坏情况来预算，而不是按收到的 prompt 来预算。

结论：按真实 tokenizer 计量的 prompt_tokens + max_new_tokens 来准入，且在 429 时始终返回 Retry-After，这样客户端会退让而不是疯狂重试。

如何计算 LLM 的 KV 缓存大小？每个 token 是 2 × num_layers × num_kv_heads × head_dim × dtype_bytes。再乘以序列长度得到单个请求的量，所有并发请求相加得到 KV 内存总用量。对于 GQA 模型要使用 num_kv_heads（而不是总 head 数），否则会高估好几倍。

当 LLM 请求过大时应该返回 429 还是 400？当单个请求超过你的硬性上下文上限时用 400——它永远不会成功，客户端必须修改它。当请求本身合法但服务器暂时耗尽了 KV 预算时用 429，并带上 Retry-After 头——等负载降下来同一个请求就能成功。

为什么 LLM 服务器测试时不 OOM，上了流量就崩？测试时通常只发一个短 prompt，几乎不占 KV 缓存。真实流量是长上下文加并发，而 KV 缓存内存随 context_length × concurrent_requests 增长。权重全程都能装下；出问题的是 KV 缓存。

按 KV token 预算来规划你的部署，而不是按权重能不能加载。在 API 边缘设置硬性上下文上限，这样一个无界请求不可能打垮整台机器；加一个并发门控用 429 和 Retry-After 卸负载；让你的推理服务器的调度器保有最后一毫秒的分配决策。如果你现在用的是 vLLM，它已经为你做了细粒度的 KV 分页和抢占——但当前置准入门控把它挡在不可能的场景之外时，它服务得最好。在租 GPU 之前算好这笔账，你的第一次部署就能在真实用户面前活下来。