Strona główna / Artykuły / Moving KV Cache Between Qwen3 Models Instead of Re-Prefilling

Artykuł opublikowany po angielsku.

Moving KV Cache Between Qwen3 Models Instead of Re-Prefilling

NVIDIA research maps Qwen3-14B attention state into Qwen3-32B with a ridge-regression bridge. See what the latency numbers mean, where it fails, and why agents care.

1817 słów

Picture a coding agent that's been running for twenty minutes. It has read through a repository, opened the docs, called a handful of tools, and piled up close to 30K tokens of context on Qwen3-14B. Then it hits the one step where a stronger model would really help.

Sending the next turn to Qwen3-32B is the easy part. Keeping the work that's already been done is the hard part. The 32B model has to chew through the whole accumulated context again before it can emit a single token. That pass is called prefill, and what it produces is the model's KV cache.

A recent NVIDIA research preprint asks a pointed question. The smaller model already built a KV cache for this conversation, so could that state be converted into the larger model's cache rather than making the larger model rebuild it from raw text? In the reported Qwen3 14B to 32B experiment, the mapper takes 277.6 ms at 32,768 tokens, against 6,975.3 ms for re-prefilling on the target model. That's roughly 25x for that one stage. It's a promising result, but it's much narrower than "25x faster LLM inference."

Why switching models wastes work

The KV cache is the attention state an autoregressive transformer builds while it reads a prompt. Once it exists, the model doesn't need to recompute attention over the whole prefix for every token it generates.

Reuse like this works naturally inside a single model. Serving systems already cache prefixes and shuttle KV state between compatible workers; NVIDIA TensorRT-LLM, for instance, supports KV-cache reuse and KV exchange for disaggregated serving.

Switching models is where it falls apart.

Qwen3-14B
   │
   │  30K-token conversation
   ▼
14B KV cache
   │
   ✕ cannot simply hand this to 32B
   │
Qwen3-32B normally re-prefills the text

Models of different sizes have different layer counts and different internal representations. Sharing a tokenizer or a model family doesn't make their caches interchangeable. NVIDIA treats this as a representation-mapping problem.

14B KV cache
     ↓
model-pair mapper
     ↓
32B-compatible KV cache
     ↓
32B starts decoding

The study is deliberately limited to within-family pairs with matched KV: source and target must agree on KV head count and on head dimension. The pairs tested come from Qwen3, Llama 3.1 and Ministral 3. Moving across families, say Qwen to Llama, is left as an open question.

A linear bridge is enough, surprisingly

The interesting finding isn't that a mapper can be trained at all. It's that a fairly plain closed-form mapper does well on several model pairs.

The paper's linear-structure analysis shows that for Qwen3 14B to 32B, the single best source layer explains about 56% of the target's key variance and 32% of its value variance. Pool a selection of source layers and those numbers climb to about 79% and 65%.

The production mapper rests on three ideas.

Pick useful source layers. A target layer doesn't just copy the source layer at the same depth. The method finds the source layers whose KV states best predict that target layer and concatenates their features.

Fit keys and values separately, per target head. The bridge itself is ridge regression, solved independently for each target layer and head, with one fit for keys and another for values. A stripped-down version of the math looks like this:

import numpy as np

def fit_ridge(x, y, lam=0.01):
    """Fit Y ≈ XW + b. Conceptual demo, not NVIDIA's implementation."""
    x_mean = x.mean(axis=0, keepdims=True)
    y_mean = y.mean(axis=0, keepdims=True)
    xc = x - x_mean
    yc = y - y_mean
    eye = np.eye(xc.shape[1], dtype=xc.dtype)
    w = np.linalg.solve(xc.T @ xc + lam * eye, xc.T @ yc)
    b = (y_mean - x_mean @ w).squeeze(0)
    return w, b

def map_cache(x, w, b):
    return x @ w + b

The real mapper's x is far bigger, since it concatenates KV features from all the selected source layers, and every target head gets its own fit. Calibration uses 500 FineWeb-Edu sequences of 1,024 tokens each, subsampled to around 128K token-level observations per target head.

Strip position out before mapping keys. This is the detail most worth remembering. RoPE rotates keys according to where a token sits, so mapping rotated keys directly risks learning the position distribution of the 1,024-token calibration data instead of the content. The method undoes the source model's RoPE rotation first, maps the position-free content representation, and then applies the target model's RoPE.

source key
    ↓
strip source RoPE
    ↓
map content representation
    ↓
apply target RoPE
    ↓
target key

Because of this, the mapper keeps working at positions it never saw during calibration, including the 32K latency test.

What the headline number covers

