RelayAttention 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 productionlevel system, vLLM, through integration with RelayAttention. The improvements are even more profound with longer system prompts.

text: Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 4945–4957 August 11-16, 2024 ©2024 Association for Computational Linguistics
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 productionlevel system, vLLM, through integration with RelayAttention. The improvements are even more profound with longer system prompts.

heading: 1 Introduction
text: After around one decade of rapid development (Sutskever et al., 2014; Vaswani et al., 2017; Radford et al., 2018; OpenAI, 2023b), we have experienced a revolution of large language models (LLMs) over the past year. LLMs like GPT-4 (OpenAI, 2023b) and Gemini (Google, 2023b) are so powerful that they can now serve as programming
† Corresponding author.
copilots (Chen et al., 2021; GitHub, 2022), universal chatbots (Google, 2023a; OpenAI, 2022), computer assistants (Microsoft, 2023a) 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 (Kwon et al., 2023). 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 (OpenAI, 2023a) 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 (Anthropic, 2023; Chen et al., 2023b; Bi et al., 2024), some emerging professional applications, such as legal analysis (Cui
4945
et al., 2023; Nay et al., 2023), healthcare applications (Steinberg et al., 2021; Rasmy et al., 2021), 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 (Kwon et al., 2023) and computations (Gim et al., 2023) 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 (Kwon et al., 2023) and FlashAttention (Dao et al., 2022; Dao, 2023), 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 (Kwon et al., 2023), 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 (Dao et al., 2022; Gu and Dao, 2023). • We propose RelayAttention, a novel approach to compute exact causal attention. It allows accessing cached hidden states of the system prompt exactly once for a batch of request tokens. We conduct an in-depth analysis of the theoretic speedup brought by RelayAttention. • We empirically verify the end-to-end efficiency improvement by integrating RelayAttention into vLLM, a production level LLM serving system, and observe non-trivial efficiency gains on several popular LLMs with different GPUs.

heading: 2 Related Works
text: 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).

heading: 2.1 Inference of Transformer-based LLMs
text: The inference of these transformer-based LLMs follows the iterative next-token-prediction paradigm. Specifically, the next token is generated in each time step by attending to all precedent tokens. The generated token is then appended to the end of the current sequence. The generation then continues until a stopping criterion (e.g., the new token is <eos>, which indicates the end of the sequence) is met. A basic approach to implementing such a generation procedure is to perform full self-attention with a casual mask over the entire up-to-present sequence at each time step, just as we do while training the model (Radford et al., 2018). This way, a single generation step takes a quadratic complexity w.r.t. the length of the up-to-present sequence. Next, we will look at how this complexity can be reduced to linear using the Key-Value Cache.

heading: 2.2 Key-Value Cache
text: Based on the observation that historical tokens are not affected by the future ones during LLM decoding, Key-Value (KV) Cache avoids repetitive computation of the hidden key-value pairs (KVs) by caching them on the fly and then reusing the cached KVs in every subsequent steps (Yu et al., 2022; Pope et al., 2023). With KV Cache, in each time step, only a single token (i.e., the latest generated one) is used as the query, and the next token is produced by attending to the cached KVs. The generation complexity thus reduces from quadratic to linear w.r.t. the up-to-date sequence length.
Some recent research further accelerates LLM inferences by pruning superfluous KV cache data (Zhang et al., 2023) or compressing it (Liu et al., 2023) to reduce key-value pairs to be cached. However, these approaches introduce algebraic discrepancies between model training and inference. Hence, they may hurt the generation quality and/or require extra finetuning efforts. In contrast, our approach maintains generation quality and is plugand-play, as it is based on a mathematical reformulation of causal attention. The acceleration of our approach comes from reducing redundant memory access of the KV cache. Therefore, it is orthogonal and complementary to prefix sharing in PagedAttention (Kwon et al., 2023), which eliminates redundant memory footprint of system prompts, and is unlike PromptCache (Gim et al., 2023), which eliminates the redundant computation of the reusable prefix KVs and thus only acceler-
ates the prompt phase (Section 3.2).

