CoDec：LLM 前缀共享解码内核
机器学习
2026-03-31 v2
摘要
多个提示之间的前缀共享为组合共享前缀的操作提供了机会，而解码阶段的注意力计算——随着上下文长度的增加成为瓶颈——是一种需要在键值（KV）缓存上进行大量内存访问的内存密集型过程。因此，本文探索了在解码阶段利用前缀共享的注意力计算潜力。然而，前缀共享机制的树状结构为高效处理共享 KV 缓存访问模式、管理复杂依赖关系和平衡不规则工作负载带来了显著挑战。为解决上述挑战，我们提出一种专用的注意力内核，将解码阶段共享前缀的内存访问合并，即 CoDec。CoDec 提供了两项关键创新：一种新颖的共享前缀注意力内核，通过优化内存层次结构并利用块内和块间并行性来提高效率；以及一种综合的工作负载平衡机制，高效估算成本、划分任务并调度执行。实验结果表明，CoDec 在解码阶段的注意力计算方面相对于最先进的 FlashDecoding 内核实现平均加速 ，内存访问降低 ，相对于 vLLM 实现每输出一个 token 的端到端时间加速 。
引用
@article{arxiv.2505.17694,
title = {CoDec: Prefix-Shared Decoding Kernel for LLMs},
author = {Zhibin Wang and Rui Ning and Chao Fang and Zhonghui Zhang and Xi Lin and Shaobo Ma and Mo Zhou and Xue Li and Zhongfeng Wang and Chengying Huan and Rong Gu and Kun Yang and Guihai Chen and Sheng Zhong and Chen Tian},
journal= {arXiv preprint arXiv:2505.17694},
year = {2026}
}