Challenges in Deploying Long-Context Transformers: A Theoretical Peak Performance Analysis

TL;DR

The paper tackles the high cost of deploying long-context transformers and identifies KV-cache size as the core bottleneck limiting concurrency and inflating prefilling and decoding costs. It introduces a concurrent programming framework to analyze end-to-end session throughput under limited GPU HBM, using a 34B Yi model with 50K context as a running example, and examines how hardware, architecture, and KV-cache compression influence performance. Key contributions include a formal decomposition into four metrics (concurrency, prefilling, decoding, context switching), an analysis of factors that affect them (context length, tensor parallelism, MoE, attention type), and a survey of lossless KV-cache compression strategies by layer, head, token, and hidden dimensions with ideas for end-to-end integration. The work aims to lay the foundation for full-stack optimizations that could reduce 1M-context inference costs to resemble 4K-context costs, enabling broader delivery of long-context AI capabilities.

Abstract

Transformer-based long context generative models power emerging AI applications like hour-long video understanding and project-level coding agent. Deploying long context transformers (e.g., 100K to 10M tokens) is prohibitively expensive compared to short context (e.g., 4K tokens) model variants. Reducing the cost of long-context transformers is becoming a pressing research and engineering challenge starting from the year of 2024. This work describes a concurrent programming framework for quantitatively analyzing the efficiency challenges in serving multiple long-context requests under limited size of GPU high-bandwidth memory (HBM) regime. We give a detailed analysis of how all additional computational costs, compared to 4K context, trace back to \textit{one single source: the large size of the KV cache}. We use a 34B GPT-3.5 level model of 50K context on A100 NVLink as a running example, and describe how its large KV cache causes four types of deployment challenges: (1) prefilling long inputs takes much longer compute time and GPU memory than short inputs; (2) after prefilling, the large KV cache residing on the GPU HBM substantially restricts the number of concurrent users being served; (3) during decoding, repeatedly reading the KV cache from HBM to SM largely increases latency; (4) when KV cache memory overflows, swapping it from HBM to DDR causes significant context switching latency. We use this framework to analyze existing works and identify possibilities of combining them to build end-to-end systems. Overall, this work offers a foundational framework for analyzing long context transformer deployment and identifies directions towards reducing the inference cost of 1M context to be as cheap as 4K.

Figures (3)

• Figure 1: A concurrent programming framework for serving multiple long-context user requests under limited GPU HBM size. There are four key factors collectively determine the overall throughput of user interaction sessions: (1) concurrency is bounded by the HBM size: the number of concurrent user being served is bounded by the size of the GPU high-bandwidth memory (HBM); (2) prefilling is compute bound: the latency to the first generated token, i.e., the prompt prefilling time, is bounded by the floating point operation per second (flops) of the GPU; (3) decoding is memory bound (under critical batch size): the latency (generated token per second) of autoregressive decoding is bounded by the bandwidth of the HBM; (4) context switching is PCIE bound: offloading user 1's KV cache to the CPU DDR and loading user 2's KV cache to the HBM is bounded by the PCIE bandwidth. Compared to the short-context setting, concurrency and context-switching challenges are much more severe under a long-context setting. All efficiency challenges from these four key factors eventually trace back to the size of the KV cache.
• Figure 2: First row: how context length changes the four key performance metrics. Increasing the context length from 4K to 50K inversely reduces concurrency, quadratically increases prefilling latency, linearly increases context switching overhead, and slightly (but also linearly) increases decoding latency. Second row: how different generations of hardware influence the metrics.
• Figure 3: Should one target prefilling or decoding? Here we show how Yi 34B young2024yi (about GPT-3.5 level) and Command R+ cohere2024command (about GPT-4 level)'s latency of prefilling and decoding (measured in seconds) on different input length and conversation rounds. The larger the model, the longer the input context, the better gain from the improved prefilling (e.g., using a 104B model to answer 5 questions of a book); the smaller the model, the more followup questions, the better gain from improved decoding (e.g., talking to a 34B AI companion for 2 hours of 100+ dialog rounds). Also note that when the context is less than 50K, changing the attention from full to linear (as many works count on linear attention to improve efficiency) does not significantly improve the overall performance.