MPCache: MPC-Friendly KV Cache Eviction for Efficient Private LLM Inference

Question 4: Compared to existing work within privacy-preserving MLaaS.

Response 4: In recent years, a large number of algorithm optimization solutions for efficient MPC inference have emerged, such as non-linear ReLU and GeLU pruning/approximation [2,3,4,6], Softmax approximation [2,3,5], etc. However, directly applying these methods to LLMs seriously damages the model accuracy. For example, directly replacing exponential in Softmax with ReLU following MPCViT [3] drastically degrades the accuracy from 78.01% to 25.04% on LLaMA-3-8B and HellaSwag dataset. Hence, these methods inevitably require expensive retraining or finetuning to recover the model accuracy, which is unrealistic and still remains unclear for LLMs. In this work, we propose MPCache, the first framework that leverages KV cache compression technique for private LLM inference. MPCache is based on post training and totally avoids the expensive cost of training or finetuning. In our experimens, we have applied our method to both SOTA 3PC protocol PUMA and 2PC protocol BumbleBee. As a result, our work can achieve more than 8x communication reduction with a sequence length of 2048 using PUMA.

References:

[1] Rathee, Deevashwer, et al. "MPC-Minimized Secure LLM Inference." arXiv preprint arXiv:2408.03561 (2024).

[2] Li, Dacheng, et al. "Mpcformer: fast, performant and private transformer inference with mpc." ICLR, 2023.

[3] Zeng, Wenxuan, et al. "Mpcvit: Searching for accurate and efficient mpc-friendly vision transformer with heterogeneous attention." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[4] Zhang, Yuke, et al. "Sal-vit: Towards latency efficient private inference on vit using selective attention search with a learnable softmax approximation." Proceedings of the IEEE/CVF International Conference on Computer Vision. 2023.

[5] Li, Fabing, et al. "Seesaw: Compensating for Nonlinear Reduction with Linear Computations for Private Inference." ICML 2024.

[6] Kundu, Souvik, et al. "Learning to linearize deep neural networks for secure and efficient private inference." ICLR 2023.

[7] Zheng, Fei, et al. "PermLLM: Private Inference of Large Language Models within 3 Seconds under WAN." arXiv preprint arXiv:2405.18744 (2024).

[8] Dong, Ye, et al. "Puma: Secure inference of llama-7b in five minutes." arXiv preprint arXiv:2307.12533 (2023).

[9] Lu, Wen-jie, et al. "Bumblebee: Secure two-party inference framework for large transformers." Cryptology ePrint Archive (2023).

[10] Pang, Qi, et al. "Bolt: Privacy-preserving, accurate and efficient inference for transformers." 2024 IEEE Symposium on Security and Privacy (SP). IEEE, 2024.

[11] Jiang, Wuxuan, et al. "Spin: An Efficient Secure Computation Framework with GPU Acceleration." arXiv preprint arXiv:2402.02320 (2024).

[12] Fan, Shengyu, et al. "Tensorfhe: Achieving practical computation on encrypted data using gpgpu." 2023 IEEE International Symposium on High-Performance Computer Architecture (HPCA). IEEE, 2023.

[13] Samardzic, Nikola, et al. "Craterlake: a hardware accelerator for efficient unbounded computation on encrypted data." Proceedings of the 49th Annual International Symposium on Computer Architecture. 2022.

We sincerely thank Reviewer 2Xee for your thoughtful feedback! Below we answer each of your comments and questions:

Question 1: Practical setup is unclear. How secret-sharing is implemented across parties in two-party computation (2PC) or three-party computation (3PC) settings.

Response 1: Thanks for your comments on the MPC settings. Here, we provide more details about the 2PC and 3PC settings.

• For 2PC setting, we follow the 2PC protocol BumbleBee where the server owns the proprietary LLM parameter and the client possesses private input data. During 2PC inference, the data is secretly shared between two parites. The linear layers are computed using homomorphic encryption (HE) and non-linear layers require interactive protocols between the two parties based on oblivious transfer (OT) and HE.
• For 3PC setting, we follow the 3PC protocol PUMA where the parameter and data are both secretly shared among three parties. During the 3PC inference process, the linear layers are computed using 2-out-of-3 replicated secret sharing (RSS) and non-linear layers are computed using OT.

