FlexBudget: Fair Comparison of KV Cache Eviction Strategies Under a Unified Compute Budget

In long-context inference, the bloating of the KV cache is an unavoidable engineering challenge. In a poster presentation at PyTorch Conference North America 2026, AMD engineer Huan Ran proposed a solution called FlexBudget, attempting to address a long-neglected methodological problem: when different teams propose various KV cache eviction rules, are we really comparing the merits of the rules themselves, or are we comparing whose CUDA kernel is written faster?

The Problem: Rules Entangled with Implementation

When a model processes ultra-long contexts, the KV cache becomes extremely large and cannot be fully retained. Researchers have therefore proposed various eviction strategies: keeping the most recent tokens, keeping tokens with the highest attention scores, or hybrid rules combining position and attention weights.

The problem is that each method usually comes with its own custom kernel, tightly binding the rule to the implementation. This means that when Method A beats Method B on some benchmark, you cannot tell whether A's eviction logic is smarter, or whether A's kernel just happens to be more friendly to the current hardware. According to Huan Ran's description in the talk, this confusion makes it difficult for both academia and industry to draw reliable conclusions.

KV Cache is a core data structure in Transformer autoregressive inference. When generating each new token, the model needs to recompute its attention with all historical tokens, and the KV cache stores the Key and Value vectors of historical tokens to avoid redundant computation. The cache size is approximately 2 × number of layers × number of attention heads × head dimension × sequence length × precision bytes. For a 70B-class model with a 32K sequence, the cache can easily reach tens of GB, even exceeding the model weights themselves. The longer the context, the more the cache grows linearly, and memory bandwidth and capacity become the main bottleneck for inference throughput. Therefore, evicting low-value KV entries becomes a key means of reducing the cost of long-context inference.

FlexBudget's Approach: Fix Everything, Only Swap the Rules

The core idea of FlexBudget is very straightforward—fix the kernel and memory budget, and only replace the eviction rule itself. The same code, the same compute quota, letting different strategies compete under completely fair conditions.

According to Huan Ran's introduction, the method has two key characteristics:

• Training-free: Can be used without any additional training or fine-tuning
• All core PyTorch: Fully implemented based on native PyTorch, without relying on custom CUDA extensions

From an engineering perspective, "all core PyTorch implementation" means FlexBudget can maintain consistency across different hardware backends, lowering the barrier to reproduction, and freeing method comparisons from hardware details.

Training-free means the method can be applied directly to an already-trained model without any gradient updates or fine-tuning. For KV cache eviction, this means the eviction decision is computed online entirely from signals visible at inference time (such as attention scores, token positions, and recency), introducing no additional parameters and not changing the original model weights. Its advantages are low deployment cost, plug-and-play capability, and no damage to the original model's capabilities; the trade-off is that rule design can only rely on inference-time statistics and cannot adaptively optimize like learning-based methods. FlexBudget chooses the training-free route, further ensuring the purity of comparisons between different rules.

All core PyTorch means the implementation only uses standard operators and APIs built into the PyTorch framework, without writing or relying on custom CUDA/C++ extensions. Custom kernels are usually manually optimized for specific GPU architectures, which can bring significant speedups, but also make methods difficult to migrate across hardware, raise the barrier to reproduction, and make it hard to separate performance gains from algorithmic contributions. Although a pure PyTorch implementation may sacrifice some peak performance, it gains hardware independence and reproducibility, which is actually more valuable in methodological comparison scenarios—it separates "rule quality" from "engineering implementation quality."

Why This Matters

KV cache optimization is a key part of reducing cost and improving efficiency in long-context inference, but the methodological foundation in this field is currently not solid. When reporting results, most papers package rule improvements and engineering optimizations together, making it hard for readers to disentangle the sources of contribution.

The value of FlexBudget is not necessarily in proposing some entirely new eviction rule, but in providing a cleaner experimental paradigm. Once the kernel and budget are unified, researchers can more clearly answer fundamental questions like "which attention signal is more worth retaining."

For engineers deploying long-context inference in real systems, this comparability is equally important—what you need is not the best number on some benchmark, but which eviction rule is truly more effective under your own hardware and budget constraints.

Huan Ran will present this work in the poster session at PyTorch Conference North America 2026 (October 20, Community area). Interested researchers can discuss the details on site.

Experimental paradigm here refers to a method comparison framework after controlling variables: fixed kernel, fixed memory budget, fixed hardware and precision, with only the eviction rule as the independent variable that changes. Traditional KV cache research often bundles rule design and kernel optimization together when publishing, making result differences impossible to attribute. A clean experimental paradigm borrows the principle of "isolating variables" from systems research, giving conclusions transferability. For academia, it allows rules from different teams to compete fairly on the same baseline; for industry, the comparability it provides is more valuable for decision-making than the best number on a single benchmark.

One-line Summary

AMD engineer proposes FlexBudget, providing a fair comparison paradigm for KV cache strategies by fixing the kernel and memory budget while only swapping the eviction rule.

Paragraph Summary

To address KV cache bloat in long-context inference, AMD engineer Huan Ran proposed FlexBudget at PyTorch Conference North America 2026. Existing KV cache eviction methods usually bind rules to custom kernels, making it impossible to determine whether performance differences come from rule quality or engineering implementation. FlexBudget fixes the kernel and memory budget and only replaces the eviction rule itself, and has two key characteristics: training-free and all core PyTorch implementation, thereby maintaining consistency across hardware backends and lowering the barrier to reproduction. Its value lies not in proposing a brand-new eviction rule, but in providing a cleaner experimental paradigm to help researchers answer which attention signal is more worth retaining, and to provide engineers deploying long-context inference with a comparable basis for decisions.