In-DRAM computing speeds up language model decoding on edge devices
Accelerating LLM Decoding with In-DRAM Parallelism
Hardware Architecture
Summary
Running large language models on small devices is slow because the processor spends a lot of time moving data between memory and itself. The authors focus on a way to do some of the calculating directly inside the memory chip (DRAM) to speed things up. They redesigned both the math steps and how data is organized so the memory can work on many parts of the problem at once, making the process faster. Their approach shows a clear speed improvement over the usual methods.
What this means in practice
- •For edge device engineers: Implement faster inference for large language models on mobile and embedded hardware using in-DRAM parallel processing.
- •For hardware architects: Design memory systems that integrate processing units to accelerate matrix operations typical in AI workloads.
Authors
Kevin Xu, Tatsuya Kubo, Daichi Tokuda, Shinya Takamaeda-Yamazaki, Issei Sato
Abstract
Autoregressive decoding of large language models (LLMs) on edge devices is dominated by memory-bound general matrix--vector multiplication (GeMV) operations. Processing-using-DRAM (PUD) is a promising computing paradigm that can alleviate this bottleneck by reducing processor--memory data transfers and enabling massively parallel computation within DRAM (e.g., across 65,536 columns). However, weight matrices in individual LLM layers (e.g., a QKV matrix in Llama3 8B with 6,144 rows) are typically too small to fully utilize parallelism in PUD. To address this challenge, we propose an algorithm--system co-design for PUD-based LLM inference. At the algorithm level, we enable a subsequent projection to be packed with the current projection by reformulating it as a computation on the current residual state plus a correction for the residual update. Layer redundancy allows this correction to be approximated using fewer bit-serial steps. At the system level, we tailor data mapping and execution to improve utilization of either columns within DRAM subarrays (SIMD) or independently controlled banks (MIMD). Overall, we achieves up to $1.41\times$ end-to-end decoding speedup over the PUD baseline, corresponding to a $2.45\times$ speedup over processor-based execution.