We describe the threat model in Appendix B and also add an overall secure inference framework in Figure 19 in Appendix H (refer to the updated Supplementary Material). The static eviction algorithm is performed during the prefill stage (refer to Section 4.2), and the dynamic selection algorithm is performed during the decoding stage and before computing the attention, which relies on the matrix multiplication, top-k, and token gathering protocols (refer to Section 4.3). All protocols in MPCache operate on ciphertext. Since our focus is on the algorithm optimization rather than protocol optimization, we put more information to the appendix, and the security of MPCache directly follows the security from PUMA and BumbleBee. We will add more detailed practical setups in our final version.

Question 2: Uncertain whether MPC is a practical solution for LLMs.

Response 2: Although MPC iccurs high communication and latency overhead, we believe MPC has strong potential for the following reasons:

• MPC provides cryptographically-strong privacy protection for both data and the LLM parameters and is also independent of the hardware platforms. Compared with HE-based inference methods, MPC inference is more accurate.
• In recent years, there is a significant amount of research on MPC inference: 1) MPC-friendly algorithm optimization [1-6]. For example, MPCFormer [2] achieves up to 5x speedup without comprimising the accuracy; 2) MPC protocol optimization [7-10]. For example, PUMA [8] achieves up to 3x speedup compared to MPCFormer; 3) hardware acceleration [11-13]. For example, [12] can achieve more than 100x speedup for the linear layers, which is computed with HE in 2PC scenario. These optimizations are orthogonal to each other and thus can be combined together to advance the feasibility of MPC inference.

While MPC for LLMs still incurrs high overhead, our work attempts to optimize the efficiency from the algorithm perspective. Specifically, MPCache considers the cost of MPC protocols to design an MPC-friendly algorithm, aiming to further improve efficiency and promote its application. More importantly, our work does not rely on specific MPC protocols and can be naturally extended to other schemes besides PUMA and BumbleBee. Our method is also orthogonal to other algorithm/protocol/hardware optimizations, and they can be combined togher to further improve the efficiency.

Question 3: Motivation of this work.

Response 3: Our work also focuses on improving the efficiency of private inference, just like existing privacy-preserving ML research. We want to clarify that the main contribution of our work is not combining MLaaS with KV cache usage. Instead, KV cache is already a fundamental component for existing generative LLMs to guarantee the inference correctness. If we do not use KV cache, the model has to re-compute all the previous tokens, significantly increasing the overhead (details can be found in Appendix A.3). Hence, we apply KV cache to private LLM inference by default, and our overall aim is to reduce the computation associated with KV cache size for better efficiency. We find the communicaiton and latency overhead mainly come from attention computation with KV cache, especially non-linear Softmax as shown in Figure 1(a). Hence, MPCache develops an MPC-friendly KV cache eviction algorithm to significantly reduce the overhead of expensive non-linear Softmax.

Dear Reviewer 2Xee,

We are writing to follow up on the rebuttal submission for our paper. We have carefully addressed all the concerns you raised in your initial review, providing additional clarifications to support our responses.

Given the significant impact of your review on the final decision, we believe our rebuttal must be fully considered. We kindly request your prompt feedback on our responses, as your feedback is essential to ensuring a fair and balanced evaluation of our work.

Please let us know if there are any additional questions you would like us to address or if further clarification is required.

Thank you for your attention to this matter!

We are very grateful for the Reviewer h8w6’s appreciation of our work and the thoughtful feedback! Below we answer each of your concerns:

Question 1: Specify the number of levels (n) selected for clustering in the experiments.

Response 1: For hierarchical clustering, we in practice choose a two-level hierarchical structure, i.e., n=2, and when the final dynamic selection ratio < 0.5, we first drop 50% clusters at the 1st hierarchical level. In our experiments on the long-context LongBench, we choose to use a cluster size of 32 at the 1st hierarchical level and 16 at the 2nd hierarchical level to trade off the efficiency and model performance. The detailed experimental settings of the hierarchical clustering can be found in Appendix F, and we will describe more details about Figure 6 in our final version.

Question 2: Lower commonality score between layers 6 and 12.

