Optimizing LLM Inference: Fluid-Guided Online Scheduling with Memory Constraints
This work addresses the challenge of efficient LLM deployment for applications requiring high computational resources, offering a novel scheduling approach that bridges operations research and machine learning.
The paper tackles the problem of optimizing LLM inference under memory constraints by formulating it as a multi-stage online scheduling problem, resulting in the WAIT and Nested WAIT algorithms that improve throughput and reduce latency compared to baselines like vLLM and Sarathi in experiments with Llama-7B on an A100 GPU.
Large Language Models (LLMs) are indispensable in today's applications, but their inference procedure -- generating responses by processing text in segments and using a memory-heavy Key-Value (KV) cache -- demands significant computational resources, particularly under memory constraints. This paper formulates LLM inference optimization as a multi-stage online scheduling problem where sequential prompt arrivals and KV cache growth render conventional scheduling ineffective. We develop a fluid dynamics approximation to provide a tractable benchmark that guides algorithm design. Building on this, we propose the Waiting for Accumulated Inference Threshold (WAIT) algorithm, which uses multiple thresholds to schedule incoming prompts optimally when output lengths are known, and extend it to Nested WAIT for cases with unknown output lengths. Theoretical analysis shows that both algorithms achieve near-optimal performance against the fluid benchmark in heavy traffic conditions, balancing throughput, latency, and Time to First Token (TTFT). Experiments with the Llama-7B model on an A100 GPU using both synthetic and real-world datasets demonstrate improved throughput and reduced latency relative to established baselines like vLLM and Sarathi. This work bridges operations research and machine learning, offering a rigorous framework for the efficient deployment of LLMs under memory constraints.