Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

UMDAM: A Unified Data Layout and DRAM Address Mapping for Heterogenous NPU-PIM

Hai Huang, Xuhong Qiang, Weisheng Zhao, Chenchen Liu

TL;DR

This work tackles memory-bound decode in edge LLM inference by co-designing data layout and DRAM address mapping to bridge NPUs and PIM. The proposed UMDAM framework uses a tile-based column-major layout and a configurable, hardware-conscious DRAM mapping to preserve NPU interleaving while enabling PIM locality, implemented with lightweight memory-controller support. Evaluations on OPT models show up to 3.0× TTFT reduction and over 2× end-to-end latency improvement, without extra storage overhead. The results demonstrate practical viability for deploying efficient NPU-PIM co-execution on commercial edge platforms.

Abstract

Large Language Models (LLMs) are increasingly deployed on edge devices with Neural Processing Units (NPUs), yet the decode phase remains memory-intensive, limiting performance. Processing-in-Memory (PIM) offers a promising solution, but co-executing NPU-PIM systems face challenges such as data layout mismatches, bandwidth loss, and redundant storage. To address these issues, we propose UMDAM, a unified memory-affinity data layout and DRAM address mapping scheme tailored for NPU-PIM co-execution. UMDAM employs a column-major, tile-based layout and a configurable DRAM mapping strategy to ensure compatibility with NPU computation while maximizing PIM efficiency -- without introducing extra memory overhead or bandwidth loss. Comprehensive evaluations on OPT models demonstrate that UMDAM reduces time-to-first-token (TTFT) by up to 3.0x and time-to-last-token (TTLT) by 2.18x, significantly improving end-to-end LLM inference efficiency on edge devices.

UMDAM: A Unified Data Layout and DRAM Address Mapping for Heterogenous NPU-PIM

TL;DR

This work tackles memory-bound decode in edge LLM inference by co-designing data layout and DRAM address mapping to bridge NPUs and PIM. The proposed UMDAM framework uses a tile-based column-major layout and a configurable, hardware-conscious DRAM mapping to preserve NPU interleaving while enabling PIM locality, implemented with lightweight memory-controller support. Evaluations on OPT models show up to 3.0× TTFT reduction and over 2× end-to-end latency improvement, without extra storage overhead. The results demonstrate practical viability for deploying efficient NPU-PIM co-execution on commercial edge platforms.

Abstract

Large Language Models (LLMs) are increasingly deployed on edge devices with Neural Processing Units (NPUs), yet the decode phase remains memory-intensive, limiting performance. Processing-in-Memory (PIM) offers a promising solution, but co-executing NPU-PIM systems face challenges such as data layout mismatches, bandwidth loss, and redundant storage. To address these issues, we propose UMDAM, a unified memory-affinity data layout and DRAM address mapping scheme tailored for NPU-PIM co-execution. UMDAM employs a column-major, tile-based layout and a configurable DRAM mapping strategy to ensure compatibility with NPU computation while maximizing PIM efficiency -- without introducing extra memory overhead or bandwidth loss. Comprehensive evaluations on OPT models demonstrate that UMDAM reduces time-to-first-token (TTFT) by up to 3.0x and time-to-last-token (TTLT) by 2.18x, significantly improving end-to-end LLM inference efficiency on edge devices.

Paper Structure

This paper contains 12 sections, 5 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Motivation
  3. Data Layout Challenges in Heterogeneous NPU-PIM
  4. Related Works on Data Layout and Address Mapping
  5. UMDAM: Unified Memory-affinity Data Layout and DRAM Address Mapping Scheme for NPU–PIM
  6. Data layout format and DRAM address mapping
  7. System Implementation
  8. Evaluation
  9. Experimental Setup
  10. Evaluation on a Single Query
  11. Deployability of UMDAM on Commercial NPU Platforms
  12. Conclusion

Figures (5)

  • Figure 1: (a) Differences between Conventional layout and PIM-optimized layout, (b) Conventional approaches.
  • Figure 2: An illustration of UMDAM: (a) Unified data layout and DRAM address mapping for NPU-PIM, (b) System overview.
  • Figure 3: TTFT speedup of UMDAM over the NPU-PIM baseline with varying prefill length and model size.
  • Figure 4: TTLT speedup of UMDAM over the NPU-PIM baseline with varying prefill-to-decode ratio.
  • Figure 5: Deployability verification of UMDAM on Ascend 310B.