Inference-Time Hyper-Scaling with KV Cache Compression

We thank Reviewer 2bQ1 for taking the time to read our paper and for providing thoughtful and constructive feedback.

W1 (a) and Q2: Broader representation of NLP tasks in the evaluation

We appreciate your suggestion to include a broader set of tasks! We note that DMS is primarily designed for compressing the KV cache during the auto-regressive decoding phase, rather than the parallel prefilling phase. Hence, the focus of our evaluation is primarily on long-generation/reasoning tasks rather than long-context tasks. Among the tasks you suggested, a prime example of the former is code generation: hence, we ran additional experiments for R1-Qwen 2.5 models on LiveCodeBench, trained with win=256 and 100 steps per CR. For parallel scaling, we report pass@all, namely whether any of the generated threads passed the tests for a given problem.

Similar to other long-generation tasks, we find that hyper-scaling with DMS yields vastly improved performance for the same compute budget, measured as cumulative memory reads.

In addition, we attach the results of Llama 3.1 8B on the QASPER benchmark (win=256, 200 steps per CR) for QA on scientific research papers. We find that DMS mostly maintains performance on this long-context NLP task.

While our focus is on long-sequence reasoning, these results indicate DMS maintains performance on other NLP tasks. We agree that broader evaluations would be valuable and are a natural extension of this work.

W1 (b) and Q1: Little hardware-level evidence for speed-ups, no reports of measurements

We provide latency measurements during generation using the Megatron-LM framework. We run our measurements on H100 GPUs in bfloat16 precision using FlashAttention2 and the flash_attn_with_kvcache routine for both vanilla and DMS. We note that we use research code for DMS that has not been optimised (e.g., through fusion of multiple operations), resulting in a visible latency overhead. In addition, we would like to note that for small to moderate token batch sizes, KV cache size does not drive the majority of latency (please refer to Appendix H and Figure 6 for details). However, a practical inference scaling scenario of serving multiple users at a time, requires substantial token batch sizes. In addition, Qwen 2.5 models have smaller KV cache sizes than the Llama 3 family or the newer Qwen 3 family, when measured in relation to the number of model parameters. Finally, we remark that our Megatron-LM framework setup is designed for fast inference, albeit we could not optimise it to the point of matching state-of-the-art inference engines due to time constraints of the response period.

For latency benchmarks, we measure the generation of reasoning chains for AIME24. The notation “Lat @ X” denotes the latency when generating a token at position X. From the results below, it emerges that even for moderate batch sizes (such as 64 and 128) and sequence lengths (>-8k), CR4 and CR8 decrease latency considerably. The reason why there is not a direct correspondence between the reduction in memory reads (due to higher CR) and reductions in latency is that their dependency emerges only at large token batch and sequence sizes. We refer to Appendix H for an explanation with heatmaps. Overall, these findings suggest that DMS results in effective latency reductions due to decreased KV cache traffic in regimes such as inference-time scaling (parallel or sequential). In turn, this enhances reasoning quality by allowing for a larger effective token budget for the same wall-clock runtime (while also reducing memory load).

We will include latency plots for all model sizes in the final version of our paper.

W2 (a) and Q3: Exploration of how the method behaves on pathological prompts

The Needle-in-a-Haystack and Variable Tracking (VT) tasks (Table 1 in the manuscript; Tables 5 and 7 in Appendix D) consist of synthetically constructed inputs, and DMS models perform well on both (often exceeding the original model’s performance). In particular, the VT task measures the ability to track long-range dependencies, and NIAH attention patterns are highly non-local. As our models are retrofitted without access to private pre-training and post-training data, we cannot confirm whether the inputs for these two tasks are OOD, although there is a high chance that this is the case. Searching for adversarial inputs is interesting, but beyond our core focus on compressing KV caches for long reasoning chains, where we assume that the prompt is well-formatted.

W2 (b) Logit distillation with sampled data raises concerns about generalization. How well the approach transfers to domains with significant distribution shifts or task variations?

