Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

IANUS: Integrated Accelerator based on NPU-PIM Unified Memory System

Minseok Seo, Xuan Truong Nguyen, Seok Joong Hwang, Yongkee Kwon, Guhyun Kim, Chanwook Park, Ilkon Kim, Jaehan Park, Jeongbin Kim, Woojae Shin, Jongsoon Won, Haerang Choi, Kyuyoung Kim, Daehan Kwon, Chunseok Jeong, Sangheon Lee, Yongseok Choi, Wooseok Byun, Seungcheol Baek, Hyuk-Jae Lee, John Kim

TL;DR

Novel PIM Access Scheduling is proposed that manages not only the scheduling of normal memory accesses and PIM computations but also workload mapping across the PIM and the NPU.

Abstract

Accelerating end-to-end inference of transformer-based large language models (LLMs) is a critical component of AI services in datacenters. However, diverse compute characteristics of end-to-end LLM inference present challenges as previously proposed accelerators only address certain operations or stages (e.g., self-attention, generation stage, etc.). To address the unique challenges of accelerating end-to-end inference, we propose IANUS -- Integrated Accelerator based on NPU-PIM Unified Memory System. IANUS is a domain-specific system architecture that combines a Neural Processing Unit (NPU) with a Processing-in-Memory (PIM) to leverage both the NPU's high computation throughput and the PIM's high effective memory bandwidth. In particular, IANUS employs a unified main memory system where the PIM memory is used both for PIM operations and for NPU's main memory. The unified main memory system ensures that memory capacity is efficiently utilized and the movement of shared data between NPU and PIM is minimized. However, it introduces new challenges since normal memory accesses and PIM computations cannot be performed simultaneously. Thus, we propose novel PIM Access Scheduling that manages normal memory accesses and PIM computations through workload mapping and scheduling across the PIM and the NPU. Our detailed simulation evaluations show that IANUS improves the performance of GPT-2 by 6.2$\times$ and 3.2$\times$, on average, compared to the NVIDIA A100 GPU and the state-of-the-art accelerator. As a proof-of-concept, we develop a prototype of IANUS with a commercial PIM, NPU, and an FPGA-based PIM controller to demonstrate the feasibility of IANUS.

IANUS: Integrated Accelerator based on NPU-PIM Unified Memory System

TL;DR

Novel PIM Access Scheduling is proposed that manages not only the scheduling of normal memory accesses and PIM computations but also workload mapping across the PIM and the NPU.

Abstract

Accelerating end-to-end inference of transformer-based large language models (LLMs) is a critical component of AI services in datacenters. However, diverse compute characteristics of end-to-end LLM inference present challenges as previously proposed accelerators only address certain operations or stages (e.g., self-attention, generation stage, etc.). To address the unique challenges of accelerating end-to-end inference, we propose IANUS -- Integrated Accelerator based on NPU-PIM Unified Memory System. IANUS is a domain-specific system architecture that combines a Neural Processing Unit (NPU) with a Processing-in-Memory (PIM) to leverage both the NPU's high computation throughput and the PIM's high effective memory bandwidth. In particular, IANUS employs a unified main memory system where the PIM memory is used both for PIM operations and for NPU's main memory. The unified main memory system ensures that memory capacity is efficiently utilized and the movement of shared data between NPU and PIM is minimized. However, it introduces new challenges since normal memory accesses and PIM computations cannot be performed simultaneously. Thus, we propose novel PIM Access Scheduling that manages normal memory accesses and PIM computations through workload mapping and scheduling across the PIM and the NPU. Our detailed simulation evaluations show that IANUS improves the performance of GPT-2 by 6.2 and 3.2, on average, compared to the NVIDIA A100 GPU and the state-of-the-art accelerator. As a proof-of-concept, we develop a prototype of IANUS with a commercial PIM, NPU, and an FPGA-based PIM controller to demonstrate the feasibility of IANUS.

Paper Structure

This paper contains 27 sections, 18 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Transformer-based LLMs
  4. Platforms for DNN Inference
  5. Motivations
  6. Diverse Computational Requirements of LLMs
  7. Partitioned vs. Unified Memory Systems in LLMs
  8. IANUS Architecture
  9. NPU & PIM Architecture
  10. Transformer-Aware NPU & PIM Microarchitecture
  11. Data Manipulation in Self-Attention
  12. Vector Operations in Vector Unit
  13. Data Allocation in PIM
  14. IANUS Microarchitecture
  15. PIM Access Scheduling
  16. ...and 12 more sections

Figures (18)

  • Figure 1: (a) Structure of GPT with vector operations marked with (V) and (b) multi-head attention mechanism shown in detail.
  • Figure 2: Generation stage of GPT-2 XL (a) Latency and FLOPs breakdown of decoders. (b) Latency breakdown of self-attention. Results are obtained using an A100 GPU.
  • Figure 3: (Left) Architecture of a core in NPU. (Middle) PIM architecture. (Right) Overall architecture of IANUS.
  • Figure 4: Data allocation and tiling scheme for a matrix-vector multiplication in PIM.
  • Figure 5: IANUS's DRAM address mapping with the mapping of tile shown in Figure .
  • ...and 13 more figures