ThAME: 3D Memory-Enabled Heterogeneous Accelerator for LLM Mixture of Experts
Hardware ArchitectureArtificial IntelligenceDistributed, Parallel, and Cluster Computing
Summary
The authors address speed and efficiency problems in running Mixture of Experts (MoE) models on regular hardware, which face limits like slow memory access and complex data traffic. They introduce ThAME, a new 3D chip setup combining different memory types and a smart way to split the work, tailored for how MoE models process data. They also create a special communication system inside the chip to handle the unpredictable flow of data better. Their tests show ThAME is much faster and uses less energy than current methods.
Authors
Pratyush Dhingra, Pramit Kumar Pal, Janardhan Rao Doppa, Partha Pratim Pande
Abstract
Mixture of Experts (MoE) architectures have emerged as a dominant paradigm for scaling Large Language Models (LLMs). However, MoE inference on conventional hardware is constrained by three fundamental bottlenecks. These encompass the massive memory bandwidth required to fetch non-contiguous expert weights, the non-deterministic scatter-gather traffic generated by input-dependent token routing, and the tail-latency dependency imposed by synchronous expert output aggregation. To address these challenges, we propose ThAME, a three-dimensional (3D) heterogeneous multi-chiplet architecture for MoE inference. ThAME employs Ferroelectric Field-Effect Transistor (FeFET)-based non-volatile and DRAM-based volatile memory chiplets with a co-designed compute mapping strategy that aligns the distinct computational profiles of attention mechanisms and expert routing. Furthermore, we design a specialized Network-on-Chip communication backbone optimized to mitigate the bottlenecks associated with non-deterministic token routing traffic across the combinatorial space of input-dependent MoE traffic patterns. Experimental results demonstrate that ThAME outperforms state-of-the-art counterparts by up to 15.7x in terms of speedup and improves energy efficiency by up to 9.8x.