Paper: ChunkAttention: Efficient Attention on KV Cache with Chunking Sharing and Batching (ICLR-2024)

Reviewer comment (Evaluation):
In Figure 4, why does the throughput decrease when the batch size becomes 128?

Author response:
- Regarding the ChunkAttn implementation, it encompasses two CUDA kernel functions: chunk_first and seq_first.  
- In the case of seq_first, optimal GPU utilization occurs when attention_head_size multiplied by batch_size exceeds the GPU's streaming multiprocessor count. For head 32, when the batch size is greater than 48, the throughput stabilizes.  
- As for chunk_first, the kernel's performance primarily hinges on matrix multiplication. In theory, larger batches should yield improved throughput. However, our matrix multiplication implementation lacks fine-tuning for batch size 128, resulting in this regression.

Reviewer comment (Evaluation):
Although 1) naive, 2) xformers, and 3) FlashAttention do not support sharing, why does the naive version outperform the others?

Author response:
- xformers and FlashAttention focus on training instead of inference, which means they lack specifical design for generation: the Q in QKV is a vector instead matrix.  
- Flash attention has been aware of this issue and published Flash-decoding. From its [blog](https://pytorch.org/blog/flash-decoding/), we also could find the flash attention outperforms naive pytorch(figure3).

Reviewer comment (Experiments):
When the model size grows on the same hardware (e.g., Lllama 13B in A100 80G), how does the result change? I think ChunkAttention can achieve more performance gain because it efficiently utilizes key-value cache memory compared to the baselines.

Author response:
- I believe ChunkAttention can yield greater performance gains due to its efficient utilization of key-value cache memory compared to the baseline methods. When the model size increases, ChunkAttention demonstrates improved performance for two reasons: 1. Enhanced utilization of L1/L2 cache, and 2. Memory savings from sharing that allow for larger acceptable batch sizes.  
- For comprehensive end-to-end testing of the entire model, more time might be required. This is due to vLLM employing numerous custom operators, necessitating the implementation of these same operators for fair comparison purposes.

Reviewer comment (Evaluation):
How long does it take to build the first prefix tree? For example, when $n_{s} = 8192$, I think it might take a non-negligible time.

Author response:
- Let's consider this scenario: with 64 requests in the tree, each having $n_s=8192$, adding a new request ($n_s=8192, s_r=0.5$) takes about 2.5ms. I find this acceptable as the cost only occurs once when a new request is added.  
- The primary overhead in building the prefix tree arises from tensor copying and transposition. Our use of Torch significantly increases the time for this process. However, by employing our custom operators, this time can be significantly reduced, as confirmed by experiments on vllm.

Reviewer comment (Experiments):
This authors did NOT compare ChunkAttention agasint prior token pruning techniques, e.g.,

Sehoon Kim, Sheng Shen, David Thorsley, Amir Gholami, Woosuk Kwon, Joseph Hassoun, and
Kurt Keutzer. Learned token pruning for transformers. In Aidong Zhang and Huzefa Rangwala
(eds.), KDD ’22: The 28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining,
Washington, DC, USA, August 14 - 18, 2022, pp. 784–794. ACM, 2022. doi: 10.1145/3534678.
3539260. URL https://doi.org/10.1145/3534678.3539260.

or KV cache compression techniques, e.g., 
Jesse Mu, Xiang Lisa Li, and Noah D. Goodman. Learning to compress prompts with gist tokens.
CoRR, abs/2304.08467, 2023. doi: 10.48550/arXiv.2304.08467. URL https://doi.org/
10.48550/arXiv.2304.08467.

Author response:
While token pruning technology also aims to improve online inference performance, it works on different aspects. You can apply token pruning together with chunk-attention. We can reference those papers but I don’t think it’s necessary to compare the results

Reviewer comment (Experiments):
There is a lack of description of the experiment setting; upon checking the source code provided in the supplementary materials, the reviewer realized that all the experiments were based on some micro-benchmark on a single transformer layer (pending confirmation if this interpretation is wrong). This makes the experimental results very problematic, as there is no evidence that the IO behavior would be the same when serving a single layer or an end-to-end model. The system should be evaluated for some real workflow instead of this micro-benchmark.

Author response:
The concern of lacking end-to-end tests on real-world LLM models is valid and well-received. We failed to provide it in the current version due to the following reasons:  
1. a comparison with existing popular inference solutions, such as vLLM, needs to be done. But these systems all have their highly customized implementations. We need to reach parity on many techniques(noise) unrelated to the optimization covered by this paper, e.g., iteration-based batching, memory offloading, and kernel fusions. Otherwise, the comparison is unfair: where the improvement comes from must be clarified. Is it a result of our optimization or some other techniques? This is also why we only publish micro-kernel level results: they are much less noisy.  
2. we are working on a highly optimized and production-ready version of this paper's solution. More time is needed before publishing comprehensive end-to-end test results. There will be an updated version of this paper after that.

Reviewer comment (Experiments):
The reviewer questions whether it is possible to provide some end-to-end evaluation for some real model over some real LLM inference traces.

Author response:
The concern of lacking end-to-end tests on real-world LLM models is valid and well-received. We failed to provide it in the current version due to the following reasons:  
1. a comparison with existing popular inference solutions, such as vLLM, needs to be done. But these systems all have their highly customized implementations. We need to reach parity on many techniques(noise) unrelated to the optimization covered by this paper, e.g., iteration-based batching, memory offloading, and kernel fusions. Otherwise, the comparison is unfair: where the improvement comes from must be clarified. Is it a result of our optimization or some other techniques? This is also why we only publish micro-kernel level results: they are much less noisy.  
2. we are working on a highly optimized and production-ready version of this paper's solution. More time is needed before publishing comprehensive end-to-end test results. There will be an updated version of this paper after that.

Reviewer comment (Writing):
In Section 3.1, the term $n_s$ is used prior to its explanation in later sections, creating a disconnect in the logical flow.

Author response:
Use of Unexplained Terms in Section 3.1: We’ll revise the structure in the next version of the paper to make sure terms are clearly defined before they are used to improve coherence.

Reviewer comment (Experiments):
The paper does not provide any references or empirical evidence to support the concept of sharable KV cache in the context of prompt engineering. This makes it challenging to gauge the validity and impact of the proposed method.

Author response:
References and Empirical Evidence for Sharable KV Cache: The concept of sharable KV cache is an observation from real-world systems, including internal LLM systems and external chatbot plugin systems such as ChatGPT. We are currently in the process of exploring ways to publish empirical data to support this observation, while ensuring we respect all necessary confidentiality constraints.

Reviewer comment (Experiments):
There is insufficient rationale behind the choice of shared prompt tokens in the experiments. The representativeness of these settings is questionable, and more clarity is needed to assess the paper's contributions effectively.

Author response:
Rationale Behind Shared Prompt Tokens: The selection of shared prompt tokens in our experiments was made with the intention to cover a broad range of cases observed in prompt systems. These cases can vary widely, with the shared prefix ranging from hundreds to tens of thousands of tokens. In the revised version of our paper, we'll provide a more detailed explanation of our experimental settings to better represent these scenarios.