OmniKV: Dynamic Context Selection for Efficient Long-Context LLMs

The dense baseline (Full) seems not to be strong enough. For Llama-3-8B model on A100, TensorRT-LLM and vLLM take less than 20ms to decode a token (sequence length = 128K, batch size = 1) . However, the dense baseline takes ~80ms, according to Figure 4.

What is the dense baseline implementation used for efficiency evaluation in Figure 4? i) what is the baseline setting for “original model running on three A100 GPUs in pipeline”

The Full (dense) Attention baseline we use here is based on Huggingface's Transformers library, employing FlashAttention2 as the backend. However, this implementation lacks kernel-level optimization, resulting in significantly slower decoding speeds compared to inference frameworks like vLLM.

As we implemented OmniKV on Lightllm, and as shown in the previous table, our method still demonstrated better latency. Because the half-precision computational power of the NVIDIA-H20 and A100 is 148 and 312 TFLOPS, the latency on NVIDIA-H20 (seq len=128k, bs=1) is higher than A100.

ii) what will be the latency for “original model” if the model is served with tensor parallelism rather than pipeline parallelism.

We have made OmniKV compatible with tensor parallelism, and OmniKV achieves superior results compared to the original model on vLLM with tensor parallelism.

Average latency of per token for each request (ms/token, tp=4):

Lightllm+OmniKV demostrates reduced decoding latency as sequence length increases. This is because the "non-filter layers" have a fixed sequence length of 2048. As the batch size decreases, the computation also decreases for "non-filter layers". But when context length decreases, the computation remains the same for "non-filter layers".

The optimization techniques proposed in this paper cannot accelerate the prefilling speed of LLM.

We have modified OmniKV to optimize the prefill stage as well. Experimental results indicate that OmniKV-prefill achieves a 1.90x reduction in latency without performance loss compared to the original model with an input length of 256K.

Specifically, during the prefill stage, we divide the sequence dimension into multiple chunks for computation. For each chunk, we also use the "filter layer" to select important tokens. Since the sequence length of the query chunk is no longer 1, we employ the uniform context selector introduced in Section 4.2 to select important tokens in the "filter layer".

Here, "r" denotes the size of the selected tokens, with the ratio of token size to chunk size being "r". Thus, the larger the value of r, the less information is lost.

The results indicate that r=2 has achieved a fine balance between efficiency and performance.

Prefill latency (s) on a single NVIDIA-H20:

Quality preserving results on LongBench:

In Figure 4, it seems that OmniKV (w/ offload) is slower than the dense baseline until the sequence length reaches 128K. For OmniKV w/o offload, the speed also lags behind the dense baseline when sequence length < 32K.

Since we have integrated OmniKV with Lightllm, we have further reduced decoding latency. Our experiments demonstrate that decoding acceleration can be achieved when the sequence length is ≥16k. It is important to note that the half-precision computational power of the H20 and A100 is 148 and 312 TFLOPS, respectively.

Decoding latency (ms/token) on a single NVIDIA-H20:

Moreover, these overheads can be entirely avoided by dynamically adjusting computation method based on input length. For instance, when the input length is 8K, the original model can be utilized, whereas for a length of 32K, OmniKV can be employed.

Thank you for your thoughtful follow-up questions and for taking the time to review our responses in such detail. We appreciate the opportunity to further clarify and address your concerns.

I interpret this to mean that "r" represents the number of chunks into which the context has been divided.

There might be some misunderstanding here. The parameter $r$ determines the size of selected important tokens, which is calculated as $r \times$ chunk_size. Here, chunk_size is a hyperparameter that represents the size of a single chunk. Please note that OmniKV-prefill does not select $r$ chunks, but $r \times$ chunk_size tokens. Token level selection offers better flexibility.

Specifically, as described in Section 4.2, OmniKV can use a local observation window to identify important tokens. We have extended this process to the prefill stage. During the prefill stage, splitting the input into chunks does not impact the correctness of the computation [1]. OmniKV-prefill treats each chunk as a local observation window to filter out important tokens from full context within the "filter layers". And the total number of selected tokens is $r \times$ chunk_size. Then, the "non-filter layers" use the important tokens identified by the "filter layers" as the context. For these "non-filter layers," the inputs for attention calculations consist of two parts: 1) The KV cache containing the selected important tokens. 2) The hidden states of the current chunk. Thus, the sequence length for these layers becomes $(r+1) \times$ chunk_size. We have further added an illustration to more clearly demonstrate the process OmniKV-prefill.

