DeepStack: Facilitating Co-Design Exploration of 3D DRAM-Stacked Accelerators for Distributed LLM Inference

Abstract

Advances in hybrid bonding and packaging have driven growing interest in 3D DRAM-stacked AI accelerators. As large language models (LLMs) scale to hundreds of billions or trillions of parameters, distributed inference across multiple 3D chips has become essential for AI serving. This trend makes cross-stack co-design critical because system-level parallelization and scheduling choices are tightly coupled with hardware characteristics such as memory organization, interconnects, and thermal constraints. We present DeepStack, an accurate performance model and efficient design space exploration (DSE) framework for distributed 3D-stacked LLM inference. At the hardware level, DeepStack captures transaction-aware memory bandwidth, bank activation constraints, buffering limitations, and thermal and power behavior. At the system level, it incorporates comprehensive parallelization strategies and execution scheduling. Through a dual-stage network abstraction and tile-level compute-communication overlap modeling, DeepStack achieves up to 100,000x faster evaluation than state-of-the-art simulators at comparable accuracy. We cross-validate DeepStack against our in-house 3D designs, an NS-3 backend with 2.12% error, and vLLM serving on eight B200 GPUs with 12.92% error. Combined with hierarchical search, DeepStack efficiently explores about 2.5 x 10^14 design points spanning the number of stacked DRAM layers, DRAM vertical connectivity, interconnects, compute-memory allocation, and distributed scheduling under thermal and area constraints. A search-space ablation shows that restricted DSE baselines can miss up to 9.5x modeled throughput. Beyond modeling and DSE, DeepStack derives design implications for distributed 3D AI systems and guides performance optimization across the stack. Source code and artifacts are available at https://github.com/tile-ai/DeepStack/tree/ae.

Update MICRO ver