很多推理场景的 prompt 前缀是固定的系统提示词、角色设定、Few-shot 示例。PrefixCaching 把前缀的 KV Cache 缓存起来相同前缀的请求直接复用省掉 Prefill 计算。CANN 的 ATB 从 8.5 版本开始原生支持 PrefixCaching。原理普通推理请求1: [系统提示 500 token 用户输入 50 token] → Prefill 550 token → Decode 请求2: [系统提示 500 token 用户输入 80 token] → Prefill 580 token → Decode重复计算前 500 tokenPrefixCaching请求1: [系统提示 500 token 用户输入 50 token] → Prefill 550 token → 缓存前 500 token 的 KV 请求2: [系统提示 500 token 用户输入 80 token] → 只 Prefill 后 80 token复用前 500 token 的 KV前缀命中时Prefill 计算量从 580 token 降到 80 tokenTTFT 降 86%。ATB 的 PrefixCaching 配置fromatbimportLLM,PrefixCachingConfig modelLLM(meta-llama/Llama-2-7b-hf,devicenpu:0,prefix_cachingPrefixCachingConfig(enableTrue,max_prefix_length4096,# 最大缓存前缀长度cache_capacity48,# KV Cache 显存上限GB))显式指定前缀最可靠的方式是显式告诉 ATB 哪部分是前缀system_prompt你是一个专业的AI助手请用中文回答问题。以下是几个示例\n示例1: ...\n示例2: ...# 把系统提示注册为前缀prefix_idmodel.register_prefix(system_prompt)# 推理时指定前缀 IDresultmodel.generate(请解释什么是昇腾NPU,prefix_idprefix_id,# 自动复用系统提示的 KV Cache)register_prefix只做一次 Prefill后续所有请求复用。自动前缀检测不显式指定前缀时ATB 自动检测公共前缀modelLLM(meta-llama/Llama-2-7b-hf,devicenpu:0,prefix_cachingPrefixCachingConfig(enableTrue,auto_detectTrue,# 自动检测公共前缀))# 这两个请求会自动复用公共前缀result1model.generate(系统你是AI助手\n用户什么是CANN)result2model.generate(系统你是AI助手\n用户什么是Ascend C)# 系统你是AI助手\n 是公共前缀只计算一次自动检测的原理对 prompt 做 token-level 的 hash 比对找到最长公共前缀。KV Cache 生命周期缓存的 KV Cache 需要管理生命周期# 注册前缀永久缓存直到手动释放prefix_idmodel.register_prefix(system_prompt)# 使用前缀resultmodel.generate(问题,prefix_idprefix_id)# 释放前缀显存不够时model.release_prefix(prefix_id)# LRU 策略自动淘汰最久未用的前缀modelLLM(model_id,prefix_cachingPrefixCachingConfig(enableTrue,eviction_policylru,# LRU 淘汰策略max_cached_prefixes64,# 最多缓存 64 个前缀))多轮对话的前缀复用多轮对话天然适合 PrefixCaching——每一轮都是在前一轮基础上追加轮次1: [系统提示 500 用户输入 50] → 缓存 550 token 轮次2: [系统提示 500 用户输入 50 模型回复 200 用户输入 60] → 复用前 550新增 260 轮次3: [系统提示 500 用户输入 50 模型回复 200 用户输入 60 模型回复 300 用户输入 40] → 复用前 810新增 340# ATB 自动处理多轮对话的前缀复用messages[{role:system,content:system_prompt},{role:user,content:什么是CANN},]# 每轮对话自动复用之前的 KV Cacheresult1model.chat(messages)messages.append({role:assistant,content:result1})messages.append({role:user,content:它和CUDA是什么关系})result2model.chat(messages)# 自动复用前缀实测效果Llama2-7BAtlas 800I A2系统提示 500 token场景TTFT (ms)Prefill 计算 (token)显存节省无 PrefixCaching1805500PrefixCaching命中255089%PrefixCaching未命中1805500500 token 前缀命中时TTFT 从 180ms 降到 25ms7 倍提速。显存预算PrefixCaching 的 KV Cache 占用显存。7B 模型4096 token 前缀KV Cache 大小 2 (KV) × num_layers × hidden_dim × seq_len × dtype_size 2 × 32 × 4096 × 4096 × 2 (bf16) 2 GB64GB 显存中缓存 10 个不同前缀 20GB。剩余 44GB 给 decode KV Cache还能服务约 20 个并发请求。PrefixCaching 是在线对话服务的关键优化。相同系统提示的请求只 Prefill 一次后续直接复用 KV Cache。ATB 8.5 原生支持TTFT 降 7 倍。显存预算要提前规划。仓库在这里https://atomgit.com/cann/ATB