[1] Agrawal, Amey, et al. "Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills." arXiv preprint arXiv:2308.16369 (2023).

What is the difference between "r=1" and the baseline?

Taking the actual experimental setup of OmniKV-prefill as an example, chunk_size=4096, r=1 means that OmniKV-prefill selects $1\times 4096$ important tokens in the "filter layers". The actual sequence length of the hidden state inputs accepted by the "non-filter layers" is thus $4096$ (important tokens)+ $4096$ (current chunk) = $8192$ tokens. When chunk_size=4096, r=2, the actual sequence length of the hidden state inputs accepted by the "non-filter layers" is $2 \times 4096$ (important tokens) + $4096$ (current chunk) = $12288$ tokens. This greatly reduces the input sequence length involved in the computation of the "non-filter layers", reducing the complexity of "non-filter layers" from $n^2$ to $(r+1)n$, thus accelerating the prefill stage.

The paper mentions that OmniKV “perform full attention" in "all layers" in prefilling stage, why is it faster than the baseline? Do you evict tokens in the prefilling stage in the current version?

In the prefill stage, the original OmniKV performs full attention in all layers. Since prefill acceleration is also important, we extended OmniKV to the prefill stage. OmniKV-prefill only executes full attention in the "filter layers" and layers before the first "filter layer" in both prefill and decode stages. For other "non-filter layers", OmniKV-prefill performs sparse attention, which is why it's faster than the baseline. OmniKV-prefill still does not discard any tokens to maintain the model's ability in multi-step reasoning scenarios.

We hope our responses address your concerns, but if there are still any aspects that remain unclear, we would be more than happy to further discuss and provide additional clarifications.

We have further created a visual comparison to showcase the differences between OmniKV during the decode stage and OmniKV-prefill during the prefill stage. You can view the illustration via this link: Comparison. This diagram also demonstrates the efficiency enhancements introduced in the prefill stage with OmniKV-prefill. Please let us know if further clarification is needed.

Thank you for your question regarding the OmniKV integration within the lightllm framework as presented in our results. In the table labeled "Decoding latency (ms/token) on a single NVIDIA-H20," the version of OmniKV we demonstrated did not have offloading activated. Lightllm utilizes complex memory management mechanisms and currently lacks support for offloading. We are currently in the process of implementing a version that supports offloading. We will update our results and share improvements once this version is available. Thank you for your interest in our work.

Thank you for your valuable feedback. Offloading does bring some overhead, but based on our simulation, in the scenario of long context, Lightllm + OmniKV w/ offload can still achieve lower latency than Lightllm + origin model and greatly expand the longest supported context. The experimental setup consists of an NVIDIA-H20 and a 16-core CPU.

We consider the overhead brought by Offload from the prefill and decode stages.

The prefill stage is computationally intensive, so OmniKV can use computation to cover the cost of transferring kv cache offload to CPU memory. Therefore, the additional overhead brought by the offload can be ignored, which is shown in Figure 4 where the prefill time for OmniKV with offload and the original model is almost the same.

For the decode stage, there are mainly two operations that bring overhead:

1. Use the index of important tokens obtained from the "filter layer" to index the full KV cache in CPU memory.

2. Load the indexed sub KV cache to the GPU. The amount of this part of the transmission is the size of sub KV cache num_important_tokens*num_layers_on_cpu*num_key_value_heas*head_dim*dtype_size*2. Due to the consistency of important tokens between OmniKV layers, we can start the load of consecutive layers at the "filter layer".

We measured the speed of CPU to GPU data transfer on server with NVIDIA H20 to be 17.7GB/s. As shown in Figure 5, num_important_tokens=2048 is enough to achieve good results. When num_important_tokens=2048, the total transmission volume for outputting a token is about 2048*24*8*128*2*2/1e9=0.2GB, which takes about 11.3ms, and the indexing operation takes about 3.5ms, so the total time is 14.8ms. It should be noted that the cpu-to-gpu bandwidth of server with A100 in Figure 4 of the paper is about 10.2GB/s, which is slower than the H20 server.

When the context length is 256K (the longest context supported by an NVIDIA H20), the total serial delay is 26.4+14.8=41.2ms. But because computation can cover transfer, the time consumed by Lightllm + OmniKV w/ offload should be less than 41.2ms, which is still faster than the 47.1ms of Lightllm + origin model.

Since the transmission volume is fixed, the latency is independent of the context length. When the GPU memory can accommodate all KV caches, OmniKV can avoid offloading. In addition, the number of layers that perform offloading can be dynamically determined based on the context length. For example, when the context length slightly exceeds the maximum that can be accommodated (e.g. 300K), we can only offload the KV cache of one layer to further reduce the impact of offloading on efficiency.

