盘古大模型团队开源推理组件,Omni Cache以内存为中心的KV管理与推理加速实践
本文整理自 AICon 上海 2026 Ken Zhang 分享《Omni Cache以内存为中心的KV管理与推理加速》通过音视频转录总结工具Ai好记进行视频转文字整理以下为精炼整理后的内容。推理时代的显存瓶颈大模型推理的性能瓶颈除了算力之外最核心的限制就是显存HBM。当前推理引擎普遍采用静态预分配方式——引擎启动后就把 KV Block Pool 预分配好每一层需要的显存空间提前固定。这种方案的第一个问题是HBM 真的用到了极致吗从 PagedAttention 起源的内存分块机制把显存分配成一个个 block虽然带来了灵活性但 block 之间跳跃寻址的额外开销不可忽视。第二个问题是非连续块寻址导致算子开发复杂度高。Ken Zhang 的团队在盘古大模型的实际运维中发现算子开发中最常遇到的问题不是算法逻辑本身而是 HBM 的管理问题——地址越界、tensor stride 不匹配、内存碎片。PagedAttention 的 block 机制使这些问题更加突出。架构翻转DRAM 为骨干HBM 为缓存传统的推理引擎以 HBM 为中心调度请求进入后先分配 HBM block资源在生命周期内持续占用即使实际并不需要一直占着这块空间。Ken Zhang 团队的方案做了一个架构翻转——把 Block Pool 移到 DRAM 上分配HBM 只作为缓存来用仅缓存当前层计算所需的 KV。在某型号 H20 或 B200 上单卡 HBM 约 100GB8 张卡约 1TB。但 HBM 的成本是 DRAM 的 5 倍以上。把 Block Pool 全部移到 DRAM 分配后HBM 上只需要保留连续的 KV 缓冲区多层复用同一个 buffer。利用层间计算的时间每层约 30ms60 层每层约 500μs做 H2D/D2H 搬运时间是够用的。如果单层 buffer 不够还可以用多 buffer 策略——单数层用一个 buffer双数层用另一个甚至可以扩展更多 buffer。性能收益Prefill 和 Decode 双端提升架构翻转带来的核心改变是Block Pool 移至 DRAMHBM 留出连续缓冲区算子从 PagedAttention 替换为 ReshapedAttention。连续寻址带来的优势很明显HBM 寻址效率大幅提高L2/L3 cache 命中率提升整体性能收益显著。FlashAttention 与 PagedAttention 相比15%-20% 的性能差异大部分来自 Block 与 Slot Mapping 的额外开销。用 ReshapedAttention 替代后这部分开销就省掉了。**Prefill 阶段吞吐大幅提升。**在 decode 侧引入稀疏 attention 策略——把 KV 分成四个区管理实现平稳的 HBM 消耗曲线。DeepSeek V4 等新模型面临 top-k attention 问题有多种窗口与选择机制TBO、SWA 滑动窗口、SFA 选择区、Recurrent虽然每次计算不需要全量 KV但全量 KV 仍需保存。分区控制后Decode 阶段的 batch size 提升了 4-8 倍。对于 DataInfer 等时延不敏感场景不关心单次时延关心系统吞吐架构翻转带来的收益更突出端到端性能提升 3 到 5 倍。P 侧与 D 侧的 KV 搬运策略在处理 Prefill 到 Decode 的 KV 转移时传统方案采用「先 P 后 D」策略——Prefill 完成后才决定 KV 分配给哪个 Decode 节点。这在以对话为主的时代是合理的因为无法提前预估 D 侧负载。但进入 AI 时代后Ken Zhang 的团队逐步转向「先 D 后 P」策略先确定 Decode 节点再做 Prefill 和新 KV 搬运。这样做的好处是实现轴层级的 KV 拉取调度更精准。同时在 Prefill 内部同一个 step 内可以预取下一层 KV。当前这个推理组件已经在管理数万张生成卡被 GM、千问、盘古等多个模型使用最近即将开源。总结Omni Cache 的核心思路并不复杂把昂贵的 HBM 从「主力存储」变成「智能缓存」让 DRAM 承担存储主力角色。通过架构翻转减轻 HBM 压力、连续寻址提升计算效率、分区管理优化 Decode 吞吐——这一套组合拳为大规模推理服务提供了实际的性能提升路径。以上内容由Ai好记转录整理。Ai好记是一款音视频转图文笔记的AI音视频总结工具支持解析B站、抖音、小红书小宇宙等平台链接及本地/网盘音视频文件转录后自动生成精华速览、思维导图和结构化笔记帮助你把几小时的视频内容变成可搜索、可复习的图文笔记。