SwiftKV: Fast Prefill-Optimized Inference with Knowledge-Preserving Model Transformation

Thank you for your detailed review of our paper and for your feedback.

On more serving systems and lossy optimization: We clarify that vLLM is not a comparison in our study but it is used as our inference backend, as vLLM is the most widely used inference engine in production today. Although SwiftKV produces models with some quality loss, the quality loss does not depend on the inference system and can be evaluated independently. The knowledge retention and total compute reduction does not change across inference systems.

We chose vLLM to demonstrate that it is possible to achieve further improvements over a strong baseline system. Supporting SwiftKV for vLLM was part of the contribution of the paper. Other systems do not support SwiftKV either, and would require implementation for each which is beyond the scope of our paper.

On integration with other techniques: In Sec 5.4 and 5.5, we showed that SwiftKV can be integrated with KV cache quantization and early exit for decoding tokens. In our comment for all reviewers, we also showed results for Llama-3.1-405B-Instruct quantized with W8A8. The model quality gap for Llama-405B is smaller than the other models we tested, which indicates that our method may be more robust for larger models.

While we agree that integration between SwiftKV and other techniques is important, it has a very large space of combinations and is infeasible for us to implement and evaluate. However, we shall add more substantial discussions to our paper. In short, SwiftKV is fully compatible with speculative decoding because we still execute all layers for decode tokens and can provide output logits to verify draft tokens with. For sparse attention, SwiftKV is agnostic to attention sparsity and, in principle, can be used in combination to achieve a compound reduction in total inference compute.

Comparison with pruning and early-exit:

For early-exit, one recent work is https://arxiv.org/pdf/2404.16710. However, the accuracy gap is much larger than ours even with continued pre-training on 500B tokens and/or training from scratch using 28B tokens (Figure 8 / Table 2), while our training budget is orders of magnitude smaller (160M x 2 epoch = 320M tokens). Also, even if they are able to early-exit one token, if any later token requires all layers, then all the missing KV-Cache belonging to the exited tokens still need to be computed. As such, it may speed up latency-sensitive use cases but it may not reduce the total computation.

For pruning, two recent works are Llama-3.2 [1] and Minitron [2]. Although they have more substantial data curation (both use a much larger model to generate synthetic data, and both use orders of magnitude more data) than we do, our method has much better accuracy (e.g., Minitron-4B has <61% on MMLU, but ours is 67.3% for 50% SingleInputKV). On the other hand, our method only applies to prefill tokens (making it suitable for input-heavy workloads) while they apply to both prefill and decode.

We will add discussions about this as well as the suggested related works in revising our paper.

To answer your specific questions:

1. Open source release: We will certainly make the code open source after the review process and anonymity requirements are lifted. We will include the distillation code, inference code, as well as SwiftKV models.
2. Improvements to TPOT: This improvement is due to mixing prefill and decode tokens in each minibatch, which is widely adopted by inference systems today, including vLLM. The decode tokens in each minibatch also benefit from the prefill tokens being processed faster. As you correctly pointed out, this may not apply to prefill-decode disaggregation. However, the Throughput/TTFT improvements are still valid. We will add this discussion to our paper.
3. Compared with KV cache reduction: Our key contribution is prefill compute reduction (SingleInputKV) combined with KV cache compression, while prior works only address KV cache compression. We refer to our top-level comment for all reviewers for a more detailed discussion. In it, we compared 50% SingleInputKV with Merge-all-Layers, which represents an idealized version of MiniCache.
4. Diverse set of datasets: The datasets used for training the model and evaluating the model (mmlu, gsm8k etc) are not the same as the workload we use to evaluate SwiftKV’s throughput and latency. For the latter we use a traffic pattern from our production where the requests have a median of 2K+ input tokens, with about 256 output tokens. This observation that inputs are much larger than outputs in production was the motivation for SwiftKV. However, if a different traffic pattern was used where generation is larger than prefill, then SwiftKV would have limited benefits.

Your response did not fully address my concerns.

First, I kindly request to see the acceleration effects of this method demonstrated on other systems (not just vllm), or a comparison with the optimization effects of other systems. It would be helpful to understand whether your method can effectively integrate with optimizations from other systems or if it demonstrates superior performance compared to them.

Secondly, I would like to ask for ablation studies specifically focusing on the kvcache compression component, as this would help substantiate the necessity of adopting AcrossKV in this approach. While I appreciate the additional experiments you provided, they only validated the necessity of early exit.

Thirdly, the work suffers from the lack of evaluations on datasets with varying prefill and decoding characteristics to showcase the method’s practical effectiveness. In fact, the proposed approach appears primarily beneficial for optimizing the prefill stage. Therefore, the usage of TPOT might mislead readers into overestimating the contribution of this paper

