MIST: A Co-Design Framework for Heterogeneous, Multi-Stage LLM Inference

Fuente: arXiv
Salvato in:
Dettagli Bibliografici
Autori principali: Bambhaniya, Abhimanyu Rajeshkumar, Wu, Hanjiang, Subramanian, Suvinay, Srinivasan, Sudarshan, Kundu, Souvik, Yazdanbakhsh, Amir, Elavazhagan, Midhilesh, Kumar, Madhu, Yu, Minlan, Raychowdhury, Arijit, Krishna, Tushar
Natura: Preprint
Pubblicazione: 2025
Soggetti:
Accesso online:
Tags: Aggiungi Tag
Nessun Tag, puoi essere il primo ad aggiungerne!!
_version_ 1866914493737467904
author Bambhaniya, Abhimanyu Rajeshkumar
Wu, Hanjiang
Subramanian, Suvinay
Srinivasan, Sudarshan
Kundu, Souvik
Yazdanbakhsh, Amir
Elavazhagan, Midhilesh
Kumar, Madhu
Yu, Minlan
Raychowdhury, Arijit
Krishna, Tushar
author_facet Bambhaniya, Abhimanyu Rajeshkumar
Wu, Hanjiang
Subramanian, Suvinay
Srinivasan, Sudarshan
Kundu, Souvik
Yazdanbakhsh, Amir
Elavazhagan, Midhilesh
Kumar, Madhu
Yu, Minlan
Raychowdhury, Arijit
Krishna, Tushar
contents Modern LLM serving now spans multi-stage pipelines including RAG retrieval and KV cache reuse, each with distinct compute, memory, and latency demands. Inference engines expose a large configuration space with no systematic navigation methodology, and exhaustively benchmarking configurations can exceed 40K in cloud costs. Simultaneously, the hardware landscape is rapidly diversifying across AMD GPUs, TPUs, and custom ASICs, while cross-vendor prefill-decode (PD) disaggregated configurations lack unified software stacks for end-to-end evaluation today. To address this gap, we present MIST, a Heterogeneous Multi-stage LLM inference Execution Simulator. MIST models diverse request stages; including RAG, KV retrieval, reasoning, prefill, and decode across complex hardware hierarchies. MIST 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, MIST 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. MIST 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.
format Preprint
id arxiv_https___arxiv_org_abs_2504_09775
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle MIST: A Co-Design Framework for Heterogeneous, Multi-Stage LLM Inference
Bambhaniya, Abhimanyu Rajeshkumar
Wu, Hanjiang
Subramanian, Suvinay
Srinivasan, Sudarshan
Kundu, Souvik
Yazdanbakhsh, Amir
Elavazhagan, Midhilesh
Kumar, Madhu
Yu, Minlan
Raychowdhury, Arijit
Krishna, Tushar
Hardware Architecture
Artificial Intelligence
Distributed, Parallel, and Cluster Computing
Machine Learning
Modern LLM serving now spans multi-stage pipelines including RAG retrieval and KV cache reuse, each with distinct compute, memory, and latency demands. Inference engines expose a large configuration space with no systematic navigation methodology, and exhaustively benchmarking configurations can exceed 40K in cloud costs. Simultaneously, the hardware landscape is rapidly diversifying across AMD GPUs, TPUs, and custom ASICs, while cross-vendor prefill-decode (PD) disaggregated configurations lack unified software stacks for end-to-end evaluation today. To address this gap, we present MIST, a Heterogeneous Multi-stage LLM inference Execution Simulator. MIST models diverse request stages; including RAG, KV retrieval, reasoning, prefill, and decode across complex hardware hierarchies. MIST 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, MIST 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. MIST 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.
title MIST: A Co-Design Framework for Heterogeneous, Multi-Stage LLM Inference
topic Hardware Architecture
Artificial Intelligence
Distributed, Parallel, and Cluster Computing
Machine Learning
url https://arxiv.org/abs/2504.09775