HeTraX: Energy Efficient 3D Heterogeneous Manycore Architecture for Transformer Acceleration

TL;DR

HeTraX is a novel 3D heterogeneous hardware architecture designed to efficiently accelerate end-to-end transformer models, achieving significant speedup and energy efficiency improvements while maintaining thermal feasibility.

Contribution

This paper introduces HeTraX, a specialized 3D heterogeneous architecture optimized for transformer kernels, addressing limitations of existing accelerators in performance and thermal management.

Findings

HeTraX achieves up to 5.6x speedup over state-of-the-art accelerators.

HeTraX improves energy-delay product (EDP) by 14.5x.

HeTraX maintains thermal feasibility during operation.

Abstract

Transformers have revolutionized deep learning and generative modeling to enable unprecedented advancements in natural language processing tasks and beyond. However, designing hardware accelerators for executing transformer models is challenging due to the wide variety of computing kernels involved in the transformer architecture. Existing accelerators are either inadequate to accelerate end-to-end transformer models or suffer notable thermal limitations. In this paper, we propose the design of a three-dimensional heterogeneous architecture referred to as HeTraX specifically optimized to accelerate end-to-end transformer models. HeTraX employs hardware resources aligned with the computational kernels of transformers and optimizes both performance and energy. Experimental results show that HeTraX outperforms existing state-of-the-art by up to 5.6x in speedup and improves EDP by 14.5x while ensuring thermally feasibility.