To address the computational challenges posed by rapidly evolving intelligent applications, leading processor manufacturers are accelerating the integration of novel computing units into core processor architectures. Represented by Intel’s Advanced Matrix Extensions (AMX), the new generation of matrix computing units achieves a leap in computational performance compared with traditional vector units, such as Intel’s AVX-512, enabling single-chip peak performance to reach the order of hundreds of TFLOPS. These advanced computing units exhibit two notable characteristics: first, their efficient operation highly depends on specific data layouts; second, their demand for memory bandwidth is an order of magnitude higher than that of vector units. This architectural evolution renders traditional vector-unit-based optimization methods ineffective, highlighting the urgent need to explore new optimization strategies compatible with these units.
As a current research hotspot in the field of artificial intelligence, Mixture of Experts (MoE) has drawn significant attention from both academia and industry, making the acceleration of its inference a central focus. Matrix operations serve as the core computational components of MoE inference, with their performance directly determining overall inference efficiency. Specifically, the Grouped Query Attention (GQA) mechanism and the expert Feed-Forward Network (FFN) constitute the most critical performance bottlenecks in MoE model inference. Existing studies have conducted in-depth optimizations centered on these two components.
The choice to optimize matrix operations for CPUs equipped with AMX units is mainly based on the following three considerations. First, compared with GPUs, CPUs provide a more cost-effective solution for MoE inference. CPUs offer larger memory capacity, enabling the deployment of MoE models on a single node without relying on multiple GPUs. At the same time, CPU compute power and memory bandwidth are continuously improving, better meeting the performance requirements of MoE inference. Second, CPU memory bandwidth is significantly higher than the data transfer bandwidth between CPUs and GPUs. In inference scenarios with limited GPU resources, executing most matrix computations directly on the CPU often yields better performance. For example, the LIA framework demonstrates that when running large language model inference on a node equipped with only a single GPU, placing the primary computations on a CPU with AMX units significantly improves performance, especially when the input batch size is only a few hundred. Finally, compared with CPUs equipped only with vector units, AMX substantially enhances matrix operation performance. Moreover, although matrix optimization techniques for vector units are already well developed, optimization research for AMX still requires further exploration.
To fully exploit the performance of AMX, FlashMatrix was proposed, specifically designed to accelerate matrix operations in MoE. Its core consisted of two synergistic optimization strategies. First, it departed from the traditional oneDNN approach of performing layout transformations on weight matrices, instead applying layout transformations only to the significantly smaller input matrices in MoE, with these transformations executed within the cache and registers. This strategy greatly reduced the overhead of layout conversion while ensuring that weight matrices consistently maintained general layouts, such as row-major or column-major order. Second, a micro-kernel with the highest compute-to-memory-access ratio was designed to fully exploit AMX registers.
The experiments were conducted on Intel Gold 6430 and Intel Platinum 8468V, covering three core operators: single matrix multiplication, GQA, and expert FFN. The results show that FlashMatrix outperforms the state-of-the-art oneDNN library on all tested CPU platforms. In particular, for the performance of expert FFN, FlashMatrix achieves an average speedup of 2.5×. For end-to-end inference performance, FlashMatrix achieves a speedup of approximately 1.2×.
The two optimization strategies incorporated in FlashMatrix, namely the efficient layout transformation strategy and the micro-kernel with the highest compute-to-memory-access ratio, significantly enhance the performance of matrix operations in MoE. These optimizations provide efficient foundational operator support for deploying MoE models on CPUs.
京公网安备11010802044758号