Optimizing Layer-Fused Scheduling of Transformer Networks on Multi-accelerator Platforms

TL;DR

The paper tackles efficient transformer scheduling on diverse multi-accelerator hardware by extending the Stream framework to support transformer-specific layers, including transpose and softmax, and enabling layer-fused execution. After validating the extended framework against a GAP8-based transformer with 8-bit quantization, it demonstrates latency predictions within ~8–9% of measured values, establishing reliable cost estimation for exploration. The study then reveals that layer fusion can substantially reduce the active feature memory footprint without increasing latency, with gains depending on the attention-head input geometry: for M < N, a memory reduction up to about 1/3 is achievable with alpha ≈ (2N+M)/(3N); for M > N, gains can reach roughly alpha ≈ 3N/(2N+M) < 1 and can be as high as ~0.3 in examples. The results, along with an open-source release, show practical methodology for transformer optimization on multi-core and multi-accelerator platforms, enabling more memory- and energy-efficient deployments of transformer workloads.

Abstract

The impact of transformer networks is booming, yet, they come with significant computational complexity. It is therefore essential to understand how to optimally map and execute these networks on modern neural processor hardware. So far, literature on transformer scheduling optimization has been focusing on deployment on GPU and specific ASICs. This work enables extensive hardware/mapping exploration by extending the DSE framework Stream towards support for transformers across a wide variety of hardware architectures and different execution schedules. After validation, we explore the optimal schedule for transformer layers/attention heads and investigate whether layer fusion is beneficial to improve latency, energy or memory requirements. Our study shows that the memory requirements for active feature data can be drastically reduced, by adapting the execution schedule based on the size of the input of the attention head.

Figures (6)

• Figure 1: Operators of a transformer network, highlighting one attention head containing 5 matrix-matrix multiplications (3x features$\times$weights, 2x features$\times$features), a transpose and a softmax layer. Stream symons2022towards is extended with these additional layer types.
• Figure 2: Illustration of Stream symons2022towards, with indicated in red the blocks where adaptations have been made.
• Figure 3: Illustration of dependencies in different layer types.
• Figure 4: Validation results of the framework with two different sized CCT-like networks on GAP8.
• Figure 5: Scheduling and total active features memory usage across time for memory-optimized layer-by-layer and layer-fused execution, with best possible latency. For the layer-fused execution, we make a distinction between $M < N$ and $M > N$ as this leads to different optimal schedulings. For each scheduling, both $M > N$ as $M = N$ as $M < N$ are shown to illustrate why an optimal scheduling for one input size is not necessarily the best one for another scheduling. The dot at the end of the plots indicates that the output of the attention head should remain active to consume as the input of the next attention head. For a), the computations of $V$ and $Q.K^T$ can be swapped, this doesn't change latency or maximal memory requirement. For c), it is also an option to fuse $Q$, $Q.K^T$ and $Q.K^T.V$, instead of doing first $Q$ completely and then only fuse $Q.K^T$ and $Q.K^T.V$.