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

## Abstract

Abstract LLM inference for popular enterprise use cases, such as summarization, RAG, and code-generation, typically observes orders of magnitude longer prompt lengths than generation lengths.
This characteristic leads to high cost of prefill and increased response latency.
In this paper, we present SwiftKV, a novel model transformation and distillation procedure specifically designed to reduce the time and cost of processing prompt tokens while preserving high quality of generated tokens.
SwiftKV combines three key mechanisms: i) SingleInputKV, which prefills later layers’ KV cache using a much earlier layer’s output, allowing prompt tokens to skip much of the model computation,
ii) AcrossKV, which merges the KV caches of neighboring layers to reduce the memory footprint and support larger batch size for higher throughput,
and iii) a knowledge-preserving distillation procedure that can adapt existing LLMs for SwiftKV with minimal accuracy impact and low compute and data requirement.
For Llama-3.1-8B and 70B, SwiftKV reduces the compute requirement of prefill by 50% and the memory requirement of the KV cache by 62.5% while incurring minimum quality degradation across a wide range of tasks.
In the end-to-end inference serving using an optimized vLLM implementation, SwiftKV realizes up to 2 × 2\times 2 × higher aggregate throughput and 60% lower time per output token.
It can achieve a staggering 560 TFlops/GPU of normalized inference throughput, which translates to 16K tokens/s for Llama-3.1-70B in 16-bit precision on 4 × \times × H100 GPUs.

## 1 Introduction

Large Language Models (LLMs) are quickly becoming an integral enabler of enterprise applications and offerings, including code and data co-pilots (Chen et al., ; Pourreza & Rafiei, ), retrieval augmented generation (RAG) (Lewis et al., ; Lin et al., ), summarization (Pu et al., ; Zhang et al., ), and agentic workflows (Wang et al., ; Schick et al., ). While it is clear that LLMs can add value to these applications, the cost and speed of inference determine their practicality. Therefore, improving the aggregate throughput and reducing latency of LLM inference has become an increasingly important topic of interest, with various efforts (Sec. 2) tackling the problem from multiple angles.

In this paper, we take a unique approach to improving LLM inference for enterprise applications based on the key observation that typical enterprise workloads process many more input tokens than output tokens. For example, tasks like code completion, text-to-SQL, summarization and RAG each submit long prompts but produce a small number of generated tokens, and a majority of enterprise LLM use cases in Snowflake incur a 10:1 ratio between prompt and generated tokens.

Based on this observation, we designed *SwiftKV*, which improves throughput and latency by: i) reducing the computation required to pre-fill the KV cache for input tokens, and ii) enabling memory savings to support larger batch sizes needed to serve LLMs more cost effectively (Sheng et al., ; Pope et al., ; Yu et al., ). SwiftKV (Fig. 1) consists of four key components:

Figure: Figure 1: Illustration of SwiftKV with 50% SingleInputKV and 2-way AcrossKV. After distillation, the KV cache of layers 5–8 can all be populated using the hidden state outputs of layer 4. For prefill tokens, layers 5–8 may be skipped altogether, while decode tokens complete all layers. Existing models may be efficiently adapted for SwiftKV by distilling from the original unmodified model using a small dataset. Model knowledge is preserved by keeping the trainable parameters limited to the Q, K, and V projections of the layers affected by SingleInputKV.
Refer to caption: x1.png

SingleInputKV. SingleInputKV rewires an existing model so that the pre-fill stage during inference can skip a number of later layers in the network, and their KV cache are computed by a single earlier layer. SingleInputKV is motivated by the observation that the output hidden states of the later transformer layers do not change significantly (see Sec. 3.2, also independently discovered by Liu et al. ()). With SingleInputKV, the computation required for pre-fill is reduced by approximately the number of layers skipped. We found that it is possible to skip at least 50% of the layers without significantly impacting the model quality (Sec. 4.2), which translates to a 50% reduction of the pre-fill computation in inference.

AcrossKV. While SingleInputKV reduces the pre-fill computation, it does not reduce the KV cache memory requirement. AcrossKV combines the KV projections from multiple adjacent layers into a single one, and share the KV cache across these layers to reduce its size in memory. AcrossKV allows significant memory savings, which unlocks higher throughput by enabling larger batches during inference. In Sec. 3, we show that AcrossKV can reduce the KV cache size by 25% with less than a 1% quality gap. We also show in ablation studies (Sec. 5) that AcrossKV is compatible with existing KV cache quantization methods, which combine to unlock 62.5% reduction in KV cache size.

Knowledge Recovery. Although SingleInputKV and AcrossKV can be applied to existing LLMs with minimal changes to their architectures, we found that the resulting model parameters should still be adapted to the new architecture to recover their original prediction quality. This can be done via distillation from the original model. With SwiftKV, we found that a lightweight distillation is sufficient, with only a fraction of the model (Q, K, and V projections of the affected layers) trained on 680M tokens, which takes less than 3 hours on a single $8\times$ H100 node for Llama-3.1-8B.

SwiftKV Optimized Inference. To realize the computation and memory reductions of SingleInputKV and AcrossKV into end-to-end throughput and latency improvements, we implemented SwiftKV in vLLM (Kwon et al., ). Our implementation includes several additional optimizations, including fusing all KV-projections beyond layer $l$ into a single GEMM operation, and integrated memory management needed to lower the KV cache memory footprint achievable via AcrossKV.

SwiftKV increases the aggregate throughput of enterprise workloads by up to $2\times$, while reducing time-to-first-token (TTFT) and time-per-output-token (TPOT) by up to 50% and 60%, respectively. In fact, for Llama-3.1-70B, SwiftKV can achieve a normalized throughput of 560 TFLops/GPU(^1^11Normalized throughput and MFU is based on number of floating point operations in the baseline model.). This is an unprecedented 56.6% MFU utilization for inference (Sec. 4.3). We show that SwiftKV incurs minimal quality degradation averaged across a wide range of tasks (Sec. 4.2), including ARC-Challenge (Clark et al., ), Winogrande (Sakaguchi et al., ), HellaSwag (Zellers et al., ), TruthfulQA (Lin et al., ), MMLU (Hendrycks et al., ), and GSM8K (Cobbe et al., ).

In addition to these main results, in Sec. 5 we discuss the impact of distillation, datasets, choice of trainable parameters for training SwiftKV. We also present our analysis of the hidden state similarities, and how AcrossKV can be extended and combined with other KV cache compression works. Additionally, we also discuss how SingleInputKV can enable compute savings not just during pre-fill but also during decoding phase.

## 2 Related Work

