TranSAGE boosts transformer model speed and energy use with spatial pipelines
TranSAGE: Automating Spatial Accelerator Generation for Large-Scale Transformer Models
Hardware Architecture
Summary
Large Transformer models are powerful but need a lot of computing power and memory, making them hard to run efficiently. The authors created TranSAGE, a system that automatically designs special hardware to run these models faster and use less energy by cleverly organizing multiple operations to work in a pipeline. TranSAGE uses smart algorithms to find the best way to split and schedule the model's tasks, then generates hardware designs automatically. Tests show TranSAGE can speed up processing and save energy much better than existing GPUs and similar tools.
What this means in practice
- •For hardware accelerator designers: Automatically generate spatial hardware accelerators optimized for large Transformer models improving throughput and energy efficiency significantly.
- •For cloud ai infrastructure teams: Deploy more efficient Transformer inference servers by integrating TranSAGE-generated accelerators to reduce operational costs and increase capacity.$Commercial implications: Enables commercial deployment of custom accelerators that outperform GPUs/NPUs in energy and speed for large language models.
Authors
Chi Zhang, Yi Sun, Weichuang Zhang, Quan Chen, Jieru Zhao
Abstract
Large-scale Transformers have achieved remarkable success across NLP, computer vision, and robotics, but their growing computational and memory demands challenge efficient inference. Spatial accelerators offer advantages in on-chip bandwidth and energy efficiency, yet designing them for large models remains costly and time-consuming. Existing accelerator generators typically focus on single-operation support, overlooking the benefits of multi-operation pipelining, which limits throughput and energy efficiency. We propose TranSAGE, a compiler-based framework that jointly optimizes scheduling for large-scale Transformers under multi-operation pipelining and generates corresponding accelerator implementations. Our approach integrates a design space exploration (DSE) engine with multi-operation dataflow optimization and a dynamic programming-based graph partitioning algorithm to efficiently identify pipeline-friendly operation clusters and their optimal dataflow parameters. We further design spatial architecture primitives and lowering passes to automatically generate efficient RTL implementations. Evaluations on four representative LLMs demonstrate that TranSAGE achieves up to 2.41x higher throughput and 13.04x higher energy efficiency than commercial GPU/NPU baselines. Compared with state-of-the-art accelerator generators, TranSAGE achieves up to 11.33x higher throughput and 6.07x higher energy efficiency. Our code is available at https://github.com/sjtu-zhao-lab/TranSAGE.