FastTPS: An Optimized Method for LLM Token Phase for AI accelerators

📄 arXiv: 2607.11211v1 📥 PDF

作者: Wenzong Yang, Danyang Zhang, Kun Cao, Tejus Siddagangaiah, Rajeev Patwari, Zhanxing Pu, Siyin Kong, Zijiang Yang, Hao Zhu, Varun Sharma, Yue Gao, Tianping Li, Fan Yang, Jicheng Chen, Yushan Chen, Fennian Zhao, Aaron Ng, Elliott Delaye, Ashish Sirasao, Sudip Nag

分类: cs.LG

发布日期: 2026-07-13

备注: 16 pages, 8 figures, 7 tables

💡 一句话要点

提出FastTPS以解决LLM推理中的内存瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 推理优化 内存管理 AI加速器 高效计算

📋 核心要点

  1. 现有的LLM推理方法在处理长序列输入时，因令牌阶段的顺序处理导致低并行性，造成吞吐量下降和内存利用不佳。
  2. FastTPS通过无重载的KV缓存连接、高效的RoPE注意力机制和高度融合的MLP设计，旨在提升LLM推理的性能和内存效率。
  3. 实验结果显示，FastTPS在AMD Ryzen AI 300系列NPU上实现了6倍的速度提升，同时保持高达93%的内存带宽利用率。

📝 摘要（中文）

随着大型语言模型（LLMs）的普及，对有效推理的需求不断增加。然而，在解码器仅限的LLM推理中，令牌阶段的顺序处理导致了固有的低并行性，从而降低了吞吐量，并在处理长序列输入时造成显著的内存开销。本文提出了FastTPS，这是一种高性能且低精度损失的方法，旨在加速LLM推理中的令牌阶段。该方法包括三个关键组件：无重载的KV缓存连接、基于优化的FLAT的高效RoPE注意力机制，以及高度融合的MLP与细粒度流水线调度。实验结果表明，FastTPS显著缓解了令牌阶段的内存瓶颈，在AMD Ryzen AI 300系列NPU上实现了6倍的速度提升，同时在Phi3-mini-4k-instruct推理过程中保持93%的峰值内存带宽利用率。

🔬 方法详解

问题定义：本文旨在解决大型语言模型推理中令牌阶段的内存瓶颈和低吞吐量问题，现有方法在处理长序列时表现不佳，导致计算资源未得到充分利用。

核心思路：FastTPS的核心思路是通过优化内存访问和计算过程，减少令牌处理的延迟和内存开销，从而提升推理速度和效率。

技术框架：FastTPS的整体架构包括三个主要模块：无重载的KV缓存连接模块、基于FLAT的RoPE注意力模块和高度融合的MLP模块。每个模块都针对特定的性能瓶颈进行了优化。

关键创新：FastTPS的关键创新在于其无重载的KV缓存连接和高效的RoPE注意力机制，这些设计使得注意力计算能够完全融合，显著提高了内存带宽的利用率。

关键设计：在设计中，FastTPS采用了BF16精度以减少内存占用，同时通过细粒度的流水线调度优化了MLP的执行效率，确保在高负载下仍能保持良好的性能。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，FastTPS在AMD Ryzen AI 300系列NPU上实现了6倍的速度提升，相较于无融合方法，且在Phi3-mini-4k-instruct推理过程中保持93%的峰值内存带宽利用率，显示出显著的性能优势。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理、对话系统和其他需要高效推理的大型语言模型应用。通过提升推理速度和内存利用率，FastTPS能够为实时应用提供更好的支持，推动AI加速器的广泛应用与发展。

📄 摘要（原文）

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.