AX650N - Sipeed Maix-IV (AXeraPi-Pro) NPU teardown

After spending a significant amount of time reverse-engineering the RK3588 NPU and examining Rockchip's 6 TOPS claim. The AXera AX650N SoC piqued my interest due to AXera's ambitious claims about the NPU's processing power, boasting "72T mixed precision computing power, native support for Transformer intelligent processing platform". Upon closer inspection, the AX650N delivers 72.0 TOPS@INT4 and 18.0 TOPS@INT8. Interestingly, the performance claim from INT8 to INT4 is a 4x gain, rather than the typical 2x improvement. There is an ongoing effort to port some of the smaller Transformer model to showcase its capabilities. However, given the performance claims I would have expected some larger models to be showcased, but that doesn't seem to be the case.

AX650N

The AX650N is part of AXera's vision processor lineup, featuring a unique capability for intelligent black light vision, which appears to generate color images in low-light conditions. From what I understand, this capability likely stems from its ISP (Image Signal Processor) unit utilizing RGB+IR data from the image sensor. This opens up potential for performing object detection in low-light environments. The AX650N is powered by 8 A55 cores and features an unusual configuration with two DDR controllers, each capable of addressing 16GB. This design allows for the potential to double data transfer rates, as the controllers operate independently and don't share the same address bus. Additionally, the SoC includes dual DSP cores (Tensilica Vision Q7 DSP) to further boost its vision processing capabilities. However, since it's primarily a vision processor, I suspect AXera may be positioning it to also capitalize on the growing interest in large language models (LLMs).

Maxi-IV board

AXera collaborated with Sipeed to produce a developer board, the Sipeed Maxi-IV, also known as the AXeraPi-Pro. The board is set to boot from eMMC by default and comes with a preinstalled Busybox image, however, the image is quite minimal. Fortunately there is a Ubuntu image is available for flashing to the eMMC. The main issue with the Ubuntu image is that it doesn't allocate enough space for the different mount points for the root filesystem for it to be useful. To fix this, I had to create a secondary image on a USB drive, boot from it (painfully slow) and then repartition the eMMC. The CPU fan runs continuously as soon as power is applied (with no PWM control), making it a noticeable distraction.

 

Upon booting the board, the surprising discovery is that Linux recognizes just 4GB of memory, despite the board having 8GB of onboard RAM. The remaining 4GB is reserved as CMM (Contiguous Memory Model), a large block of physical memory allocated for onboard peripherals like the ISP, video encoder/decoder, and NPU. It’s important to note that CMM differs from CMA, which is typically reserved by the Linux kernel. The ratio of CMM to Linux memory can be adjusted. The main challenge with 4GB (or less) available for Linux is the difficulty in fully utilizing all 8 CPU cores.

 

SDK support

To begin development the AX650N SDK is essential, but accessing it can be challenging if you're outside of China. For reasons unknown, Sipeed only hosts the SDK on their Baidu account. I reached out to Sipeed support with no response, email was sent on July 17, 2024 and still no response. I encountered the same lack of response from the Sipeed Telegram channel. Contacting AXera doesn't help they point you back to Sipeed. After weeks of trying I finally manage to get hold of the SDK without Sipeed.

The SDK includes a small kernel patch that gets applied against kernel 5.15.73, however the main irritation is that most peripheral drivers are provided as pre-compiled binaries, making the kernel effectively closed source. The rest of the SDK provides a set of closed source user space libraries exposing APIs to the different IP blocks. The primary drawback with a closed source approach is addressing security concerns as highlighted on Sipeeds NanoKVM. SDK documentation is primarily in Chinese with occasional English versions.

 

Transformers Support

As mentioned earlier, my main focus is the NPU architecture and performance of Transformer models. There are number of pre-built transformer models (Qwen1.5-1.8B, Qwen2-0.5B, MiniCPM-1B/2B, Phi-3-mini) available here, unfortunately again hosted on Baidu. The largest one is Phi-3 mini which requires increasing the CCM space to 5 or 6GB, due to other kernel modules reserving CCM space while the kernel boots. Executing Phi-3 mini results in a token generation rate of approximately 4.4 tokens per second.

