KVLRelay

Experiment scripts for KV-Lingo: train a directional KV-cache translator and compare a model handoff with native target re-prefill. Implements head-wise linear maps for Qwen3-0.6B ↔ Qwen3-1.7B. An independent, partial reproduction; quality and speedup remain experimental.

Run

Use Python 3.12.14, the version tested with the locked dependencies.

uv venv --python 3.12.14
uv sync --locked --extra ifeval
make run DEVICE=cuda                    # Pilot: prepare, train, evaluate
make prepare PROFILE=paper-budget
make recipe PROFILE=paper-budget DEVICE=cuda RUN=qwen-headwise

Set DEVICE=mps on Apple Silicon. The Makefile pins the original Hugging Face checkpoints, Nemotron chat data, and the official IFEval grader. make prepare, train, benchmark and ifeval also run independently. make ifeval-native DEVICE=cuda scores all 541 prompts without a translator; prepare its data with .venv/bin/python ifeval.py ifeval --prepare first. By default Qwen3-1.7B builds the cache and Qwen3-0.6B continues, with a learning rate of 1e-4. DIRECTION=s2l reverses the pair, with 1e-3. Both directions share prepared data. TRAIN_ARGS="--fit-weighting attention" weights each prefix position in the closed-form fit by the target's attention to it; recipe does not forward this option. make switch times re-reading against switching on chats joined from held-out data and cropped to each length; set them with SWITCH_ARGS="--lengths 4000 8000". make task TASK=arc and make task TASK=gsm8k HANDOFF=trace score generate-and-parse accuracy from a translator; make task-native scores native paths only. With HANDOFF=trace, the source reasons and the target writes only the answer, compared against the target reading the same trace as text. Resume with TRAIN_ARGS=--resume, IFEVAL_ARGS=--resume, TASK_ARGS=--resume or RECIPE_ARGS=--resume on the corresponding target, keeping the same run name and settings. Generated data, maps and reports live under data/, artifacts/ and results/.

Protocol

Stage	Script	Output
Prepare	paper_data.py	Disjoint natural/packed Nemotron splits; reasoning on/off; hashes
Train	train_translator.py	FP64 closed-form fit, then FP32 KL distillation; frozen models
Measure	benchmark.py	Held-out target KL, reply perplexity, native controls, four timing paths
Measure	switch_cost.py	Re-read and switch wait on joined chats, up to 8k tokens by default; each path timed apart
Generate	ifeval.py	Official strict instruction accuracy; resumable answer-level report
Generate	task_eval.py	ARC-Challenge and GSM8K accuracy, prompt or trace hand-off; resumable item-level report
Orchestrate	recipe.py	Validation-only LR selection, seeds 42/43/44, IFEval item-level mean/SE

Both profiles use BF16 models by default (DTYPE=float32 selects FP32). The adapter supports batch-one, full-attention Qwen3 with matching layer/KV-head counts, a shared tokenizer and default RoPE. Frozen model weights are unchanged; only the maps are trained. Atomic checkpoints retain optimizer/RNG/update state every 100 updates; resume rejects changed data, code, models or training settings.

Profile	Train / validation / test	Fit; updates × batch	Prefix / reply cap	IFEval
pilot	256 / 64 / 64	256; 100 × 1	1,024 / 128	32
paper-budget	40,000 / 64 / 64	400; 5,000 × 8	16,384 / 2,048	541

run executes one training cell. recipe first scores native models, searches five half-decade-spaced rates by final validation KL, extends edge winners, then trains and evaluates three seeds. Closed-form initialization is reused across cells. Local checks cover checkpoint recovery and recipe orchestration; the large recipe is unexecuted. The 40,000-example corpus and initial LR window are our choices; sampling is independent.

KL measures closeness to the native target's token distribution; perplexity scores the recorded reply. IFEval separately grades generated answers. Lower KL alone does not establish useful transfer: compare task accuracy and latency against both native models. Research runs have no default accuracy-loss gate; deployment limits are optionally selectable with IFEval's --max-accuracy-loss.

Timing compares target re-prefill, existing-source handoff, source-prefill plus handoff, and existing-target continuation through first-token logits. Both models are resident; warm handoff requires existing source pre-normalization features. Loading, tokenization and network transfer are excluded; reports retain raw trials. Failed checks return nonzero after writing JSON; make run still collects IFEval. CV ≤10% is our timing diagnostic screen.

Recorded M3 pilot

Apple M3, 16 GiB: Qwen3-0.6B→1.7B, seed 42; 32 natural fitting conversations, 100 updates, eight held-out continuations (512/64-token caps); 32 IFEval prompts, 1,024 generated-token cap. This smaller run uses neither preset's full budget.

Cache	Target KL BF16 ↓	Target KL FP32 ↓	IFEval BF16 /32 ↑
Native source	0.5260	0.5227	14
Native target	0	0	13
Closed-form	0.8145	0.8194	9
Distilled	0.3666	0.3674	6

FP32 evaluation reused BF16-trained maps. Native controls passed 6/8 BF16 and 8/8 FP32; reconstruction exactly matched ordinary split execution in both precisions. Timing stability passed 5/8 BF16 and 2/8 FP32 contexts (30 repeats, five warmups). Distillation lowered KL by 55%, but task accuracy lost 21.9 points against the target; 11/32 translated answers reached the output cap. No validated speedup.

Where this can be taken forward

From what it's obvious: self-hosted model routing, coding-agent handoffs and document QA over shared context.

Since this requires access to internal KV tensors, pair-specific quality/latency evidence and a native re-prefill fallback.

Ordinary hosted chat APIs cannot expose this transfer path through a prompt alone.

Looking for contributors for larger runs, more pairs, corrections or negative results with commands and reports.

Code: MIT. Downloaded models, data and grader retain their own licenses.

Signoff: @tensorhan