heading: 2.3 Other Optimizations for LLM Inference
text: Besides the KV Cache, several other techniques optimize LLM inference in a post-training manner. For example, network quantization techniques can also be applied to LLMs as they are architectureagnostic, even though they may need some adaptations to improve the generation stability and quality (Frantar et al., 2022; Xiao et al., 2023; Lin et al., 2023). FlashAttention (Dao et al., 2022; Dao, 2023) is another technique to optimize LLMs’ throughputs on GPUs by avoiding redundant write/read of attention probability matrix into/from DRAM. A production-level LLM serving system may also include continuous batching (Yu et al., 2022), which enables iteration-level scheduling of requests, and speculative sampling (Chen et al., 2023a; Leviathan et al., 2023), which uses a smaller model to generate a draft and then uses the large model to check and correct it. Our approach can work together with these components with no conflicts.

heading: 3 Methodology
text: 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).

heading: 3.1 Preliminary: The Latency of Operators
text: To increase the utilization of arithmetic units, modern processors use pipelining to allow concurrent memory access and computation. For a perfectly parallelized operator, which maximizes the overlap of data transfer and computation, the runtime latency is determined by the larger one between total memory access time and total computation time. Given a processor that takes tm for per-byte access, and tc for a floating operation on average, the ratio r of the total computation time over the total memory access time for an operator is:
r = tc ×#floating operations tm ×#byte access = I × tc tm , (1)
where I is the arithmetic intensity of the operator:
I = #floating operations
#byte access . (2)
When I < tmtc , r is less than 1, the operator is memory-bound. This means that the bottleneck of the operator is memory access, and we can accelerate it only if we can reduce the memory access time. The speed of modern GPUs far outpaces the bandwidth of its memory (i.e., tctm 1 ), and thus it typically requires a high arithmetic intensity to achieve full utilization of the computing capability (e.g., A100-SXM4 GPU requires at least 38.2).
For a half-precision (2 bytes/element) general matrix multiplication (GEMM) of problem size (m,n, k): C = ABT , where C ∈ Rm×n, A ∈ Rm×k, B ∈ Rn×k, the arithmetic intensity is:
Igemm = 2mnk
2(mk + nk +mn) < min{m,n, k}.
(3) When m,n, k are all large (e.g., > 128), the operation can saturate the utilization of the computing capability due to high arithmetic intensity. This is normally true for linear projection operations in LLM inference, where m is the number of tokens in a batch, k is the input hidden dimension, and n is the output hidden dimension. However, as a special case of GEMMs, the general matrix-vector product (GEMV) operation, in which there is a vector in A and B, is always memory-bound as Igemv < 1. This is the case for casual attention computation with cached KVs, as we will show in Section 3.2.

heading: 3.2 Bottleneck of LLM Services
text: Given a batch of user prompts, the LLM inference is usually divided into two phases: the prompt
phase, which computes the hidden states of the full prompts (i.e., the concatenation of system prompt and user prompts) and generates the first new tokens; and the autoregressive generation phase, which generates all subsequent tokens sequentially, one token for each request at a time step. In this work, we focus our investigation on the autoregressive generation phase as it contains the hot spot of response generation1.
In Fig. 3, we demonstrate a time step during the autoregressive generation phase, with the batch size assumed to be 2. There are two key observations: 1. The computation of attention is memory-
bound. This is because the attention computation for a request mainly involves two GEMVs (red dashed boxes in Fig. 3), with an arithmetic intensity lower than 1. It thus requires memory access reduction for acceleration. 2. There are redundant memory accesses in the typical scenarios where a shared system prompt is prepended to request-specific contexts. Specifically, the cached key-value pairs of the shared system prompt (system KVs) are read from off-chip DRAM multiple times, each for a request in the batch (red arrows in Fig. 3). Such redundancy becomes a substantial overhead when the system prompt is long. Section 3.3 proposes the core design of RelayAttention for removing the redundant memory access.

