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

arXiv:2510.25977v4 Announce Type: replace Abstract: 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.

Related

• YOCO++: Enhancing YOCO with KV Residual Connections for Efficient LLM Inference
• Adaptive Layer Selection for Layer-Wise Token Pruning in LLM Inference
• Dynamic Model Routing and Cascading for Efficient LLM Inference: A Survey
• MultiTok: Variable-Length Tokenization for Efficient LLMs Adapted from LZW Compression
• Lil: Less is Less When Applying Post-Training Sparse-Attention Algorithms in Long-Decode Stage
• Latent-Condensed Transformer for Efficient Long Context Modeling

Source: arXiv cs.CL | 2026-04-27