We hope our responses address your concerns.

Weakness 1: The memory manipulation doesn't seem to be compatible with fundamental serving mechanisms like continuous batching and prefix caching. It's questionable how such algorithm can work in real world.

OmniKV is compatible with existing inference frameworks. We have implemented OmniKV within Lightllm. As Lightllm is a lightweight inference framework which is easier to adapt compared to popular vLLM. The development on vLLM is still ongoing.

Since Lightllm supports continuous batching and prefix caching, we compared the impacts of enabling and disabling these features on inference efficiency. Experiments here were conducted on a single NVIDIA-H20.

Continuous Batching:

We first sending a 64K-length request (A), followed by a 30-second wait before sending two 4K-length requests (B and C). During the process, the observed real-time running batch size changed as follows:

Batch Size: 0 => 1 (A start) => 3 (B, C join) => 2 (A end, response) => 0 (B, C end, response)

Enabling Continuous Batching reduced the end-to-end latency for requests B and C from 28.9 seconds to 23.5 seconds. But it also led to a slight increase in the latency of request A (37.8s => 41.7s), as A, B, and C were decoding simultaneously for a certain period.

Prefix Caching:

Using the novel Harry Potter, we initially sent a request to summarize the entire book. Then, after randomly removing a portion of the latter content, we sent these sub-contexts for summarization. In this scenario, the overall average end-to-end latency was reduced from 65.8 seconds to 35.49 seconds.

Algorithmic compatibility: OmniKV utilizes the exact same decoding process as standard LLMs, with the only distinction being that in a few "filter layers". These layers mask out unimportant caches but retain all KV caches. So, OmniKV remains fully compatible with continuous batching and prefix caching.

Related to weakness 1, how does this method compare to vLLM/SGLang with FlashInfern/FlashAttention backend, in terms of throughput and latency? If slower, is OmniKV compatible with practical serving framework like vLLM?

The experimental results are as follows. All experiments were conducted on a single NVIDIA-H20. "origin" refers to using original model.

The throughput results of decoding (tp=1, batch size=1):

Due to OmniKV's reduction in computation, Lightllm+OmniKV achieved the best results when the sequence length exceeded 64K. Moreover, OmniKV achieved a 1.64x speedup for an input length of 256k.

(2) quality test on needle-in-a-haystack on 128k, 256k, 512k ctx.

Utilizing the model Llama-3-8B-Instruct-Gradient-1048k, OmniKV achieved perfect accuracy on the needle-in-a-haystack benchmark. We have integrated the results of the needle-in-a-haystack experiment with a context length of up to 512k into Figure 7 of the paper.

Could you benchmark against [1], for both speed and quality preserving?

Firstly, it is crucial to clarify that MInference is a method specifically designed to optimize prefill latency, whereas OmniKV targets the latency in the decoding phase and extends the maximum context length supported under memory constraints. These two methods are orthogonal to one another. However, we have still applied OmniKV to optimize the prefill stage, naming this approach OmniKV-Prefill.

Specifically, during the prefill stage, we divide the sequence dimension into multiple chunks for computation. For each chunk, we also use the "filter layer" to select important tokens. Since the sequence length of the query chunk is no longer 1, we employ the uniform context selector introduced in Section 4.2 to select important tokens in the "filter layer".

Here, "r" denotes the size of the selected tokens, with the ratio of token size to chunk size being "r". Thus, the larger the value of r, the less information is lost.

Prefill latency on a single NVIDIA-H20:

MInference is specifically designed for Prefill acceleration and optimizes multiple CUDA kernels. But OmniKV-prefill solely relies on the Huggingface-transformers library, consequently less efficient. However, OmniKV-prefill still achieves a 1.90x latency improvement over the original model with an input length of 256K.

Quality Preserving:

The results indicate that r=2 has achieved a fine balance between efficiency and performance.

It will be quite challenging to make the proposed method work with tensor parallel,pipeline parallel, and context parallel, which are also essential in production.

can you benchmark TP=4,8 for <10b scale models? Also, tp=8, pp=4, for any 70B models?

I'd like to see (1) latency and throughput for prompt length of 128k, 256k, 512k(probably need ctx parallel) and output length of 512

OmniKV is compatible with these three parallelism methods. Due to the simplicity of OmniKV, the adaptation process is relatively straightforward.

Tensor parallelism: We have implemented OmniKV w/ tensor parallelism in Lightllm. Experimental results indicate that with TP enabled, OmniKV achieves considerable acceleration compared to vLLM with original model.

