TL;DR
When you swap between different-sized LLMs in production (14B to 32B, for example), the receiving model has to re-run the entire prefill from scratch. This paper proposes a closed-form linear mapper that transfers the KV cache across models in the same family, achieving 2.7-25x speedup over re-prefill with 73-98% accuracy retention on most model pairs.
The Problem
Modern LLM deployments are not single-model systems. Three patterns dominate production:
Model cascading — route easy queries to a small model, hard ones to a large model. Mid-conversation switching — escalate to a stronger model when complexity increases. Dynamic routing — pick the best model per request based on classifier signals.
Every swap forces the receiving model to redo prefill from scratch. Prefill processes every input token through every layer to build the KV cache. For a RAG pipeline with thousands of context tokens, that's hundreds of milliseconds of GPU time — on every single handoff.
The reason reuse seemed impossible: source and target models have different layer counts, hidden dimensions, and attention heads. You can't just copy the cache.
How It Works
The authors discovered that KV caches within the same model family have strong linear structure. On Qwen3 14B→32B:
- A single source layer explains 56% of variance in target keys, 32% in values
- Using multiple source layers pushes this to 79% keys, 65% values
This means a simple ridge regression can map source KV caches to the target space. The pipeline has three steps:
Step 1 — Layer selection. For each target layer, select the top-k most predictive source layers (by R²).
Step 2 — RoPE stripping. Remove rotary position embeddings from keys before mapping, making the mapper position-free and reusable across all context lengths.
Step 3 — Ridge fit. Solve the closed-form ridge regression on just 500 calibration sequences:
$$W^* = (X^\top X + \lambda I)^{-1} X^\top Y$$
The mapper is fitted per attention head and reused at inference time.
Show Me The Code
import torch
from typing import List
def rotate_half(x: torch.Tensor) -> torch.Tensor:
x1 = x[..., : x.shape[-1] // 2]
x2 = x[..., x.shape[-1] // 2 :]
return torch.cat([-x2, x1], dim=-1)
def strip_rope(keys: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
"""Remove RoPE encoding to get position-free keys."""
return keys * cos - rotate_half(keys) * sin
def fit_ridge_mapper(
source_kvs: List[torch.Tensor], # [(N, D), ...] top-k source layers
target_kv: torch.Tensor, # (N, D) target layer KV
lambda_reg: float = 1e-4
) -> torch.Tensor:
"""Fit closed-form ridge regression mapper."""
X = torch.cat(source_kvs, dim=-1) # (N, k*D)
XtX = X.T @ X
XtY = X.T @ target_kv
reg = lambda_reg * torch.eye(X.shape[1], device=X.device, dtype=X.dtype)
W = torch.linalg.solve(XtX + reg, XtY) # (k*D, D)
return W
def transfer_kv(
source_kvs: List[torch.Tensor],
W: torch.Tensor,
cos: torch.Tensor = None,
sin: torch.Tensor = None,
is_key: bool = True
) -> torch.Tensor:
"""Transfer KV cache from source to target model space."""
if is_key and cos is not None:
kvs = [strip_rope(kv, cos, sin) for kv in source_kvs]
else:
kvs = source_kvs
return torch.cat(kvs, dim=-1) @ W
# --- Calibration (one-time, offline) ---
# 500 sequences x 1024 tokens from FineWeb-Edu
# source_layer_kvs: list of (N, D) tensors from top-k source layers
# target_layer_kv: (N, D) tensor from the target layer
W = fit_ridge_mapper(source_layer_kvs, target_layer_kv)
# --- Inference (fast path) ---
predicted_kv = transfer_kv(source_layer_kvs_at_runtime, W, cos=cos, sin=sin)
The calibration is a one-time offline cost. At inference time you only run the matrix multiply.
Benchmark Results
Tested across 6 model pairs in 3 LLM families:
| Model Pair | Accuracy Retention | Speedup vs Re-prefill |
|---|---|---|
| Qwen3 14B -> 32B | 98% | 25x |
| Qwen3 32B -> 72B | ~91% | ~18x |
| Other passing pairs | 73-89% | 2.7x+ |
| Failure cases (2 pairs) | Degraded | -- (MLP recovers +37pp) |
Key highlights:
- 4 out of 6 pairs hit 73-98% accuracy retention with the linear mapper alone
- 25x speedup on the best-case Qwen3 pair
- 500 calibration sequences is all you need — practically free
- Stable across multi-turn handoffs (no error accumulation)
- When linear fails, a nonlinear MLP recovers up to +37 percentage points on HellaSwag
Gotchas & Limitations
It only works within the same model family. The method requires matching KV head count and per-head dimension between source and target. Cross-family transfer (e.g., Qwen to Llama) is not supported.
Two out of six pairs fail the linear test. The paper does not yet offer a reliable way to predict which pairs will work before you try. Pre-deployment validation is mandatory.
Values are harder to transfer than keys. Single-layer R² for keys reaches 56%, but only 32% for values. The nonlinear MLP helps, but adds inference overhead.
Benchmarks are HellaSwag-centric. Performance on math reasoning, coding, or long-context retrieval tasks is not characterized.
Try It Today
The paper is fully self-contained and the method is simple enough to implement in an afternoon:
- Collect KV caches from both models on 500 calibration sequences
- Select top-k predictive source layers per target layer using R²
- Strip RoPE from keys, fit ridge regression closed-form
- Slot the mapper into your model-switching logic
If you are running vLLM or a PagedAttention-based system, this integrates naturally as a pre-processing step before inserting the transferred KV cache into the paged memory pool.
The paper is at https://arxiv.org/abs/2608.03893 — well worth reading if you work on LLM serving infrastructure.
What's your experience with KV cache optimization in production? Drop a comment.
Sources
- Heo et al. (2026). Cross-Model KV Cache Transfer in LLM Families: A Closed-Form Linear Mapping for Prefill Reuse. https://arxiv.org/abs/2608.03893
- arXiv cs.LG recent listings: https://arxiv.org/list/cs.LG/recent
Top comments (0)