The token rate appears to be lower than expected given the AX650N's TOPS claims. Phi-3 Mini seems to be using INT8 quantization, although it's unclear which specific approach, I'm guessing it's w8a8 with SmoothQuant. In comparison, Phi-3 Mini on the RK3588, also using INT8 (w8a8), reportedly runs at up to 6.46 tokens per second. The AX650N claims 18.0TOPS@int8, while the RK3588 is rated at 6 TOPS@int8. Although the performance may be lower, it's important to also consider the quality of the quantized models by evaluating each model's perplexity on both platforms, which I haven't done. Furthermore the vendor for both SOCs aren't particularly forthcoming on perplexity evaluations when releasing models. Regardless of perplexity, the expectation was that the additional 12 TOPS would lead to better performance. The next step is to take a closer look at the NPU architecture to gain a clearer understanding of what might be causing this discrepancy. Given the platform's closed-source nature, this presents several challenges.

 

To deploy models to the NPU, AXera rely on their Pulsar toolchain, which converts ONNX files into their proprietary axmodel format. At a high level. here's my understanding of how the axmodel files are processed:

 

The axmodel files contain a mix of ONNX data and an internal graph representation of the model, which is sent to the NPU kernel driver. To execute an LLM model, each layer is stored in its own axmodel file, containing descriptions of the input/output parameters and corresponding weights. The input/output parameters and weights are loaded into CCM. To run the layer, the kernel driver receives references to the input/output parameters, weights, and a list of NPU commands, all of which are included in the file.

Note, there is no api to directly interface with NPU, the only mechanism is by generating a axmodel file. For example to instruct the NPU to multiple 2 int8 matrices the SDK contains a bunch of sample which surprisingly rely on loading an axmodel file.

To efficiently run quantized models on the NPU, it ideally needs the ability to execute a Gemm operation. What's notable from the Pulsar documentation is that the ONNX Gemm operation is restricted as outlined below, suggesting that de-quantization cannot be performed while applying FMA.

alpha: Not supported yet

beta: Not supported yet

 

Power Consumption

After the board completes booting, the idle current consumption is about 450mA, equating to 5.4 watts (12V x 450mA). While running Phi-3 Mini, the consumption averages 680mA or 8.6 watts. For running a small LLM, the AX650N NPU is highly efficient, with the NPU consuming just 2.71 watts and benefiting from direct access to CPU memory. With a well-designed board, these figures could be significantly reduced.

Neutron NPU

The AXera NPU architecture, known as Neutron, has different iterations across the AX6XXX chipsets. It is primarily designed to support vision applications by executing Convolutional Neural Networks (CNNs). Documentation is sparse, but for the AX650N, there's a brief description in Pulsar along with the diagram below:

The AX650 and M76H NPUs are mainly composed of three Conv convolution cores and three groups of Vector cores. These Conv and Vector cores are allocated in a 1:1 ratio and divided into three groups of vNPUs

 

The diagram states that the NPU consists of 3 convolution cores, 3 vector cores, and 3 SDMA (System direct memory access) units. We’ll need to investigate further to confirm if this is accurate. The NPU can be set up to function either as 3 distinct Virtual NPUs (vNPUs) or as a single NPU with access to all cores.

The SDK documentation provides more details, highlighting the NPU performance as:

Max. 43.2 TOPS @INT4 and 10.8 TOPS @INT8.

This differs from the original claim of 18.0 TOPS @INT8, with only 10.8 TOPS @INT8 coming from the NPU. I speculate that the remaining 7.2 TOPS might be partially attributed to the dual DSP cores, which feature 256 MACs @INT8. However, I’m still unsure about the remaining TOPS and remain unconvinced that the performance figures are genuine. I suspect there may be some creative accounting at play!

Examination of the memory map offers additional insights into the architecture.

 

The EU0-EU12 labels denote a total of 13 Execution Units (EUs), along with 3 SDMA units, bringing the overall count to 15. The NPU also features its own local memory through OCM (On-Chip Memory), with an address map indicating 11.5MB (0xAFFFFF). Data transfers between the CMM and OCM must be managed by the SDMA units both before and after instructing the NPU cores to execute commands. There will be an overhead of moving each layers weight data to OCM especially for larger LLM models.

 

I had hoped the Execution Units were general-purpose compute cores akin to those in a GPU. Although this isn't documented, extensive debugging suggests they are fixed-function IP blocks. The 13 EUs are divided as follows, with some names and descriptions being educated guesses based on the limited information at hand.

 