Response 2: Since the middle layers also show a little lower commonality score, we compare the performance with and without applying the layer-wise sharing strategy in the middle layers under different KV cache budgets in the following table. As shown in the table, applying layer-wise sharing strategy to these layers have acceptable influence on the model performance.

We also think this is an interesting phenomenon and we will further explore this problem in our future work.

Question 4: The main techniques in Section 4.3 and 4.4 contribute to the end-to-end performance more marginally when compared with top-k parallel and static eviction.

Response 4: We agree that top-k and static eviction contribute most in our ablation study. However, for smaller models, e.g., GPT-2, the overhead proportion of the max protocol and multiplication in Equation (4) becomes larger. Below we give example cases on GPT-2 with a sequence length of 512 and 1024 to show the effectiveness of our proposed techniques. As observed, techniques, i.e., Linearization & Reordering (Section 4.3) and Layer-wise Index Sharing (Section 4.4) both contribute to the efficiency improvement significantly.

In conclusion, the techniques proposed in Section 4.3 and 4.4 are orthogonal to other optimizations, and they are all MPC-friendly and also have different impact in different situations. We will add more comprehensive experimental results and analysis in our final version.

Question 5: How should parameters (alpha, gamma, #adjacent layers, eviction ratio) be determined in practice.

Response 5: In practice, these parameters are pre-determined before inference and public to all the parties. In our experiments, $\alpha$ used in similarity approximation is empirically determined and we fix it to 0.6 (ablation study in Figure 17). The number of adjacent layers for layer-wise index sharing is fixed to 2 for better accuracy (ablation study in Figure 12). In practice, we believe the eviction ratio should be determined based on two aspects: 1) the expected sequence length and 2) the latency that can be tolerated/afforded by the server and the client. MPCache proposes a novel algorithm with configurable hyper-parameters to enable exploring the Pareto front of the LLM performance and efficiency.

Question 6: How does the eviction of the value cache influence the overall performance.

Response 6: For KV cache eviction, once key cache is pruned, the corresponding value cache is useless. Hence, key cache and value cache are simultaneously pruned correspondingly based on the selected token indices, as described in Algorithm 1. In our experiments, we prune both key and value cache.

Question 7: The difference between parallel top-k protocol and existing works such as CipherGPT.

Response 7: We follow a similar top-k protocol as CipherGPT. However, the top-k protocol in CipherGPT is designed for selecting k elements from a single vector and not optimized with parallelization, and hence, for token selection in this work, multiple attention heads need to be computed sequentially, leading to significant overhead. In our work, we improve the top-k protocol with parallelization along the attention head axis to support the head-wise token selection, which is much more efficient than the top-k protocol in CipherGPT.

References:

[1] Zhang, Yancheng, et al. "HEPrune: Fast Private Training of Deep Neural Networks With Encrypted Data Pruning." NeurIPS 2024.

[2] Li, Fabing, et al. "Seesaw: Compensating for Nonlinear Reduction with Linear Computations for Private Inference." ICML 2024.

[3] Kundu, Souvik, et al. "Learning to linearize deep neural networks for secure and efficient private inference." ICLR 2023.

[4] Cho, Minsu, et al. "Selective network linearization for efficient private inference." ICML 2022.

We sincerely thank Reviewer 782C for your thoughtful feedback! Below we answer each of your comments and questions:

Question 1: Directly multiplying the cache tokens with the one-hot vector cannot remove tokens.

Response 1: Sorry for the confusion caused by our unclear description. Here we give a more clear explanation. We first take retrieving 1 token as an example. The first step of token gathering is converting an index $i$ into an one-hot vector $o\in\mathbb R^{1\times T}$ based on comparison protocol, where $T$ denotes token number. Given a key cache $K\in\mathbb R^{T\times D}$, where $D$ denotes the hidden dimension, multiplying $o$ with $K$ (matrix-vector multiplication) can generate an output with dimension $1\times D$, which is the retrieved token. To extend the case to retrieve $m$ tokens, we concatenate $m$ one-hot vectors to form a matrix $O\in \mathbb R^{m\times T}$, and then multiply $O$ with $K$ (matrix-matrix multiplication) to generate an output with dimension $m\times D$, which is the retrieved tokens. In our updated Supplementary Material, we have improved the protocol description in Appendix B.4.

