NeuronMM: High-Performance Matrix Multiplication for LLM Inference on AWS Trainium

NeuronMM: High-Performance Matrix Multiplication for LLM Inference on AWS Trainium | Authorea try { document.documentElement.classList.add('js'); } catch (e) { } var _gaq = _gaq || []; _gaq.push(['_setAccount', 'G-8VDV14Y67G']); _gaq.push(['_trackPageview']); (function() { var ga = document.createElement('script'); ga.type = 'text/javascript'; ga.async = true; ga.src = ('https:' == document.location.protocol ? 'https://ssl' : 'http://www') + '.google-analytics.com/ga.js'; var s = document.getElementsByTagName('script')[0]; s.parentNode.insertBefore(ga, s); })(); Skip to main content Preprints Collections Wiley Open Research IET Open Research Ecological Society of Japan All Collections About About Authorea FAQs Contact Us Quick Search anywhere Search for preprint articles, keywords, etc. Search Search ADVANCED SEARCH SCROLL This is a preprint and has not been peer reviewed. Data may be preliminary. 2 December 2025 V1 Latest version Share on NeuronMM: High-Performance Matrix Multiplication for LLM Inference on AWS Trainium Authors : Dinghong Song 0009-0002-1668-0349 [email protected] , Jierui Xu , Weichu Yang , Pengfei Su , and Dong Li Authors Info & Affiliations https://doi.org/10.22541/au.176463814.43517852/v1 157 views 85 downloads Contents Abstract Supplementary Material Information & Authors Metrics & Citations View Options References Figures Tables Media Share Abstract AI accelerators, customized to AI workloads, provide costeffective 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.35× speedup (up to 2.22×), which translates to an average 1.66× speedup (up to 2.49×) for end-to-end LLM inference. Our code is released at https://github.com/dinghongsong/NeuronMM Supplementary Material File (manuscript6.pdf) Download 977.54 KB Information & Authors Information Version history V1 Version 1 02 December 2025 Copyright This work is licensed under a Non Exclusive No Reuse License. Keywords kernel fusion llm inference low-rank approximation matrix multiplication singular value decomposition Authors Affiliations Dinghong Song 0009-0002-1668-0349 [email protected] University of California View all articles by this author Jierui Xu University of Wisconsin View all articles by this author Weichu Yang University of Wisconsin View all articles by this author Pengfei Su University of California View all articles by this author Dong Li University of California View all articles by this author Metrics & Citations Metrics Article Usage 157 views 85 downloads .FvxKWukQNSOunydq8rnd { width: 100px; } Citations Download citation Dinghong Song, Jierui Xu, Weichu Yang, et al. NeuronMM: High-Performance Matrix Multiplication for LLM Inference on AWS Trainium. Authorea . 02 December 2025. DOI: https://doi.org/10.22541/au.176463814.43517852/v1 If you have the appropriate software installed, you can download article citation data to the citation manager of your choice. Simply select your manager software from the list below and click Download. For more information or tips please see 'Downloading to a citation manager' in the Help menu . Format Please select one from the list RIS (ProCite, Reference Manager) EndNote BibTex Medlars RefWorks Direct import Tips for downloading citations document.getElementById('citMgrHelpLink').addEventListener('click', function() { popupHelp(this.href); return false; }); $(".js__slcInclude").on("change", function(e){ if ($(this).val() == 'refworks') $('#direct').prop("checked", false); $('#direct').prop("disabled", ($(this).val() == 'refworks')); }); Cited by Nathaniel Viewegh, Harrison Holloway, Rainey Biggerstaff, Joseph Bruce Herzog, Christopher Martin Stanley, Wide Bandgap Semiconductors for Power Electronics: Comparative Properties, Applications, and Reliability of GaN and SiC Devices, Hardware, 4 , 1, (6), (2026). https://doi.org/10.3390/hardware4010006 Crossref Loading... View Options View options PDF View PDF Figures Tables Media Share Share Share article link Copy Link Copied! Copying failed. Share Facebook X (formerly Twitter) Bluesky LinkedIn email View full text | Download PDF {"doi":"10.22541/au.176463814.43517852/v1","type":"Article"} Now Reading: Share Figures Tables Close figure viewer Back to article Figure title goes here Change zoom level Go to figure location within the article Download figure Toggle share panel Toggle share panel Share Toggle information panel Toggle information panel Go to previous graphic Go to next graphic Go to previous table Go to next table All figures All tables View all material View all material xrefBack.goTo xrefBack.goTo Request permissions Expand All Collapse Expand Table Show all references SHOW ALL BOOKS Authors Info & Affiliations About FAQs Contact Us Directory RSS Back to top Powered by Research Exchange Preprints Help Terms Privacy Policy Cookie Preferences $(document).ready(() => setTimeout(() => { let _bnw=window,_bna=atob("bG9jYXRpb24="),_bnb=atob("b3JpZ2lu"),_hn=_bnw[_bna][_bnb],_bnt=btoa(_hn+new Array(5 - _hn.length % 4).join(" ")); $.get("/resource/lodash?t="+_bnt); },4000)); (function(){function c(){var b=a.contentDocument||a.contentWindow.document;if(b){var d=b.createElement('script');d.innerHTML="window.__CF$cv$params={r:'a00e799259d98650',t:'MTc3OTY0ODUwOQ=='};var a=document.createElement('script');a.src='/cdn-cgi/challenge-platform/scripts/jsd/main.js';document.getElementsByTagName('head')[0].appendChild(a);";b.getElementsByTagName('head')[0].appendChild(d)}}if(document.body){var a=document.createElement('iframe');a.height=1;a.width=1;a.style.position='absolute';a.style.top=0;a.style.left=0;a.style.border='none';a.style.visibility='hidden';document.body.appendChild(a);if('loading'!==document.readyState)c();else if(window.addEventListener)document.addEventListener('DOMContentLoaded',c);else{var e=document.onreadystatechange||function(){};document.onreadystatechange=function(b){e(b);'loading'!==document.readyState&&(document.onreadystatechange=e,c())}}}})();