FASTTPS: AN OPTIMIZED METHOD FOR LLM TOKEN PHASE FOR AI ACCELERATORS

Overview

FastTPS (Fast Tokens Per Second) is a dataflow optimization strategy designed to accelerate the token generation phase of decoder-only Large Language Models (LLMs). During this phase, models generate tokens sequentially, which often leads to underutilization of hardware resources due to memory bottlenecks and limited computational parallelism. FastTPS addresses these inefficiencies by optimizing how data is managed and processed on general-purpose AI accelerators.

Key Optimization Strategies

The paper introduces three primary components to enhance inference throughput:

• Global KV Cache Management (GKVC): This strategy replaces the traditional, memory-intensive 'Concat' operator with a static 3D memory management approach. By pre-allocating memory and updating key-value pairs in place, it eliminates the need for frequent memory reorganization and data transfers, ensuring continuous execution on the accelerator.
• TPSFLAT (Fused Logit Attention Tiling): Building on existing tiling concepts, this method optimizes the Rotary Positional Embedding (RoPE) and attention mechanism. It leverages the continuity provided by GKVC to fuse multiple operators, thereby increasing the operational intensity and reducing the frequency of memory access.
• Fusion MLP: This component optimizes the Multi-Layer Perceptron (MLP) block by rearranging weight layouts and implementing a new multi-op fusion schedule. By using fine-grained matrix division, it reduces I/O overhead and improves the utilization of the accelerator's computational units.

Performance and Impact

Experimental results on the AMD Ryzen AI 300 series NPU demonstrate that FastTPS significantly outperforms standard implementations. The method achieved a 6x speedup during Phi3-mini-4k-instruct inference while maintaining 93% peak memory bandwidth utilization. Furthermore, the authors report that FastTPS preserves high computational precision, achieving up to 10x higher accuracy compared to standard FlashAttention implementations, making it suitable for high-precision deployment scenarios.