Orion's Belt

Inspiration

I have been looking at recent large language model architectures, and how inference can be run on Arm edge CPUs, particularly on later ISAs like v8.2 and v9. As transformer architectures are extensively covered, I wanted to look at how Gated Delta Net layers could be handled with CPU, allowing us to dedicate the NPUs to transformer layers.

What it does

The key deliverable are three GDN kernels for Arm CPUs - Cumulative decay, Causal Conv1D, and Gated Delta Rule Scan. Though NPU and GPU solutions were also explored for GDN, the goal is to allow for heterogenous inference across the different types of compute on ARM edge devices.

How we built it

We have an existing cluster of Arm V8.2 single board computers, including Turing Pi RK3588 (A76+A55), Jetson Nano (A57) and Raspberry Pi 5 (A76) devices.

Challenges we ran into

Originally, we were looking at exploring Arm V9 via a Orion O6 (which is where the name for the repo comes from), but the hardware ended up being too difficult to acquire in the competition period. This meant that the benchmarking focused on Neon rather than any of the affordances of the Scalable Vector Extension 2 framework. It would have also been cool to explore Scalable Matrix Extension, but this does not appear to be yet supported by Edge ARM devices.

Accomplishments that we're proud of

• Getting working kernels that provide a performance improvement, at least for the prefill stage.
• Being able to benchmark against disparate ARM devices.

What we learned

This project was a valuable way to learn about both the Gated Delta Net architecture, as well as the mechanics and quirks of ARM CPU inference.

What's next for Orion's Belt

It would be valuable to explore full heterogenous inference across CPU/GPU/NPU, as well as exploring the architecture of v9 devices.

Built With

• arm
• jetson-nano
• kimi
• nemotron
• qwen
• raspberry-pi
• rk3588