arXiv:2510.25977cs.CL2025-10

通过奇异值分解压缩与分块优化,提升AWS Trainium上大模型推理效率。

NeuronMLP: Efficient LLM Inference via Singular Value Decomposition Compression and Tiling on AWS Trainium

  • 基于SVD压缩与分块技术,针对Trainium架构定制优化
  • 在压缩比0.05下实现内核级1.35倍加速,端到端推理快1.21倍
  • 适合关注AI芯片高效推理的工程师和研究者

新兴AI加速器正成为大语言模型(LLM)高效推理的新机遇。AWS最新推出的Trainium加速器凭借异构架构提供了有吸引力的推理选项,但其流水线阵列结构和特殊数据布局要求使得高性能利用面临挑战。本文提出NeuronMLP,一种基于奇异值分解(SVD)压缩与分块的高效LLM推理方法,专为Trainium设计。通过内核融合与新型缓存策略,减少软件管理内存层次中的数据移动,最大化SRAM带宽,避免昂贵的矩阵转置操作。该方法重点优化了LLM中关键的多层感知机(MLP)层。在九个数据集和六个近期LLM上的评估显示,相较于AWS官方基于Neuron Kernel Interface(NKI)的矩阵乘法内核,NeuronMLP在压缩比0.05下,内核级平均提速1.35倍,端到端推理平均提速1.21倍。

原文摘要 · Abstract (English)

Emerging AI accelerators have started to gain attention and offer new opportunities for efficient inference of large language models (LLMs). Trainium, an AI accelerator recently developed by Amazon Web Services (AWS), provides an attractive option for LLM 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 propose NeuronMLP, an efficient LLM inference method based on Singular Value Decomposition (SVD) compression and tiling on AWS 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. The proposed method is specifically optimized for multi-layer perceptron (MLP) layers in LLMs, which serve as a critical computational kernel for inference on Trainium. Evaluating on nine datasets and six recent LLMs, we show that NeuronMLP significantly outperforms the state-of-the-art Neuron Kernel Interface (NKI)-based matrix multiplication (matmul) kernel implemented by AWS on Trainium: at the kernel level, it achieves an average 1.35x speedup, which translates to an average 1.21x speedup for end-to-end LLM inference, under a compression ratio of 0.05.

大模型推理SVD压缩AWS Trainium高效计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。