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

2026-07-13

AI Summary

This work addresses the inefficiencies in large language model (LLM) inference caused by sequential token processing, which leads to low parallelism, underutilized compute units, and high memory overhead—particularly pronounced in long-sequence scenarios. To overcome these limitations, the authors propose three key techniques: a reload-free KV cache concatenation enabling fully fused attention, a block-wise optimized FLAT-based high-precision RoPE attention mechanism, and a highly fused MLP architecture with fine-grained pipeline scheduling. These innovations synergistically enhance computational efficiency and memory bandwidth utilization on AI accelerators. Evaluated on the AMD Ryzen AI 300 series NPU, the proposed approach achieves a 6× speedup over non-fused baselines while sustaining 93% of peak memory bandwidth utilization during inference with Phi-3-mini-4k-instruct.

Abstract

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.

Problem

• LLM inference
• token phase
• low parallelism
• memory overhead
• AI accelerators

Innovation

• KV Cache fusion
• RoPE attention
• MLP pipeline scheduling
• token-phase acceleration
• memory bandwidth optimization