UniFormer: Unified and Efficient Transformer for Reasoning Across General and Custom Computing

TL;DR

UniFormer addresses the cross-platform deployment bottlenecks of Transformer models by introducing a dual-branch attention architecture that fuses global linear attention with local block-wise attention under a GEMM-centric design. The approach enables high parallelism and compute–storage fusion, employing a Triton-based kernel for the global branch on GPUs and FPGA-friendly block computations for hardware pipelines. Empirical results show substantial latency-throughput gains on GPUs (1.79×–2.36× throughput, ~82% Top-1 accuracy on ImageNet) and large speedups (180×–470×) plus energy efficiency benefits on FPGAs, establishing cross-platform efficiency without sacrificing accuracy. This cross-platform unification advances deployability and real-time performance of efficient Transformers on both general-purpose and customised hardware, with implications for scalable, low-power AI systems.

Abstract

The success of neural networks such as convolutional neural networks (CNNs) has been largely attributed to their effective and widespread deployment on customised computing platforms, including field-programmable gate arrays (FPGAs) and application-specific integrated circuits (ASICs). In the current era, Transformer-based architectures underpin the majority of state-of-the-art (SOTA) larger models that are also increasingly deployed on customised computing hardware for low-power and real-time applications. However, the fundamentally different parallel computation paradigms between general-purpose and customised computing often lead to compromises in model transfer and deployability, which typically come at the cost of complexity, efficiency or accuracy. Moreover, many cross-platform optimisation principles have also remained underexplored in existing studies. This paper introduces UniFormer, a unified and efficient Transformer architecture for both general-purpose and customised computing platforms. By enabling higher parallelism and compute-storage fusion, UniFormer achieves state-of-the-art (SOTA) accuracy and latency on GPUs while exhibiting strong adaptability on FPGAs. To the best of our knowledge, this paper is the first efficient Transformer work that jointly considers both general-purpose and customised computing architectures.

Figures (7)

• Figure 1: Problem setting (left): Recent efficient Transformer works leverage customised GPU kernels to maximise parallelism and memory utilisation, while FPGA designs can already enable fine-grained dataflows and pipelined execution by high-level synthesis under significantly reduced computational complexity. Moreover, both general and customised computing architectures seek more advanced compute–storage fusion mechanisms to alleviate the common bandwidth bottleneck. Existing works (right): Representative efficient Transformer implementations with target platforms.
• Figure 2: Matrix multiplication on parallel computing architectures, such as FPGAs and GPUs, can be further decomposed into fine-grained parallel blocks and multiply-add operations to enable customised and hardware-aware optimisations. Favourable trade-offs between design complexity and efficiency make matrix multiplication an efficient computational primitive across both architectures.
• Figure 3: The proposed dual-branch Transformers. The input features are split into more parallel streams and can be fused with Triton kernels with inner-loop and outer-loop scheduling strategies to maximise data and memory usage (e.g., SRAM and HBM in GPUs, BRAM and DRAM in FPGAs).
• Figure 4: Throughput Analysis. Throughput performance on an H800 GPU (batch size = 64, head = 16, dimension = 64) under varying sequence lengths. The plot illustrates the latency and throughput characteristics across model configurations. The more results are provided in Appendices A.1 & A.2.
• Figure 5: Latency vs. input context length ($N_{\text{CTX}}$) under five different attention on 4070 Ti Super.