首token延迟(TTFT, Time To First Token)是大模型推理体验的关键指标。用户可以容忍生成速度慢一些,但等待3秒才出第一个字是不可接受的。Prefix Cache是降低TTFT最有效的技术之一——对于共享system prompt的场景,可以将TTFT降低50%-80%。本文深入解析这项技术。

一、为什么首token延迟高

1.1 Prefill阶段的计算瓶颈

大模型生成回复的第一步是处理输入prompt(prefill阶段)。对于1000 token的输入,prefill需要一次性计算所有token的KV Cache。

与生成阶段(每次只处理1个token)不同,prefill阶段是计算密集型的:

  • 输入1000 token,需要计算1000×1000的注意力矩阵
  • 计算量与输入长度的平方成正比
  • 在A100上,1000 token prefill约需300-500ms

1.2 重复计算的浪费

在实际应用中,大量请求共享相同的前缀:

请求1: [System Prompt 500 tokens][用户问题A 20 tokens]
请求2: [System Prompt 500 tokens][用户问题B 30 tokens]
请求3: [System Prompt 500 tokens][用户问题C 15 tokens]

每个请求都要重新计算那500 token的KV Cache——完全相同的计算重复了3次。如果有1000个请求,就是1000次重复计算。

Prefix Cache的核心思想:缓存共享前缀的KV Cache,后续请求直接复用

二、Prefix Cache的技术原理

2.1 KV Cache回顾

Transformer推理中,每个token在每一层都会产生Key和Value向量,存储在KV Cache中供后续token使用:

Layer 0: K=[k0, k1, ..., kn], V=[v0, v1, ..., vn]
Layer 1: K=[k0, k1, ..., kn], V=[v0, v1, ..., vn]
...
Layer L: K=[k0, k1, ..., kn], V=[v0, v1, ..., vn]

Prefix Cache就是在请求完成后,不丢弃这些KV Cache,而是按前缀哈希存储,供后续请求复用。

2.2 前缀匹配与复用

class PrefixCache:
    def __init__(self, max_size_gb=4):
        self.cache = {}  # prefix_hash → KV Cache blocks
        self.max_size = max_size_gb
        self.current_size = 0
    
    def compute_prefix_hash(self, token_ids):
        """计算token序列的哈希"""
        return hash(tuple(token_ids))
    
    def find_longest_prefix(self, token_ids):
        """找到最长可复用的前缀"""
        for length in range(len(token_ids), 0, -1):
            prefix_hash = self.compute_prefix_hash(token_ids[:length])
            if prefix_hash in self.cache:
                return length, self.cache[prefix_hash]
        return 0, None
    
    def store(self, token_ids, kv_cache):
        """存储KV Cache"""
        prefix_hash = self.compute_prefix_hash(token_ids)
        size = self.estimate_size(kv_cache)
        
        # LRU淘汰
        while self.current_size + size > self.max_size and self.cache:
            self.evict_oldest()
        
        self.cache[prefix_hash] = kv_cache
        self.current_size += size
    
    def estimate_size(self, kv_cache):
        """估算KV Cache内存大小"""
        # 每层每个token: 2(K和V) × num_heads × head_dim × 2bytes(FP16)
        num_tokens = kv_cache[0].shape[1]
        per_token = 2 * self.num_heads * self.head_dim * 2 * self.num_layers
        return num_tokens * per_token

2.3 增量哈希优化

直接对整个前缀计算哈希效率低(O(n²))。使用增量哈希可以在O(n)时间内计算所有前缀哈希:

class IncrementalHash:
    def __init__(self):
        self.hash = 0
    
    def update(self, token_id):
        """增量更新哈希"""
        self.hash = (self.hash * 31 + token_id) % (2**64)
        return self.hash
    
    def all_prefix_hashes(self, token_ids):
        """一次遍历计算所有前缀哈希"""
        hashes = []
        h = 0
        for tid in token_ids:
            h = (h * 31 + tid) % (2**64)
            hashes.append(h)
        return hashes

三、在vLLM中的实现

vLLM从0.4.x版本开始原生支持Prefix Cache(称为Automatic Prefix Caching, APC):

3.1 开启方式

from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-2-13b",
    enable_prefix_caching=True,  # 开启Prefix Cache
)

3.2 工作原理

vLLM的Prefix Cache建立在PagedAttention之上:

1. 请求到达,vLLM计算输入token的前缀哈希序列
2. 在BlockTable中查找匹配的物理块
3. 命中的块直接引用(不复制)
4. 未命中的部分正常计算prefill
5. 计算完成后,新块注册到全局缓存
请求A: [shared_prefix_500_tokens][unique_A]
物理块: [block_0][block_1]...[block_31][block_32_A]