Therefore, I will maintain my previous rating.

Thank you for your detailed review of our paper and for your feedback. We will correct all of the copy editing suggestions in revising our paper.

On end-to-end results: We would like to clarify the results presented in Section 4.3 (Figures 3 and 4) showcase the end-to-end performance of our work. Specifically, these results highlight the performance of our final trained and transformed SwiftKV models, fully integrated into the vLLM system, which is the most widely used open-source serving system in industry. The measurements reflect the end-to-end throughput and latency improvements observed in our production system under simulated traffic.

If the "end-to-end results" being requested refer to a different context or additional metrics beyond what is already presented, we kindly ask for clarification so we can address the concern more effectively.

On optimizing 5% of the total model: To clarify, the 5% refers to the K and V projections of the later layers skipped by SingleInputKV that still need to be run. For example, with 50% SingleInputKV, we only need to run the KV projections of the last 50% of layers (skipping the Q, O, attention, and MLP operations), which amount to <5% of the later layers. We provide a full breakdown of the compute reductions in our top-level comment for all reviewers.

On Section 3.4, Knowledge Recovery: When AcrossKV is used, only the Wk and Wv of the initial layer in each AcrossKV block is trained. We will clarify this in Sec 3.4.

Why 4-way caching is better than 2-way caching: As stated in Appendix C.1, we did not search for the best hyper-parameters for different model sizes and/or different settings (e.g., the 4-way and 2-way caching) due to the limited resource and high cost to do so. As such, all the results we got are likely sub-optimal. Similar observation can be found in Table 1: the accuracy of 50% SingleKV + 4-way is higher than 2-way on GSM-8K.

On 16K tokens/s and 560 TFLOPs: You are correct that the 70B model requires much more (approx 8-9x) floating point operations than the 8B model, and that is why the per-GPU tokens/s is much higher for the 8B model. The compute performance for the 70B model improves primarily due to the individual matrix multiplication operations being larger than for the 8B model, which is more suitable for the parallel GPU architecture. Even though the compute for 70B is distributed, it is still on the same node, which has very fast communication between GPUs.

Full model vs partial model fine-tuning: We believe one reason that full fine-tuning performs worse is it can cause the model to overfit to the distillation dataset, which in SwiftKV’s case is very small compared to the scale of pretraining (or even instruction tuning). As stated in Appendix C.1, the total number of training tokens we have is 160M x 2 (epoch) = 320M. In Sec 5.2, we also mentioned prior works that hypothesize that the core knowledge of transformer models are stored in their MLP layers instead of Attention layers. Combined with the fact that SwiftKV only modifies the attention-related parts of the model, it makes sense that we should only train the Q, K, and V weights. However, if we are able to do the full scale pre-training and post-training (i.e., multi-iterations of SFT / DPO etc), the full model training might be better.

On dataset selection: There is a very large search space of dataset mixtures that can lead to better accuracy, that would be infeasible to fully explore in the scope of our paper due to limited resources and high cost. Instead, we simply chose two very popular instruction datasets. Even though we could have optimized the dataset mixture to better fit the distribution of our evaluation tasks, we believe it can still be informative to readers to know the effect of dataset selection when it does not align well with the evaluation tasks. Therefore, we chose to use two basic datasets in our main evaluation, and provide insights into better dataset selection in our ablations (Sec 5.3). Nevertheless, it is both surprising and promising for SwiftKV that even a basic selection of data (with only 160M tokens) can lead to reasonable results.

Thank you for your detailed review of our paper and for your helpful feedback.

On evaluating a wider range of models: We provide the accuracy scores for two additional models, Mistral-Small-Instruct-2409 and Llama-3.1-405B-Instruct (FP8) in our top-level comment for all reviewers. For Mistral-Small, we observe similar impacts to model quality and throughput as Llama-8B. For Llama-3.1-405B-Instruct with W8A8 quantization, we additionally observe a smaller model quality gap, indicating that our method may be more robust in larger models. We are happy to include these results and their corresponding system performance results in revising our paper.

On computation cost reduction: First, we would like to clarify that with SingleInputKV, only the decode tokens, not prefill tokens, need to calculate the attention outputs for all layers. For prefill tokens, if using 50% SingleInputKV, the last half of layers only need to calculate the key and value projections to populate the KV cache. This is done based on the output of the middle layer, not the immediate preceding layer, which allows the query projections, attention outputs, and MLPs to be skipped entirely for the last half of layers.

As suggested, we provide a more detailed breakdown of the compute in our top-level comment for all reviewers, which we are happy to add to our revised paper.