heading: 3.3 LLM Serving with RelayAttention
text: The key idea of RelayAttention is to group multiple matrix-vector multiplications between the batched queries and the cached KVs into single matrix-
1Besides, the prompt phase can effectively saturate GPU utilization as it involves large matrix multiplications.
matrix multiplications, as shown in Fig. 4, allowing system KVs to be read from DRAM exactly once per batch. Algorithm 1 summarizes the algorithm in Pytorch-like pseudo code. It divides the computation of a causal attention layer into three steps: system attention step, context attention step, and relay fusion step. In the system attention and context attention steps, we compute two intermediate attention outputs as if the LLM is prompted by the shared system prompt / request-specific context only. In the relay fusion step, we compute the final output as a convex combination of the two intermediate outputs. Next, we show that RelayAttention is computing a mathematical reformulation of casual attention.
Without loss of generality, we consider a single sequence in the batch and a single attention head. Formally, given an on-the-fly sequence R at generation step t, we divide it into three segments (in order): (1) the system prompt of length s, (2) the user prompt of length u, and (3) the response generated by the LLM of length t− 1. Let ki,vi ∈ Rd denote the hidden key, value embedding of the token at position i ≤ l = s + u + t, and qt ∈ Rd denotes the hidden query embedding in the current step. The casual attention output ot is defined as:
ot = Attention(qt, {ki}li=1, {vi}li=1)
= l∑
j=1
exp(qtkTj )
σ1→lt vj ,
(4)
where σb→et = ∑e j=b exp(qtk T j ) is the sum-exp between the start position b and end position e > b, associated with qt. By splitting the summation in Eq. (4) at position s, which is the end system
Algorithm 1 Pseudocode for RelayAttention.
# INPUT: # q: query tensor for new inputs , (b, m, h, d) # k: key tensor for new inputs , (b, m, h, d) # v: value tensor for new inputs , (b, m, h, d) # kv_cache: context KVs , (N, 2, b, l-s, h, d) # sys_kv_cache: sys. KVs , (N, 2, 1, s, h, d) # layer_id: the index of current layer , int # l_cache: the length of cached key -value , int # OUTPUT: # o: the output of causal attention
# note: (1) we modified the interface of multi -head # attention to return the log -sum -exp (lse); # (2) the order of context attention and system # attention doesn't matter because of no dependency
# context attention , as if there is no system prompt # k.size (1) = 1 in autoregressive generation phase # k.size (1) > 1 in prompt phase l_new = l_cache + k.size (1) kv_cache[layer_id , 0, l_cache:l_new , ...] = k kv_cache[layer_id , 1, l_cache:l_new , ...] = v o, lse = multihead_attention(
q, k_cache[layer_id , 0, :l_new , ...], v_cache[layer_id , 1, :l_new , ...], casual_mask=True)
# system attention bsz , len , nhead , dim = q.size() q1 = q.view(1, bsz*len , nhead , dim) k_sys , v_sys = sys_kv_cache[layer_id ]. unbind(dim =0) o_sys , lse_sys = multihead_attention(
q1, k_sys , v_sys) o_sys = o_sys.view(bsz , len , nhead , dim) lse_sys = lse_sys.view(bsz , len , nhead , 1) # relay fusion alpha_sys = 1 / (1 + exp(lse - lse_sys )) alpha_usr = 1 - alpha_sys o = o * alpha_usr + o_sys * alpha_sys
prompt, we have:
ot = s∑
j=1
exp(qtkTj )
σ1→lt vj +
l∑
j=s+1
exp(qtkTj )
σ1→lt vj .
(5)
Consider the first term on the right side of Eq. (5). As it is close to the Attention(·, ·, ·) operation in Eq. (4), with only a difference in the numerator, it can be rewritten as a rescaled attention:
σ1→st σ1→lt s∑
j=1
exp(qtkTj )
σ1→st vj . (6)
This rescaling trick (Milakov and Gimelshein, 2018; Rabe and Staats, 2021) can also be applied
to the second term on the right side of Eq. (5), and thus ot is a convex combination of two scaled attention terms:
ot =α sys t Attention(qt, {ki}si=1, {vi}si=1)+ αctxt Attention(qt, {ki}li=s+1, {vi}li=s+1),
(7)
where αsyst = σ1→st σ1→lt and αctxt = σs+1→lt σ1→lt
= 1 − αsyst are the combination coefficients. Discussion. Back to the view of a batch, the first term in Eq. (7) for all concurrent requests, namely system attention, can be grouped to use large matrix multiplications. This essentially eliminates the redundant access of system KVs as shown in Fig. 4. In practice, as the sum of exponentials σb→et is not numerically stable to compute directly, we use the log-sum-exp trick to return log(σb→et ) in attention computation, and the computation of αsyst is reformulated accordingly in Algorithm 1. While reformulating the casual attention, we did not assume step t 6= 1. This means that RelayAttention is also applicable to the prompt phase, where the input of a request is not a single token generated in the last step but contains multiple tokens from the user prompt, as reflected in Algorithm 1. Peripheral adaptations. There are two major adaptations needed to make RelayAttention work better within existing inference systems. First, instead of using a single KV cache for both the system prompt and the request-specific context, we use a separate system KV cache to store system KVs and fill it offline before model serving. This can be viewed as a combination of prefix sharing in PagedAttention, which eliminates redundant memory footprint of system KVs, and PromptCache (Gim et al., 2023), which eliminates redundant computation in the prompt phase. Second, as the system
KVs are already computed offline, we add an offset of s (i.e., the length of the system prompt) in the position of those request-specific context tokens to make sure of correct position embedding. Fig. 5 summarizes the key modifications to integrate RelayAttention into an existing LLM serving system (e.g., vLLM (Kwon et al., 2023) or TensorRT-LLM (Nvidia, 2023)).

