KVBridge

Memory-bounded, failure-safe cross-model KV-cache transfer for LLM families.

KVBridge is a production-oriented implementation and systems extension of NVIDIA's August 2026 paper, Cross-Model KV Cache Transfer in LLM Families. It maps a source model's attention cache into a compatible target model so the target can continue without replaying the entire prompt.

The core algorithm is paper-faithful: target-layer-specific source selection, independent K/V ridge maps, cross-head features, and RoPE removal/reapplication. The systems work around it is original to this project: cancellation-resistant centered statistics, bounded-memory fitting, out-of-core SafeTensors shards, resource planning, revision/tokenizer gates, tamper-evident artifacts, guarded fallback, structured telemetry, a Hugging Face handoff adapter, and a reproducible validation suite.

Evidence boundary: T0-T1 pass locally. A real Tesla-T4 T2 run on Qwen3 0.6B→1.7B completed and failed both preregistered quality gates despite a measured mapping-speed advantage; its raw rows and hash-bound manifests are checked in. No 14B→32B quality or H100/H200 result is claimed. That remains a partner/high-memory experiment.

What ships

Capability	Implementation
Paper mapper	Per-target-layer, per-target-head ridge maps for K and V
Layer selection	Head-averaged single-source R², then top-k cross-layer features
Position handling	Exact inverse/forward RoPE using factors emitted by each model
Modest-host fitting	Target-layer blocks, deterministic token stride, CPU/CUDA statistics
Distributed path	Mergeable sufficient statistics plus torch.distributed all-reduce
Data plane	SafeTensors calibration shards; no pickle deserialization
Artifact plane	Atomic writes, SHA-256, model fingerprints, measured FP32/BF16 storage
Runtime safety	Shape/architecture/numerical/latency/logit-KL gates, explicit fallback
Integration	Qwen/Llama-style Hugging Face DynamicCache handoff adapter
Operations	JSON resource planner, structured handoff events, CI, package build
Research	Attention-output metrics, pinned T2 jobs, local results, multi-lab protocol

Thirty-second proof

python -m pip install -e .
kvbridge demo --output artifacts/demo
kvbridge inspect artifacts/demo
kvbridge plan configs/qwen3_14b_to_32b.paper.json

The demo downloads nothing. It creates two synthetic model-family members with different depths and RoPE bases, hides a known cross-layer affine map, fits only from cache pairs, and validates on an unseen sequence. The checked-in run recovered every true source layer and achieved holdout R² above 0.999998.

System design

flowchart LR
    A["Source model prefill"] --> B["Capture K/V + source RoPE"]
    B --> C["Strip source RoPE"]
    C --> D["Top-k layer features"]
    D --> E["Per-head K/V ridge maps"]
    E --> F["Apply target RoPE"]
    F --> G{"Runtime gates"}
    G -->|accept| H["Target decode from mapped cache"]
    G -->|reject| I["Visible full-prefill fallback"]
    J["Versioned SafeTensors artifact"] --> E
    K["Telemetry + canary probe"] --> G

Loading

Fitting is a two-stage, multi-pass process. The calibration input can be a Python sequence or a factory that reopens shards. Only centered sufficient statistics are retained:

[ W=(X^T X + \lambda I)^{-1}X^T Y,\qquad b=\bar{Y}-\bar{X}W ]

Target-layer blocking bounds peak memory. Batchwise Chan/Welford updates accumulate centered covariance directly, avoiding cancellation from subtracting large raw moments after a scale run. For Qwen3 14B→32B at k=8, the planner reports a 1,073,872,896-parameter (4.0005 GiB FP32) artifact and about 0.56 GiB of fit statistics per one-layer block. The cost is more sequential passes over calibration shards; this is intentional and configurable.

Local evidence

The committed local results were generated on an 8-thread CPU-only PyTorch 2.8 runtime. Each mapping measurement uses 10 warmups and 100 timed iterations.

Case	Cache geometry	Holdout R²	Exact layer recovery	Median map latency
Micro	3→2 layers, 2 heads × 4 dim, 64 tokens	0.99999873	Yes	0.339 ms
Small	5→4 layers, 2 heads × 8 dim, 128 tokens	0.99999895	Yes	0.781 ms
Medium	8→6 layers, 4 heads × 16 dim, 256 tokens	0.99999892	Yes	2.844 ms

These tests validate shape semantics, layer selection, ridge recovery, RoPE round trips, serialization, and scaling behavior. They do not estimate Qwen/Llama downstream accuracy or GPU speedup.

Real-model T2 evidence

The checked-in Qwen3 0.6B→1.7B T2 evidence was executed on a free Tesla T4 at commit e69e139. It used 16×512-token FineWeb-Edu calibration sequences (2,048 stride-sampled observations), eight held-out 256-token sequences, FP16 models, CUDA FP32 accumulation, and a 58,777,600-parameter BF16 mapper.

Metric	Result	Preregistered gate
Cache R² mean	0.2347	diagnostic only
Attention-output cosine mean / minimum	0.6568 / 0.2937	minimum ≥ 0.90 — failed
One-token logit KL mean / p95	1.6145 / 4.5903	p95 ≤ 0.20 — failed
Next-token agreement	25.0%	diagnostic only
Transfer median / p95	46.09 / 54.23 ms	hardware-local
Target prefix-prefill median	87.29 ms	hardware-local
Per-case prefill/transfer ratio median	3.074×	speed only; not acceptance