On showing the independent effect of AcrossKV: We emphasize that the key contribution of SwiftKV is prefill compute reduction (SingleInputKV) combined with KV cache memory reduction, and refer to our top-level comment for a detailed discussion. As demonstrated, even Merge-all-Layers, which can be considered an idealized version of AcrossKV, has limited throughput benefits if applied without SingleInputKV in compute-bound scenarios. Due to this fact, and the fact that KV cache memory compression has been explored thoroughly in prior works, we do not believe an in-depth evaluation of AcrossKV-only is crucial to our paper.

Thank you for your detailed review of our paper and for your feedback.

Response to W1 and W2: Indeed, several independent works including ours made observations regarding the similarity of activations across transformer layers. However, we are not aware of any work that leverages this insight to reduce the amount of pre-fill computation to speed up inference. InfiniGen reduces the KV cache memory needed during the decoding phase to reduce data transfer between CPU and GPU to speed up heterogeneous inference with CPU offloading, however it does not reduce the prefill computation required.

As we shared in our comment for all reviewers, the key novel contribution of SwiftKV is not KV cache compression alone but rather prefill compute reduction (SingleInput KV) combined with KV cache compression. Approaches such as InfiniGen and PyramidKV that only reduce KV cache memory have very little impact on the throughput in production environments where there is ample fast GPU memory. In contrast, the compute reduction from SingleInputKV significantly improves throughput.

We appreciate the pointers to the additional related works on KV cache reduction and will add a more thorough discussion to our related works in revising our paper.

Response to W3: Unlike MInference or ALISA that focuses on reducing only the quadratic attention computation and KV cache memory, SingleInputKV reduces the total computation in the prefill stage of the model, not just the attention computation. SingleInputKV is agnostic to the sparsity pattern used in attention and as such SingleInputKV is complimentary and can be combined with MInference, ALISA, or any other sparse attention mechanism to achieve a compound reduction in total inference compute. For instance, if MInference or ALISA can reduce overall inference compute by 2x, then combining it with a 2x reduction from SwiftKV will result in a compounded reduction of 4x.

We also would like to point out that we showed SwiftKV can be combined with KV cache quantization in Sec 5.4. However, combining our method with all existing methods and showing the improvement is beyond the scope of a research work.

Response to W4: We provide the accuracy scores for two additional models, Mistral-Small-Instruct-2409 and Llama-3.1-405B-Instruct (FP8) in our top-level comment for all reviewers. For Mistral-Small, we observe similar impacts to model quality and throughput as Llama-8B. For Llama-3.1-405B-Instruct with W8A8 quantization, we additionally observe a smaller model quality gap, indicating that our method may be more robust in larger models. We are happy to include these results and their corresponding system performance results in revising our paper.

Answer to Q1: Indeed, GSM8K appears to be the task we evaluated that has the largest gap, which may be due to the small size of the model or higher demand for reasoning. In Sec 5.3, we explored another likely explanation, which is simply that our distillation can benefit from more math and coding related data. We found that using an additional 83K math and coding examples (only 7% of our existing dataset) improved the GSM8K score by 0.53 points. Although it is not in the scope of our paper to identify the optimal distillation dataset (also as stated in Sec 5.3), we believe future work in this direction can help to bridge the gap for GSM8K. Possible directions included better post-training selection, advanced training approaches (like DPO), and pre-training enhancement which are also stated in our Limitation and Future Direction section.

Also, as the question itself stated: there is a noticeable gap of LLaMa-3.1-8B but the gap reduces to ~1.5 points for LLaMa-3.1-70B with the same training data and hyper-parameters. This indicates a larger model may be more resilient to compression, and this is also illustrated by the 405B cases in the above table, where the gap further reduces to 0.9 points.

Answer to Q2: Given the main goal of SwiftKV is to reduce the prefill compute rather than KV cache memory, we felt that including AcrossKV was not critical to illustrate the effect of our distillation procedure on compute reduction. However, we are happy to take the suggestion and include a row for AcrossKV in Table 3 in revising our paper.

Answer to Q3: “Our fine-tuned model” here is not a SwiftKV model, but a model fine-tuned on (not distilled from) Llama-3.1-8B. The purpose of the “Our fine-tuned model” row serves only to illustrate the difference in data quality between what was used to distill SwiftKV and what Meta (presumably) used to create Llama-3.1-8B-Instruct.

Our message is that, since there is a large quality gap between the model we fine-tuned ourselves versus Meta’s instruct model, there is likely a large quality gap in the datasets used to fine-tune them. This further illustrates (as in our answer to Q1) that improving the distillation dataset quality is likely a promising path to improving SwiftKV in future work.