Logit distillation (model matching) mitigates domain mismatch more effectively than language modelling loss (data matching) when retrofitting LLMs [1, 2]. In addition, DMS performs well under short retrofitting runs (e.g., 100 steps per CR in the LiveCodeBench experiments above), which involve considerably less data than used for pre-training. Combined with small learning rates, this suggests minimal risk of overfitting to the distillation data and altering the model behaviour.

[1] S.T. Sreenivas, S. Muralidharan, R. Joshi, M. Chochowski, A.S. Mahabaleshwarkar, G. Shen, J. Zeng, Z. Chen, Y. Suhara, S. Diao, C. Yu, W.-C. Chen, H. Ross, O. Olabiyi, A. Aithal, O. Kuchaiev, D. Korzekwa, P. Molchanov, M. Patwary, M. Shoeybi, J. Kautz, B. Catanzaro. LLM Pruning and Distillation in Practice: The Minitron Approach, 2024. [2] B. Minixhofer, I. Vulić, E. M. Ponti. Universal Cross-Tokenizer Distillation via Approximate Likelihood Matching. arXiv preprint arXiv:2503.20083, 2025.

We appreciate Reviewer eA55’s time and effort in reviewing our work and for the helpful suggestions they offered.

W1: Which L-W-CR setups yield the best results and why?

In our manuscript, we hint at the fact that sequential scaling rather than parallel scaling brings more gains (L240); however, we also remark that often hybrid strategies employing parallel and sequential scaling lie on the Pareto frontier. For CR, the key trade-off is whether the accuracy gains from scale outweigh accuracy losses due to sparsification. In our main experiments, CR 4x consistently achieves the best balance. Post-submission, we improved performance further using win=256, shifting the optimal CR to 8x - 10x. Hence, the optimal CR could be pushed even further by developing better retrofitting strategies and model configurations that do not degrade the model’s quality.

W2: Does the method fail with fewer than 1B training tokens?

DMS works even on small numbers of training tokens, as we demonstrate below in the answer to Question 2. However, longer retrofitting may reduce the variability of the results. In our experiments after the submission, we found that enlarging the sliding window also helps to stabilise the training and reduce variability, which can be seen in Figure 4 (right). With win=256, we were able to improve the results of R1-Qwen models shown in Figure 5 with only 100 steps for every CR increment instead of 3000 steps. With 100 steps for every CR increment at 8k context and batch size 1024, retrofitting requires 839M tokens per CR, and we have not attempted to optimise these numbers further. Please note that DMC uses 24B–168B tokens depending on CR; their 1B number applies to the neuron decay phase, not full adaptation.

Q1: Figures 2 and 3 are a bit crowded with the L-W-CR labels

We agree and will revise the plots to improve clarity: adjusting label positions, plot layout, font size, and potentially removing less important labels from non-Pareto points.

Q2: How would the right subfigure of Figure 4 look like with a 1B training token budget (as in the DMC paper?)

Reducing the amount of training data increases the variability of results, especially at small sliding window and model sizes. We have run additional experiments retrofitting Llama 3.2 1B Instruct with only 50 steps per CR, achieving CR 4x with ~630M training tokens. We ran 3 seeds for different window sizes (16, 64, 256). Considering that the vanilla model achieves 47.0 on GSM8K, the additional data-efficient GSM8K scores are:

While DMS incurs some degradation due to the small model size and the extremely fast training regime, we note that a larger window size significantly stabilises training and mostly preserves the original performance.

A discussion of limitations is missing

We discuss the limitations of our work in Appendix A. We will also expand this in the Related Work section to highlight which limitations are inherited from DMC and which choices in the DMS training procedure were able to overcome some of the limitations of DMC, such as being data-hungry, more cumbersome to implement, being more prone to lead to performance degradation etc.

We appreciate Reviewer BS8X’s careful reading and helpful comments, which we address below.

W1: Improvements in writing

We will revise the manuscript to correct the title casing, resolve issues with the notation, and reorder the introduction of $q[0]$ (see below). We will also double-check the format instructions to rule out any other inconsistencies.

