FastTPS加速LLM推理中的令牌阶段，提升AI加速器利用率。
FastTPS: An Optimized Method for LLM Token Phase for AI accelerators
- 通过无重载的KV缓存拼接减少内存开销，融合注意力计算
- 采用分块优化的RoPE注意力与细粒度流水线，实现6倍提速
- 适合长序列推理场景，尤其适用于NPU等AI加速器
大语言模型（LLMs）的普及带来了对高效推理的持续需求。然而，解码器仅有的令牌阶段存在固有的低并行性，导致在处理长序列输入时，内存开销大，计算单元利用率低。为此，本文提出FastTPS，一种针对通用AI加速器的高性能、低精度损失的令牌阶段加速方法，包含三个核心组件：(1) 基于AI加速器的无需重载的KV缓存拼接，降低内存访问开销并实现注意力完全融合；(2) 基于分块优化的FLAT的高效高精度RoPE注意力；(3) 细粒度流水调度的高融合MLP。实验结果表明，FastTPS显著缓解了令牌阶段的内存瓶颈，在AMD Ryzen AI 300系列NPU上以BF16精度实现相比无融合方案6倍速度提升，同时在Phi3-mini-4k-instruct推理中保持93%峰值内存带宽利用率。
原文摘要 · Abstract (English)
The popularity of large language models (LLMs) escalates an ongoing demand for effective inference. However, due to the sequential processing of tokens during the token phase in decoder-only LLMs inference, the inherent low parallelism leads to reduced throughput and suboptimal utilization of the computing units on artificial intelligence (AI) accelerators, particularly when handling long-sequence inputs that impose significant memory overhead. Recently, many reported methods have been developed as potential solutions, since they emerge with numeric deviation. This paper presents FastTPS, a high performance and low-precision loss method for accelerating the token-phase in LLM inference on general AI accelerators which includes three key components: (1) AI accelerator-enabled reloading-free KV Cache concatenation which decreases memory access overhead as well as enables full fusion of Attention, (2) high-efficiency and high-accuracy 'RoPE' attention based on the tiling optimized FLAT, and (3) highly-fused MLP with fine-grain pipeline scheduling. Our results confirm that FastTPS significantly alleviates memory bottlenecks in the token phase, delivering a 6x speed improvement (compared to none-fusion) on an AMD Ryzen AI 300 series NPU with BF16 precision while sustaining 93% peak memory bandwidth utilization during Phi3-mini-4k-instruct inference.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。