Understanding and Optimizing Multi-Stage AI Inference Pipelines

TL;DR

MIST presents an event-driven, end-to-end simulator for multi-stage AI inference that jointly models diverse pipeline stages (RAG, KV retrieval, reasoning, prefill, decode) across heterogeneous hardware. By integrating workload traces, routing policies, batching strategies, and multi-level memory hierarchies, it achieves end-to-end fidelity within $6\%$ of real-system performance and enables principled exploration of deployment and architectural decisions, including significant tokens-per-dollar gains in representative workloads (e.g., up to $4.7\times$ in some cases). The framework’s modular HW/SW abstractions and validation against real hardware establish a practical foundation for hardware-software co-design and future hardware design exploration for LLM inference. Overall, MIST enables cost-efficient optimization of current deployments and systematic design of next-generation AI infrastructure, reducing trial-and-error in hardware selection, batching strategies, and KV-cache architectures.

Abstract

The rapid evolution of Large Language Models (LLMs) has driven the need for increasingly sophisticated inference pipelines and hardware platforms. Modern LLM serving extends beyond traditional prefill-decode workflows, incorporating multi-stage processes such as Retrieval Augmented Generation (RAG), key-value (KV) cache retrieval, dynamic model routing, and multi step reasoning. These stages exhibit diverse computational demands, requiring distributed systems that integrate GPUs, ASICs, CPUs, and memory-centric architectures. However, existing simulators lack the fidelity to model these heterogeneous, multi-engine workflows, limiting their ability to inform architectural decisions. To address this gap, we introduce HERMES, a Heterogeneous Multi-stage LLM inference Execution Simulator. HERMES models diverse request stages; including RAG, KV retrieval, reasoning, prefill, and decode across complex hardware hierarchies. HERMES supports heterogeneous clients executing multiple models concurrently unlike prior frameworks while incorporating advanced batching strategies and multi-level memory hierarchies. By integrating real hardware traces with analytical modeling, HERMES captures critical trade-offs such as memory bandwidth contention, inter-cluster communication latency, and batching efficiency in hybrid CPU-accelerator deployments. Through case studies, we explore the impact of reasoning stages on end-to-end latency, optimal batching strategies for hybrid pipelines, and the architectural implications of remote KV cache retrieval. HERMES empowers system designers to navigate the evolving landscape of LLM inference, providing actionable insights into optimizing hardware-software co-design for next-generation AI workloads.

Figures (13)

• Figure 1: (a) LLM inference request types: Question-answering (Standard); News search (RAG pipeline) gao2024ragsurvey; Code generation (KV cache reuse) jiang2024codegen; Chat support (RAG + KV cache) dam2024chatbotsurvey; and Reasoning Math (Multi-turn reasoning + Reward Model) ahn2024llmmaths. (b) Scheduling three News Search requests across 1 CPU and 3 GPU nodes.
• Figure 2: Understanding the LLM inference serving stack and how MIST models the stack. MIST simulates collection of clients. Each client consist of a scheduler which issues steps(e.g. tokens for decode, chunks for prefill, Rerank/Retreival for RAG) to a HW cluster (e.g., Nvidia HGX, AMD MI300X, CPU host with offloading memory instance, ..). The HW cluster is a collection of multiple HW Nodes (e.g. NPUs, memory and CPUs).
• Figure 3: Batching mechanisms and their latency impact on the prefill and decode phases.
• Figure 4: Steps in (a) RAG & (b) KV Retrieval.
• Figure 5: End-to-end runtime comparison of vLLM real HW runtime and MIST simulated runtime for different parallelization with HGX:H100x8 running Llama3.1-70B. For each hardware configuration, we vary the context length, number of requests, and chunk size.