Hardware and System Optimizations. Lower-precision quantization techniques like FP8 (Kuzmin et al., ) can enable the use of tensor-cores to accelerate inference (Luo et al., ). System approaches like PagedAttention (Kwon et al., ), Tensor-Parallelism(Shoeybi et al., ), Split-Fuse (Holmes et al., ; Agrawal et al., ), Flash Attention (Dao et al., ), and their optimized implementations in TensorRT (NVIDIA, ), FasterTransformer (NVIDIA, ), vLLM (Kwon et al., ), and DeepSpeed-Inference (Aminabadi et al., ) enable better parallelization, batching, and scheduling to eliminate performance overheads and achieve better hardware peak utilization without impacting model quality.

Memory Compression. A wide range of techniques have been developed to reduce the memory footprint for inference. Lower-precision quantization techniques like FP8/FP4 can reduce the memory footprint for both KV cache and parameters (Hooper et al., ). Attention optimization techniques like MQA (Shazeer, ), GQA (Ainslie et al., ), low-rank attention (Chang et al., ) also reduce the KV Cache. All of these approaches are complementary to SwiftKV.

Similar to our work AcrossKV, MiniCache (Liu et al. ()) reduces KV Cache by consolidating adjacent layers of KV cache. While both consolidate KV cache across adjacent layers, AcrossKV enables consolidating more than just two layers, allowing for higher level of compression. Furthermore, due to the distillation based fine-tuning, AcrossKV does not require any token retention strategy where distinct KV caches are stored for special tokens even for consolidated layers. The simplicity of AcrossKV makes it easier to translate the KV cache savings into real world performance improvements.

Other Inference Optimizations.
Speculative Decoding based approaches (Cai et al. (); Xia et al. ()) reduce the number of decoding steps to speed up the generation phase of inference, but it does not reduce prefill computation. Model distillation techniques (e.g., Sreenivas et al. ()) transfer the knowledge of a large teacher model into a smaller student model, reducing inference computation by decreasing the number of parameters. However, these methods require significant retraining over long token horizons, and typically incur a significant quality degradation from the teacher to the student model.
In contrast, SwiftKV reduces the computation of the model, not its number of parameters, by approximating the KV cache of later layers, and aims to fully recover the model performance using distillation. Moreover, it requires minimal training resources including compute and data.

## 3 SwiftKV: Design and Implementation

Figure: Figure 2: LEFT: the input similarity of several smaller scale models (Fig. A.1 in the Appendix shows a similar observation for larger models). RIGHT: The time per forward pass of Llama-3.1-8B. SingleInputKV effectively reduces the forward pass processing time across a range of batch sizes.
Refer to caption: x2.png

### 3.1 Preliminaries

In transformer models (Vaswani et al., ), attention enables each token to focus on other tokens by comparing *queries* ($Q$) with *keys* ($K$) and using *values* ($V$) to compute the final representation. For a sequence of input tokens $x^{(1)},\dots,x^{(n)}$, the projections are defined as follows: $Q=XW_{Q},\quad K=XW_{K},\quad V=XW_{V}$, where $X\in\mathbb{R}^{n\times d}$ are the input embeddings, and $W_{Q}\in\mathbb{R}^{d\times d_{k}}$ and $W_{K},W_{V}\in\mathbb{R}^{d\times d_{g}}$ are trained model parameters with $d_{g}|d_{k}$. Hereafter, we may also refer to $W_{K}$ and $W_{V}$ as a single matrix $W_{KV}\in\mathbb{R}^{d\times 2d_{k}}$.

During the *prefill phase* of inference, the model processes the entire input sequence at once, computing $K$ and $V$ for all tokens in parallel (or in chunks in the case of Split-Fuse (Holmes et al., ; Agrawal et al., )). This typically occurs when the model handles an initial prompt or context.

During the *decoding phase* of inference, new tokens are generated one at a time. When predicting the next token, only the query ($Q^{(t+1)}$) for the new token needs to be computed, while the model must attend to the keys and values ($K^{(1)},\dots,K^{(t)}$, $V^{(1)},\dots,V^{(t)}$) of all previously processed tokens.

To optimize efficiency in the decoding phase, *KV caching* is employed. After processing each token $t$, the newly computed $K^{(t)}$ and $V^{(t)}$ are stored in a cache. For the next token $t+1$, only the new query $Q^{(t+1)}$, key $K^{(t+1)}$, and value $V^{(t+1)}$ are computed. The attention computation will then utilize the cached $K$ and $V$ from all prior tokens, allowing for reduced computational overhead during inference.

### 3.2 SingleInputKV: Project KV cache from a single layer

Assume the input of $l$-th layer is ${\bf x}_{l}$, and its $i$-th token is ${\bf x}_{l}^{(i)}$.
Prior studies (Liu et al., ; Gromov et al., ) showed that ${\bf x}_{l}$ becomes more similar as the depth grows.
Here, we conduct a similar study.

We compute the average input similarity between $l$-th layer’s input and all remaining layers’ input, i.e.,

$$ $\mathrm{SimScore}({\bf x}_{l})=\frac{\sum_{j=l+1}^{L}\mathrm{Similarity}({\bf x }_{l},{\bf x}_{j})}{L-l},$ (1) $$

where $L$ is the number of layers in the model and $\mathrm{Similarity({\bf x}_{l},{\bf x}_{j})}$ is the average cosine similarity between all ${\bf x}_{l}^{(i)}$ and ${\bf x}_{j}^{(i)}$ tokens.

We use 50 random training examples from HuggingFaceH4/ultrachat_200k to estimate $\mathrm{SimScore}({\bf x}_{l})$, and
the results of Llama-3.1-8B-Instruct, Mistral-7B-Instruct-v0.1, and Qwen2.5-7B-Instruct are shown in the left of Fig. 2.
As the layers get deeper, $\mathrm{SimScore}({\bf x}_{l})$ gets higher.
Particularly, around half of the depth, the average similarity of ${\bf x}_{l}$ with ${\bf x}_{>l}$ is above $0.5$ for all models, which shows that the difference of input hidden states are small in deeper layers.

Based on this observation, the first key component of SwiftKV is to use $l$-th layer’s output ${\bf x}_{l+1}$ to compute the KV cache for all remaining layers. More specifically, SwiftKV retains the standard transformer architecture up to and including the $l$-th layer, but the KV cache for all remaining layers are computed immediately using ${\bf x}_{l+1}$, i.e.

$$ ${\mathbf{KV}}_{j}={\mathbf{W}}_{KV}^{j}{\bf x}_{l+1},~{}~{}~{}~{}\text{for all $j>l$,}$ (2) $$

where ${\mathbf{KV}}_{j}$ is the KV cache for $j$-th layer and ${\mathbf{W}}_{KV}^{j}$ is its KV projection weight matrix.