W2: Does DMS require independent retrofitting for different compression ratios?

Fortunately, this is not the case! A single run with a gradually increasing CR yields multiple checkpoints for various compression levels. All our experiments are based on this setup. This enables a more efficient use of compute and flexible continuation for higher compression ratios. We will clarify this better if the paper is accepted.

Below we specify resource estimates for retrofitting the models used in the manuscript at 100 steps/CR to CR 8x. Please see the response to Q5 for details.

W3: Practical evaluation on efficiency

We provide latency measurements during generation using the Megatron-LM framework. We run our measurements on H100 GPUs in bfloat16 precision using FlashAttention2 and the flash_attn_with_kvcache routine for both vanilla and DMS. We note that we use research code for DMS that has not been optimised (e.g., through fusion of multiple operations), resulting in a visible latency overhead. Also, for small to moderate token batch sizes, KV cache size does not drive the majority of latency (see Appendix H and Figure 6 for details). However, a practical inference scaling scenario of serving multiple users at a time and multiple rollouts per query (parallel scaling) requires substantial batch sizes. In addition, Qwen 2.5 models have smaller KV cache sizes than the Llama 3 family or the newer Qwen 3 family, when measured in relation to the number of model parameters. Finally, we remark that our Megatron-LM framework setup is designed for fast inference, albeit we could not optimise it to the point of matching state-of-the-art inference engines due to time constraints of the response period.

For latency benchmarks, we measure the generation of reasoning chains for AIME24. The notation “Lat @ X” denotes the latency when generating a token at position X. From the results below, it emerges that even for moderate batch sizes (such as 64 and 128) and sequence lengths (>-8k), CR4 and CR8 decrease latency considerably. The reason why there is not a direct correspondence between the reduction in memory reads (due to higher CR) and reductions in latency is that their dependency emerges only at large token batch and sequence sizes. We refer to Appendix H for an explanation with heatmaps. Overall, these findings suggest that DMS results in effective latency reductions due to decreased KV cache traffic in regimes such as inference-time scaling (parallel or sequential). In turn, this enhances reasoning quality by allowing for a larger effective token budget for the same wall-clock runtime (while also reducing memory load).

We will include latency plots for all model sizes in the final version of our paper.

Q1: Extraction of the decision variable from $q[0]$

We agree with the Reviewer that the role of $q[0]$ in predicting the eviction decision should be better clarified in Section 3, we will adopt this suggestion in the next version of our paper. The reason why it was only introduced in the experimental setup is that this is not part of the “core” design of our method. In principle, alternative approaches exist, such as introducing a trainable vector $w$ and computing the decision variable as $x^Tw$. Reusing q[0] or k[0] is primarily a form of performance optimisation.

Q2: Ambiguous L, H, T symbols on L149

Thank you for bringing this to our attention. We will revise the formula using $|L|$, $|H|$, and $|T|$ to denote the number of layers, attention heads, and sequence length, respectively.

Q3: Avoiding materialisation of the attention mask

As shown in Figure 1(b), the attention mask can be reconstructed from a vector $\alpha = [log(1 - \alpha_1), \log(1 - \alpha_2), …, \log(1 - \alpha_n)]^T$ , as these values are replicated as columns of the mask, with the diagonal and sub-diagonals zeroed to account for the sliding window. This regular structure can be leveraged for a more efficient implementation of self-attention. A detailed discussion of a similarly compressed representation for a family of structured masks—which our mask belongs to—is given in [1], along with source code based on FlashAttention 2. We will convey this better in the text and reference this work in the final version of our manuscript.

[1] G. Wang, J. Zeng, X. Xiao, S. Wu, J. Yang, L. Zheng, Z. Chen, J. Bian, D. Yu, H. Wang, FlashMask: Efficient and Rich Mask Extension of FlashAttention, 2025.

Q4: Only one q[0] is retrofitted for a GQA group, so are the eviction decisions consistent for all heads within a group?