heading: 3.4 Theoretical Speedup
text: In this section, to derive the theoretical speedup of RelayAttention by the memory access reduction, we analyze the memory access during the attention computation of an autoregressive generation step.
Without RelayAttention, given a batch of b request tokens, the number of elements n to transfer between DRAM and SRAM is:
n = bd︸︷︷︸ queries + b(s+ c)d︸ ︷︷ ︸ cached KVs + bd︸︷︷︸ outputs , (8)
where d is the embedding dimension, s is the length of the shared system prompt, and c is the length
of request-specific context. With RelayAttention enabled, the number of elements to access n′ is:
n′ = (bd+ sd+ bd)︸ ︷︷ ︸ system attention +(bd+ bcd+ bd)︸ ︷︷ ︸ context attention + 3bd︸︷︷︸ relay fusion . (9) Therefore, the speedup p is:
p = n
n′ =
s+ c+ 2
s/b+ c+ 7 . (10)
In Fig. 6, we plot the speedup brought by using RelayAttention. The small gaps between the practical and corresponding theoretical curves verify our analysis.
Though the speedup of standalone RelayAttention can be analyzed, it is still a question of how an end-to-end LLM serving system can benefit from RelayAttention. In Section 4, we provide an empirical study to answer this question.

heading: 4 Experiments
text: 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 (Touvron et al., 2023) for evaluation unless stated otherwise. We demonstrate the improvement for more models in Section 4.4.

