大模型推理的显存大头不是权重，是 KV Cache。Llama2-70B 的权重约 140GB，但 32 个并发请求、序列长度 4096 的 KV Cache 就要 80GB。ATB 的 KV Cache 管理模块决定了你的推理服务能同时服务多少个请求。

KV Cache 的显存需求

单层 KV Cache 大小计算：

KV Cache = 2 (K+V) × num_kv_heads × head_dim × seq_len × batch_size × dtype_size

Llama2-7B (GQA, 32 kv_heads, head_dim=128, fp16):
  单层 = 2 × 32 × 128 × 4096 × batch × 2 bytes
       = 64MB × batch

  32 层 = 2GB × batch

batch=32 时 KV Cache 占 64GB。权重才 14GB，KV Cache 是权重的 4.5 倍。

Paged KV Cache

标准实现预分配连续显存给每个请求的 KV Cache。问题：序列长度不确定，预分配太浪费。

ATB 用 Paged KV Cache（类似 vLLM 的 PagedAttention）：

KV Cache 切成固定大小的 Block：
  block_size = 16 tokens
  每个 block = 2 × num_kv_heads × head_dim × 16 × 2 bytes = 256KB

按需分配：
  请求 A 生成到 token 48 → 分配 3 个 block
  请求 B 刚开始 → 分配 1 个 block
  请求 C 结束 → 释放 block

显存利用率从 40-60%（静态预分配）提升到 90%+

BlockManager 的调度

ATB 的 BlockManager 维护一个 block 池：

# BlockManager 伪代码
class BlockManager:
    def __init__(self, total_blocks):
        self.free_blocks = list(range(total_blocks))
        self.request_blocks = {}  # request_id → [block_ids]

    def allocate(self, request_id, num_tokens):
        num_blocks = ceil(num_tokens / block_size)
        blocks = self.free_blocks[:num_blocks]
        self.free_blocks = self.free_blocks[num_blocks:]
        self.request_blocks[request_id] = blocks

    def free(self, request_id):
        blocks = self.request_blocks.pop(request_id)
        self.free_blocks.extend(blocks)

当一个请求生成完毕（遇到 EOS 或达到 max_tokens），它的 block 被释放回池子，新请求可以复用。这比静态预分配节省 30-50% 的显存。

Prefix Caching

多个请求共享相同前缀（比如 system prompt）时，KV Cache 的前缀部分可以共享：

请求 A: [system prompt] + "你好"
请求 B: [system prompt] + "再见"

system prompt 的 KV Cache 只算一次，两个请求共享

ATB 的 Prefix Caching 实现：

from atb import LLM

model = LLM("model_id", device="npu:0",
            enable_prefix_caching=True)  # 开启前缀缓存

# 相同 system prompt 的请求自动共享 KV Cache
results = model.generate([
    "System: 你是助手\nUser: 你好",
    "System: 你是助手\nUser: 再见"
])
# system prompt 部分的 KV Cache 只计算一次

Llama2-7B + 1K token 的 system prompt：不开 prefix caching 时每个请求要 1GB 的 KV Cache 做前缀；开了之后只有第一个请求算一遍，后续请求零开销。

显存碎片问题

Paged KV Cache 解决了"预分配浪费"的问题，但引入了新问题：block 碎片。

一个请求的 KV Cache 分散在多个不连续的 block 里。Attention 计算时需要读这些 block，如果 block 不连续，DMA 读取效率低。

ATB 的做法：每次分配 block 时尽量选择连续的 block。不够连续时才散开。在 block 池维护一个连续区间的空闲列表：

# 简化的连续分配逻辑
free_ranges = [(0, 100), (150, 200), ...]  # 连续空闲区间

def allocate(num_blocks):
    # 优先从最大连续区间分配
    for start, end in sorted(free_ranges, key=lambda r: r[1]-r[0], reverse=True):
        if end - start >= num_blocks:
            return list(range(start, start + num_blocks))
    # 没有足够大的连续区间，分散分配
    return allocate_scattered(num_blocks)

实际测试中，连续分配比例在 80-90%。剩下的 10-20% 分散 block 对性能影响约 3-5%。

性能数据

Atlas 800I A2，Llama2-7B，KV Cache 对比：

配置	最大并发数	显存利用率	首 token 延迟
静态预分配	12	45%	65ms
Paged KV Cache	24	88%	70ms
Paged + Prefix Caching	30	92%	55ms

首 token 延迟略高 5ms（Paged 的间接寻址开销），但并发数翻倍。Prefix Caching 反而降低了首 token 延迟（共享前缀不重算）。

---

KV Cache 管理是推理服务的命脉。ATB 的 Paged KV Cache 让显存利用率翻倍，Prefix Caching 让共享前缀零开销。如果你的推理服务并发上不去，先看 KV Cache 的分配策略。仓库在这里：

https://atomgit.com/cann/ATB