Fast Inference for Augmented Large Language Models

TL;DR

LAMPS minimizes request completion time through a unified scheduling approach that considers the total length of requests and their handling strategies during API calls, and predicts the strategy that minimizes memory waste of a request during its API calls, aligning with but improving upon existing approaches.

Abstract

Augmented Large Language Models (LLMs) enhance the capabilities of standalone LLMs by integrating external data sources through API calls. In interactive LLM applications, efficient scheduling is crucial for maintaining low request completion times, directly impacting user engagement. However, these augmentations introduce scheduling challenges due to the need to manage limited memory for cached information (KV caches). As a result, traditional size-based scheduling algorithms, such as Shortest Job First (SJF), become less effective at minimizing completion times. Existing work focuses only on handling requests during API calls by preserving, discarding, or swapping memory without considering how to schedule requests with API calls. In this paper, we propose LAMPS, a novel LLM inference framework for augmented LLMs. LAMPS minimizes request completion time through a unified scheduling approach that considers the total length of requests and their handling strategies during API calls. Recognizing that LLM inference is memory-bound, our approach ranks requests based on their consumption of memory over time, which depends on both the output sizes and how a request is managed during its API calls. To implement our scheduling, LAMPS predicts the strategy that minimizes memory waste of a request during its API calls, aligning with but improving upon existing approaches. We also propose starvation prevention techniques and optimizations to mitigate the overhead of our scheduling. We implement LAMPS on top of vLLM and evaluate its performance against baseline LLM inference systems, demonstrating improvements in end-to-end latency by 27%-85% and reductions in TTFT by 4%-96% compared to the existing augmented-LLM system, with even greater gains over vLLM.

Figures (11)

• Figure 1: Illustration of an augmented-LLM request. The API fetches detailed information about the 2024 Nobel Prize.
• Figure 2: Impact of including API calls: using a subset of INFERCEPT abhyankarinfercept, we compare two variations of the dataset—one with API calls and one without. (a) KV cache usage (%) over time when all API calls are handled using Preserve. (b) Number of completed requests over time using Preserve. (c) Number of completed requests over time using Discard.
• Figure 3: Comparison of scheduling policies for API-augmented requests with a memory budget of $6$ unit, requests lengths and API duration are summarized in Table \ref{['tab:example_durations']}. (a) FCFS: yielding an average request completion time of 11.66 units. (b) SJF: achieving an average request completion time of 10.33 units. (c) SJF by Total Length: Orders requests by total length, achieving an average request completion time of 11 units (d) Optimized: Integrates length and API handling, achieving an average request completion time of 10 units.
• Figure 4: Memory consumption over time for a request with one API call using three memory management strategies: (1) Preserve, (2) Discard and Recompute, and (3) Swap. The highlighted area represents memory waste for one request.
• Figure 5: LAMPS architecture. LAMPS minimizes completion time for API-augmented requests through three steps: predicting pre-API output length and API properties, determining request handling strategies (preserve, discard, or swap) to minimize memory waste, and implementing a scheduling policy based on the handling method and request output length.