Dissecting and Re-architecting 3D NAND Flash PIM Arrays for Efficient Single-Batch Token Generation in LLMs

TL;DR

This paper proposes a novel 3D NAND flash PIM architecture optimized for large language model token generation, achieving significant speedup and efficient integration with minimal area overhead.

Contribution

It introduces a re-architected 3D NAND flash PIM array with optimized configurations and mapping methods for LLMs, enabling efficient single-batch token generation.

Findings

2.4x speedup over RTX4090 with vLLM

Comparable performance to four A100 GPUs

Minimal 4.9% latency overhead and small die area

Abstract

The advancement of large language models has led to models with billions of parameters, significantly increasing memory and compute demands. Serving such models on conventional hardware is challenging due to limited DRAM capacity and high GPU costs. Thus, in this work, we propose offloading the single-batch token generation to a 3D NAND flash processing-in-memory (PIM) device, leveraging its high storage density to overcome the DRAM capacity wall. We explore 3D NAND flash configurations and present a re-architected PIM array with an H-tree network for optimal latency and cell density. Along with the well-chosen PIM array size, we develop operation tiling and mapping methods for LLM layers, achieving a 2.4x speedup over four RTX4090 with vLLM and comparable performance to four A100 with only 4.9% latency overhead. Our detailed area analysis reveals that the proposed 3D NAND flash PIM architecture can be integrated within a 4.98mm2 die area under the memory array, without extra area overhead.