KV Cache 推理引擎
KV Cache Inference Engine
高效的自回归推理引擎,含 KV 缓存管理和工具调用状态机
子问题
1.缓存预分配与位置追踪
2.prefill-decode 分离
3.多 sample 并行解码
4.工具调用 token 注入
5.FA3/SDPA 运行时自动检测与 API 统一
6.滑动窗口 pattern 字符串到层级窗口大小的映射
7.安全表达式求值的沙箱化与超时保护
8.forced token 与 sampled token 的 mask 区分
各项目的解法1 solutions
Signals
横向对比
| 维度 | nanochat |
|---|---|
| 缓存预分配 | 固定形状 (n_layers,B,T,H,D) 一次性 zeros 分配,decode 阶段按 hint 预估总长 |
| prefill-decode分离 | batch=1 prefill 后 KVCache.prefill() 克隆到 N sample 并行解码 |
| 注意力后端 | SimpleNamespace 统一 API,运行时 sm90 检测自动 FA3/SDPA 切换 |
| 滑动窗口 | 字符串 pattern SSSL 平铺到层,最后一层强制全上下文 |
| 工具调用注入 | RowState.forced_tokens deque + mask 标记区分采样/强制 token |
| 流式生成 | yield (token_column, token_masks) 逐步输出,支持多 sample 并行 |
最佳实践
1.batch=1 prefill 后克隆缓存到多 sample
2.prefill 缓存精确分配 seq_len=len(tokens) 避免浪费,decode 缓存按 hint 一次分配
3.logits.expand() 替代 repeat() 创建视图节省显存
4.位置推进集中在最后一层 forward 确保所有层读写一致
5.SimpleNamespace 导出统一 API 实现注意力后端零成本切换