NeuronMLP：基于奇异值分解压缩和 AWS Trainium 上的瓷砖实现的高效 LLM 推理
计算与语言
2026-04-27 v4
摘要
新兴的 AI 加速器正在开始获得关注并为高效的大型语言模型（LLM）推理提供新机会。Trainium 是亚马逊云服务（AWS）最近开发的 AI 加速器，为 LLM 推理提供了有吸引力的选项。然而，由于其 systolic array 架构和数据布局的特殊要求，利用 Trainium 架构实现高性能可能具有挑战性。在本文中，我们提出了 NeuronMLP，这是一种基于奇异值分解（SVD）压缩和 AWS Trainium 上的瓷砖的高效 LLM 推理方法。我们引入了一系列针对 Trainium 定制的技术，基于 kernel 融合和新颖的缓存策略以减少软件受控内存层次结构中的数据移动，最大化 SRAM 带宽，并避免昂贵的矩阵转置。该方法特别针对 LLM 中的多层感知器（MLP）层进行了优化，这些层是 Trainium 上推理的关键计算内核。评估在九个数据集和六个最新 LLM 上进行，我们表明 NeuronMLP 在 AWS 在 Trainium 上实现的基于 NKI 的矩阵乘法（matmul）内核的最先进方法方面显著优于：在内核级别，平均加快 1.35 倍，这转化为端到端 LLM 推理平均加快 1.21 倍，在压缩比为 0.05 的情况下。
引用
@article{arxiv.2510.25977,
title = {NeuronMLP: Efficient LLM Inference via Singular Value Decomposition Compression and Tiling on AWS Trainium},
author = {Dinghong Song and Jierui Xu and Weichu Yang and Pengfei Su and Dong Li},
journal= {arXiv preprint arXiv:2510.25977},
year = {2026}
}
备注
12 pages, 8 figures