KRI is a family of training-free algorithms for picking which KV blocks
each request should attend to. Compute a per-block score offline, write
it as a .pt prior, load it at serve time
through the routing substrate, take the top K. No model changes; no
custom inference; the same model that produced the prior serves the
request. Knlp leverages
vLLM cartridges as the routing accelerator,
so swapping algorithms and sweeping K becomes decode-only iteration,
not prefill-per-experiment.
KRI stands for KV Routing Index — a family of training-free per-block scorers. Each variant answers the same question with a different signal: how much does each KV block contribute to the quality of the output, conditional on the rest of the cache? Variants split on whether the score is computed from attention weights, from logit shifts under leave-one-block-out, from gradients, from query-aware geometry, or from spectral structure on the block graph.
Bake the prior once per cartridge. Ship the .pt
file alongside the cartridge. Load it through the
CartridgeKRIProvider on the
routing branch; the
connector takes care of the top-K selection per request and the
sparse inject into paged cache.
Where each algorithm lands on the production benchmark. The bar is “beat random middle at 5σ.” KRI-D-kv-sum holds the line; xa25 is the cross-attention repair overlay applied on top of any selector (it implements the CacheBlend algorithm, arxiv 2405.16444, at a fixed 25% HKVD refresh ratio). Numbers are the validated 2026-04 sweep; lower-K is harder, higher-K closes the gap to full attention.
The A<a>R<r>K<k> tag
describes how the K-block budget is split across three regions of
the prefix. The cartridge prefix is divided into anchor (head),
mid, and recent (tail); each tag pins specific counts to each
region.
# 128-block cartridge prefix, K=16 budget, A1R2K13 split: block: 0 1 2 3 4 5 6 7 ... 120 121 122 123 124 125 126 127 [A] [ -- pick best 13 by KRI score in the middle 125 blocks -- ] [R][R] ^ anchor=1 ^^^ recent=2
Examples: A0R0K16 is pure KRI selection
with no structural pinning. A4R4K8
pins half the budget structurally and lets KRI pick the other
half. A1R2K13 is the production
structural choice — minimal structural anchoring, KRI does
most of the work.
| # | algorithm | score | vs random | signal | regime |
|---|---|---|---|---|---|
| 1 | KRI-D-kv-sum + xa25 | 26.5% | +5.0pp / 5σ | hidden L2 leave-one-block-out, sum across layers | multi-needle long-context |
| 2 | A1R2K13 + xa25 | ~24% | +2.5pp | structural (anchor + recent) + KRI-D mid blocks | regime-robust; absorbs bad probes |
| 3 | Hybrid (anchor + recency) | ~22% | +0.5pp | position-only policy | zero-bake fallback |
| 4 | attention_prefill prior + xa25 | ~22% | flat | layer-mean prefill attention probability | recency-biased on causal attention |
| 5 | eigenvector centrality + xa25 | pending | — | dominant eigenvector of block-attention graph | addresses KRI-A recency bias and KRI-D marginal-MI |
| — | Random middle (baseline) | 21.5% | 0.0pp | uniformly sampled mid-prefix blocks | baseline; xa25 stays flat (-0.5pp) |
| — | Recency-only (last K) | lower | − | last K blocks only | end-weighted tasks; collapses on beginning-weighted |
| variant | per-block score | where computed | strength | weakness |
|---|---|---|---|---|
| KRI-D-kv-sum | Sum across layers of L2 between full hidden and leave-block-out hidden | offline; one prefill per block | robust leader on multi-needle long-context | marginal-MI redundancy when needles co-occur |
| KRI-D | KL(logits_full || logits_leave-block-out) | offline; one prefill per block | direct logit-level signal | weaker than kv-sum on multi-block ablation |
| KRI-A | Mean per-layer attention probability mass per block during prefill | one prefill | baseline; cheapest after prefill | recency bias on causal attention |
| KRI-G | L2 of gradient of LM-head loss w.r.t. block-mean K | one backward | tracks loss curvature | numerically sensitive on low-precision |
| KRI-Q | Q·K agreement against block-mean K (per query) | online | query-aware, free at serve time | flat on untrained block-mean K |
| KRI-T (TriAttention) | Pre-RoPE Q/K trigonometric factorization: content-stationary centers + position-modulated trig series | online | query-aware, position-aware, zero bake cost | requires multi-needle evaluation under cross-task drift |
| kmeans_keys | K-means cluster centroids over block-mean K; rank by cluster size | offline | captures structural redundancy | k-means hyper-sensitivity to init |
| attention_prefill | Layer-mean prefill attention (post-softmax) per block | one prefill | functional, prompt-conditioned signal | recency-skewed, untrained-LM-head failure mode |
| eigenvector centrality | Dominant eigenvector of the block-level prefill attention graph | one prefill + power iteration | addresses both KRI-A recency bias and KRI-D marginal-MI redundancy | head-pooling weights load-bearing; needs synthetic-Q variant |
CartridgeConnector with
CartridgeKRIProvider loading the prior.
A cartridge is a pre-computed KV cache shipped alongside the model. The cartridge tokens have already been pushed through the transformer at bake time; the resulting K and V tensors are saved to disk as fixed-size blocks. At serve time vLLM loads those blocks straight into paged cache — no prefill runs over the cartridge prefix. The model sees the cache as if it had just finished a long prompt it never actually processed.
CartridgeConnector alone — a
.pt file on local disk,
CartridgeStore loading it into
paged cache, and the routing-prior plugin doing block selection.
That is the entire stack you need to ship and compare a new
KRI variant.
Single-process, single-GPU, single-node. Cartridge lives on local disk; the connector loads it directly. Zero external dependencies beyond vLLM and the routing branch. This is what every result on this page was measured with.
Add LMCache only when you want cross-process tiering,
persistent shared storage, or cross-node serving. The routing
branch ships a native adapter with sparse
retrieve_chunks on the
use_native switch —
deployment plumbing, not a research dependency.
# Once per cartridge (offline, minutes-to-hours): $ python routing/bake_kri_yours.py \ --cartridge /data/carts/patient_00.pt \ --model meta-llama/Llama-3.2-3B-Instruct \ --out /data/priors/patient_00_kri_yours.pt # Then, for every routing config you want to try (~71 ms per request): $ vllm serve meta-llama/Llama-3.2-3B-Instruct \ --kv-transfer-config '{ "kv_connector": "CartridgeConnector", "kv_connector_extra_config": { "cartridges": {"patient_00": "/data/carts/patient_00.pt"}, "routing": { "enabled": true, "K": 16, "priors": {"patient_00": "/data/priors/patient_00_kri_yours.pt"} } }, "kv_role": "kv_both" }' # Same cartridge, same prior, every request: decode only. Sweep K, sweep priors, # swap algorithms by pointing routing.priors at a different .pt file.
cartridge.pt + model
↓
KRI baker (one of: bake_kri_d_kv_prior.py, bake_kri_t_prior.py, ...)
↓
prior.pt = {
block_affinities: Tensor[num_layers, num_kv_heads, num_blocks], # KRI-D-kv-sum
# or:
kmeans_blocks_perK: dict[int, list[int]], # per-K block lists
# or:
kmeans_blocks: list[int], # legacy single list
}
CartridgeConnector.__init__:
routing.priors[cart_id] = path/to/prior.pt
CartridgeKRIProvider.register_prior(prefix_hash=cart_id, prior_path=...)
CartridgeConnector.update_state_after_alloc(request, ...):
manifest = CartridgeKRIProvider.get_manifest(
prefix_hash=cart_id,
query_hash=None, # KRI-G path; query_hash=req for KRI-Q/KRI-T
K=routing.K,
)
_request_selected_block_ids[req_id] = set(manifest.block_indices)
CartridgeConnector.get_block_skip_list(req_id, num_logical_blocks):
return sorted([b for b in range(num_logical_blocks)
if b not in _request_selected_block_ids[req_id]])
# Scheduler turns the skip list into null_block via
# KVCacheManager.null_block_positions.
# Attention backend skips null_block at non-selected positions.
Read KRI as content selection. Earlier evaluation on NIAH-copy showed that, with position-matched random blocks of equal count, the measured quality matches KRI — on that workload, what matters is placement (which positions the model can see), not content (which blocks). KRI's placement profile carries most of its value at K small. On multi-needle long-context, KRI's content selection adds measurable signal: LongHealth K=16 + xa25 lifts KRI-D-kv-sum 5pp over random middle at 5σ (n=200).
LongHealth at K=16 with xa25. Content selection contributes distinctively beyond placement.
Position-matched random blocks match KRI. Placement carries the value; content selection adds nothing.
KRI-D's leave-one-block-out signal underweights blocks whose information appears in multiple other blocks (marginal-MI failure). Use kv-sum or eigenvector variants here.
Your idea becomes a routing algorithm in four steps. Each step has one concrete file to write or one concrete config knob to set. If you can produce a tensor that scores cartridge blocks, you can ship a KRI.
One Python script. Input: a cartridge + a model. Output: a tensor of per-block scores.
routing/bake_kri_<you>.py
.pt
Pick one of three on-disk schemas. The simplest is
block_affinities: a 3-D tensor and
you are done.
<cart_id>_kri_<you>.pt
One line in extra_config. No vLLM
code change. CartridgeKRIProvider
handles loading and top-K resolution.
routing.priors
Beat the production stack on LongHealth K=16 or document the negative. Both outcomes ship.
+>0pp vs KRI-D-kv-sum
Compute one score per block. The signal can be anything — attention probabilities, hidden-state distance under leave-one-block-out, query-aware similarity, eigenvector centrality on the block graph, your own idea. The shape is what matters.
# Required output shape: # block_affinities: Tensor[num_layers, num_kv_heads, num_blocks] # The connector mean-pools across layers/heads and takes top-K per request. import torch from transformers import AutoModelForCausalLM def bake(cartridge_path: str, model_id: str) -> torch.Tensor: cart = torch.load(cartridge_path) # your cartridge model = AutoModelForCausalLM.from_pretrained(model_id).eval().cuda() num_layers, num_heads, num_blocks = shapes_from(cart) scores = torch.zeros(num_layers, num_heads, num_blocks) for b in range(num_blocks): scores[:, :, b] = your_signal(model, cart, block_idx=b) return scores if __name__ == "__main__": scores = bake("patient_00.pt", "meta-llama/Llama-3.2-3B-Instruct") torch.save({"block_affinities": scores}, "patient_00_kri_yours.pt")
cartridge.num_tokens // block_size.
Cartridge tokens divide evenly into fixed-size blocks; that count
is your last dimension.
Three on-disk shapes are accepted. The first is the easiest and covers most variants; only fall back to the others if your algorithm fundamentally produces a different shape.
| key | shape / type | when to use |
|---|---|---|
block_affinities |
Tensor[L, H, N] |
default. Continuous per-block scores. Connector picks top-K. |
kmeans_blocks_perK |
dict[int, list[int]] |
Your algorithm decides differently per K (e.g. k-means with re-clustering per budget). |
kmeans_blocks |
list[int] |
Legacy single list, ignored beyond first K entries. |
No vLLM code change. The connector loads any
.pt in the supported schema and the
CartridgeKRIProvider handles the rest.
Add your file path under
extra_config.routing.priors keyed by
cartridge id.
{
"kv_connector": "CartridgeConnector",
"kv_connector_extra_config": {
"cartridges": {"patient_00": "/data/carts/patient_00.pt"},
"routing": {
"enabled": true,
"K": 16,
"priors": {
"patient_00": "/data/priors/patient_00_kri_yours.pt"
}
}
}
}
CartridgeKRIProvider.get_manifest(prefix_hash, query_hash, K=16).
That returns a BlockManifest with
block_indices. The scheduler turns the
complement into null_block; attention
skips those positions. Sparse inject writes only the selected blocks
into paged cache.
The production stack is the target. Run your prior through the same eval harness used to validate KRI-D-kv-sum. Two outcomes both ship: a win goes into the routing-prior table on this page; a documented negative goes into the kill log so the next person does not retry the same idea.
Production stack: KRI-D-kv-sum + A1R2K13 + xa25 at K=16. 26.5% vs random middle 21.5%, 5σ at n=200.
The bar for KRI variants is competitive, not crushing. A clean +0.5pp at the same K with no latency regression is a credit-worthy result.
You have one Python file in
routing/, one
.pt artifact per cartridge, one new
row in the variants table on this page, and a serving config that
other people can drop into vllm serve.
No vLLM patch. No model fine-tuning. No custom inference. Your
algorithm just runs.
KRI is the algorithm; the routing branch is the substrate. The substrate also accepts non-KRI algorithms through the same interface (hybrid, recency, future ScoutAttention and RetrievalAttention plug-ins).
KRI is training-free; the KRI-FT PEFT vehicle fine-tunes a model under a routing mask so it tolerates more aggressive K. Use either independently. The two compose in a 2×2 factorial (base / KRI-FT × dense / routed).
SPF is a separate scheduler-side prefetch experiment that did not clear its survival gate. KRI and the routing branch are independent of SPF; the substrate that originally hosted KRI under SPF now lives on the routing branch.