Question 2: Paper lacks description of the whole secure inference process, and more information is revealed (e.g., how many tokens in the cache are used).

Response 2: Thanks for your kind advice! We add an overall framework of the whole secure inference process as shown in Figure 19 in Appendix H (refer to the updated Supplementary Material), and we will add it to our final version. Take 2PC inference as an example as illustrated in Figure 19, the server owns the proprietary LLM parameter and the client possesses private input data. During 2PC inference, the data is secretly shared between two parites. Following BumbleBee, linear layers are computed using homomorphic encryption (HE) and non-linear layers require interactive protocols between the two parties based on oblivious transfer (OT) and HE. The static eviction algorithm is performed during the prefill stage (refer to Section 4.2), and the dynamic selection algorithm is performed during the decoding stage and before computing the attention, which relies on the matrix multiplication, top-k, and token gathering protocols (refer to Section 4.3).

MPCache follows a common threat model where all the parties are semi-honest, i.e., the parties follow the designed protocols but are curious and attempt to learn extra information. In our threat model, we assume all the parties are aware of the LLM architecture and eviction ratios, which is consistent with HEPrune [1], Seesaw [2], SENet [3], SNL [4], etc. We argue that this information does not compromise the client’s data or inference results, nor does it enable the client to access the model’s parameters.

In practice, we believe the static and dynamic eviction ratio should be determined based on two aspects: 1) the expected sequence length and 2) the latency that can be tolerated/afforded by the server and the client. Hence, eviction ratio does not reveal the information of data and model parameters. Our framework proposes a novel algorithm with configurable hyper-parameters to enable exploring the Pareto front of the LLM performance and efficiency.

Question 3: Unclear citations.

Response 3: Thanks for your kind suggestion! We have revised the unclear citations in our updated version, and we will organize the citations more carefully to enhance the paper presentation in our final version.

Dear Reviewer 782C,

We are writing to follow up on the rebuttal submission for our paper. We have carefully addressed all the concerns you raised in your initial review, providing additional data and clarifications to support our responses.

Given the significant impact of your review on the final decision, we believe our rebuttal must be fully considered. We kindly request your prompt feedback on our responses, as your feedback is essential to ensuring a fair and balanced evaluation of our work.

Please let us know if there are any additional questions you would like us to address or if further clarification is required.

Thank you for your attention to this matter!

Thanks to the authors for their detailed response.

Many of my concerns are addressed. The updated description about the token gathering protocol for one-token retrieving in Appendix B.4 and Algorithm 2 is well-organized to me. Additionally, I encourage the authors to include the discussion about the security threats as well as hyperparameter setting to your final version. It is challenging for the client and server to pre-determine the hyperparameter before inference. The client has limited knowledge about the proprietary model of the server, it would be difficult for the client to see the impact of different choices of parameters like $\alpha$ and #adjacent layers. The pre-defined parameters such as $\alpha=0.6$ and #adjacent layers $=2$ might not be robust for inputs of different lengths and from different datasets. I believe acknowledging this challenge can help the readers to better understand how MPCache can be deployed in practice.

However, some of my concerns remain unsolved.

• The protocols in MPCache. I understand that MPCache involves the MatMult, TopK and Gather protocols as shown in Figure 19. However, I have concerns about how these protocols are used in MPCache.

  • For the Gather protocol, to generate the one-hot vector, equality protocol should be used instead of comparison (Algorithm 2 line 2). Comparison protocols returns 0 if the $id$ is smaller than $i$ and 1 otherwise (or vice versa, as shown in Figure 3 in Iron).
  • For the TopK protocol, the authors mentioned that they follow the design of CipherGPT, except for the parallelization across the attention heads. The TopK protocol in CipherGPT is shuffle-based, because the comparison results need to be revealed for efficient quick sorting (Algorithm 2 line 8 in CipherGPT). With the shuffle-based TopK protocol, the relative order of tokens and KV caches cannot be preserved. However, as shown in Appendix D of Bolt (S&P 2024), the relative order of tokens is important during pruning.

