TLDR: This project introduces RISC-V microkernel support into IREE, an MLIR-based machine learning compiler, to accelerate GenAI workloads. By enabling `linalg.mmt4d` lowering for RISC-V64 and developing optimized `mmt4d` microkernels for prefill and decode phases, the team achieved significant performance gains (up to 50x in decode for single-threaded) for the Llama-3.2-1B-Instruct model compared to upstream IREE, addressing a critical performance gap for RISC-V in AI hardware.
In the rapidly evolving landscape of artificial intelligence, the efficient execution of Generative AI (GenAI) models is paramount. A recent project by 10xEngineers, detailed in their paper Accelerating GenAI Workloads by Enabling RISC-V Microkernel Support in IREE, addresses a significant challenge in this area: optimizing GenAI performance on RISC-V hardware.
IREE, which stands for Intermediate Representation Execution Environment, is a powerful compiler and runtime built on MLIR. Its primary function is to take machine learning workloads, apply various optimizations like operator fusion and tiling, and then generate highly optimized binaries for a wide array of platforms, including CPUs, GPUs, and specialized accelerators. While IREE generally performs well, custom-optimized kernels are often superior for mixed-precision calculations, which are common in GenAI.
Currently, IREE includes a library of these optimized microkernels for popular architectures like x86 and ARM CPUs. However, despite the growing presence and importance of RISC-V in the AI hardware sector, dedicated RISC-V microkernels have been conspicuously absent. This gap has led to suboptimal performance of GenAI models when run on RISC-V-based systems. The core objective of this project was to bridge this gap by integrating RISC-V microkernel support directly into IREE.
Understanding the Core Mechanism
At the heart of many GenAI workloads is matrix multiplication, a computationally intensive operation. IREE handles this using MLIR’s linalg dialect contraction operations. To achieve optimal performance, especially with tiled matrix multiplication, data needs to be arranged efficiently to minimize cache misses. IREE employs a clever strategy involving `tensor.pack` to rearrange 2-D matrices into a tiled 4-D format where data is stored contiguously. This packed data is then processed by `linalg.mmt4d` for optimized matrix multiplication. Finally, `tensor.unpack` converts the 4-D result back to its original 2-D layout.
While IREE’s architecture is designed to be extensible for integrating custom microkernels, the existing library only provided these specialized kernels for x86 and ARM64 architectures. This project specifically focused on implementing the crucial `mmt4d` microkernels for the RISC-V target, enabling the same level of optimization for this emerging hardware.
The Implementation Approach
The team at 10xEngineers adopted a two-pronged methodology to achieve their goal. The first step involved modifying IREE’s compilation pipeline to correctly transform linalg contraction operations into `tensor.pack`, `tensor.unpack`, and `linalg.mmt4d` operations specifically for the RISC-V64 target. This required adjusting the `iree-codegen-materialize-device-encoding` pass, which is responsible for determining optimal tile sizes for matrix dimensions based on the target architecture. A key innovation here was the introduction of VLEN-aware tiling for RISC-V64, with specific tile sizes chosen for the ‘prefill’ and ‘decode’ phases of LLM inference (6, VLEN/8, 1 for prefill and 1, VLEN/4, 1 for decode). These sizes were carefully selected to maximize hardware register utilization without causing performance degradation due to register spills.
The second, equally vital step was the actual implementation of the microkernel functions. The `mmt4d` microkernels were developed for the `f16xf16->f32` mixed-precision case, meaning inputs are 16-bit floating-point numbers, and the output is a 32-bit floating-point number. Recognizing the distinct computational patterns of LLM phases, separate `mmt4d` microkernels were crafted for the ‘prefill’ phase (which typically involves General Matrix-Matrix Multiplication, or GEMM) and the ‘decode’ phase (which often involves General Matrix-Vector Multiplication, or GEMV).
Also Read:
- Streamlining Chiplet Design with Multi-Agent AI Guidance
- Making AI Smart and Small: Quantized Neural Networks for Microcontrollers
Impressive Performance Gains
To validate their work, the researchers rigorously tested the newly implemented microkernels. They evaluated the Llama-3.2-1B-Instruct model, compiled with their enhanced 10x-IREE framework, using the LM-Evaluation-Harness. The accuracy results were identical to the Huggingface version of the model, confirming the correctness of their implementation.
The performance benchmarking, conducted on a MILK-V Jupiter board equipped with 1.66GHz x 8 RISC-V vector cores, revealed substantial improvements. For single-threaded operations, the decode phase showed an astounding 50x performance gain compared to the upstream IREE. In multi-threaded scenarios, the prefill phase saw a 2x performance increase, while the decode phase experienced a remarkable 17x boost. These figures, reported in tokens per second, clearly demonstrate the significant acceleration achieved by integrating RISC-V microkernel support into IREE, making GenAI workloads much more efficient on RISC-V hardware.
