Hardware-based Heterogeneous Memory Management for Large Language Model Inference

TL;DR

This paper tackles the memory capacity and bandwidth bottlenecks of large language model inference by introducing H2M2, an asymmetric memory system that couples capacity-centric LPDDR with bandwidth-centric HBM and attaches computation to both memories. It develops a dynamic kernel-memory mapping policy and hardware memory abstraction to adapt to the evolving KV-cache footprint during generation, reducing data movement and improving accelerator utilization. The key contributions are a head-aware kernel-memory mapping scheme, a runtime dynamic mapping algorithm with low overhead, and a hardware-assisted memory abstraction that enables flexible KV-cache management. Empirical results on GPT3-175B, Chinchilla-70B, and Llama2-70B show substantial speedups over LPDDR baselines (up to 2.94x) with modest overheads, demonstrating H2M2’s practicality for scalable, cost-effective LLM serving.

Abstract

A large language model (LLM) is one of the most important emerging machine learning applications nowadays. However, due to its huge model size and runtime increase of the memory footprint, LLM inferences suffer from the lack of memory capacity in conventional systems consisting of multiple GPUs with a modest amount of high bandwidth memory. Moreover, since LLM contains many bandwidthintensive kernels, only focusing on the memory capacity without considering the bandwidth incurs a serious performance degradation. To handle such conflicting memory capacity and bandwidth demands in a cost-effective way, this study investigates the potential of heterogeneous memory systems, proposing H2M2. It uses an asymmetric memory architecture consisting of capacity-centric and bandwidthcentric memory with computation units attached to each memory device. With the asymmetric memory, we first analyze the effect of kernel-memory mapping for the asymmetric memory. Second, we propose a dynamic runtime algorithm that finds a mapping solution considering the characteristics of LLM operations and the change of footprint during LLM inference. Third, we advocate the need for memory abstraction for the efficient management of the asymmetric memory. H2M2 outperforms the conventional homogeneous memory system with LPDDR by 1.46x, 1.55x, and 2.94x speedup in GPT3-175B, Chinchilla-70B, and Llama2-70B, respectively.

Figures (19)

• Figure 1: Accelerator substrate used by H2M2.
• Figure 2: The topology of the decoder layer in GPT3. Orange boxes with W labels represent the weight parameter tensors, blue boxes represent the input activation tensors, and gray boxes represent the KV cache tensors.
• Figure 3: Footprint breakdown for batch size 32, sequence length growing from 256 to 2048 in GPT3-175B. The portion of attention increases due to the increase of KV cache size.
• Figure 4: Four possible configurations of memory systems applicable for LLM inference acceleration.
• Figure 5: Techniques for supporting optimal mapping granularity with asymmetric memory, for a single decoder layer. For clarity in explanation, a simplified view of a decoder layer is presented, visualizing GEMM/GEMV kernels mainly.