• The secure computation framework is still unclear. The standard secure inference process of Transformers, such as that in Iron and BumbleBee, is clear. However, how MPCache is integrated in this process still remains unclear. Figure 19 does not correctly reflect the workflow of MPCache (in Algorithm 3). Since the KV Eviction is performed before attention in Figure 19, I assume that Figure 19 only illustrates the decoding step and the dynamic selection. There are multiple issues. For example, the query $q$ is needed to calculate the similarity, but is not passed to the KV Eviction module. Also, $Q$ is used to denote the queries in the prefilling phase. The static eviction is performed after the attention, which is not reflected in Figure 19 or Appendix H.

• I understand that both key cache and value cache are evicted (line 19 and 20 in Algorithm 3). During dynamic selection, if the similarity is calculated using the value cache, will the final performance be the same?

There are still some minor issues about the citation. For example, there are two citations for LazzyLLM (line 594 and 598). I encourage the authors to check their cited papers to make sure the references are correct.

Dear Reviewer 782C,

As the deadline for the rebuttal period is approaching, we would like to know if your concerns about this paper have been resolved. If there are any other questions or concerns, please let us know.

Thank you again for your time and effort in reviewing our work!

We appreciate Reviewer crCT's constructive feedback, which helps us improve our work! Below we answer each of your comments and questions:

Question 1: Ability of these techniques to generalize to other models. How the method would be useful in more realistic settings where both the model and the data are private.

Response 1: We believe our proposed static-dynamic eviction algorithm has good generalization capability. This is because of two reasons: 1) redundant tokens that do not impact the decoding of any downstream tokens generally exists in different input sentences, which is observed in many previous works. These tokens can thus be statically evicted without impacting the LLM performance. In our experimentals, we use the static eviction algorithm to prune ~50% tokens, which is very different from [1,2] that only rely on static eviction and only prune once to the specific ratio; 2) the effectiveness and generalization ability of dynamic eviction originates from the Softmax attention. As the Softmax function normalizes the summation of the attention scores to 1, it limits the maximum number of tokens that impact each LLM decoding step. We observe the sum of attention scores of the top-100 ranked tokens is usually larger than 0.95, indicating majority of the impact comes from these top-100 ranked tokens. This observation makes us confident to use a high eviction ratio during the dynamic eviction phase without significant accuracy degradation.

In practice, we believe the static and dynamic eviction ratio should be determined based on two aspects: 1) the expected sequence length and 2) the latency that can be tolerated/afforded by the server and the client. Our framework proposes a novel algorithm with configurable hyper-parameters to enable exploring the Pareto front of the LLM performance and efficiency. The framework has been validated across different benchmarks and different LLMs, which demonstrates its value and potential for practical usecases.

In our revised version, we will provide more discussions on the generalization capability.

Question 2: Exact protocols used for performing the dynamic eviction.

Response 2: The protocols used for dynamic eviction can be divided into the following three parts: 1) Similarity approximation. As shown in Equation (6), we invoke the matrix multiplication protocol to measure the relationship between the query token and previous tokens. 2) Top-k selection. We then invoke top-k protocol to obtain the indices of important tokens. The protocol is the same as that in CipherGPT and we optimize the protocol with parallelization for fast head-wise token selection. 3) Token gathering. This protocol consists of two sub-protocols: a) comparison protocol to securely convert integer indices into one-hot vectors; and b) matrix multiplication protocol of the KV cache and one-hot vectors to obtain the selected tokens. We describe the details in Appendix B.4.

In this work, our focus is on the algorithm optimization for efficient inference rather than protocol optimization, so we directly follow the protocols proposed by 3PC PUMA and 2PC BumbleBee, which are implemented in the SecretFlow framework. We will provide detailed protocol descriptions more clearly in our revised version.

Question 3: Some citations are inaccurate or missing.

Response 3: Thanks for your kind suggestion! We have revised the incorrect citations and added missing citations in our updataed version. We will further check the citations in our final version.

References:

[1] Li, Yuhong, et al. "Snapkv: Llm knows what you are looking for before generation." arXiv preprint arXiv:2404.14469 (2024).

[2] Cai, Zefan, et al. "Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling." arXiv preprint arXiv:2406.02069 (2024).