Tasa: Thermal-aware 3D-Stacked Architecture Design with Bandwidth Sharing for LLM Inference

TL;DR

This work targets the memory bandwidth bottleneck in autoregressive LLM inference by analyzing thermal challenges in 3D-stacked architectures and proposing Tasa, a thermal-aware heterogeneous 3D design. Tasa combines high-performance P-cores for GEMM with high-efficiency E-cores for GEMV and introduces a bandwidth-sharing scheduling scheme to better utilize available bandwidth across cores, all while flattening heat distribution. Thermal analyses show 3D stacks suffer from heat concentration and poorer scalability compared to 2D, motivating architectural interventions; experiments with Llama-65B and GPT-3 66B demonstrate up to 5.55°C peak-temperature reductions and up to 2.85× end-to-end speedups over GPU baselines, plus favorable energy efficiency. Overall, Tasa extends the latency–temperature Pareto frontier for LLM inference on 3D-stacked hardware, enabling scalable, thermally balanced accelerators with practical performance gains.

Abstract

The autoregressive decoding in LLMs is the major inference bottleneck due to the memory-intensive operations and limited hardware bandwidth. 3D-stacked architecture is a promising solution with significantly improved memory bandwidth, which vertically stacked multi DRAM dies on top of logic die. However, our experiments also show the 3D-stacked architecture faces severer thermal issues compared to 2D architecture, in terms of thermal temperature, gradient and scalability. To better exploit the potential of 3D-stacked architecture, we present Tasa, a heterogeneous architecture with cross-stack thermal optimizations to balance the temperature distribution and maximize the performance under the thermal constraints. High-performance core is designed for compute-intensive operations, while high-efficiency core is used for memory-intensive operators, e.g. attention layers. Furthermore, we propose a bandwidth sharing scheduling to improve the bandwidth utilization in such heterogeneous architecture. Extensive thermal experiments show that our Tasa architecture demonstrates greater scalability compared with the homogeneous 3D-stacked architecture, i.e. up to 5.55 $\tccentigrade$, 9.37 $\tccentigrade$, and 7.91 $\tccentigrade$ peak temperature reduction for 48, 60, and 72 core configurations. Our experimental for Llama-65B and GPT-3 66B inferences also demonstrate 2.85x and 2.21x speedup are obtained over the GPU baselines and state-of-the-art heterogeneous PIM-based LLM accelerator

Figures (14)

• Figure 1: 3D-stacked architecture with DRAMs on logic die.
• Figure 2: (a) The simulated heat map of 3D-stacked architecture for LLM inference, and (b) the temperature distribution across vertical dies.
• Figure 3: (a) The horizontal temperature distribution of the logic die in 3D architecture and (b) conventional 2D architecture.
• Figure 4: Simulated temperatures with varying cores for 2D/3D architectures.
• Figure 5: Thermal analysis with varying distance between cores.