1. Convolution Unit (3x) - Capable of performing convolution operations ( Depthwise/Group Conv, Dilation, and ConvTranspose)
   
2. Computer Vision Unit (3x) - Image normalization, reszing, cliping & CV remap/warp.
3. Tensor Unit (3x) - Operations to support activation functions, pooling & elementwise calculations & reduction calcuations.
   
4. MAU (Matrix Arithmetic Unit)  (x1) - Multiply two vectors, int8/int16 inputs and fp16/f32 outputs, topN (N <32) outputs

For LLMs, I assume the workload is distributed across the Convolution and Tensor Execution Units (EU), possibly involving the Matrix Arithmetic Unit (MAU). If the MAU is used, it could create a bottleneck due to the single instance. Additionally, I suspect the EUs can't always operate in parallel—for example, the Convolution EU likely needs to complete its output before the Tensor EU can use it as input. Moreover, the limited OCM memory might prevent EUs of the same type from running in parallel when dealing with large weight data.

Wrap Up

The AX650N is an intriguing SoC for image and vision applications, though its suitability for LLMs feels like an afterthought or marketing hype. The performance claims are hard to verify, and my testing doesn't seem to support the TOPs figures. Additionally, the closed-source nature of the software makes it challenging, if not impossible to develop for. This raises doubts in my mind about its suitability for use in a commercial product. The lack of support from Sipeed for the Maix-IV only exacerbates these issues. Ideally, the Maix-IV should have 16GB of memory, evenly split between Linux and CCM. This would enable running larger LLMs while fully utilizing all 8 CPU cores. In its current incarnation its very limited for this purpose.

I welcome AXera to reach out and clarify or address the points raised in this post.

Fleetwood was kind enough to donate a board for me to review as we both shared an interest in validating the performance claims for LLMs.

RK3588 - Reverse engineering the RKNN - Running llama2.c with TinyStories

I've finally reached a point with reverse engineering where we can start evaluating the usefulness of the NPU for LMMs. I've crafted a basic library (rk3588-npu) with sufficient features for initial integration. A good reference application for integration is llama2.c because its a single C file, the code structure is straightforward to follow and more importantly modify. We will use TinyStories (stories110m) for testing since the models are relatively small, making it easier to troubleshoot when the outputs go a stray. Credit to karpathy for providing llama2.c

Note, I've set the cpu and npu cores to max clk speeds.

We'll start with running run.c against stories110m, this is the fp32 version using the cpu with a single thread (single core). As we see roughly 9.7 tokens/second.
 

 

 

Next I converted run.c to use fp16 (_Float16) along with the weights. We see a slight drop in performance to roughly 9.3 tokens/s as arithmetic operations require a conversion back to fp32.

As with the fp32 version a single core run at a 100% as the code is single thread.

 

The next step was to offload all FP16 multiplications to the NPU. With a vocabulary size of 32,000, the largest multiplication is 768 x 32,000, with others being 768 x 768, 768 x 2048, and 2048 x 768. For efficient execution, the model weights need to reside entirely in memory, accessible to the NPU. This requires them to be within the 4GB address space, which can be problematic for larger models. In our case, the weights are roughly 256MB, requiring an expansion of the kernel CMA memory allocation to 512MB. Additionally, the weights needed conversion to the NPU format.

The changes result in a additional uplift of roughly 21 to 23 tokens/s depending on the length of the output as per the video below. Conservatively we could say a doubling.

 

CPU fluctuates between 30-60% for the single. The CPU is still critical for a number of reason:

1. We're still having to rely on memory copies to send/receive the remaining data for the multiplication to occur on the NPU.

2. Invocation of the NPU kernel driver requires CPU cycles.

3. The rest of the llam2.c code stills runs of the CPU.

Although the results look promising we need to bear mind that TinyStories is very small model as per the architecture. Furthermore its fortunate that the converted weights can fit in memory without having to shuffle weights between userspace and physical memory. In additional fp16 format would further limit the possibility for larger models to run efficiently. So conclusion so far there is some uplift but mileage will vary depending on model size and number of layers.

RK3588 - Reverse engineering the RKNN (Rockchip Neural Processing Unit)