Throughput results of decoding (tp=4), using Llama-3-8B-Instruct-262K:

Average latency of per token for each request (ms/token, tp=4):

Throughput results of decoding (tp=8), using Llama-3-8B-Instruct-262K:

Average latency of per token for each request (ms/token, tp=8):

Lightllm+OmniKV experiences reduced decoding latency as sequence length increases. This is because the "non-filter layers" have a fixed sequence length of 2048. As the batch size decreases, the computation also decreases for "non-filter layers". But when context length decreases, the computation remains the same for "non-filter layers".

Technical detail of tp: The attention heads is distributed across different GPUs in tensor parallelism, and the outputs are aggregated using all-reduce, making inter-GPU communication a key factor in efficiency.

For $\text{tp}=n$ , OmniKV maintains $n$ indices of important tokens for each card. Each index of important tokens is computed solely by the card's own attention heads, thereby eliminating the need for inter-GPU communication. Since different attention heads may have distinct indices, the system intuitively exhibits greater flexibility. We have also substantiated this through experiments.

Pipeline parallelism: We have implemented OmniKV w/ tensor parallelism in transformers (huggingface). The experiments of Llama-3.1-70B shown in Table 1 is conducted using pipeline parallelism. And the support for Lightllm and vLLM is ongoing.

To implement OmniKV in pipeline parallelism, we only need to load the corresponding subset of KV cache onto the correct GPU device. The overhead introduced is quite small, usually batch_size * k_size * tp_num, with k_size usually set to 2048.

Context parallelism: OmniKV remains fully compatible. We typically require only around 2048 tokens to achieve excellent results. For this short sequence length, the communication latency in multi-GPU parallelism often exceeds the computational latency of a single GPU. Hence, we can employ context parallelism only in dense (full) attention layers, while utilizing a single GPU for sparse attention layers. This means we simply need to configure which layers will implement context parallelism. We can allocate samples to different GPUs when the batch size > 1, to prevent resource wastage.

Regarding the 70B model setup with tp=8 and pp=4, we currently lack of GPUs to run the experiments. We are actively reaching out to acquire the resources and will provide the results once available.

Dear Reviewer eZxa,

Thank you for taking the time to provide us with your thoughtful and constructive feedback. We greatly appreciate your insights and have prepared our rebuttal to address each of your concerns with the utmost respect.

We kindly ask if you could confirm whether our responses have sufficiently resolved your concerns. Please feel free to reach out with any further questions or suggestions, as we remain open to continued discussion.

Best regards, The Authors

Dear Reviewer eZxa,

We hope this message finds you well. Thank you once again for your insightful review and the valuable feedback provided. We have carefully addressed the concerns you raised in our rebuttal, particularly focusing on the deployment challenges of OmniKV in practical applications and comparison with MInference.

We demonstrated that OmniKV, when integrated with Lightllm (a lightweight inference acceleration framework), enhances efficiency, showing that Lightllm+OmniKV outperforms vLLM+origin and Lightllm+origin configurations. Additionally, we have extended OmniKV to the prefill stages, referring to this adaptation as OmniKV-prefill. This adaptation has achieved a significant speed improvement of 1.90x compared to the original model, while maintaining performance quality on par with the MInference benchmarks.

We sincerely hope that our responses and the additional benchmarks provided have addressed your concerns effectively. We are keen to receive your feedback to confirm whether the issues have been resolved to your satisfaction or if further clarification is needed.

Thank you for your continued engagement and support in refining our work. We look forward to your response.

Static KV selection methods like H2O can not only save inference time but also save total memory without CPU offloading. If this is solved, I will raise my score.

The static KV selection methods like H2O are susceptible to discarding tokens that may not have garnered sufficiently high attention at earlier stages, but these tokens could prove pivotal in later phases of generation. During multi-step inference (e.g., in CoT scenarios), tokens receiving high attention scores shift as generation evolves as shown in Figure 1b. Therefore, OmniKV retains the entire KV cache and dynamically selects during the decoding process.

The comparison between OmniKV and H2O is presented in Tables 1 and 2. OmniKV maintains higher quality than H2O, particularly in Chain-of-Thought scenarios, as illustrated in Figure 3 and Table 6.

Inference-time pipeline parallelism, especially in the same node, is uncommon. A tensor parallelism baseline is expected in this case, even if I can understand it is very hard to beat.

