arXiv:2502.04563cs.LGcs.AI2025-02被引 31

专为晶圆级芯片设计的LLM推理系统,性能超传统GPU数百倍。

WaferLLM: Large Language Model Inference at Wafer Scale

  • 基于新型PLMR模型,实现晶圆级并行计算优化
  • 在Cerebras WSE2上实现GEMV速度提升606倍、能效高16倍
  • 适合追求极致算力与能效的晶圆级AI硬件开发者

新兴AI加速器越来越多采用晶圆级制造技术,将数十万颗AI核心集成于网格架构中,配备总容量达数十GB的分布式片上内存和每秒数十PB的超高带宽。然而,现有LLM推理系统针对共享内存架构(如GPU)优化,难以充分发挥此类加速器潜力。本文提出首个晶圆级LLM推理系统WaferLLM,其基于新提出的PLMR模型(读作“Plummer”),精准捕捉晶圆级架构的独特硬件特性。借助该模型,WaferLLM开创性地实现了晶圆级并行,优化了数十万片上核心的利用率。同时,提出MeshGEMM和MeshGEMV,是首个可有效扩展至晶圆级加速器的GEMM与GEMV实现。评估表明,WaferLLM的加速器利用率比当前最优方法高出200倍。在Cerebras WSE2晶圆级加速器上,其GEMV运算速度较NVIDIA A100 GPU快606倍,能效提升16倍;全量LLM推理相比运行SGLang和vLLM的A100 GPU集群提速10-20倍。随着晶圆级AI模型、软件与硬件的持续成熟,这些优势将进一步放大。项目已开源:https://github.com/MeshInfra/WaferLLM。

原文摘要 · Abstract (English)

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.

LLM推理晶圆级高性能计算GEMM

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