heading: 4.1 Experimental Setup
text: Data. Two datasets are used in our evaluation: ShareGPTv3 (ShareGPT, 2023) and MMLU (Hendrycks et al., 2021). SharedGPTv3 (ShareGPT, 2023) contains 53k real conversions between users and ChatGPT (OpenAI, 2022). MMLU is a benchmark for measuring massive multitask language understanding in few-shot settings. It consists of 57 tasks covering various subjects and domains, such as mathematics, history, law, and medicine. Each subject/task contains a series of single-choice questions, and 5 extra questions with answers (as few-shot examples). The statistics of the benchmarking data are summarized in Table 1.
Hardware. Our experiments involve three GPUs: A40, A100-PCIE-40GB, and A100-SXM4-80GB. However, A40 is used unless stated otherwise. The hardware specifications are listed in Table 2.
Three Approaches used for comparison:
• vLLM 2: a state-of-the-art open-source LLM inference system designed for high throughput LLM serving. We use the v0.2.6 release. Note that the core component of vLLM, PagedAttention (Kwon et al., 2023), allows storing the shared system KVs exactly once by the prefix sharing technique mentioned in their paper, but this technique is not included in the vLLM v0.2.6 release. Considering the importance to save memory for a higher concurrency, we implement a stronger baseline, vLLM-PS as specified below. • vLLM-PS: the augmented version of vLLM implemented by us. It integrates not only prefix sharing but also PromptCache (Gim et al., 2023), which precomputes the system KVs and reuses them to mitigate the burden of the prompt phase. Therefore, vLLM-PS eliminates both redundant memory footprint and unnecessary computations of system KVs. • vLLM-RA (ours): the modfied vLLM with our RelayAttention integrated. Compared with vLLM-PS, this version further eliminates the redundant memory accesses of system KVs, as discussed in Section 3.3. Note that, vLLM and the two variants mentioned above use dynamic batch size: requests are dynamically batched with varied batch sizes in {1, 2, 4, 8, 16, 24, 32, 40, ..., 256} by a scheduler, depending on the available memory and the number of pending requests in the queue. vLLM would schedule as many requests as possible each time to maximize the hardware utilization.

heading: 4.2 Noninteractive Batch Processing
text: For the non-interactive batch processing scenarios where users just submit their jobs to the LLM services and harvest the processing results later, we consider the throughput (number of tokens per
2https://github.com/vllm-project/vllm
second) and processing time as the key metrics. We plot the throughputs w.r.t. the length of system prompt for processing ShareGPTv3 on the three GPUs in Fig. 7. For vLLM, the throughputs degrade heavily as the system prompt becomes long for two reasons: (1) the system prompt occupies too much memory, and thus heavily limits the batch size/concurrency of decoding; (2) it takes too much time to handle long system prompts during causal attention computation. With prefixing sharing, vLLM-PS solves the first problem and achieves up to 108% improvement on the throughput. Our vLLM-RL further solves the second problem and increases the throughput from 1.06× to 4.36× of vLLM. Table 3 shows results of the few-shot test on MMLU. We can see that using a long system prompt to include more examples is crucial for improving accuracy. In the case of the 5-shot test, our vLLM-RA provides a 76% reduction of processing time on both A40 and A100-SXM4-80GB GPUs.

heading: 4.3 Interactive Serving
text: An important LLM application is chatbots (OpenAI, 2022; Google, 2023a), in which interactive LLM services are typically provided. Unlike the noninteractive scenario, though we still expect a high throughput for good hardware utilization, we also care about the normalized latency (i.e., average per-token latency), which is crucial for user experience. Following PagedAttention (Kwon et al., 2023), we sample 1000 requests from the ShareG-
PTv3 dataset to benchmark the efficiency. The request arrival times are generated using Poisson distribution with different request rates.
As shown in Fig. 8, as the request rate increases, the throughput grows gradually until reaching a maximum. In contrast, the latency remains low at the beginning and then goes up steeply when the highest throughput is achieved. Around the latency of 0.5s/token, where the user experience and hardware utilization is balanced, vLLM-RA sustains higher request rates than both vLLM and vLLM-PS with clear margins (up to ∼ 2.2× when the system prompt length is 2048).

heading: 4.4 The Improvement for More Models
text: To verify the efficiency improvement for more models, we choose several other popular LLMs such as Llama2-13B, Llama-30B, Phi-2 (Microsoft, 2023b), and Mistral-7B (Jiang et al., 2023) to run the noninteractive batch processing of ShareGPTv3. As shown in Table 4, vLLM-RA also provides consistent improvements for these LLMs.

