NeuronMM: High-Performance Matrix Multiplication for LLM Inference on AWS Trainium
This work addresses performance bottlenecks in LLM inference on a specific AI accelerator, offering incremental improvements over existing implementations.
The paper tackled the challenge of achieving high-performance matrix multiplication for LLM inference on AWS Trainium accelerators by introducing kernel fusion and novel caching strategies, resulting in an average 1.35x speedup for matmul kernels and 1.66x for end-to-end inference.
AI accelerators, customized to AI workloads, provide cost-effective and high-performance solutions for training and inference. Trainium, an AI accelerator recently developed by Amazon Web Services (AWS), provides an attractive option for LLM training and inference through its heterogeneous architecture. However, leveraging Trainium architecture for high performance can be challenging because of its systolic array architecture and special requirement on data layout. In this paper, we design high-performance matrix multiplication (matmul), a critical compute kernel, for LLM inference on Trainium. We introduce a series of techniques customized to Trainium based on kernel fusion and novel caching strategies to reduce data movement across the software-managed memory hierarchy, maximize SRAM bandwidth, and avoid expensive matrix transpose. Evaluating with nine datasets and four recent LLMs, we show that our system largely outperforms the state-of-the-art matmul implemented by AWS on Trainium: at the level of matmul kernel, it achieves an average 1.35x speedup (up to 2.22x), which translates to an average 1.66x speedup (up to 2.49x) for end-to-end LLM inference.