The internal operations and capabilities of the RK3588 NPUs are mainly concealed within a closed-source SDK known as RKNPU2. Given the huge interest in Large Language Models (LLMs) and the quest for optimal matrix multiplications for transformer models. I was curious to understand the implementation of the newly introduced matrix multiplication API (rknn_matmul_run) to the sdk . A thorough examination of the RKNN section in the TRM (Technical Reference Manual) reveals no native mechanism for matrix multiplication, especially for vectors.

To grasp whats going on, the initial step was to understand how the NPU functioned. While the TRM furnished a detailed list of registers and a brief overview of the core units constituting the NPU. It notably lacked essential information on programming the registers for executing operations. For example there were no specifics about deriving or calculating register values based on factors such as data formats (e.g., int8 vs. float16) or the size of input data or weights. Furthermore there was no information on how construct a pipeline for the NPU to execute. Fortunately, I had a slight advantage from a previous reverse engineering attempt on the V831 NPU. Nevertheless, even armed with this knowledge, it has still required several months of trial and error, extensive analysis of data streams, encountering a few dead ends, and numerous attempts at reverse engineering. Finally, I managed to understand how to activate the NPU and get it to execute simple operations.

The RK3588 NPU seems to be distant cousin of the NVDLA architecture in that the some of the terminology is similar and the core units has similar functions and pipe lines to NVDLA although they have been named differently. One of primarily differences is that we can give the NPU a list of tasks (RKNN terminology) to execute and then wait for completion. For example if I have simple neural network consisting of 3 layers and each layer consists of convolution + bias then it is possible to feed 3 tasks (each performing convolution + bias) to the NPU along with the necessary input, weight and bias values. Subsequently we just wait for the NPU to notify when its complete.

The image presented above is extracted from the TRM and has been altered because the description provided in the TRM doesn't entirely align with their diagram, and, more crucially, the register naming convention. Here is my interpretation, each NPU comprises of three distinct units:

• CNA - Convolution Network Accelerator (include CORE rectangle). In the TRM it refers to the Neural Network Accelerating Engine, CNA isn't described.
  
• DPU - Data Processing Unit
• PPU - Planar Processing Unit

Based on the above, the NPU is primarily designed for running conventional Convolutional Neural Networks (CNNs). This is attributed to the CNA core feature, which revolves around executing convolutions by inputting image or feature data along with the corresponding weights. The emphasis on CNNs is further evident by the majority of RKNPU2 samples provided, such as YOLOX, Mobilenet, and ResNet. The CNA output can be directed to the DPU, where element-wise operations such as addition, multiplication, and RELU can be carried out. Subsequently, the DPU's output can be channeled to the PPU, where operations like min, max, and average pooling are executed. Additionally, there is the option to directly feed data to the DPU or PPU without necessitating a convolution step.

To execute convolutions efficiently, the CNA employs multiply-accumulate (MAC) operations. The performance of a CNA is partially determined by the number of MAC units used. According to the TRM, for a single NPU core the count of MAC operations depends on the input data type:

• 1024 int8 MAC operations per cycle
• 512 float16 MAC operations per cycle

Each MAC cell caches 1x1x16 weight bytes, for int 8 its 16 values whilst for float16 it  reduces to 8. We require 2 MAC cells to perform float 16 hence the reduction in operations per cycle. Internally feature and weight data must conform to Rockchips NC1HWC2 format where C2 is the aforementioned value. One 1x1x16 cube of feature data is then shared by all MAC cells to calculate partial sums which are then sent to the accumulator. At higher level the CNA appears to execute a block operation, as observed in my tests where, for instance, the MAC caches 32 channels of weight data for fp16. Hence the requirement to layout weights in kernel groups each with 32 channels.

 

Performance is also affected by the access time to input and weight data, the CNA incorporates a second level cache known as convolution buffer (cbuf). In the above diagram the 384KB onboard memory is partly for that purpose. Importantly the numbers of MAC units plus the cbuf influence how large of a convolution can be completed in one task.

