We prototype sparsity-awareness from zero on existing AMD neural processing unit (NPU) XDNA 1-2 fabric for I/O performance gain. By purely alleviating memory I/O, we show up to ~1.59x improvement in FFN throughput which translates to ~1.25x increase in end-to-end token generation throughput. Our approach is generalizable to existing models such as Meta AI OPT-6.7B, Mixtral-8x22B, ReluFalcon-40B, and scalable to upcoming AMD NPU architectures optimized toward LLM workloads.

Billion-parameter Large Language Models (LLMs) have revolutionized various AI applications. However, the computational demand for inference presents significant challenges for deployment on resource-limited devices, e.g. commodity workstations. Therefore, there has been great interest in ReLU-fication as it offers minimal impact on model training convergence and online accuracy performance, while introducing up to 95% exploitable activation sparsity. Now, we aim to translate the newfound activation sparsity into tangible performance gain in the form of high token generation throughput.

Many works aim to reduce the computational (FLOPs) requirements during inference. However, it is known that LLM inference is memory-bound. Therefore, we orthogonally focus on I/O reduction. Layerwise in the ReLU-fied architecture, the FFN features up projection clipped by ReLU and fed into down projection, resulting in a sparse-vector GeMV. Given a DMA buffer descriptor encoding a weight matrix transfer, we decompose and selectively load rows that correspond to activated elements in the sparse up projection to realize I/O performance gain.

We work toward integrating this functionality on existing NPU dataflow in three phases. In Phase 0, we program the AIE array using IRON toolchain to identify an upper-bound speedup ratio, which is near-linear. In Phase 1, we co-design our application testbench with internal task token synchronization to integrate the enhancement for alternating sparsity patterns into the existing XDNA 1-2 fabric. (This is the worst-case pattern as rows cannot be packed into fewer BDs). Directly extendable, Phase 2 is left as future work and identifies the end-to-end latency of online BD generation and execution. Our main contribution lies in the layer-wise upper- and lower-bound token generation performance gains afforded by Phase 0 and Phase 1, and thus scales with model depth. Even with sparse descriptor (under) utilization under the worst-case activation pattern, we demonstrate a ~1.59x improvement in FFN throughput, which translates to ~1.25x increase in end-to-end token generation throughput according to our in-house hardware dataflow-model architecture mapper. Moreover, internal benchmarks reveal that DMA descriptor execution and hence our identified performance gains are bottlenecked by the single control processor, which is expected to change in future NPU iterations.

Currently, we underutilize DMA descriptors as we statically created BDs for each row but discarded them during runtime if they correspond to an inactivated element, thus preventing reuse within the current batch of activated rows. Moreover, the activation vector pattern remains fixed. Nevertheless, this gives us a lower-bound performance gain, and it shows our enhancement is feasible to integrate in next-generation architectures such as XDNA 4+ which features per-AIE-column control processors. Our approach is independent, modular and compatible with speculative decoding, W4A8KV4 and accumulator aware quantization, and many more! Therefore, we expect speedup ratios on XDNA 4+ to be closer to that of Phase 0 (i.e. near-linear).

The most salient application is accelerating foundation model inference on commodity workstations equipped with Ryzen AI. Looking to agentic workflow platforms by Microsoft Research, this has great potential for Mixture of Expert (MoE) serving as experts are hierarchical MoEs and use FFNs at the core. Lower-level weight fetching is on-demand, and there is little model gating or routing in lower levels, so weights are frequently not preloaded.