BridgeKV prefills with the big model and decodes with the small one: translating the prompt cache makes requests 3.0× faster on a single A40
Synopsis
BridgeKV lets a large model in a model family run prefill once, then a lightweight KV-stitching module translates the large model's prompt cache into the small model's native cache layout, so every output token is generated by the small model alone. Across six model pairs and five datasets, it scores above fine-tuned small-model baselines with matched trainable parameter budgets in all 30 settings, closing up to 80% of the quality gap to running the large model for both phases. On one A40, Llama-3.1-Instruct-8B→Llama-3.2-Instruct-1B serves a 1K-token prompt with 256 output tokens in 2.96 s, close to the 2.80 s of the small model alone and 3.0× faster than the 8.83 s of the large model alone.
Interpretation
Within one request, prefill goes to the large model and per-token decoding goes to a smaller model of the same family, with no large-model blocks in the decode loop. Earlier prefill-decode disaggregation systems ran the same model on both sides of the transfer, and cross-model cache reuse covered only fine-tuned variants with identical layer counts and KV layouts. On one A40 with a 1K-token prompt and 256 output tokens, Llama-3.1-Instruct-8B→Llama-3.2-Instruct-1B takes 2.96 s, against 2.80 s for the small model alone and 8.83 s for the large model alone.
For each layer of the small model, the stitching module learns a weighted combination of a few nearby layers of the large model, then applies separate residual nonlinear maps to keys and values. Keys are stored after rotary position embedding, so a fixed map on stored keys would depend on token position; the module removes the large model's rotation from source keys before mapping and applies the small model's rotation afterward. For Llama-3.1-Instruct-8B→Llama-3.2-Instruct-1B the module has 118M parameters, 12% of the small model's non-embedding parameters, and runs once per request at the handoff.
Both models stay frozen while only the stitching module is trained, so the small model reading the stitched cache matches the large model's next-token distribution while the stitched cache stays close to the small model's native cache. Prior cross-model cache translators fit the receiving model's own cache or outputs rather than explicitly transferring the larger model's prediction behavior. In 12 of the 30 settings only the first training stage is used, so the small model decodes with exactly its fine-tuned baseline weights, and BridgeKV still scores higher in all 12.
Perspective
Hardware measurements cover one A40 at batch size 1 with prompt lengths from 512 to 8K tokens, for the four pairs whose decode anchor is 1B or 0.5B; the 30-setting quality results come from the Llama and Qwen instruction-tuned families on five datasets. The 33.86× and 27.5× figures come from a simulator that models request arrivals, queues and cache transfer, not from hardware measurement.
One stitching module is trained per model pair and dataset, and the two models must share a tokenizer; richer policies that adapt the handoff during generation are left as future work. The size of the quality gain varies by task and model pair, with the weakest case a 0.2% relative improvement on SAMSum for Llama-3.1-Instruct-8B→Llama-3.2-Instruct-3B.