This is a useful negative result: the bridge was faster on this setup, but the mapped state was not safe to serve. The runtime policy therefore requires full-prefill fallback for this pair/configuration. Recompute the lightweight publication record with:

python experiments/validate_published_evidence.py \
  configs/qwen3_0.6b_to_1.7b.t2-smoke.json \
  results/t2/qwen3-0.6b-to-1.7b-smoke

Fitting from out-of-core shards

from kvbridge.config import FitConfig, ModelSignature
from kvbridge.fit import fit_mapper
from kvbridge.io import calibration_shard_factory

source = ModelSignature(...)
target = ModelSignature(...)

mapper = fit_mapper(
    calibration_shard_factory("data/calibration"),
    source,
    target,
    FitConfig(
        top_k=8,
        ridge_alpha=0.01,
        accumulation_dtype="float32",
        selection_target_layer_block_size=8,
        target_layer_block_size=1,
    ),
)
mapper.save("artifacts/qwen3-14b-to-32b")

The factory is re-iterable because memory-bounded fitting makes several deterministic passes. Increase block sizes on high-memory nodes to reduce I/O.

For GPU fitting, set accumulation_device="cuda". Artifact compression and runtime compute precision are separate: a BF16 artifact can be loaded once with mapper.to("cuda", dtype=torch.float32) on hardware without native BF16 execution. The serving hot path then reuses resident weights instead of transferring them per request.

Live Hugging Face handoff

from kvbridge.huggingface import greedy_handoff_generate
from kvbridge.mapper import CrossModelKVMapper

mapper = CrossModelKVMapper.load("artifacts/qwen3-14b-to-32b")
tokens = greedy_handoff_generate(
    source_model=source_model,
    target_model=target_model,
    mapper=mapper,
    input_ids=input_ids,
    max_new_tokens=64,
    eos_token_id=tokenizer.eos_token_id,
)

The adapter withholds the final prompt token, maps the preceding prefix cache, and lets the target consume that final token to produce its first logits. The adapter is deliberately marked experimental until each target Transformers release/model revision passes the integration matrix in the experiment protocol.

Failure is a first-class outcome

GuardedTransferEngine rejects non-finite or unbounded caches, enforces an optional latency budget, invokes an application quality probe, and emits a structured event on both acceptance and fallback. It never silently substitutes a failed bridge.

Shadow probes can be deterministically sampled by request ID, so retries receive the same oracle decision and sampling remains auditable. Batch size and token count are bounded before mapping; the serving layer still owns request-level concurrency and memory limits.

Production rollout should progress through:

1. Offline reconstruction and attention-output diagnostics.
2. Shadow traffic with full target prefill as the oracle.
3. Canary traffic with automatic quality and latency fallback.
4. Pair-specific benchmark gates and rollback thresholds.
5. Broader traffic only after revision-pinned evidence passes.

Reproduction tiers

Tier	Hardware	Purpose	Status
T0	Any CPU	Unit, corruption, fallback, artifact, planner tests	Passing
T1	Any CPU	Synthetic structural/scale sweep	Completed
T2	1 capable GPU	Tiny/small same-family end-to-end integration	Executed; 0.6B→1.7B gates rejected
T3	Multi-GPU lab	Qwen3 14B→32B paper reproduction	Partnership-ready config
T4	Multi-lab	Long-context, quantized, distribution-shift study	Proposed protocol

See architecture, experiment protocol, Lightning H100/H200 runbook, collaboration tracks, threat model, and production checklist.

Free-GPU real-model path

The pinned configs/qwen3_0.6b_to_1.7b.t2-smoke.json job is sized for a free 16 GB-class notebook GPU. It captures revision-pinned FineWeb-Edu shards, fits on CUDA, saves a BF16 artifact, and evaluates real target-query attention cosine plus shadow-prefill logit KL. On Kaggle:

git clone https://github.com/souvikDevloper/kvbridge.git
cd kvbridge
bash scripts/kaggle_t2_smoke.sh

The driver is stage-resumable by default. It reuses calibration, mapper, or evaluation output only after validating config/model provenance, SafeTensors structure, per-shard SHA-256 records, mapper checksums, finite standard-JSON metrics, and recomputed aggregates and bootstrap confidence intervals. Set KVBRIDGE_RESUME=0 for a clean fail-closed run into empty paths.

See the free-GPU runbook. The raw result exists, but the repository correctly records the run as rejected because its preregistered gates did not pass.

Development

python -m pip install -e ".[dev]"
ruff check src tests experiments
pytest
python experiments/run_local_scale.py --repeats 100 --warmup 10
python -m build

Scope and limitations

• v0.1 deliberately gates on a shared tokenizer, matched KV heads/dimensions, and dense full attention: the regime validated by the NVIDIA paper.
• Reconstruction R² is a debugging metric, not a deployment acceptance metric. The paper finds attention-output cosine more predictive of downstream retention.
• A mapper is directional and model-revision-specific. Updating either checkpoint invalidates the artifact fingerprint and requires recalibration.
• Full Qwen/Llama evaluation requires gated model access and substantial accelerator memory. The repository contains the protocol and planner, not a false claim that those runs occurred here.
• No API stability is promised before v1.0.

Attribution

KVBridge is an independent implementation and systems extension, not an NVIDIA product. Algorithmic credit belongs to Heo et al., arXiv:2608.03893. Related approaches include LatentAlign, Cache-to-Cache, IAM, and DroidSpeak.

License

Apache-2.0. See LICENSE.