KV Graft Steering: Concept Transfer via KV States

The full-vs-cached difference was small immediately before the sampled trajectories split.

At the first A/B divergence in the four runs:

• the argmax token was still the same,
• the top-25 token IDs still overlapped 25/25,
• the maximum absolute logit difference was roughly 0.14–0.19,
• the total-variation distance between the post-top-p distributions was roughly 0.016–0.030.

But torch.multinomial selected different tokens, and from that point onward the autoregressive histories were different.

That seems compatible with a small execution-path numerical difference being amplified by sampling rather than with an immediately different semantic state.

PyTorch explicitly notes in its numerical-accuracy documentation that mathematically identical floating-point computations are not guaranteed to be bitwise identical; operation order, implementation, platform, and precision can matter.

So I would be careful about interpreting a small-seed sampled-output difference before checking the deterministic path.

In particular, I would not use the four-seed check above to estimate a “success rate”; four seeds are nowhere near enough for that. The useful observation is the much narrower one:

cached visible == exact-token split KV

for every tested 400-token trajectory, while the full-recompute path was the one that separated.

There is another small but important confound in the current scripts.

Conceptually the visible arm is equivalent to tokenizing:

tokenizer(REF_PROMPT + " " + PROMPT)

whereas the current KV path in gen_kvgraft.py tokenizes the pieces separately:

tokenizer(REF_PROMPT + " ")
tokenizer(PROMPT)

For the exact Qwen tokenizer revision I tested, concatenating those separately encoded pieces did not reproduce the joint encoding.

For the list example:

joint visible:       60 tokens
current split KV:    29 + 33 = 62 tokens

and the zorp case also differed.

A split that did preserve the exact joint sequence was:

joint = tokenizer(REF_PROMPT + " " + PROMPT).input_ids

ghost = tokenizer(REF_PROMPT).input_ids
live  = tokenizer(" " + PROMPT).input_ids

assert joint == ghost + live

So I think an explicit token-ID assertion would be a useful low-cost control anywhere the claim is “same text, different representation/channel.”

The general issue is that a character/string concatenation point is not necessarily a tokenizer boundary. There is a useful production-oriented example in vLLM’s current incremental prompt encoding RFC: its proposed implementation explicitly says not to trust token alignment at the append point, backs up to a safe boundary, re-encodes the tail, and verifies overlapping token IDs before splicing.

Once I used the exact split above, ordinary cached visible prompting and split-KV prompting became indistinguishable in the four long sampled trajectories I tested.

The repo already has a FORCE_CONCAT fair-control path, which is also useful here: it helps separate the split-token sequence itself from the subsequent cache execution path.

I think there are several interesting claims here, but they do not all need to stand or fall together.

1. Hidden/precompiled prefix

This seems like the strongest current interpretation.

The guide is converted to model-specific KV state ahead of the live prompt, and the live interaction does not need to contain the guide text visibly.

That is useful even if:

visible cached prefix == frozen exact-token prefix

because the application/deployment properties can still be different.

For example:

• keep tool grammar or formatting instructions out of the visible user transcript,
• reuse a fixed guide across requests,
• swap prepared guide state without reconstructing the visible prompt at the application layer,
• potentially persist or transport prepared prefix state,
• use the plain prefix as a known-good baseline before trying layer/K/V interventions.

Transformers’ prefix-cache prefill/reuse mechanism is probably the cleanest ordinary baseline for this part. Its more general caching explanation also makes the key contract clear: past K/V are reused rather than recomputed, and custom generation loops have to preserve the appropriate cache/mask semantics.

2. A stronger behavioral channel

This is the part I think still needs isolation.

The posted 1/3 visible vs 3/3 KV observation is an observation of those runs, but the comparison currently changes tokenization and execution path at the same time.

The three-path control above suggests that, once those are separated, the behavioral difference can follow cached vs full-recompute execution rather than visible vs hidden prefix representation.

That makes me hesitant to interpret the list result as evidence that KV itself makes the instruction “execute faster.”

3. Steering by actually modifying the KV state

This is a separate and potentially more interesting direction.

Once the stored cache is altered — selected layers, K/V scaling, replacement, mixing, learned transformations, etc. — it is no longer just ordinary prefix-cache equivalence.

There is existing work near that branch:

• KV Cache Steering for Controlling Frozen LLMs applies a one-shot intervention directly to the cache.
• Memory Inception uses text-derived KV banks at selected layers for hidden steering; it is quite close in spirit, although the mechanism is not the same as storing the ordinary full text prefix unchanged.
• Prompt Cache is closer to the operational/reuse side: it precomputes attention states for recurring prompt modules and reuses them later.

So I would probably use plain exact-token KV reuse as the clean baseline, then treat actual cache manipulation as the steering experiment.

That keeps the operational idea and the intervention idea from being conflated.

I would also separate the current start-prefix case from a stronger interpretation of “graft anywhere.”

For a normal causal transformer, a guide at the beginning of the sequence does not need information from the later live prompt. Prefilling it separately is therefore the natural cache case.

A KV chunk created independently and then inserted after some other context is different.

That chunk was originally computed without attending to the new preceding context, so it is generally not equivalent to the KV state that a normal full forward would have produced there.

This exact problem appears in CacheBlend: independently cached text chunks that are no longer prefixes lack their interactions with newly preceding text, so CacheBlend selectively recomputes part of the cached representation rather than assuming direct reuse is equivalent.

So I would make this boundary explicit:

prefix graft
    → ordinary causal prefix-cache semantics can apply

independently prepared KV inserted later in the sequence
    → different problem; preceding-context dependence matters

That second case may still be useful — it just deserves its own evaluation rather than inheriting the prefix result.

I would also be cautious with “same-width” as a compatibility rule.

Even before asking whether two models’ states mean the same thing, hidden width by itself does not generally guarantee that their KV tensors even have compatible geometry. The Transformers cache interface stores K/V with dimensions including the number of KV heads and per-head dimension, not merely the model hidden size; see the Transformers cache description.

And recent cross-model work imposes substantially stronger conditions.

Cross-Model KV Cache Transfer in LLM Families, for example, studies matched-KV model pairs sharing KV-head count and per-head dimension. Even then it does not simply copy the state: it removes RoPE from keys before fitting per-head mappings, and some model pairs still degrade substantially.

Likewise, DroidSpeak studies reuse between fine-tuned models with the same architecture, yet still selectively recomputes some layers to preserve quality.

So I would probably keep:

same-model frozen prefix

as the known-good baseline, and treat:

cross-model KV reuse

as a separate experiment with explicit checks for:

• layer count / mapping,
• KV-head count,
• head dimension,
• positional encoding treatment,
• cache format,
• dtype,
• and, ultimately, downstream quality.

That direction is interesting in its own right, but it is much stronger than what the same-model examples need to establish.

These do not affect the basic hidden-prefix idea, but they may matter if the repo grows into a more mechanistic KV-ablation tool.

GHOST_DROP

In the current gen_kvgraft.py, the GHOST_DROP intervention appears to zero the K tensor for the selected layer while retaining V.

I would therefore think of that as a K-only ablation, rather than a literal no ghost control.

With ordinary attention:

softmax(Q K^T) V

zeroing K does not remove those prefix positions from the softmax, and the retained V values can still contribute.

Even zeroing both K and V is not necessarily identical to removing/masking the prefix positions, because zero-valued positions can still participate in the attention normalization.

So if the goal becomes mechanistic attribution, a useful small matrix might be:

K-only
V-only
K+V
true prefix removal / masking

rather than treating one of those as synonymous with “no ghost.”

Saved GHOST_PACK

There also seems to be a distinction between the pack-format idea and the current load path.

The conceptual goal of a pack is attractive: prepare a prefix once and later restore its KV tensors without doing the guide prefill again.

But in the current gen_kvgraft.py, the load path appears to construct a fresh cache using the guide IDs before replacing its stored K/V tensors.

If that reading is right, the .pt file is already useful as a state container, but loading it is not yet equivalent to “zero guide forward.”

That looks like an implementation detail rather than a limitation of the idea. If eliminating prefill becomes important, I would make direct cache restoration from stored tensors its own benchmark and version the pack against at least:

model/revision
tokenizer/revision
exact guide token IDs
Transformers/cache format
dtype
positional/cache configuration

The reason I would keep this explicit is that the cache is not just an opaque semantic blob; it has an implementation contract and tensor layout that can change independently of the visible guide text.