WaferLLM: Large Language Model Inference at Wafer Scale

2025/02/06 by He, Congjie, Huang, Yeqi, Mu, Pei +5 · 12 citations
#Artificial Intelligence (cs.AI) #Distributed #Emerging Technologies (cs.ET) #FOS: Computer and information sciences #Hardware Architecture (cs.AR) #Machine Learning (cs.LG) #Parallel #and Cluster Computing (cs.DC)

paper · doi:10.48550/arxiv.2502.04563

Abstract

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× higher accelerator utilization than state-of-the-art methods. Leveraging a wafer-scale accelerator (Cerebras WSE2), WaferLLM delivers GEMV operations 606× faster and 16× more energy-efficient than on an NVIDIA A100 GPU. For full LLM inference, WaferLLM achieves 10-20× 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.

Cited by

• PADE: A Predictor-Free Sparse Attention Accelerator via Unified Execution and Stage Fusion
• TEMP: A Memory Efficient Physical-aware Tensor Partition-Mapping Framework on Wafer-scale Chips
• Design in Tiles: Automating GEMM Deployment on Tile-Based Many-PE Accelerators
• FlashFuser: Expanding the Scale of Kernel Fusion for Compute-Intensive Operators via Inter-Core Connection
• WATOS: Efficient LLM Training Strategies and Architecture Co-exploration for Wafer-scale Chip
• ContextPilot: Fast Long-Context Inference via Context Reuse
• Agentic AI Home Energy Management System: A Large Language Model Framework for Residential Load Scheduling
• From Principles to Practice: A Systematic Study of LLM Serving on Multi-core NPUs
• Patterns behind Chaos: Forecasting Data Movement for Efficient Large-Scale MoE LLM Inference
• Are We Scaling the Right Thing? A System Perspective on Test-Time Scaling
• SmartGen: Seamless Disaggregated LLM Inference with Selective KV Cache Transfer
• PCCL: Photonic circuit-switched collective communication for distributed ML

Related

• AI Accelerators for Large Language Model Inference: Architecture Analysis and Scaling Strategies
• SAIL: SRAM-Accelerated LLM Inference System with Lookup-Table-based GEMV
• LLM-Inference-Bench: Inference Benchmarking of Large Language Models on AI Accelerators
• HPIM: Heterogeneous Processing-In-Memory-based Accelerator for Large Language Models Inference
• A Hardware Evaluation Framework for Large Language Model Inference