Prefill Compute Reduction. SingleInputKV can enable significant reduction in prefill computation during LLM inference. Originally, all input tokens need to be processed by all transformer layers in order to generate the first new token.
Using SingleInputKV, after layer $l$, since all the required KV cache are computed, for all remaining layers ($>l$), we only need to process the last token in the input. When prefill computation dominates generated token computation, this reduces the total inference computation to approximately $l/L$. See Appendix B for more details.

### 3.3 AcrossKV: Sharing KV cache for consecutive layers

GQA (Ainslie et al., ), one of the most widely adopted KV cache compression methods, showed that the KV cache can be easily shared within a transformer layer. Later, Liu et al. () showed that the KV cache can be merged for certain pairs of adjacent layers.
AcrossKV extends the ideas to cross-layer KV cache sharing.

Particularly, instead of computing KV cache for all of the remaining layers as shown in equation 2, AcrossKV selectively chooses one layer to compute the KV cache for several consecutive layers and share it within the small group. The key idea is shown in Fig. 1. As AcrossKV can combine multiple layers’ KV caches into a single one rather than just two adjacent layers, it offers higher potential reduction ratio compared to Liu et al. () while simplifying the implementation to realize the benefits of the KV cache reduction. (See Sec.  2 for more detailed comparison with Liu et al. ()).

### 3.4 Knowledge Recovery

While SingleInputKV preserves all the original parameters, it re-wires the architecture so that the KV cache projections may receive different inputs. We found that this re-wiring (and AcrossKV) requires fine-tuning to recover the original capabilities from the modified model. As we only change the computation of the attention part for layer $>l$, this can be achieved by fine-tune just the ${\mathbf{W}}_{QKV}$ weight matrices from the $(l+1)$-th layer onwards. However, instead of directly fine-tuning these parameters using standard LM loss, we find that distilling using the output logits of the original model allows for better knowledge recovery (see Sec. 5 for more details).

Implementing the Distillation. Since only a few ${\mathbf{W}}_{QKV}$ parameters need fine-tuning, we are able to do a memory efficient parameter-sharing based distillation. More specifically, we keep a single copy of the original model weights in memory that are frozen during training, and add an extra trainable copy of the ${\mathbf{W}}_{QKV}$ parameters for layers $>l$ initialized using the original model (See Fig. 1).

During the training, we create two forward modes for the later layers $>l$, one with original frozen parameters using original architecture, and another with the SwiftKV re-wiring using new QKV projections i.e.,

$$ ${\bf y}_{teacher}={\mathbf{M}}({\bf x},SwiftKV=False),~{}~{}\text{and}~{}~{}~{ }{\bf y}_{student}={\mathbf{M}}({\bf x},SwiftKV=True),$ (3) $$

where ${\bf y}_{\cdot}$ is the final logits, ${\mathbf{M}}$ is the model, and ${\bf x}$ is the input. Afterwards, we apply the standard distillation loss ($L$) upon the outputs with temperature ($\tau$) using (Hinton et al., ).
After the distillation, the original KV projection layers $>l$ are discarded during inference.

This method allows us to perform the distillation for Llama-3.1-8B-Instruct on 680M tokens of data in 3 hours using 8 H100 GPUs, and Llama-3.1-70B-Instruct in 5 hours using 32 H100 GPUs across 4 nodes.

### 3.5 Optimized Implementation for Inference

LLM serving systems can be complex and incorporate many simultaneous optimizations at multiple layers of the stack, such as PagedAttention (Kwon et al., ), Speculative Decoding (Leviathan et al., ), SplitFuse (Holmes et al., ; Agrawal et al., ), and more. One benefit of SwiftKV is that it makes minimal changes to the model architecture, limited to only a few linear projection layers. This means that SwiftKV can easily be integrated into existing serving systems without implementing new kernels (e.g. for custom attention operations or sparse computation) or novel inference procedures.

Implementation in vLLM. To realize the performance benefits of SwiftKV, we integrated it with vLLM . Our implementation is compatible with vLLM’s chunked prefill, which processes prefill tokens in chunks and may mix prefills and decodes in each minibatch. During each forward pass, after completing layer $l$, the KV-cache for the remaining layers ($>l$) are immediately computed, and only the decode tokens are propagated through the rest of the model layers.

GEMM and Memory Optimizations. Upon this basic implementation, we implemented two additional optimizations. First, *SingleInputKV fusion*: instead of computing the KV cache ${\mathbf{KV}}_{j}$ for each layer $j>l$ one at a time, we fused all ${\mathbf{W}}_{KV}^{j}$ into one large weight matrix ${\mathbf{W}}_{KV}^{j>l}$ so that their KV cache can be computed with a single efficient GEMM operation. Second, *AcrossKV reduction*: we modified vLLM to only allocate one layer’s KV-cache for each group of merged layers, which realizes the memory gains of AcrossKV.

## 4 Main Results

### 4.1 Setup

