Saved in:
| Main Authors: | , , , , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | https://arxiv.org/abs/2505.17694 |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866915895796826112 |
|---|---|
| author | Wang, Zhibin Ning, Rui Fang, Chao Zhang, Zhonghui Lin, Xi Ma, Shaobo Zhou, Mo Li, Xue Wang, Zhongfeng Huan, Chengying Gu, Rong Yang, Kun Chen, Guihai Zhong, Sheng Tian, Chen |
| author_facet | Wang, Zhibin Ning, Rui Fang, Chao Zhang, Zhonghui Lin, Xi Ma, Shaobo Zhou, Mo Li, Xue Wang, Zhongfeng Huan, Chengying Gu, Rong Yang, Kun Chen, Guihai Zhong, Sheng Tian, Chen |
| contents | 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. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2505_17694 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | CoDec: Prefix-Shared Decoding Kernel for LLMs Wang, Zhibin Ning, Rui Fang, Chao Zhang, Zhonghui Lin, Xi Ma, Shaobo Zhou, Mo Li, Xue Wang, Zhongfeng Huan, Chengying Gu, Rong Yang, Kun Chen, Guihai Zhong, Sheng Tian, Chen Machine Learning 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. |
| title | CoDec: Prefix-Shared Decoding Kernel for LLMs |
| topic | Machine Learning |
| url | https://arxiv.org/abs/2505.17694 |