v02_kv_cache
本版本引入什么
本版本加入简单的 KV 缓存,并把生成过程分为预填充和解码两个阶段。
预填充处理提示词并保存 K/V;解码阶段每个请求只处理一个新词元,并关注缓存中的历史状态。
为什么引入
完整重算会浪费工作,因为旧词元每次产生的 K/V 都相同。KV 缓存保存这些历史状态,让解码只关注新词元。
这是路线图中第一个面向推理服务形态的优化。
核心原则
缓存按层、按请求保存:
K: [num_kv_heads, cached_seq_len, head_dim]
V: [num_kv_heads, cached_seq_len, head_dim]
解码时,模型为新词元投影 K/V,将它们追加到缓存,再计算新查询对已缓存前缀的注意力。
建议对比的文件
kvcache.py:第一版缓存结构。forward_params.py:缓存和阶段元数据。layer/gqa.py:预填充/解码分支。llm.py:生成流程变化。
保留的权衡
缓存结构简单并由请求独占,容易阅读,但尚未解决调度、固定解码形状、内存碎片或图捕获。