Training and Evaluation. We use Llama-3.1-8B-Instruct and Llama-3.1-70B-Instruct as our base models for SwiftKV.
Our training datasets include a mixture of the full supervised training data from HuggingFaceH4/ultrachat_200k (Ding et al., ) and teknium/OpenHermes-2.5 (Teknium, ). We evaluated model quality using a modified LM-Eval-Harness (Gao et al., )(^2^22[https://github.com/neuralmagic/lm-evaluation-harness/tree/llama_3.1_instruct](https://github.com/neuralmagic/lm-evaluation-harness/tree/llama_3.1_instruct)) due to its support for the custom prompt format of Llama-3.1, particularly for MMLU and MMLU-CoT (Hendrycks et al., ), GSM8K (Cobbe et al., ), and Arc-Challenge (Clark et al., ). For more details, please see Appendix C.

Compression Metrics. For prefill computation, we report the approximate reduction as $(L-l)/L$ due to SingleInputKV, and for KV cache, we report the exact memory reduction due to AcrossKV. For example, SwiftKV with SingleInputKV ($l=L/2$) and 4-way AcrossKV is reported as 50% prefill compute reduction and 37.5% KV cache memory reduction. We further study how these theoretical compute and memory reductions translate into end-to-end inference improvements in Sec. 4.3.

Inference Performance. In our inference evaluation, we focus on two common scenarios: *batch-inference* for cost sensitive scenarios and *interactive-inference* for latency sensitive scenario.

*Batch-Inference* When processing text in bulk or when serving a model under high usage demand, it is important to achieve high *combined throughput* in terms of input + output tokens processed per second. For bulk processing, the combined throughput determines the time it takes to finish the job. For interactive use, it determines the volume of concurrent users and requests that can be supported per unit of hardware. In both scenarios, the combined throughput is a key determinant of the cost of serving the model.

*Interactive-Inference* In interactive scenarios (e.g., chatbots, copilots), not only the combined throughput is important, but also metrics that define the end-user experience. Chief upon them are the time-to-first-token (TTFT) and time-per-output-token (TPOT). TTFT is the time between the user sending a message and receiving the first token in the response. TPOT is the time between each output token after the first token has been received. Low TTFT and TPOT are desirable by interactive applications to deliver smooth usage experiences.

For all experiments on Llama-3.1-8B-Instruct, we use 1 NVIDIA H100 GPU with 80GB of memory, and for all experiments on Llama-3.1-70B-Instruct, we use 4 NVIDIA H100 GPUs running the model with 4-way tensor parallelism. We provide the full hardware and vLLM configurations in Appendix C.2

### 4.2 Model Quality with Compression

Llama-3.1-8B-Instruct.
The top rows of  Table 1 show that SwiftKV can preserve model quality well until 50% prefill reduction using SingleInputKV. For 25% prefill reduction, the accuracy degradation is only about 0.12 points and for 50% reduction, the gap is about 1 point (^3^33Note that we did not try to find the best training recipe, regarding to either training data (e.g., we did not include any math or coding datasets) or training pipeline (e.g., we did not include reinforce-learning rated steps, like DPO/RLHF). Yet, the quality of SwiftKV is close to the original base Llama-3.1-8/70B-Instruct models. In Sec. 5.3, we show that better data recipe could boost the model performance and close the quality gap further.). When we push to 62.5% reduction (i.e. SingleInputKV with $l=12$ and $L=32$), the accuracy drops to 66.09 points, which is significantly lower than the baseline. This can be explained by the drop in activation similarity from 0.61 to 0.51 between layer 16 to layer 12 (Fig. 2).

The bottom rows of Table 1 show the model quality when adding AcrossKV to 50% SingleInputKV. From pure SingleInputKV to 2-way AcrossKV, the accuracy drops about 0.9 points with 25% KV cache reduction.
The accuracy drops by another 0.32, going from 2-way to 8-way sharing, and 0.62 when going all the way to 16-way sharing.
Particularly, for the extreme case, i.e., using a single KV cache for all remaining layers, the accuracy is only about 2.5 points lower than pure SingleInputKV, and could be useful for more memory constrained cases, e.g., embedding and/or mobile devices.

Furthermore, the design of AcrossKV is complementary to many existing KV cache compression methods. In Sec. 5.4, we show that AcrossKV can be combined with quantization to achieve 62.5% reduction in KV cache memory with only a 1-point accuracy gap compared to SingleInputKV only.

Table: Table 1: Llama-3.1-8B-Instruct

Table: Table 2: Llama-3.1-70B-Instruct

Llama-3.1-70B-Instruct.
Table 2 shows that with 50% prefill reduction using SingleInputKV, Llama-3.1-70B-Instruct incurs a 1.3 point drop in accuracy which is slightly higher than the results of Llama-3.1-8B-Instruct. However, Llama-3.1-70B-Instruct is more resilient to AcrossKV, incurring less than a 0.35 point drop in accuracy even for 4-way sharing across layers.

### 4.3 Inference Performance

Batch Inference Performance.
Fig. 3 shows the results of Llama-3.1-8B and Llama-3.1-70B across several workloads with a range of input lengths. SwiftKV achieves higher combined throughput than the baseline model across all the workloads we evaluated.

For Llama-3.1-8B-Instruct, with 2K input tokens per prompt, SwiftKV achieves $1.2-1.3\times$ higher combined throughput than the baseline model, and our benefits increase further to $1.8-1.9\times$ higher combined throughput with 128K inputs. Note that for an input length of 8K tokens, SwiftKV achieves a staggering 30K tokens/sec/GPU (480 TFLOPS/GPU). For Llama-3.1-70B with 2K input tokens per prompt, SwiftKV achieves $1.4-1.5\times$ higher combined throughput than the baseline model, which improves to $1.8-2.0\times$ better combined throughput for 128K inputs.(^4^44While the total compute savings is roughly $2\times$, the end-to-end speedup is lower due to two main reasons: i) the performance improvement is limited to the decoding computation which needs the output activation of all the layers. Fig. 2 (right) shows the max possible speedup for Llama-3.1-8B-Instruct during model forward pass despite the decoding overhead, and ii) due to additional vLLM overheads outside of the model forward pass, such as sampling, optimizing which is beyond the scope of the paper.) As expected, SwiftKV provides greater improvements when the inputs are long.

We also observe AcrossKV can further improve the combined throughput due to its ability to reduce the memory usage for the KV-cache and supporting larger batch sizes. For sequence length of 8K, Llama-3.1-70B-Instruct with SwiftKV achieves a combined throughput of over 16K toks/sec over 4xH100 GPUs which corresponds to 560 TFLOPS/GPU of bf16 performance when normalized to baseline. This is an unprecedented throughput for BF16 inference workloads.

Figure: Figure 3: Combined input and output throughput for Llama-3.1-8B (left) and Llama-3.1-70B (right) across a range of input lengths (bottom). For each experiment, a number of requests are sent to vLLM at once such that the total number of tokens is roughly 15M. Each request generates 256 output tokens.
Refer to caption: x4.png

Interactive-Inference Performance.
Fig. 4 shows the TTFT and TPOT of Llama-3.1-70B-Instruct across a range of request arrival rates and input lengths. When the arrival rate is too high, the TTFT explodes due to the request queue accumulating faster than they can be processed by the system. However, SwiftKV can sustain $1.5-2.0\times$ higher arrival rates before experiencing such TTFT explosion. When the arrival rate is low, SwiftKV can reduce the TTFT by up to 50% for workloads with longer input lengths. In terms of TPOT, SwiftKV achieves significant reductions for all but the lowest arrival rates, up to 60% for certain settings. A similar story unfolds for Llama-3.1-8B, which can be found in Fig. C.1 in the Appendix.

Figure: Figure 4: Time to first token (TTFT, top) and time per output token (TPOT, bottom) for input lengths 2000 (left), 8000 (middle), and 32000 (right) for Llama-3.1-70B. For each experiment, a range of different request arrival rates is simulated. Each request generates 256 output tokens.
Refer to caption: x5.png

## 5 Ablation and Discussion

Table: Table 3: Impact of Distillation and Full/Partial Model Finetuning on Llama-3.1-8B-Instruct

### 5.1 The impact of distillation

To demonstrate the effectiveness of our distillation, we train Llama-3.1-8B-Instruct with 50% SingleInputKV and no AcrossKV using the standard language model loss, and compare it with our distillation based approach discussed in  Sec. 3.4.
The results are shown in Table 3 (a).
As we can see, the model trained with distillation has a 2.64 point higher average.
Particularly, for generative tasks, i.e., MMLU-Cot and GSM-8K, the performance improvement is 4.13 and 6.74, respectively.

### 5.2 Full model training vs. partial model training.