heading: 5 Conclusion
text: 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 RelayAttention 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.
Limitations
The limitations of RelayAttention can be reflected by the theoretical speedup (Eq. (10)). First, it helps batched inference (b > 1). The larger the batch size, the more efficient RelayAttention is. When there is only one request, which is the typical case on device-side applications, RelayAttention does not help. Therefore, RelayAttention is suitable for cloud-serving scenarios. Second, when the requestspecific contexts (including user prompts and responses) are long (e.g., 2× longer than the shared system prompt), the computation time is dominated by the processing of them; thus the efficiency gain will diminish. However, as the context length has a long-tailed distribution in many applications (e.g., chatbots), where the majority of user prompts and responses are short, the efficiency gain brought by RelayAttention is still considerable.
Broader Impact Statement
It is revealed by RelayAttention that the inference cost caused by a shared prefix can be amortized by a batch of requests. As a result, the shared prefix is much cheaper than request-specific contexts. This can reduce the hosting cost of existing LLM applications. It may also encourage LLM customization with longer system prompts or prefix-tuning (Li and Liang, 2021) in practice.

heading: B More Information of The Datasets
text: The ShareGPTv3 dataset contains both user prompts and LLM responses. The distributions of the length are plotted on the top of Fig. 9. We use synthesized system prompts during benchmarking with this dataset.
For the MMLU dataset, we use the provided fewshot examples as system prompts and the questions as user prompts. The generation length is set to 32 and we extract the answer in A, B, C, D as the first capital letter in the responses. The length distributions of system prompts and user prompts are shown in Fig. 9 bottom.

heading: C Benchmark with Synthetic Workloads
text: In the section, we benchmark the efficiency with synthetic workloads, where the user prompt length and the generation length are both fixed for all requests. Though this is far from real-world scenarios, it is useful to test the limit of an LLM serving system because such perfectly length-aligned requests eliminate the burden of scheduling. We adopt three combinations of user prompt length and generation length, (64, 128), (128, 256), and (256, 512) for benchmarking, and plot the trend of throughput w.r.t. the system prompt lenth in Fig. 10. Notably, in the most challenging case where the request-specific contexts have a length of 256 + 512 = 768, RelayAttention still provides an up to 2.2× speedup when the system prompt length is 2048.

heading: D Extension to Multi-Application Hosting
text: RelayAttention assumes that incoming requests share the same system prompt, which implies the serving process provides only one application. Although this is reasonable for applications (e.g., ChatGPT) with wide usage, deploying multiple applications in a single serving process is more economical in some scenarios. In such cases, it just requires some engineering efforts to support: for
example, tagging each request with an application ID and implementing a scheduler to batch requests with the same application ID each time.

heading: Future Work
text: Based on the observation that historical tokens are not affected by the future ones during LLM decoding, Key-Value (KV) Cache avoids repetitive computation of the hidden key-value pairs (KVs) by caching them on the fly and then reusing the cached KVs in every subsequent steps (Yu et al. , 2022; Pope et al. , 2023). With KV Cache, in each time step, only a single token (i. e. , the latest generated one) is used as the query, and the next token is produced by attending to the cached KVs. The generation complexity thus reduces from quadratic to linear w. r. t. the up-to-date sequence length. Some recent research further accelerates LLM inferences by pruning superfluous KV cache data (Zhang et al. , 2023) or compressing it (Liu et al. , 2023) to reduce key-value pairs to be cached. However, these approaches introduce algebraic discrepancies between model training and inference. Hence, they may hurt the generation quality and/or require extra finetuning efforts. In contrast, our approach maintains generation quality and is plugand-play, as it is based on a mathematical reformulation of causal attention. The acceleration of our approach comes from reducing redundant memory access of the KV cache. Therefore, it is orthogonal and complementary to prefix sharing in PagedAttention (Kwon et al. , 2023), which eliminates redundant memory footprint of system prompts, and is unlike PromptCache (Gim et al. , 2023), which eliminates the redundant computation of the reusable prefix KVs and thus only acceler-
ates the prompt phase (Section 3. 2)