请求B到达: [shared_prefix_500_tokens][unique_B]
物理块: [block_0][block_1]...[block_31][block_32_B]
                            直接复用A的block_0到block_31!
                            只需计算block_32_B

3.3 块级缓存粒度

vLLM以block为单位缓存(通常16 token/block),这意味着前缀匹配精度是16 token。即使前缀不完全相同,只要前16 token匹配就能部分复用。

四、性能实测

4.1 测试环境

  • GPU: A100-80GB
  • 模型: Llama-2-13B
  • System Prompt: 500 tokens
  • 用户问题: 20-50 tokens

4.2 结果对比

场景TTFT (无Cache)TTFT (有Cache)提升
500 token前缀420ms85ms80%
1000 token前缀780ms90ms88%
2000 token前缀1450ms95ms93%
无共享前缀50ms50ms0%

关键发现:前缀越长,Prefix Cache的收益越大。对于2000 token的system prompt,TTFT从1.45秒降到95ms——从"用户感知延迟"到"即时响应"。

4.3 吞吐量影响

Prefix Cache不仅降低延迟,还提升吞吐量:

并发数无Cache吞吐有Cache吞吐提升
10850 tok/s1450 tok/s71%
502100 tok/s3800 tok/s81%
1003100 tok/s5200 tok/s68%

五、适用场景与局限

5.1 高收益场景

  1. 统一System Prompt:所有用户共享同一个system prompt,命中率接近100%
  2. Few-shot模板:固定的few-shot示例作为前缀,不同问题共享
  3. 多轮对话:之前的对话历史作为前缀,新增一轮只需计算增量
  4. 文档RAG:检索到的文档作为共享前缀

5.2 低收益场景

  1. 无共享前缀:每个请求的输入完全不同 → 命中率为0
  2. 动态前缀:前缀中有时间戳等变化内容 → 需要设计为"变化部分放在后面"
  3. 长尾输入:请求量少且分散 → 缓存被淘汰后才能命中

5.3 前缀设计建议

# ❌ 不好:时间戳在前缀开头,导致每次都不同
prompt = f"当前时间: {datetime.now()}\n系统提示: {system_prompt}\n用户问题: {question}"

# ✅ 好:固定部分在前,变化部分在后
prompt = f"{system_prompt}\n用户问题: {question}\n当前时间: {datetime.now()}"

六、进阶优化

6.1 多级缓存

L1: GPU显存缓存 (4GB)  → 命中延迟: ~0ms
L2: CPU内存缓存 (16GB) → 命中延迟: ~10ms (需PCIe传输)
L3: 磁盘缓存 (100GB)   → 命中延迟: ~100ms (需读磁盘)

L1未命中时查L2,L2未命中时查L3。对于热门system prompt,L1命中率可达90%+。

6.2 预热缓存

在服务启动时主动计算热门前缀的KV Cache:

async def warmup_prefix_cache():
    """服务启动时预热热门前缀"""
    hot_prefixes = await get_hot_system_prompts()  # 从日志分析热门prompt
    
    for prefix in hot_prefixes:
        token_ids = tokenizer.encode(prefix)
        kv_cache = await model.prefill(token_ids)
        prefix_cache.store(token_ids, kv_cache)
    
    logger.info(f"预热了{len(hot_prefixes)}个前缀缓存")

6.3 智能淘汰策略

单纯的LRU可能不够。结合访问频率和前缀长度:

def evict_score(self, entry):
    """综合评分决定淘汰优先级"""
    frequency = entry.access_count
    recency = time.time() - entry.last_access
    size = entry.token_length  # 越长越值钱
    
    # 评分越低越先淘汰
    return frequency * size / (recency + 1)

七、与Session Cache的区别

Session Cache(会话缓存)是多轮对话场景下的特化:

多轮对话:
  Round 1: [System][Q1] → 生成A1,缓存 [System][Q1][A1] 的KV
  Round 2: [System][Q1][A1][Q2] → 复用前缀,只计算 [Q2]
  Round 3: [System][Q1][A1][Q2][A2][Q3] → 复用更多前缀

Prefix Cache是更通用的方案,Session Cache是其子集。vLLM的APC自动处理了这种情况。

结语

Prefix Cache是投入产出比极高的推理优化——开启一个配置项就能获得50-90%的TTFT降低。对于使用统一system prompt的Agent应用,这几乎是必选项。但要注意前缀设计——将固定内容放在prompt开头,动态内容放在后面,才能最大化缓存命中率。推理优化的艺术,往往就藏在这些细节之中。

本文同步发布于 硅基AGI论坛