Our models use GQA, where keys are shared among queries. Extracting the decision variable from a single query for each KV group impacts only one query-key pair, whereas extracting the variable from k[0] would impact all query-key pairs in which this key is used. We emphasise that all queries within the GQA group share the same KV cache after eviction.

Q5: Method not being robust in Figure 4 right for win=16

The win=16 result highlights that DMS can evict even very recent tokens effectively—contrary to prior assumptions—when delayed eviction is used. That said, a larger window (e.g., win=256) yields better data efficiency and makes the model operate more like vanilla, which results in overall better accuracy (as seen in Figure 4, right). Importantly, since CR is enforced over the entire context (e.g., 4k or 8k), increasing the window size should not increase memory use. We ran a series of post-submission experiments confirming that enlarging the window to 256 tokens does improve the results in reasoning models, and allows for reducing the number of retrofitting steps by an additional factor of 30x from 3000/CR to 100/CR. We will clarify this new practical insight and update the results in Figures 2 and 3. Below we show an excerpt from the results for R1-Qwen 7B in these two setups: (win=16, 3000 steps/CR) and (win=256, 100 steps/CR). AIME24 scores and KV budgets have been averaged over 3 runs with different random seeds due to the small dataset size.

The table shows that win=256, even with 30x shorter retrofitting, yields much better models.

As for the results in Figure 4 right for win=16, Llama 3.2 1B Instruct model proved brittle during retrofitting, with GSM8K scores degrading quickly (e.g., DMC at CR 3x and 4x in Table 1). We use "robust" to mean the model still performs reasonably well and does not collapse the score under difficult settings (e.g., win=16 and low token budgets), but also maintains strong results with larger windows like win=256.

L1: Multiple runs per CR and data efficiency

Please see our response to W2 and Q5, which hopefully assuages the Reviewer’s concerns, since only a single run is needed for all CRs. In addition, we pushed data efficiency to 60x the previous SOTA trainable KV cache compression (DMC)

L2: Showing latency measurements gains instead of proxy metrics

We refer to our latency measurements reported in our response for W3. In a nutshell, DMS translates into effective latency reduction in regimes with large batch sizes and sequence lengths, which are common in inference-time scaling..

We appreciate Reviewer 1K9q’s taking the time to comment on our paper and would like to thank them for their questions and remarks.

Q1: Can different retrofitted setups (context lengths and CRs) share checkpoints during retrofitting?

Yes indeed! A single DMS retrofitting run produces a sequence of checkpoints with increasing compression ratios. The target CR increases linearly from 1x to the preset maximum, and intermediate checkpoints (e.g., 2x, 3x) are saved along the way. The DMS models in Figures 2 and 3 were generated from single retrofitting runs per model size.

The answer for context length has an additional nuance. For Figures 2 and 3, we trained R1-Qwen models with a fixed 8k context and evaluated them by generating sequences beyond this length, such as 16k and 32k. Hence, the DMS retrofitting procedure did not affect length extrapolation (possibly due to taking only a few steps and teaching to evict recent tokens).

Q2: In which scenarios training-free KV cache compression methods would be more attractive?

Apart from the main advantage, which is the lack of need for extra training, training-free methods might offer advantages in particular edge cases. Sliding-window methods like TOVA maintain full KV cache fidelity until the window is exceeded, preserving vanilla accuracy on short inputs. This stands in contrast with DMS, which compresses all sequences to some extent, albeit it compresses shorter sequences to a ratio inferior to the target (and longer ones to a ratio superior to the target). For the same reason, methods like TOVA allow for pre-specifying the memory cap, whereas the KV cache in DMS grows as a function of the sequence length.

Nonetheless, for most intents and purposes, retrofitting is superior to training-free alternatives, as these do not offer a statistical guarantee of preserving performance across tasks and often lead to significant degradations [1]. We will include all these considerations in the Limitations section.

[1] P. Nawrot, R. Li, R. Huang, S. Ruder, K. Marchisio, E.M. Ponti, The Sparse Frontier: Sparse Attention Trade-offs in Transformer LLMs, 2025.