WaferLLM: Large Language Model Inference at Wafer Scale

📄 arXiv: 2502.04563v3 📥 PDF

作者: Congjie He, Yeqi Huang, Pei Mu, Ziming Miao, Jilong Xue, Lingxiao Ma, Fan Yang, Luo Mai

分类: cs.LG, cs.AI, cs.AR, cs.DC, cs.ET

发布日期: 2025-02-06 (更新: 2025-05-30)

🔗 代码/项目: GITHUB

💡 一句话要点

WaferLLM：晶圆级大语言模型推理系统，充分利用晶圆级加速器的算力。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 晶圆级加速器 大语言模型推理 并行计算 PLMR模型 MeshGEMM MeshGEMV 高性能计算 AI加速

📋 核心要点

  1. 现有LLM推理系统针对GPU等共享内存架构优化，无法充分利用晶圆级加速器的大规模并行计算能力和片上内存带宽。
  2. WaferLLM基于PLMR模型，针对晶圆级架构特性进行优化，实现了晶圆级LLM并行，并设计了MeshGEMM和MeshGEMV。
  3. 实验表明，WaferLLM在加速器利用率、GEMV运算速度和能效方面均显著优于现有方法，并实现了LLM推理的加速。

📝 摘要（中文）

本文提出了WaferLLM，首个晶圆级大语言模型推理系统。该系统基于一种名为PLMR的新型模型，该模型能够捕获晶圆级架构独特的硬件特性。WaferLLM率先实现了晶圆级LLM并行，优化了数十万片上核心的利用率。此外，WaferLLM还引入了MeshGEMM和MeshGEMV，这是首个专为在晶圆级加速器上有效扩展而设计的GEMM和GEMV实现。评估结果表明，WaferLLM的加速器利用率比最先进的方法高出200倍。在晶圆级加速器（Cerebras WSE2）上，WaferLLM的GEMV运算速度比NVIDIA A100 GPU快606倍，能效高16倍。对于完整的LLM推理，WaferLLM比运行SGLang和vLLM的A100 GPU集群实现了10-20倍的加速。随着晶圆级AI模型、软件和硬件的不断成熟，这些优势预计将继续增长。WaferLLM已在https://github.com/MeshInfra/WaferLLM开源。

🔬 方法详解

问题定义：现有的大语言模型推理系统主要针对GPU等共享内存架构设计，无法有效利用晶圆级加速器所提供的数十万个核心、数十GB的片上内存以及PB/s级别的片上内存带宽。这导致晶圆级加速器的利用率低下，无法充分发挥其潜力。

核心思路：WaferLLM的核心思路是针对晶圆级加速器的硬件特性，设计一种新的并行计算模型和优化算法，从而最大化利用晶圆级加速器的计算资源和内存带宽。通过将LLM推理任务分解成多个子任务，并将这些子任务分配到晶圆上的各个核心进行并行计算，从而实现加速。

技术框架：WaferLLM的整体框架包括PLMR模型、晶圆级LLM并行策略、MeshGEMM和MeshGEMV等关键组件。PLMR模型用于描述晶圆级加速器的硬件特性，指导并行策略的设计。晶圆级LLM并行策略负责将LLM推理任务分解成多个子任务，并分配到晶圆上的各个核心。MeshGEMM和MeshGEMV是针对晶圆级架构优化的GEMM和GEMV实现，用于加速矩阵乘法等核心计算。

关键创新：WaferLLM的关键创新在于：1) 提出了PLMR模型，能够准确描述晶圆级加速器的硬件特性；2) 设计了晶圆级LLM并行策略，能够充分利用晶圆级加速器的计算资源；3) 实现了MeshGEMM和MeshGEMV，能够高效地执行矩阵乘法等核心计算。与现有方法相比，WaferLLM能够显著提高加速器利用率和推理速度。

关键设计：PLMR模型的具体参数设置未知。MeshGEMM和MeshGEMV的关键设计在于如何将矩阵乘法任务分解成多个子任务，并分配到晶圆上的各个核心进行并行计算，同时最小化数据传输开销。具体的分解和分配策略以及数据传输优化方法未知。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

WaferLLM在Cerebras WSE2晶圆级加速器上进行了评估，结果表明其加速器利用率比现有方法高出200倍。GEMV运算速度比NVIDIA A100 GPU快606倍，能效高16倍。对于完整的LLM推理，WaferLLM比运行SGLang和vLLM的A100 GPU集群实现了10-20倍的加速。这些结果表明WaferLLM在晶圆级加速器上具有显著的性能优势。

🎯 应用场景

WaferLLM的潜在应用领域包括大规模语言模型的快速推理、AI驱动的科学计算、以及需要高性能计算的各种应用。通过充分利用晶圆级加速器的算力，WaferLLM可以加速这些应用的开发和部署，并为未来的AI发展提供更强大的计算平台。

📄 摘要（原文）

Emerging AI accelerators increasingly adopt wafer-scale manufacturing technologies, integrating hundreds of thousands of AI cores in a mesh architecture with large distributed on-chip memory (tens of GB in total) and ultra-high on-chip memory bandwidth (tens of PB/s). However, current LLM inference systems, optimized for shared memory architectures like GPUs, fail to exploit these accelerators fully. We introduce WaferLLM, the first wafer-scale LLM inference system. WaferLLM is guided by a novel PLMR model (pronounced as "Plummer") that captures the unique hardware characteristics of wafer-scale architectures. Leveraging this model, WaferLLM pioneers wafer-scale LLM parallelism, optimizing the utilization of hundreds of thousands of on-chip cores. It also introduces MeshGEMM and MeshGEMV, the first GEMM and GEMV implementations designed to scale effectively on wafer-scale accelerators. Evaluations show that WaferLLM achieves up to 200$\times$ higher accelerator utilization than state-of-the-art methods. Leveraging a wafer-scale accelerator (Cerebras WSE2), WaferLLM delivers GEMV operations 606$\times$ faster and 16$\times$ more energy-efficient than on an NVIDIA A100 GPU. For full LLM inference, WaferLLM achieves 10-20$\times$ speedups over A100 GPU clusters running SGLang and vLLM. These advantages are expected to grow as wafer-scale AI models, software, and hardware continue to mature. WaferLLM is open-sourced at https://github.com/MeshInfra/WaferLLM.