Sorry for the confusing description. We employed pipeline parallelism solely to ensure a fair comparison in computation with the original model. We then implement OmniKV with tensor parallelism on Lightllm, and test the throughput of decoding shown below. Here "origin" means using original model.

Throughput results of decoding (tp = 4), measured in tokens per second, using Llama-3-8B:

It is worth noting that the half-precision compute power of the NVIDIA H20 is 148 TFLOPS, which is lower than the 312 TFLOPS of the A100. If you have the throughput results for the A100 for comparison.

The experimental results demonstrate that OmniKV is fully compatible with tensor parallelism and yields superior performance compared to the original model on popular vLLM.

This paper uses cross-layer KV cache selection similarity. However, this idea was introduced by InfiniGen in a previous paper. No proper discussion is seen in this paper. If this is solved, I will raise my score.

InfiniGen introduced cross-layer similarity between consecutive two layers, leveraging this characteristic to pre-select the critical KV cache for layer $i+1$ during the processing of layer $i$.

However, OmniKV has uncovered that certain "filter layers" can identify tokens with consistently high attention scores across multiple consecutive layers. Although InfiniGen loads the KV cache for layer $i+1$ while processing layer $i$, the loading time may still exceed the computation time of layer $i$, resulting in GPU idling. In contrast, OmniKV can pre-load the cache for multiple layers $[i+2, i+n],(n>2)$ during the "filter layer" $i$, thus significantly reducing GPU idle time. This is also why we designate layer $i+1$ to perform full attention when layer $i$ is "filter layer".

The baselines shown in this paper are not proper. Only InfLLM is a dynamic KV selection method, but it is mainly used for extending context windows. I will ask for comparisons with Quest or Loki, as they are also dynamic KV selection methods.

Thank you for your advice. We conducted experiments with Quest on LongBench. We utilized the model LongChat-v1.5-7b-32k, as Quest does not support GQA. In these experiments, both Quest and OmniKV were assigned a cache budget corresponding to 30% of the total sequence length.

Performance of longchat-v1.5-7b-32k on LongBench

The experimental results indicate that OmniKV achieves performance comparable to Quest, while additionally reducing KV cache memory usage by 70%.

Decoding efficiency (batch size = 1), measured in milliseconds per token:

The experimental setup was conducted on a single NVIDIA-H20, utilizing the LongChat-v1.5-7b-32k model. The results indicate that OmniKV has comparable efficiency with Quest in short-text scenarios. However, OmniKV is capable of supporting significantly longer texts.

The performance of OmniKV in short-text scenarios still holds potential for further improvement. Quest is implemented using customized kernels, whereas OmniKV lacks deeply optimized CUDA kernels.

I do not think offloading the KV cache to the CPU can be seen as saving memory. I suggest the authors to correct this. Offloading is more like a system technique rather than actually reducing memory. Other than this, the authors cite and discuss InfiniGen in related works and provide a comparison with Quest, which is a more state-of-the-art and related baseline. This paper also seems to have a strong system implementation. I raise my score to 6.

Is it introduces some hyper-parameters that could be non-trivial to tune?

Firstly, OmniKV is insensitivity to hyperparameters, as shown in Table 8 of the paper, where the majority of parameter configurations yield similar results. The only exception is when the fourth layer is used as a "filter layer," which leads to relatively poor performance. However, this phenomenon aligns with the lower Inter-Layer Attention Similarity of layer 4 depicted in the left graph of Figure 6a.

Furthermore, as shown in Figures 5 and 6, we find that the "filter layers" are an inherent property of the model and do not vary with changes in tasks. This implies that for a given model, we only need to set one hyperparameter that can be applied across different tasks, greatly facilitating OmniKV's application in real-world scenarios.

Why attention score matrix could be similar across different layers?

The generalizability to bigger LLMs such as Llama 3.1 405B.

Deja Vu[1] have noted the existence of hidden state similarity between consecutive two layers, and posited that this may primarily arise from residual connections.

We intuitively believe that the crucial information carried by certain tokens should be progressively deepened and refined across layers, which may cause multiple layers to focus on a repeated subset of tokens. Thus, further experiments explored the similarity of attention maps between arbitrary layers. As shown in Figure 8, even between layers that are far apart, there exists a high degree of Inter-Layer Attention Similarity, providing some validation for this hypothesis.

In Llama 3.1 405B, as illustrated in Figure 9, the Inter-Layer Attention Similarity validates our intuition. We then ensure this through experiments on 2 tasks from LongBench.

[1] Liu, Zichang, et al. "Deja vu: Contextual sparsity for efficient llms at inference time." International Conference on Machine Learning. PMLR, 2023.