Our distillation method only fine-tuned the ${\mathbf{W}}_{QKV}$ parameters, as discussed in Sec. 3.4, with the hypothesis that it preserves the knowledge from the original models compared to fine-tuning the entire model. This hypothesis aligns with Meng et al. (), Geva et al. (), and Elhage et al. (), which suggest that MLP layers player a more prominent role in storing knowledge.

To validate this, we fine-tuned a model with 50% SingleInputKV on Llama-3.1-8B-Instruct where all parameters in the latter 50% of layers are trained.
The results are shown in Table 3 (b).
Note that the model quality of full model distillation is about 4.5 points lower than our proposed partial model distillation.

### 5.3 The impact of fine-tuning datasets

Note that in Sec. 4, we did not try to maximize the performance of SwiftKV from the data recipe perspective since the search space is very large and outside the scope of our paper. However, we want to share some initial findings about the dataset recipe.

How good is the data used to train SwiftKV?
We chose the datasets to train SwiftKV due to their popular adoption and broad domain and task coverage.
However, as compared to other high-quality domain specific fine-tuning datasets, they may have weaknesses.
To measure the quality of these two datasets, we directly fine-tuned a model using the Llama-3.1-8B base model, and compared this trained model with the Llama-3.1-8B-Instruct model released by Meta.

The results are shown in Table 4 (a).
The original Llama-3.1-8B-Instruct has a average score of 73.71 but the model trained using our two datasets only achieved 65.77.
This indicates the training data used for SwiftKV is not optimal and there may be opportunities to further improve the results we reported in Sec. 4 as discussed next.

Does more math/coding data help GSM-8K?
From Table 1, the main degradation among 7 tasks for 50% SingleInputKV is GSM-8K.
This may be due to the lack of math and coding examples in the two datasets we picked to train the model.
To verify this, we distilled SwiftKV using one extra math-related dataset, gretelai/synthetic-gsm8k-reflection-405b (GretelAI, ), and one extra coding dataset, ise-uiuc/Magicoder-OSS-Instruct-75K (Wei et al., ), in total about $8\text{K}+75\text{K}=83\text{K}$ samples, and about 16M tokens.

The results are reported in Table 4 (b).
The performance of all tasks except Winogrande are slightly improved, with the average score being 0.23 higher.
Particularly, GSM-8K improves the most, with a 0.53% improvement. This is expected since we added extra math and coding datasets.
Considering the small amount of new data (83k vs. 1.2M), the improvement is remarkable.

This study indicates that improvements in distillation data is potentially an important direction for future work, particularly domain-specific datasets to reduce the quality gap compared to the original model when using SwiftKV.

Table: Table 4: The impact of datasets on Llama-3.1-8B-Instruct.

### 5.4 Combining with other KV cache compression methods

SwiftKV explores an orthogonal design space from many existing KV cache compression methods, which means that it can be easily combined with them, e.g., sliding window (Jiang et al., ), token-level pruning (Liu et al., ), quantization (Hooper et al., ) etc.
In this section, we show the combined effect of SwiftKV with per-token KV cache FP8 quantization  (Yao et al., ) using PyTorch’s natively supported float8_e4m3fn.  Table 5 shows the accuracy degradation is within 0.4 points for all cases, even though we applied post-training quantization with no quantization-aware finetuning.

Appendix D explores a second, potentially interesting, trade-off between AcrossKV (inter-layer) vs GQA (intra-layer) KV cache sharing.

Table: Table 5: Llama-3.1-8B-Instruct KV cache quantization results.

Table: Figure 5: Density of early exit probabilities and alignment of early exit vs final logits.

### 5.5 Simple Early Exit for Decoding Tokens

SingleInputKV allows all the KV cache needed for generating future tokens to be computed without having to forward-pass though the entire LLM. This means that even the decoding phase could exit earlier without worrying about missing KV cache for subsequent tokens.

To test the feasibility, we added an early exit language modeling head. We then used the input to SingleInputKV layer to calculate output logits, and incorporated them as part of the distillation training. Our results are preliminary and requires further evaluation, but we found that the alignment between early exit logits and the final output logits to be over 66% when the largest probability from the early exit logits is over 95% (Fig. 5). We used this as our early exit criteria (i.e., decoding tokens exit early if it predicts an output token with 95%+ probability), and  Table 6 shows a sample result. See Appendix E for more details.

## 6 Conclusions

In this paper, we presented SwiftKV, a novel model transformation for reducing inference cost for prompt-dominant workloads, combined with a KV cache reduction strategy to reduce memory footprint, and a light-weight distillation procedure to preserve model accuracy. While we presented strong results on the effectiveness of SwiftKV, exploration of parameter-preserving model transformations for inference optimization is still in its early stages. We have identified both limitations as well as areas of improvement. Given the simplicity and effectiveness of SwiftKV, we hope that this will spark further exploration from the AI community which we hope will continue lower the latency and cost of inference.

## Limitations and Future Work

It is important for every work to acknowledge its limitations and suggest future directions, particularly for LLM-related works.
In our work, we did not aim to optimize the training data selection though we provide potential ways in Sec. 5.3.
Additionally, we did not include a detailed benchmark analysis for our method.
However, as shown in Sec. 5.3, we ensured that our datasets were not cherry-picked to overfit the reported tasks.
Furthermore, we did not finetune our model with advanced post-training approaches, like DPO and RLHF, which we leave for future work.
Finally, we hypothesize that our method can work even better when combined with pretraining or continued-pretraining,
but due to resources constraints, we did not explore this direction.
We hope to revisit these ideas in the future.

## Appendix A Additional Motivation

Figure: Figure A.1: The input similarity of small scale models (left) and large scale models (right).
Refer to caption: x7.png

## Appendix B Calculating Computation Reduction

For a vanilla transformer model Vaswani et al. () with a hidden dimension size $d$, the total amount of computation per token per layer is roughly given by $24d^{2}$, where $4d^{2}$ goes to KV projections, another $4d^{2}$ goes to Q projection and context GEMM, and the remaining $16d^{2}$ goes to MLP computation. For popular GQA (Ainslie et al., ) based architectures like Llama (Touvron et al., ), Mistral (Jiang et al., ), Qwen (Bai et al., ), and many others, the KV projection is reduced to $4d^{2}/G$, where $G$ is the group size and is $\geq 4$. As such, the compute cost of the KV projection is negligible ($\frac{4d^{2}/G}{20d^{2}+4d^{2}/G}$ which is less than 5% when $G\geq 4$), compared to the rest of the computation.

With SingleInputKV, pre-fill tokens only require computing this negligible KV projection for layers $>l$. Therefore, it reduces the total pre-fill computation to approximately $l/L$.
Fig. 2 shows significant end-to-end reduction in model forward time for different value of $l$ compared to $L$. In Sec. 4.3, we show that such compute reduction translates to considerably higher aggregate serving throughput and lower latency, reducing both serving cost, and improving end-user experience.