Some of you may have already deduced that the matrix multiplication API is essentially executed through a 2D convolution. For instance, let's consider matrix A as [M x K] and matrix B as [K x N]. Matrix A represents the feature data arranged in an Mx1xK (hwc) format, while matrix B denotes the weight data organized in a 1x1xNxK (hwck) format. Consequently, the resulting matrix C [M x N] is arranged as Mx1xN. I'm at the point where I have a simple test running which asks the NPU to perform a matrix multiplication. I'm using matrices data derived from a GGML testcase (test-mul-mat.cpp) to verify the output is correct. To run the test check out my repo and build, sadly I'm still testing against a kernel 5.10 on a Rock-5b. If the test runs output should be as below & screenshot above.

rock@rock-5b:~/rk3588-npu/build$ ./matmul_4_36_16
drm name is rknpu - 20220829 - RKNPU driver
input dma is ffff8000, output dma is ffffa000, weights dma is ffff9000
Size of npu_regs 112
RKNPU_SUBMIT returned 0
=========================================================================================================
 1224.0 1023.0 1158.0 1259.0 1359.0 1194.0 1535.0 1247.0 1185.0 1029.0  889.0 1182.0  955.0 1179.0 1147.0
 1216.0 1087.0 1239.0 1361.0 1392.0 1260.0 1247.0 1563.0 1167.0 1052.0  942.0 1214.0 1045.0 1134.0 1264.0
 1125.0  966.0 1079.0 1333.0 1287.0 1101.0 1185.0 1167.0 1368.0  990.0  967.0 1121.0  971.0 1086.0 1130.0
  999.0  902.0 1020.0 1056.0 1076.0  929.0 1029.0 1052.0  990.0 1108.0  823.0  989.0  759.0 1041.0 1003.0
=========================================================================================================

Regarding reverse engineering, I've reached a stage where I understand the majority of register settings that impact convolution when dealing with feature data as input. The primary uncertainty lies in determining the bank sizes for feature/weight data, however I'm hopeful that this can be deduced. After dedicating a significant amount of time to analyzing the NPU, here is a list of key areas that you should be aware of:

1. All data pointers within the NPU (e.g., input, weights, outputs, task lists) are 32-bit and must reference physical memory. Consequently, this restricts the memory range to 4GB, making it impractical to leverage a board with 16/32GB memory for the NPU to use. Moreover, it potentially imposes limitations on the types of models that can be executed on the NPU.

2. The claim of 6 TOPS should be approached with caution. While each NPU core is rated at 2 TOPS, there are registers that could potentially enable convolution across all 3 cores. However, after analyzing the data streams generated by the SDK, it appears that this feature is never utilized. Additionally, there doesn't seem to be a similar capability available for the DPU/PPU units, which would restrict its usability. In my view, the most effective approach is to treat them as individual cores and execute models on each one, taking into account the memory constraints mentioned earlier.

3. The SDK matrix multiplication API, in certain aspects, represents an inefficient utilization of the NPU. There is the overhead of memory allocation, a kernel call, and instructing the NPU to execute a single convolution. Ideally, the NPU should be tasked with executing multiple operations and providing all the supplied data for those operations. Typically this is how the NPU is utilized when running a CNN model (ie YOLOvX). The caveat here is that the converted model is limited to contains layers where the operations are supported by the NPU.

4. Initial bench marking for the multiplication of two fp16 [512 x 512] matrices suggests that I could achieve completion in a respectable time of around 1ms. Please note, this involves sending 2 tasks to the NPU, as mentioned earlier due to the cbuf limitation. Unfortunately, this is only part of the story when it comes using vectors data as input. The costly operations involve converting the matrices to feature and weight data formats, and vice versa for the output, if done at runtime. I made an effort to create a highly optimized conversion routine for vector to feature data conversion. According to my benchmarks, this process takes approximately 2ms for fp16 [512 x 512] matrices. I would estimate 12-15ms to perform all the conversions for the matrices mentioned above. Ideally, the matrix for the weight data should be converted ahead of time to reduce conversion overhead and, if possible, persisted for reuse.

5. I was hoping there was the capability to use a programmable core to perform custom operations. Unfortunately this isn't case and your left with using OpenCL as the alternative. This brings it own challenges if you need to shuffle data between OpenCL and the NPU.

There is still more to discover about the other units (DPU/NPU) and I'll spend time doing that. Lastly TRM v1.0 contains numerous gaps and inconsistencies for RKNN, if anyone has later version it would be greatly appreciated.