arXiv:2511.15503cs.ARcs.DC2025-11被引 1

DCC编译器协同优化数据重排与计算代码,显著提升内存计算架构的机器学习性能。

DCC: Data-Centric Compilation of Machine Learning Kernels for Processing-In-Memory Architectures

  • 以数据为中心设计联合优化框架,统一处理数据重排与计算调度。
  • 在HBM-PIM上平均提速2.21倍,AttAcc上平均提速3.92倍,最高达13.17倍。
  • 支持多种PIM设备,适合需高效利用内存带宽的大型模型推理场景。

高性能主机处理器可集成处理内存(PIM)设备,通过利用PIM核心处的大内存带宽,加速机器学习(ML)模型中内存密集型核函数,包括大语言模型(LLMs)。然而,主机需连续元素分布在不同DRAM银行,而PIM核心要求其本地银行内元素连续,这导致执行时需进行数据重排,带来显著性能与可编程性挑战,尤其在支持多种PIM设备时更为突出。现有编译方法缺乏对多样ML核函数与多PIM设备的系统性优化,且在计算代码优化阶段常忽略数据重排开销。我们证明了数据重排与计算代码优化相互依赖,需在调优过程中联合优化。因此,我们设计了DCC——首个面向PIM系统的数据中心编译器,通过统一调优流程联合优化数据重排与计算代码。DCC集成多层PIM抽象以支持多种后端。它能有效协同数据分块策略与计算循环分块方案,应用PIM特定代码优化,并借助快速准确的性能预测模型,为给定核函数在目标PIM架构上选择最优代码调度。评估显示,各独立ML核函数在HBM-PIM上最高提速7.68倍(平均2.21倍),在AttAcc PIM上最高提速13.17倍(平均3.92倍),相比纯GPU执行。在端到端大语言模型推理中,DCC在AttAcc上使GPT-3和LLaMA-2平均提速4.52倍(最大7.71倍于LLaMA-2)。DCC已开源:https://github.com/SPIN-Research-Group/DCC。

原文摘要 · Abstract (English)

High-performance Host processors can integrate Processing-In-Memory (PIM) devices, which can accelerate memory-intensive kernels of Machine Learning (ML) models, including Large Language Models (LLMs), by leveraging the large memory bandwidth available at PIM cores. However, Host processor needs consecutive elements distributed across DRAM banks, while PIM cores need consecutive elements within their local banks. This necessitates data rearrangements in ML kernel execution that pose significant performance and programmability challenges, further exacerbated by the need to support diverse PIM devices. Current compilation approaches lack systematic optimization for diverse ML kernels and multiple PIM devices, and may largely ignore data rearrangement costs during the compute code optimization step. We show that data rearrangements and compute code optimization are interdependent, and need to be jointly optimized during the tuning process. Therefore, we design DCC, the first data-centric ML compiler for PIM systems that jointly co-optimizes data rearrangements and compute code in a unified tuning process. DCC integrates a multi-layer PIM abstraction to support multiple PIM backends. DCC enables effective co-optimization of data partitioning strategies with compute loop partitioning schemes. DCC applies PIM-specific code optimizations, and leverages a fast and accurate performance prediction model to select the bestperforming code schedule for a given kernel on a target PIM architecture. Our evaluations in various individual ML kernels show that DCC achieves up to 7.68x speedup (2.21x average) on HBM-PIM, and up to 13.17x speedup (3.92x average) on AttAcc PIM, over GPU-only execution. In end-to-end LLM inference, DCC on AttAcc accelerates GPT-3 and LLaMA-2 by 4.52x average (up to 7.71x in LLaMA-2) over GPU. DCC is open-sourced at https://github.com/SPIN-Research-Group/DCC.

PIM编译器大模型推理数据重排

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。