Demystifying AI Platform Design for Distributed Inference of Next-Generation LLM models

TL;DR

GenZ introduces an analytical framework that links LLM architectures, inference optimizations, and hardware platforms to end-to-end distributed inference performance. It combines a model profiler, NPU characterizer, and platform characterizer to predict compute, memory, and interconnect requirements, and validates predictions against multiple real hardware platforms with a geomean error of $5.82\%$. Through four case studies—from scaling HW characteristics to extreme-scale AI assistants—GenZ provides actionable insights into platform design choices, network topologies, and microarchitectural tradeoffs for next-generation LLM inference. The work offers open-source tooling and a web app to support design-space exploration, aiming to guide AI engineers and computer architects in building efficient, scalable AI hardware platforms. Overall, GenZ fills a critical gap by enabling rapid, high-fidelity exploration of how model architectures and software optimizations interact with distributed hardware to meet SLOs for diverse LLM use cases.

Abstract

Large language models (LLMs) have shown remarkable performance across a wide range of applications, often outperforming human experts. However, deploying these gigantic models efficiently for diverse inference use cases requires carefully designed hardware platforms with ample computing, memory, and network resources. With constant innovation in LLM serving optimizations and model architecture evolving at breakneck speed, the hardware requirements to meet Service Level Objectives (SLOs) remain an open research question. To answer the question, we present an analytical tool, GenZ, to efficiently navigate the relationship between diverse LLM model architectures(Dense, GQA, MoE, Mamba), LLM serving optimizations(Chunking, Speculative decoding, quanitization), and AI platform design parameters. Our tool estimates LLM inference performance metrics for the given scenario. We have validated against real hardware platforms running various different LLM models, achieving a max geomean error of 5.82.We use GenZ to identify compute, memory capacity, memory bandwidth, network latency, and network bandwidth requirements across diverse LLM inference use cases. We also study diverse architectural choices in use today (inspired by LLM serving platforms from several vendors) to help inform computer architects designing next-generation AI hardware accelerators and platforms. The trends and insights derived from GenZ can guide AI engineers deploying LLMs as well as computer architects designing next-generation hardware accelerators and platforms. Ultimately, this work sheds light on the platform design considerations for unlocking the full potential of large language models across a spectrum of applications. The source code is available at https://github.com/abhibambhaniya/GenZ-LLM-Analyzer . Users can also be tried it on at https://genz-llm-analyzer.streamlit.app/ without any setup on your web browser.

Figures (20)

• Figure 1: Platform requirements for two workloads.
• Figure 2: An overview of GenZ framework.
• Figure 3: Typical LLM Model Architecture. Each Layer has multiple parallel heads. For MoE models, there are multiple parallel MLP layers out of which 'k' are activated.
• Figure 4: Various parallelization strategies for neural network training and inference. Each colored box represents an accelerator (NPU), and the numbers correspond to model layers.
• Figure 5: GenZ LLM inference modeling against vLLM inference in prefill and decode stage for varying batch size on various platforms with different $\tau_{p}$. Line color blue, red, and green represent LLaMA2-7B, LLaMA2-13B, and OPT-175B.