arXiv:2505.17694cs.LG2025-05被引 8

提出CoDec,让大模型解码时共享前缀,显著提速并省内存。

CoDec: Prefix-Shared Decoding Kernel for LLMs

  • 设计专用注意力内核,合并共享前缀的内存访问
  • 相比SOTA快1.9倍,内存访问减少120.9倍
  • 适合长上下文、多请求并发的推理场景

多个提示共享前缀为融合操作提供了机会,但解码阶段的注意力计算因上下文增长成为瓶颈,且需频繁访问前缀的键值(KV)缓存,带来高内存开销。本文探索在解码阶段利用前缀共享潜力。然而,前缀共享的树形结构给注意力计算带来挑战:如何高效处理共享KV缓存访问模式,同时管理复杂依赖并平衡不规则工作负载。为此,我们提出专用注意力内核CoDec。其核心创新包括:一种新型共享前缀注意力内核,优化内存层次结构并利用块内与块间并行性;以及一套完整的负载均衡机制,可高效估算代价、划分任务并调度执行。实验表明,相较于当前最优的FlashDecoding内核,CoDec在解码阶段注意力计算上实现平均1.9倍加速和120.9倍内存访问减少;相比vLLM,端到端每输出一个词元时间降低3.8倍。

原文摘要 · Abstract (English)

Prefix-sharing among multiple prompts presents opportunities to combine the operations of the shared prefix, while attention computation in the decode stage, which becomes a critical bottleneck with increasing context lengths, is a memory-intensive process requiring heavy memory access on the key-value (KV) cache of the prefixes. Therefore, in this paper, we explore the potential of prefix-sharing in the attention computation of the decode stage. However, the tree structure of the prefix-sharing mechanism presents significant challenges for attention computation in efficiently processing shared KV cache access patterns while managing complex dependencies and balancing irregular workloads. To address the above challenges, we propose a dedicated attention kernel to combine the memory access of shared prefixes in the decoding stage, namely CoDec. CoDec delivers two key innovations: a novel shared-prefix attention kernel that optimizes memory hierarchy and exploits both intra-block and inter-block parallelism, and a comprehensive workload balancing mechanism that efficiently estimates cost, divides tasks, and schedules execution. Experimental results show that CoDec achieves an average $1.9\times$ speedup and $120.9\times$ memory access reduction compared to the state-of-the-art FlashDecoding kernel regarding attention computation in the decode stage and $3.8\times$ end-to-end time per output token compared to the vLLM.

大模型推理注意力优化内存效率

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。