Here are the latencies reported for Qwen3 14B to 32B on an 8xH100 node:

  • 64 tokens: 14.0 ms for mapping versus 61.7 ms for re-prefill, about 4x faster.
  • 8K tokens: 67.8 ms versus 1,154.8 ms, about 17x faster.
  • 32K tokens: 277.6 ms versus 6,975.3 ms, about 25x faster.

Read these carefully. They compare the time to run the mapper with the time the receiving model needs to re-prefill. They are not a 25x speedup of generation overall, because decoding still has to happen afterward. The paper also flags a deployment caveat: a real multi-process serving setup might have to ship the mapped cache over to the target process, and the mapper benchmark doesn't turn that into a network-inclusive number.

Quality is good on some pairs and poor on others. Qwen3 14B to 32B keeps 97.6% of the 32B target's average standalone accuracy over the five reported benchmarks. Qwen3 8B to 32B keeps 87.5%, Llama 3.1 8B to 70B keeps 72.8%, and Ministral 3B to 8B keeps 76.2%. Then Ministral 3B to 14B and 8B to 14B drop to just 44.2% and 41.6% average retention with the linear mapper. In this study, matching KV dimensions is a prerequisite, not a promise of a clean handoff.

Where this helps agents

Long-running agents are exactly the place where repeated prefill hurts.

Here's what a routing policy built on this could look like:

"""Illustrative serving API. The paper does not provide this interface."""
state = agent.run(model="qwen3-14b", task=task)

if state.needs_stronger_reasoning:
    mapped_kv = kv_bridge.convert(
        source_model="qwen3-14b",
        target_model="qwen3-32b",
        kv_cache=state.kv_cache,
    )
    state = agent.resume(
        model="qwen3-32b",
        kv_cache=mapped_kv,
    )

Routing to a bigger model isn't new; that already works today. The shift is that the computation state could travel with the conversation. A research agent might explore cheaply on a small family member, hand off to a larger one for a hard stretch of reasoning, then move back down for routine work. Mappings are directional, so 14B to 32B and 32B to 14B need separate mappers, and the paper motivates both quality escalation and cost reduction.

The multi-turn experiment alternated Qwen3 14B and 32B across 100 CoQA conversations. By turn 10 the small-to-large gap had widened by about 1.7 percentage points, and large-to-small drift grew by roughly 0.33 points per turn. That's decent evidence for conversational switching, though the researchers themselves warn that linear drift could build up over much longer sessions.

Limits that matter in production

The paper is good enough to change how you think about model routing. It's not something you can switch on in a serving stack yet.

  • The compatibility envelope is narrow. All reported experiments stay inside model families, use dense full-attention models, and run six pairs with identical KV head counts and head dimensions. This method hasn't been tried across families or on mismatched KV shapes.
  • Every pair needs its own validation. Two Ministral mappings fail badly under ridge regression even though their KV architecture lines up. A nonlinear MLP recovers HellaSwag retention by about 24 to 37 percentage points on those failing pairs, which is encouraging, but it shows the simple linear story doesn't hold everywhere.
  • The mapper isn't small. The evaluated mappers hold roughly 1.01B to 3.36B parameters and take 4 to 12 GB per direction. Qwen3 14B to 32B comes in near 1.07B parameters and 4 GB. Fitting needs about 47 to 87 minutes per pair on a single 8xH100 node.

Since mappings are directional, a router across many models piles up pair-specific state. At the paper's average mapper size, that works out to 39 GB with three models, then 79 GB with four, and 131 GB once you reach five. The mappers can sit in host memory or storage instead of VRAM, but they're still infrastructure you have to run.

One more thing: there's no official implementation attached to the paper today. The arXiv record has the paper and its source, but no NVIDIA repository for this mapper, so the snippets above are explanations rather than drop-in code. For the serving-side context, see the TensorRT-LLM docs on KV cache reuse and the paper's arXiv record.

What to keep in mind

People have mostly treated KV caches as state that belongs to one model. If cross-model transfer gets reliable enough, that assumption goes away.

For agent infrastructure, the interesting future isn't just a router that picks a model per request. It's a router that moves a developed conversation state between compatible models without paying the full context cost a second time.

None of this makes KV cache portable in general. It doesn't remove decode cost, it doesn't solve cross-family routing, and the weakest pairs show why testing each pair matters. But the Qwen3 14B to 32B result shows that the work stored in one model's attention state can be reused by another model for less than recomputing it from text.

If that carries over into real serving systems, cascading between models gets far more interesting for long-running agents. The question stops being only which model should answer the next turn. It also becomes how much of the previous model's computation can come along.