Fast Transformer Inference on ARM-Based HMPSoCs

Abstract: Transformer models have set new performance standards for ML tasks. However, their resource-intensive deployment on resource-constrained edge devices for cloud-free, on-chip transformer inference remains challenging. The ARM Compute Library (ARM-CL) framework provides low-latency CNN inference on ARM-based edge devices but lacks support for transformer inference. In this work, we implement several new transformer kernels in ARM-CL to support native transformer execution. Our extended ARM-CL achieves up to three times faster transformer inference compared to state-of-the-art CPU/GPU implementations on an ARM-based embedded board. Furthermore, heterogeneous multi-processor system-on-chips (HMPSoCs) powering edge devices provide both embedded CPUs and GPUs. We introduce cooperative CPU-GPU transformer inference, which executes memory-intensive operations on the CPU while utilizing the GPU for highly parallelizable, compute-intensive operations. This cooperative execution, implemented with minimal overhead, further reduces transformer inference latency by up to 15.72% compared to the best single-processor inference on ARM-CL.

• The paper extends ARM Compute Library to enable native, low-latency transformer inference on ARM-based HMPSoCs.
• It employs a cooperative CPU-GPU layer switching strategy to optimally assign transformer layers based on memory and compute demands.
• Experimental results show up to 15.72% latency reduction compared to single-processor setups, enhancing real-time edge machine learning.

Efficient Transformer Inference on ARM-Based Heterogeneous Multi-Processor SoCs

Introduction and Motivation

Transformer architectures dominate contemporary machine learning applications across NLP, vision, and multimodal domains due to their generalized attention mechanism and high representational capacity. However, deploying transformer-based models on embedded ARM-based edge devices presents distinct challenges due to the significant computational and memory bandwidth requirements. Traditional inference frameworks (TVM, PyTorch, ARM Compute Library) either lack optimal support for embedded ARM-Mali GPU heterogeneity or omit transformer-specific primitives. This work addresses these deficiencies by proposing extensions to ARM Compute Library (ARM-CL) that realize native, low-latency transformer inference for ARM-based HMPSoCs. The focus is on exploiting processor heterogeneity to minimize inference latency for on-chip transformer execution.

ARM-CL Extension for On-Chip Transformer Inference

The ARM-CL framework originally provides highly optimized kernels for CNN operations on ARM CPUs and GPUs, but lacks support for transformer-specific layers, including Multi-Head Attention (MHA), Scaled Dot-Product Attention (SDPA), and layer normalization. The paper implements ARM-CL extensions featuring kernels for embedding, MHA, SDPA, and normalization, fully leveraging NEON intrinsics for CPUs and OpenCL for GPUs.

Figure 1: ARM-CL implementation enables CPU-GPU layer-switched transformer inference with model-specific kernel scheduling and shared tensor communication for layer switching.

The extended ARM-CL achieves computational equivalence with PyTorch kernels, with tailored optimizations for edge hardware. Models supported include BERT-base, DistilBERT, MobileBERT, SqueezeBERT, and GPT-2, compiled with runtime control for granular execution profiling and minimal overhead. These extensions are released as open-source under the MIT license, promoting reproducibility and further research.

Layer-Wise Performance Characterization

Micro-benchmarks systematically characterize inference latency across transformer layers (Embedding, Attention Linear, SDPA, FF, AddNorm) for varying model depth ($d_{model}$) and token sequence length ($\mathcal{L}$). The analysis demonstrates pronounced performance heterogeneity between CPU and GPU across the fundamental transformer layers.

CPU outperforms GPU for memory-bound layers (Embedding, SDPA, AddNorm) due to superior serial access, while GPU achieves lower latency on compute-intensive layers (Attention Linear, FF) owing to parallel MMULs. However, GPU advantages diminish as MMUL workloads exceed the limited L2 cache size (e.g., 128 KB on evaluated ARM Mali G52), resulting in degraded bandwidth efficiency for larger $\mathcal{L}$ or $d_{model}$.

This layer-wise heterogeneity establishes the necessity for a hybrid inference strategy—assigning each layer to the most suitable processor based on the workload’s memory-compute profile.

Cooperative CPU-GPU Layer-Switched Transformer Inference

To optimally exploit processor heterogeneity, the work develops cooperative CPU-GPU layer-switched inference within ARM-CL. The approach leverages ARM-based HMPSoC’s shared memory for inter-processor communication, eliminating the need for explicit memcpy operations encountered in previous graph partitioning methods.

Figure 2: CPU-GPU interactions facilitate seamless layer switching by shared tensors and direct memory access, minimizing communication overhead during transformer inference.

The implementation schedules layers for execution on the processor delivering the lowest latency, with shared tensors allocated in GPU-addressable memory regions accessible by both CPU and GPU via OpenCL. This upscaling method mitigates switching overhead and maintains dataflow integrity despite transformer recurrence and residual connections.

Experimental Evaluation and Numerical Results

Evaluations are performed on the Khadas VIM 3 BASIC (Amlogic A331D HMPSoC) platform with ARM Cortex-A73 CPUs and ARM Mali G52 GPUs under constrained DRAM and cache sizes. The extended ARM-CL achieves significant reductions in inference latency compared to TVM and ExecuTorch on both CPU and GPU.

• Comparison against TVM: ARM-CL delivers on average $2.34\times$ (CPU) and $2.23\times$ (GPU) lower latency than TVM for transformer inference.
• Comparison against ExecuTorch: ARM-CL matches ExecuTorch performance on CPU and uniquely enables GPU acceleration on ARM-Mali.

Single-processor results show GPU outperforms CPU for BERT-base, DistilBERT, and GPT-2, whereas CPU is preferable for SqueezeBERT and MobileBERT due to their structural workloads. The proposed CPU-GPU layer-switched inference consistently outperforms both single-processor methods across all models, with up to 15.72% latency reduction (averaging 10.95%) when compared to the best single-processor configuration.

Practical and Theoretical Implications

This research provides robust evidence for processor and layer-wise computational heterogeneity in transformer inference on embedded ARM platforms. The proposed CPU-GPU layer switching exploits these heterogeneities, yielding practical benefits for privacy-preserving, cloud-free, real-time ML applications on mobile and IoT devices.

The theoretical implications extend to optimal workload partitioning in deep learning compilers targeting heterogenous embedded hardware. The study’s insights generalize to other HMPSoC platforms, contingent on shared memory support and cache architecture.

Future Directions

Opportunities include further optimization of cache-aware MMUL kernels, adaptive graph partitioning for runtime workload balancing, and leveraging hardware accelerators (e.g., NPUs) with transformer-specific support. Exploring auto-tuning strategies to dynamically assign inference workloads based on current hardware utilization and memory constraints could further enhance responsiveness, especially in concurrently multi-tasking edge devices.

Conclusion

The paper delineates architectural and kernel-level extensions to ARM Compute Library, enabling efficient transformer inference for ARM-based HMPSoCs. By characterizing and exploiting performance heterogeneity at the layer and processor level, the cooperative CPU-GPU layer-switched inference strategy achieves substantial latency reduction. This work advances the practical deployment of transformer-based AI models on resource-constrained edge devices and informs future research in hardware-aware ML compiler designs (2606.02836).