## Appendix C Experimental Details

### C.1 Training and evaluation details

We directly use the Huggingface LLama-3.1 checkpoints, particularly, “meta-llama/Meta-Llama-3.1-8B-Instruct” and “meta-llama/Meta-Llama-3.1-70B-Instruct”.
For datasets, we use the supervised finetuning datasets from “HuggingFaceH4/ultrachat_200k” and “teknium/OpenHermes-2.5”, which in total is about 1.2M samples, and about 160M tokens.
We set training epochs to be 2, learning rate to be 3e-4, weight decay to be 0.05, warm up ratio to be 5%, maximum sequence length to be 8192 with attention separated sequence packing, the distillation temperature to be 2.0, and the training batch size to be 32 for both Llama-3.1-8B-Instruct and Llama-3.1-70B-Instruct.

Our evaluation follows [https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8](https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8) using the github repository [https://github.com/neuralmagic/lm-evaluation-harness/tree/llama_3.1_instruct](https://github.com/neuralmagic/lm-evaluation-harness/tree/llama_3.1_instruct).
The main reason behind this is that the implemention from the repository aligns with original Llama-3.1 evaluation, which has superme scores over the original Lm-eval-harness repository.
One issue we found in the provided commands is the one used to run MMLU-5-shots.
Directly using the command does not give us desired accuracy.
Therefore, we added both --apply_chat_template and --fewshot_as_multiturn, and the accuracy is even slightly higher than what they reported.

For all tasks, we follow the same number of few shots and/or chain of thoughts as the provided commands.
We present the number of shots and metrics used in the paper in Table C.1.

Table: Table C.1: The setting for different tasks

### C.2 Inference Speedup Evaluation Details

Figure: Figure C.1: Time to first token (TTFT, top) and time per output token (TPOT, bottom) for input lengths 2000 (left), 8000 (middle), and 32000 (right) for Llama-3.1-8B. For each experiment, a range of different request arrival rates is simulated. Each request generates 256 output tokens.
Refer to caption: x9.png

#### Hardware Details.

We ran all inference speedup experiments on a AWS p5.48xlarge instance, with 8 NVIDIA H100 GPUs, 192 vCPUs, and 2TB memory. Llama-3.1-8B-Instruct experiments are run using 1 of the 8 GPUs, and Llama-3.1-70B-Instruct experiments are run using 4 of the 8 GPUs.

#### vLLM Configuration.

We ran all experiments with enforce_eager and chunked prefill enabled with max_num_batched_tokens set to 2048. To run each benchmark, we instantiated vLLM’s AsyncLLMEngine and submitted requests using its generate method according to each benchmark setting. For each request, the inputs are tokenized before being submitted, and the outputs are forced to a fixed length of 256.

#### Llama-3.1-8B Latency Evaluation.

See Fig. C.1.

## Appendix D Inter-layer AcrossKV vs Intra-Layer KV cache Reduction

In this section, we share different design choices of AcrossKV, which considers the tradeoff between GQA (Ainslie et al., ) and the acorss layer sharing into the design.
Particularly, when $\text{AcrossKV}\geq 2$, we can either use GQA and AcrossKV together or we can simply use AcrossKV to get all savings.
For instance, when $\text{AcrossKV}=4$ (a.k.a., the second row of the final session in Table 1), we have KV cache reduction from both GQA and AcrossKV.
However, we can either do multi-query attention (MQA) for all 16 layers or do multi-head attention (MHA) but share the KV cache for all 16 layers.

Table: Table D.1: Llama-3.1-8B-Instruct AcrossKV design

We present the $50\%$ SingleInputKV reduction with MQA, GQA plus AcrossKV, and GQA plus MHA in Table D.1, that all have the same KV cache reduction, 37.5%.
AcrossKV-GQA actually provides the best performance.
One thing to notice is that the AcrossKV-MHA is actually worse than the result of AcrossKV-16x from from Table 1 even though AcrossKV-MHA has larger KV cache than AcrossKV-16x.
We hypothesis that this might be related to hyper-paramter tuning but did not invest deeper.
Also, note that pure MQA leads to worst performance, which is about 17 points lower than AcrossKV-GQA

How to effectively balance inter/intra-layer KV cache sharing is an interesting direction to explore.
We hope that our initial experiments here shed some light for future research.

## Appendix E Early Exit Details

Thanks to SingleInputKV, there is no need to go through the entire network to compute all KV cache for later tokens generation.
This also brings another unique opportunity as compared to standard transformer architecture design: it allows the model to perform early exit to speedup both the prefill and generation phases without worrying about missing KV cache.

To add early exit upon SwiftKV, we add an extra language modeling head and use the input to calculate SingleInputKV to directly compute the output logits.
We apply the same distillation loss of this early exit logits with the original teacher’s logits and directly sum the two distillation losses together as our final optimizing loss.

After training, we first look the alignment between the early exit logits and the final logits.
The figure is shown in Fig. 5.
The left axis plot the density of early exit largest probabilities.
And the right axis shows the probabilities when the maximum of early exit logits aligns with the final logits for each bin.
As can be seen, when the maximum logits is larger than $0.95$, the alignment likelihood is larger than 66%.
Note that this is achieved within 160M token training and longer training should help the alignment.

How to use early exit is always an interesting directly and research topic.
Fully exploring this is out of the scope of this paper.
Here we adopt a simple heuristic approach.
When the largest probability is larger than 0.95, we directly use the early exit prediction.
Otherwise, we still process the remaining layers to get the final prediction.
We apply this logic to some questions we selected from Alpaca (Taori et al., ) and the results are shown in Table 6 and Appendix E.1.

### E.1 Early Exit Examples

Here we provide more examples of early exit from Alpaca dataset.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 4
questions: Q1. In Table 1, there is a significant drop in performance for 8B model on the math dataset GSK-8K. I suppose this is the harder case, meaning that the proposed method may not work well for the case of small models on tasks demanding more logic and reasoning. An analysis on why this performance drop occurs would be interesting. 

Q2. In Table 3, to show the impact of distillation, AcrossKV is disabled. However, to reduce KV cache, AcrossKV should be enabled for inference serving, right?

Q3. In Table 4, the performance of "our fine-tuned model" is significantly inferior to the base model. The result seems to be negative to the usefulness of your techniques. I don't quite understand the logic here.
soundness: 3
strengths: S1. Inference optimization for in Transformer-based LLMs is an important topic which has been extensively studied in recent years.

S2. Several key components have been proposed in this paper, with their usefulness showcased in the evaluation. 

S3. The proposed method is orthogonal to many existing optimizations and they can be used jointly to further optimize the performance.
summary: This paper studies the optimization of inference in Transformer-based LLMs. It presents SwiftKV, a solution to reducing the KV cache and inference time to long contexts up to 128K. The proposed method features three parts: SingleInputKV, AcrossKV, and knowledge recovery. Experiments show the effectiveness of the proposed method and the usefulness of its components, as well as how they work jointly with other optimization techniques.
weaknesses: W1. SingleInputKV borrows observations and ideas from previous works, as stated in the submission (such observation has also been utilized in the InfiniGen paper published at OSDI 2024).

W2. A core technique in the proposed method is cross-layer KV cache compression. The comparison/discussion with state-of-the-art KV cache compression/merging/cross-layer works is missing, e.g., PyramidKV and infini-attention. It is encouraged to discuss the difference and the novelty compared to existing KV cache compression techniques in the related work. Some surveys can be found here:
https://github.com/October2001/Awesome-KV-Cache-Compression

W3. Whereas the paper discusses long inputs, it lacks discussions with recent works on long contexts (see the above link), such as MInference, which optimize the prefilling of long contexts. Some of them exploit the sparsity to reduce KV cache and speed up inference, e.g., ALISA. 

W4. The proposed method is not training free, yet only Llama-3.1 models are evaluated. It is unclear if the performance (and its optimal parameter settings) also translates to models. Extension to other open models, such as Mistral, would be beneficial to understanding the contributions of this work.

---

confidence: 4
contribution: 2
initial_rating: 6
presentation: 2
questions: Please refer to W1 and W2.
soundness: 2
strengths: S1. The paper proposes two techniques derived from insights from prior research, demonstrating their efficacy in reducing computational and memory costs during LLM inference.

S2. The authors show that fine-tuning can alleviate the decline in benchmark scores, emphasizing the practicality of the proposed methods without notably sacrificing model performance.
summary: This paper proposes two methods to reduce the cost of the prefill stage during LLM inference. 
The first method, called SingleInputKV, reuses the output hidden state vector from the i-th layer in attention layers as the input vector to generate key and value (KV) vectors of the subsequent layers. In previous methods, the j-th layer used the output hidden state vector from the (j-1)-th layer as input.
The second method, called AcrossKV, enables the KV vectors generated by the i-th layer to be reused by the following layers. In previous methods, each layer generates its own KV vectors by multiplying the input vector with its weight.
These techniques reduce computational costs by reusing the input and KV vectors of earlier layers for later layers. They also decrease the number of KV vectors that need to be cached for the decode stage. The proposed methods build on prior work [1], which showed minimal differences in the values of input vectors across layers as the number of layers increases in transformers.
The authors implemented these techniques in Llama-3.1-8B and Llama-3.1-70B models,  showing that while the performance on the LLM benchmark remains largely unaffected, both time and memory usage in the prefill stage are reduced by almost two times.

[1] Songwei Liu, Chao Zeng, Lianqiang Li, Chenqian Yan, Lean Fu, Xing Mei, and Fangmin Chen. Foldgpt: Simple and effective large language model compression scheme, 2024c. URL https://arxiv.org/abs/2407.00928.
weaknesses: W1. The experiments in the paper are somewhat limited.
- The authors evaluate the proposed techniques only on Llama-3.1 models. Testing a wider variety of models would strengthen the results. If the proposed methods could demonstrate their benefits across transformer models with different attention mechanisms (e.g., sparse attention, low-rank attention), scaling approaches (e.g., wide scaling, deep scaling, sparse scaling), and sizes (Llama-3.2-1B, Llama-3.2-3B, Llama-3.2-8B,  Llama-3.2.-11B, Llama-2-13b, and Llama-3.2-70B, Llama-3.2-90B, and Llama-3.1-405B), it would enhance the paper’s contribution. 
- The authors need to demonstrate whether applying SwiftKV to larger models yields more significant results compared to small models. Incorporate models such as  Llama-2-13B and Llama-2-7B. If applying the proposed methods to Llama-2-13B yields better results than Llama-2-7B in terms of both cost and the benchmark scores, it would strengthen the contribution of the paper.
- There is no experiment to show the independent effect of AcrossKV without the presence of SingleInputKV, leaving the isolated impact of AcrossKV unexplored. The authors need to compare a baseline model to one with only AcrossKV applied.

W2. The justification for the claimed reduction in computational cost is insufficient. The paper needs to specify which operations are being skipped by SingleInputKV clearly by providing a detailed breakdown of the computational costs for each component of a Transformer model, comparing the baseline to SwiftKV. In Figure 1, SingleInputKV still appears to need to generate the output hidden state vector of every attention layer, which is a primary computational task in Transformer models. This is because the proposed method needs to generate the query vector for each attention layer in Figure 1, and the output hidden state vector from the (i-1)th layer is required to compute the query vector for the i-th layer.

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 2
questions: * Section 3.4, Knowledge Recovery: The description suggests that the distillation is done for every of the later layers, but Figure 1 suggests that at least W_K and W_V are only trained for the initial layer of each AcrossKV block.
* Why are the results more or less consistently better for 4-way caching compared to 2-way caching for the 70B model? That seems kind of counterintuitive.
* footnote 5, page 7: What are the end-to-end results?
* Section 4.3: "a combined throughput of over 16K toks/sec over 4xH100 GPUs which corresponds to 560 TFLOPS/GPU"
  * So that is around 4k token per second for each GPU compared to 30k tokens/sec/GPU for 8B Llama model. But because the 70B model is much more complex, there are more floating point operations necessary?
  * Any notion why the pure compute performance increases despite a more "distributed" setting (multiple GPUs)?
* Any notion why the full model fine-tuning performs so much worse than the partial model fine-tuning?
* Section 5.3: "This may be due to the lack of math and coding examples in the two datasets we picked to train the model."
  * Why did you choose these datasets, if at least the coding use case is serving as a motivational example?
* Doesn't the discussion in Appendix B the whole point of the paper, i.e. trying to optimize a part than accounts for less than 5% of the total compute time?
soundness: 3
strengths: The ideas for the various optimizations are presented reasonable clearly and they seem novel as well, especially their combination. The evaluation on a number of models and datasets/benchmarks supports their performance claims and a reasonable ablation study is provided as well.
summary: SwiftKV proposes several techniques to improve to reduce computation and memory footprint for LLM inference while maintaining a similar level of accuracy. In particular, they propose to skip the pre-fill stage of later layers by rewiring the model and rely instead on intermediate computation results from an earlier layer, leading to a reduced amount of computation. The authors also reduce the memory footprint of later layers by sharing a KV cache across multiple subsequent layers. Additional they use a distillation/fine-tuning process of the affected model part to reduce the difference in accuracy compared to the original model. They show in their evaluation computation improvements for throughput as well as latency.
weaknesses: For me, the biggest issue is that end-to-end results are missing, which makes it hard for me to put the presented inference results (throughput, latency) into context, which also makes me question how useful the presented numbers are.

* apart from SingleInputKV, all the other optimizations are not properly motivated regarding the reasoning why they should work (some form of microbenchmark)
* end to end results are missing, especially since some of their writing, if I am not mistaken, suggests that they target a part of the pipeline that only compromises 5% of the "runtime"
* line 314: "The accuracy drops by another 0.32, going from 2-way to 8-way" - it is actually 1.32 according to the table
  * similarly the number for 16-way is wrong as well
* The authors might consider removing 5.5 to have more space for presenting the other content/result in more detail.
* not clear why the used benchmarks are representative for the use cases mentioned as motivation in the introduction

detailed copy editing comments:
* related work:
  * "their optimized implementations in TensorRT (NVIDIA, 2019)" - all previously mentioned techniques were published after 2019
* Figure 2, right side:
  * The parameters used in the legend are not explained at all in the caption. It is possible to understand after reading the subsequent text.
  * The subfigure is never actually referenced in the text, except in the appendix, as "proof" for some statement later and in a footnote.
* Figure 4: Artifacts in the layering of the curves, sometimes a dot is at the top for one datapoint and and then further down for other datapoints. But maybe that was intentional?
* Table 3: There is a horizontal line missing after "(a) The effect of distillation".
* line 447: "which suggest that MLP layers player a more prominent role" - missing verb, probably it should be "play" instead of "player"
* Figure 5 is not readable
* minor issues:
  * typos:
    * line 119: "Tensor-Parallelism(Shoeybi et al., 2020)" - missing space
    * lines 130 to 138: additional brackets around the year for the citations
    * line 157/158: "(Holmes et al., 2024; Agrawal et al., 2024))" - additional bracket at the end 
    * line 360: "toks/sec" - probably "token/sec"
    * line 859: "superme" - perhaps "supreme"?
    * line 923: "hyper-paramter" - missing e
    * line 923: "but did not invest deeper" - probably "investigated"
  * references:
    * Clark et al. 2018: cited differently than the other arXiv papers
    * Cobbe et al. 2021: misses place, where the paper was published
    * Dao et al. 2024:
      * year states 2024, but conference abbreviation suggests 2022
      * conference abbreviation is nowadays NeurIPS
    * Ding et al. 2021: misses place, where the paper was published
    * Elhage et al. 2021: url not clickable
    * GretelAI 2024: url not clickable
    * Hendrycks et al. 2021: cited differently than the other ICLR papers
    * Hinton et al. 2015: cited differently than the other arXiv papers
    * Kuzim et al. 2024:
      * year states 2024, but conference abbreviation suggests 2022
      * conference abbreviation is nowadays NeurIPS
    * Lewis et al. 2020: conference abbreviation is nowadays NeurIPS
    * Liu et al. 2024a:
      * cited differently than the other arXiv papers
      * cited twice (2024b)
    * Liu et al. 2024d:
      * year states 2024, but conference abbreviation suggests 2023
      * conference abbreviation is nowadays NeurIPS
    * Meng et al. 2024:
      * year states 2024, but conference abbreviation suggests 2022
      * conference abbreviation is nowadays NeurIPS
    * Pourreza and Rafiei 2024:
      * year states 2024, but conference abbreviation suggests 2023
      * conference abbreviation is nowadays NeurIPS
    * Sakaguchi et al. 2019: cited differently than the other arXiv papers
    * Wei et al. 2023: cited differently than the other arXiv papers

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. The code link appears to be invalid. Could you make the code open-source to enhance reproducibility?
2. SwiftKV focuses primarily on optimization during the prefill stage. How should we interpret the decrease in TPOT shown in the performance results?
3. Could you provide results comparing the performance of SwiftKV with more competitive baselines, such as Minicache, as mentioned in your paper? Could you clarify the connections and differences between your method and existing work, including its strengths and weaknesses? Could you demonstrate whether your method can be integrated with other approaches? Additionally, could you outline the potential application scenarios for your method?
4. As I understand, most datasets used in your paper consist of multiple-choice questions, leading to longer prefill times and shorter decoding times. I’m interested in seeing SwiftKV's performance on more diverse datasets.
soundness: 2
strengths: 1. The writing is good.  
2. Experimental results indicate that the proposed method effectively reduces latency while preserving knowledge.
summary: This paper proposes SwiftKV, a method that reduces LLM inference latency while preserving knowledge. SwiftKV combines Early Exit, KV cache compression, and Knowledge Distillation techniques, demonstrating latency improvements in performance evaluation.
weaknesses: 1. Limited Efficient Experiments: VLLM serves as the only baseline in the performance results, which limits the demonstration of this work’s necessity and effectiveness. The authors’ method is a lossy optimization approach, and they should compare it with more serving systems to demonstrate respective performance improvements and knowledge retention. Although other methods may not conflict with the authors' approach, they may not be easily integrated (e.g. the strategy of Early Exit can hardly apply to Speculative Decoding[1], or combined with certain sparse attention methods like PowerInfer[2], quantization method like GPTQ[3] may result in significant performance degradation.). If the authors cannot demonstrate the effectiveness of their method compared to others, or show that it can integrate with other methods for added benefits, the significance of this work is greatly diminished.

2. Lack of Key Assumptions: Some critical assumptions are missing, such as noting that latency-sensitive servers often adopt disaggregated systems to handle the prefill and decode stages separately. This omission could impact the reported TTFT and TPOT performance results, because in the disaggregated systems, TPOT will hardly be influenced due to improvements in the prefill stage.

[1] Cai, Tianle, et al. "Medusa: Simple llm inference acceleration framework with multiple decoding heads." arXiv preprint arXiv:2401.10774 (2024).

[2] Song, Yixin, Zeyu Mi, Haotong Xie, and Haibo Chen. "Powerinfer: Fast large language model serving with a consumer-grade gpu." arXiv preprint arXiv:2312.12456 (2023).

[3] Frantar, Elias, Saleh Ashkboos, Torsten Hoefler, and Dan Alistarh. "Gptq: Accurate post-training quantization for generative pre-trained transformers." arXiv preprint arXiv:2210.17323 (2022).

Meta-review: While SwiftKV presents a solid engineering approach to optimizing LLM inference through KV cache optimization and model transformation (which I totally appreciate the effort especially the opensourced release), it (as the reviewers suggested) falls slightly below the bar for ICLR. Though the work demonstrates meaningful performance improvements, several limitations prevent acceptance: the core techniques largely build upon existing methods, the experimental evaluation is limited, and there's insufficient comparison with state-of-the-art KV cache compression techniques.

Decision: Reject