R elay A ttention for Efficient Large Language Model Serving with Long System Prompts (2024)

Abstract: A practical large language model (LLM) service may involve a long system prompt, which specifies the instructions, examples, and knowledge documents of the task and is reused across requests. However, the long system prompt causes throughput/latency bottlenecks as the cost of generating the next token grows w.r.t the sequence length. This paper aims to improve the efficiency of LLM services that involve long system prompts. Our key observation is that handling these system prompts requires heavily redundant memory accesses in existing causal attention computation algorithms. Specifically, for batched requests, the cached hidden states (i.e., key-value pairs) of system prompts are transferred from off-chip DRAM to on-chip SRAM multiple times, each corresponding to an individual request. To eliminate such a redundancy, we propose RelayAttention, an attention algorithm that allows reading these hidden states from DRAM exactly once for a batch of input tokens. RelayAttention is a free lunch: it maintains the generation quality while requiring no model retraining, as it is based on a mathematical reformulation of causal attention. We have observed significant performance improvements to a production-level system, vLLM, through integration with RelayAttention. The improvements are even more profound with longer system prompts.

## introduction after around one decade of rapid development @xcite @xcite , we have experienced a revolution of large language models (llms) over the past year. llms like @xcite and gemini @xcite are so powerful that they can now serve as programming figure 1 : llama-30b attention inference latency w.r.t. system prompt length (a40 gpu, batch size 32). we set the length of (request-specific) contexts, which include user prompts and previously generated tokens, to 128. copilots @xcite , universal chatbots @xcite , computer assistants @xcite and other roles that penetrate our daily life. however, the high inference cost of these large models has become a substantial obstacle to serving more people @xcite . it is therefore important to improve the hardware utilization so that llms can have a higher throughput within a fixed hardware budget. llm services commonly use an applicationspecific system prompt @xcite to specify the task's instructions. the system prompt is concatenated with the user prompt as the full input to the llm for response generation and is shared by all requests to a service. the system prompt becomes long if the service provider wants to provide detailed guidelines and examples for better response quality or apply more restrictions/policies for ethical safety. as the sequence length that llms can process grows @xcite @xcite , some emerging professional applications, such as legal analysis @xcite ti (16gb) may be a good choice. it has 16 gb gddr6 memory, which supports running 7b llms locally in half-precision. it is now on sale at only @xmath0500. <@xmath1. for customer inquires, provide suggestions faithfully according to the documents provided here. when providing shopping suggestions, also attach links to the suggested items. available items are listed below: <@xmath2, status, price, link [item 1 specs] [item 2 specs] .... rtx 4060ti (16gb), on sale, $399, https://xxxxx/yyy .... </@xmath3: <@xmath4... </@xmath5, we use the shopping assistant as an example application. @xcite , healthcare applications @xcite , and the shopping assistant example shown in fig. 2 , may include one or more knowledge documents to provide domain-specific knowledge, resulting in even longer system prompts. although long system prompts are beneficial to improving the generation quality or enabling new applications, they also pose a challenge to the llm service: the inference throughput and latency of the service can be heavily degraded, thus increasing the per-request cost. this is inherently caused by the causal attention, in which each new token is generated by "looking at" all precedent ones. in this paper, we propose a novel approach to mitigate the efficiency problem of using long system prompts in llm services. our key observation is that there are not only redundant memory footprint @xcite and computations @xcite corresponding to the system prompt, but also unnecessary memory accesses during causal attention computation. specifically, while the system prompt is shared by all requests, its hidden states (i.e., key-value pairs) are read from dram multiple times by existing attention algorithms such as pagedattention @xcite and flashattention @xcite , each for an individual request in the batch. this severely slows down llm inferences, which are known to be memory-bound (section 3.2). to eliminate such redundant memory access, we pro-pose relayattention, an exact algorithm to compute causal attention based on a mathematical reformulation of it. the key idea of relayattention is to group the matrix-vector multiplications corresponding to the system prompt into matrix-matrix multiplications, which allow loading the hidden states of the system prompt from dram exactly once for all request tokens in a batch (section 3.3). as a result, the attention inference latency grows much slower than pagedattention w.r.t. the length of system prompt, as shown in fig. 1 . we provide an in-depth analysis of the theoretic speedup of the standalone attention based on the io redundancy reduction (section 3.4). our empirical results for end-to-end serving further verify the efficiency: integrating relayattention into vllm @xcite , an already highly optimized productionlevel llm serving system, we still observe up to 2.2× sustainable request rate and 2.0× throughput with the llama2-7b model for a chatbot workload. similar efficiency improvements are also observed for several other popular llms and are consistent on several data center gpus. the efficiency gains continue growing with longer system prompts. our key contributions can be summarized as: • we have identified a llm service bottleneck that has not been studied by existing works: there are highly redundant memory accesses caused by long system prompts. we anticipate that our analysis will inspire more works on deep architectures with io-awareness @xcite . ## related works our approach aims to improve the inference efficiency of transformer-based llms (section 2.1). it is based on extending the widely used key-value cache mechanism (section 2.2). we also briefly review other techniques for accelerating llm inference, which may complement ours (section 2.3). ## methodology in this section, we elaborate on the proposed approach. we begin with a brief preliminary of the hardware utilization of operators in section 3.1, followed by an analysis of the bottleneck in llm serving in section 3.2, which shows that the redundant memory access slows down the inference especially when the system prompt is long. we then introduce relayattention, a novel algorithm to compute exact causal attention that allows the elimination of the redundancy in section 3.3. finally, we analyze the theoretical acceleration of relayattention over existing approaches from the perspective of io-awareness (section 3.4). ## experiments in this section, we conduct experiments to answer the question of how much our approach can help an end-to-end llm serving system. we provide the experimental setup in section 4.1. our major evaluation is conducted with consideration of two scenarios: noninteractive batch processing (section 4.2) and interactive service (section 4.3). we use the llama2-7b model @xcite for evaluation unless stated otherwise. we demonstrate the improvement for more models in section 4.4. ## conclusion in this paper, we have identified a bottleneck of using long system prompts in llm services: there are highly redundant memory accesses corresponding to those system kvs. we have proposed relay-attention to compute exact causal attention while removing the redundant memory accesses. an analysis of the theoretical speedup of relayattention is provided. extensive experiments over different gpus, models, and datasets empirically verify the efficiency gains brought by relayattention.