大模型的推理加速综述
大模型的推理加速综述
大模型的推理加速是当前AI领域研究的热点问题。本文从技术角度出发，详细介绍了大模型推理过程中的三个关键技术点：Prefill & Decode、KV Cache以及并行推理。这些技术对于提高大模型的推理效率和降低计算成本具有重要意义。
大模型的推理加速综述
在CCF Chip会议上，作者听到了许多关于大模型的研究报告，其中一些有趣的技术点值得关注。本文将重点阐述大模型推理加速时的细节问题，主要集中在Prefill & Decode、KV Cache以及并行推理这三个关键技术上。
Prefill & Decode
研究者们发现，大模型的推理过程可以分为两个阶段：Prompt phase和Token generation phase。前者需要处理用户输入的Prompt，后者则是逐个生成输出Token。这两个阶段在访存和计算特性上存在显著差异：
- Prompt phase是典型的计算受限场景，计算瓶颈占据主导。
- Token generation phase则是典型的访存受限场景，访存瓶颈占据主导。
这种区分有助于通过合理分配计算与访存资源来实现高效的大模型推理。
从访存角度来说，Token phase的访存需求量远大于Prompt phase。在Batch Tokens接近(10^2)时，Token phase的访存量达到了接近Prompt phase的两倍。
KV Cache
KV Cache是大模型推理加速中的重要技术，其原理基于注意力机制的公式：
[Attention(Q,K,V)=softmax(\frac{QK^T}{\sqrt{d_k}})V ]
在Decode阶段，可以复用已经生成的Key和Value向量，从而减少计算开销。这种缓存技术可以显著优化大模型的推理效率。
并行推理
并行推理主要分为三种方法：数据并行、模型并行和流水线并行。
- 数据并行：将大数据集切分为小batch，分散到多个节点上独立运行模型，再汇总结果。
- 模型并行：将大型模型层或模块分散到多个节点运行，以减少单个节点的容量需求。
- 流水线并行：将模型的不同层次分散到不同节点，以流水线形式进行并行工作。
在实际应用中，这三种方法往往需要灵活组合使用。
总结
大模型的推理加速技术对存算架构设计有重要指导意义。例如，Prefill&Decode两阶段的划分要求架构具备计算与访存的资源分配能力；KV Cache需要专用的芯片结构支持；并行推理则要求芯片在数据流分配与控制上具有灵活性。
参考文献
- Patel, P., Choukse, E., Zhang, C., Shah, A., Goiri, Í., Maleki, S., & Bianchini, R. (2024). Splitwise: Efficient generative LLM inference using phase splitting.https://arxiv.org/abs/2311.18677
- Yuan, J., Liu, H., Shaochen, Zhong, Chuang, Y.-N., Li, S., Wang, G., Le, D., Jin, H., Chaudhary, V., Xu, Z., Liu, Z., & Hu, X. (2024). KV Cache Compression, But What Must We Give in Return? A Comprehensive Benchmark of Long Context Capable Approaches.https://arxiv.org/abs/2407.01527
- Li, S., Zhao, Y., Varma, R., Salpekar, O., Noordhuis, P., Li, T., Paszke, A., Smith, J., Vaughan, B., Damania, P., & Chintala, S. (2020). PyTorch Distributed: Experiences on Accelerating Data Parallel Training.https://arxiv.org/abs/2006.15704
- Shoeybi, M., Patwary, M., Puri, R., LeGresley, P., Casper, J., & Catanzaro, B. (2020). Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism.https://arxiv.org/abs/1909.08053
- Huang, Y., Cheng, Y., Bapna, A., Firat, O., Chen, M. X., Chen, D., Lee, H., Ngiam, J., Le, Q. V., Wu, Y., & Chen, Z. (2019). GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism.https://arxiv.org/abs/1811.06965