国防科技大学学报2026，Vol.48Issue(3)：357-367,11.DOI:10.11887/j.issn.1001-2486.25090058
面向AMX单元的矩阵算子优化方法
Matrix operator optimization method for AMX unit
摘要
Abstract
In the inference process of mixture of experts models,matrix operators constitute the primary performance bottleneck,with those in the attention module and expert computation being particularly time-consuming.Although existing approaches have extensively optimized matrix operators on GPUs,the substantial differences between GPU and CPU architectures in memory hierarchy and compute units make these optimizations difficult to transfer directly to CPU platforms.To address this limitation,FlashMatrix was introduced as a matrix-operator optimization scheme tailored for CPU equipped with advanced matrix extensions.FlashMatrix incorporates an efficient data layout transformation strategy that avoids additional memory-access overhead caused by layout conversions,and employs a carefully designed micro-kernel for matrix multiplication that achieves an optimal compute-to-memory ratio through effective register reuse.Experimental results show that,compared with the state-of-the-art CPU matrix-computation library oneDNN,FlashMatrix delivers an average 2.5 × speedup.For end-to-end inference performance,FlashMatrix achieves a speedup of approximately 1.2 ×.关键词
AMX单元/矩阵算子/数据布局/性能优化Key words
AMX unit/matrix operators/data layout/performance optimization分类
信息技术与安全科学引用本文复制引用
杨维铃,方建滨,董德尊..面向AMX单元的矩阵算子优化方法[J].国防科技大学学报,2026,48(3):357-367,11.基金项目
国家自然科学基金委员会联合基金资助项目(U24B20151) （U24B20151）
