generate() produces wrong output when past_key_values is an externally-built DynamicCache

# Add `DynamicCache.from_kv_tensors()` — safe reconstruction from prefill KV

## Problem

KV cache compression methods (SnapKV, H2O, KIVI, and user implementations) commonly follow this pattern:

```python

# Prefill

out = model(**inputs)

kc = out.past_key_values

# Compress

new_cache = DynamicCache()

for li in range(n_layers):

k, v = compress(kc[li].keys, kc[li].values)

new_cache.update(k, v, li)

# Generate — UNRELIABLE

model.generate(input_ids=gen_last, past_key_values=new_cache)

```

This silently produces wrong output. We reproduced on Qwen3-VL-8B (MMBench-Video, 1747 samples): **10.7pp accuracy loss with bitwise-identical KV tensors**. Passing explicit `attention_mask` and `position_ids` restores quality to within 0.2pp of baseline.

## Root Cause

`generate()` cannot infer the correct attention mask when `past_key_values` is an externally-built `DynamicCache`. The cache stores KV tensor values but the positional context (sequence length, per-layer position tracking used internally by `generate()`) is not properly propagated.

## Proposed Fix

Add `DynamicCache.from_kv_tensors()` classmethod:

```python

@classmethodclassmethodclassmethodclassmethodclassmethodclassmethodclassmethodclassmethod

def from_kv_tensors(cls, kv_pairs, seq_length):

“”"Create a DynamicCache from pre-existing KV tensors.

Args:

kv_pairs: list of (k, v) tuples, one per layer.

k, v have shape [batch, n_kv_heads, seq_len, head_dim]

seq_length: total sequence length (prefill tokens)

Returns:

DynamicCache with K/V set and proper internal state

“”"

cache = cls(num_hidden_layers=len(kv_pairs))

for li, (k, v) in enumerate(kv_pairs):

cache.update(k, v, li)

return cache

```

The `num_hidden_layers` parameter (optional, new) pre-allocates the internal list and enables proper initialization. It’s backward compatible.

## Minimal Reproduction

```python

# Identity reconstruction test

out = model(**inputs)

kc_orig = out.past_key_values

# OLD (wrong):

new_cache = DynamicCache()

for li in range(n_layers):

new_cache.update(kc_orig[li].keys, kc_orig[li].values, li)

model.generate(input_ids=last_token, past_key_values=new_cache)

# → output diverges from baseline

# NEW (correct):

new_cache = DynamicCache.from_kv_tensors(

[(kc_orig[li].keys, kc_orig[li].values) for li in range(n_layers)],

seq_length

)

# Generate with explicit args:

mask = torch.ones(1, seq_length + 1, device=device, dtype=torch.long)

pos = torch.tensor([[seq_length]], device=device)

model.generate(input_ids=last_token, attention_mask=mask,

position_ids=pos, past_key_values=new_cache)

# → output matches baseline

```

## Testing

- [ ] Identity reconstruction test on Qwen3-VL-8B (≤0.2pp difference)

- [ ] Identity reconstruction test on Qwen2.5-VL-7B (no crash)

- [ ] Backward compatibility: existing usage with `past_key_values` from prefill works unchanged

- [ ] Text-only LLM (Llama) identity reconstruction test

## Related

All KV compression methods that follow the “extract → compress → rebuild DynamicCache → generate()” pattern are affected. Methods that operate in-forward (SnapKV’s hooks, H2O’s in-place eviction) are immune.

For now, after trying Qwen3-VL 2B in Colab, it looked like this path had quite a few behavioral quirks:


My current reading is that the 10.7-point degradation and the recovery with explicit generation inputs are important observations, but the current example does not yet isolate the effect of DynamicCache.from_kv_tensors() itself.

In the “old” and “new” paths, at least three things change together:

  1. how the cache object is constructed;
  2. whether a full-length attention_mask is supplied;
  3. whether position_ids are supplied explicitly.

The proposed factory also receives seq_length, but the sample implementation does not use it to change any cache state. It still creates a DynamicCache and calls update() for each layer.

I therefore tried a smaller identity-reconstruction control with Qwen/Qwen3-VL-2B-Instruct. In that control:

  • the original cache and several externally rebuilt caches produced identical first-step logits under identical arguments;
  • generate() failed for all cache variants when given only the one-token suffix without a full attention mask;
  • adding the full mask was sufficient to make generate() work;
  • adding position_ids without the full mask was not sufficient;
  • Qwen3-VL also had important multimodal positional state outside the K/V tensors.

So the narrowest explanation supported by this control is not yet “externally built DynamicCache loses its internal sequence state.” It looks more like:

Resuming generate() from a non-empty cache requires an unambiguous contract between the cached prefix, the unprocessed suffix, the full attention mask, and any model-specific positional state.

A K/V constructor could still be useful for validation and ergonomics, but it may not be the component that restored generation in the example above.

The continuation path that worked

The safest default path I found was:

  1. Build the cache from the already processed prefix only.
  2. Pass only the not-yet-processed suffix as input_ids.
  3. Pass an attention mask covering:
past_kv_length + new_tokens_length
  1. Compare the original and rebuilt caches using exactly the same generation arguments.
  2. For Qwen3-VL, preserve the model-side multimodal positional state as well as the K/V tensors if the cache moves between model instances or processes.

This matches the general contract described in the Transformers caching documentation: the attention mask used with a cache should cover both the past K/V and the current input.

A compact form of the boundary is:

# The last token is intentionally excluded from the prefill.
prefix_ids = input_ids[:, :-1]
current_ids = input_ids[:, -1:]

prefix_out = model(
    input_ids=prefix_ids,
    attention_mask=full_attention_mask[:, :-1],
    pixel_values=pixel_values,
    image_grid_thw=image_grid_thw,
    mm_token_type_ids=mm_token_type_ids[:, :-1],
    use_cache=True,
)

cache = prefix_out.past_key_values

# Full mask = cached prefix + current unprocessed token.
output = model.generate(
    input_ids=current_ids,
    past_key_values=cache,
    attention_mask=full_attention_mask,
    max_new_tokens=1,
    do_sample=False,
)

The highest-information controls appear to be:

Result Likely branch
Original and rebuilt caches both fail without the full mask Resume-input / mask contract, not reconstruction-specific
Original and rebuilt caches both work with mask_only Explicit position_ids were probably not independently required
Only the rebuilt cache fails under identical arguments Compare layer types, cache lengths, config-derived metadata, and model-side state
Same-instance continuation works, fresh-instance continuation fails State exists outside the K/V tensors
Identity reconstruction works, actual compression fails Logical positions, retained-token mapping, per-layer lengths, or backend behavior
Small Qwen3-VL control: setup and cache-equivalence results

Setup

I used:

  • Qwen/Qwen3-VL-2B-Instruct
  • model revision 89644892e4d85e24eaac8bacfd4f463576704203
  • Transformers 5.14.1
  • Torch 2.11.0+cu128
  • Tesla T4
  • fp16
  • SDPA
  • one synthetic 224Ă—224 image
  • batch size 1
  • one-token continuation
  • full multimodal sequence length: 82
  • cached prefix length: 81

This was not an attempt to reproduce the 8B MMBench-Video result. It was only an identity-reconstruction and generation-resume sanity check.

Cache variants

I compared:

# A. Original cache
original = copy.deepcopy(prefill_cache)
# B. The pattern used in the post
rebuilt_empty = DynamicCache()

for layer_idx, layer in enumerate(prefill_cache.layers):
    rebuilt_empty.update(
        layer.keys.detach().clone(),
        layer.values.detach().clone(),
        layer_idx,
    )
# C. Config-aware cache, then update()
rebuilt_config = DynamicCache(
    config=model.config.text_config,
)

for layer_idx, layer in enumerate(prefill_cache.layers):
    rebuilt_config.update(
        layer.keys.detach().clone(),
        layer.values.detach().clone(),
        layer_idx,
    )
# D. Existing tensor-data constructor path in the tested version
kv_pairs = [
    (
        layer.keys.detach().clone(),
        layer.values.detach().clone(),
    )
    for layer in prefill_cache.layers
]

rebuilt_ctor = DynamicCache(
    ddp_cache_data=kv_pairs,
    config=model.config.text_config,
)

In the tested version, all four variants had:

  • 28 initialized layers;
  • DynamicLayer for all 28 layers;
  • reported cache length 81;
  • identical K/V shapes;
  • bitwise-identical K/V values.

Across the 15 valid original-versus-rebuilt first-step comparisons, the logits were also exactly identical.

That does not prove that all models, cache layer types, Transformers versions, or compressed caches behave identically. It does indicate that, for this Qwen3-VL identity case, DynamicCache().update() did not itself introduce a generation difference.

The current DynamicLayer.get_seq_length() implementation obtains its physical sequence length from the K tensor’s sequence dimension. Other cache layer types can have different semantics—for example, sliding layers may track cumulative processed length separately—so an API intended for real compression probably needs to distinguish those concepts. The relevant implementations are in cache_utils.py.

Mask/position matrix and the first model call made by generate()

The 2Ă—2 result

For each cache variant, I tried:

cases = {
    "neither": {},
    "mask_only": {
        "attention_mask": full_attention_mask,
    },
    "position_only": {
        "position_ids": current_position_ids,
    },
    "both": {
        "attention_mask": full_attention_mask,
        "position_ids": current_position_ids,
    },
}

The generate() result was:

Generation arguments Original cache Rebuilt caches
Neither Failed Failed
position_ids only Failed Failed
Full attention_mask only Worked Worked
Full mask + position IDs Worked Worked

The important part was what reached the first actual model call.

Without the full mask

Input supplied to generate():

past cache length: 81
input suffix length: 1

First model call:

input_ids shape:       [1, 0]
attention_mask shape:  [1, 1]
past cache length:     81

The only unprocessed token had been sliced away before the model forward pass.

With explicit position IDs, but no full mask

The same thing happened:

input_ids shape:       [1, 0]
attention_mask shape:  [1, 1]
position_ids:          explicitly supplied
past cache length:     81

The position tensor could not restore an input token that had already been removed by the generation input-preparation path.

With the full mask only

First model call:

input_ids shape:       [1, 1]
attention_mask shape:  [1, 82]
past cache length:     81
position for current token: inferred as 81

Generation succeeded.

With the full mask and explicit positions

This also succeeded, but it did not show an additional recovery over mask_only in this control.

In the tested generation/utils.py, generation prepares a next_sequence_length and then slices input_ids to that length before the model call. It also creates a default attention mask from the currently supplied input if no full mask is available, and attention-mask-derived position IDs are created using the cumulative mask.

That creates an ambiguity when all of these are true:

  • past_key_values is already non-empty;
  • input_ids contains only the unprocessed suffix;
  • the suffix has length 1;
  • no full-length attention mask is supplied.

The generated one-token mask can look like the mask for a complete one-token input rather than the mask accompanying a one-token suffix after a long cached prefix.

This is conceptually close to earlier cache-continuation problems such as Transformers issue #36151, where generation inferred the cache/input position incorrectly when reusing externally supplied past_key_values. It is not necessarily the same implementation bug or the same Transformers version, but it is a nearby boundary.

For reproducibility, I captured the first real model call with a forward pre-hook rather than replacing prepare_inputs_for_generation():

captured = []

def capture_first_call(_module, args, kwargs):
    if not captured:
        captured.append({
            "input_ids_shape": (
                tuple(kwargs["input_ids"].shape)
                if kwargs.get("input_ids") is not None
                else None
            ),
            "attention_mask_shape": (
                tuple(kwargs["attention_mask"].shape)
                if kwargs.get("attention_mask") is not None
                else None
            ),
            "position_ids_shape": (
                tuple(kwargs["position_ids"].shape)
                if kwargs.get("position_ids") is not None
                else None
            ),
            "past_length": (
                kwargs["past_key_values"].get_seq_length()
                if kwargs.get("past_key_values") is not None
                else None
            ),
        })

handle = model.register_forward_pre_hook(
    capture_first_call,
    with_kwargs=True,
)

One additional quirk: direct one-token forward() and generate() did not have exactly the same convenience behavior for Qwen3-VL. A direct forward() with a full-length mask but no explicitly sliced current-token 3D position IDs produced a shape mismatch, whereas generate() prepared the current-token positions correctly. So a manual forward loop and GenerationMixin should probably be treated as two separate contracts.

Qwen3-VL has important state outside the K/V tensors

Qwen3-VL’s multimodal position handling is not represented solely by the K/V tensors.

In Transformers 5.14.1, the model stores rope_deltas on the model instance. The source describes it as the difference between sequence length and multimodal RoPE position, and explicitly caches it on the model:

The tested prefix prefill produced:

rope_deltas = -42

I kept the K/V tensors and full attention mask unchanged and changed only that model-side state.

Model-side state Max absolute logit difference Cosine similarity Top-1 result
Correct saved rope_deltas Reference Reference left
rope_deltas = None 17.1328 0.30597 The
rope_deltas = 0 17.1328 0.30597 The
Correct value + 1 1.21875 0.99833 left

The None and zero cases were identical in this control, and neither was a small numerical perturbation. The top-5 overlap with the correct state was also zero.

This matters especially for:

  • serializing a cache and restoring it later;
  • transferring a cache to another process;
  • transferring a cache to a fresh model instance;
  • distributed serving or worker migration;
  • reconstructing K/V tensors without replaying the multimodal prefill.

It may not explain the reported benchmark degradation if the exact same model instance performs the prefill and immediately performs generation, because the model-side rope_deltas may still be present.

So I would keep this as a separate branch:

same model instance works
fresh model instance fails
    → inspect model-side multimodal positional state

rather than treating it as the confirmed cause of the 10.7-point difference.

It also means that a generic K/V-only factory cannot, by itself, define a completely portable Qwen3-VL resume state. The broader object may need to include model-specific metadata, or the model may need a documented hook for exporting and restoring that metadata.

The exact public attribute may evolve—the tested source already notes deprecation of the output-side attribute in favor of model-side state—so the durable concept is “model-side multimodal positional state,” not necessarily one permanent field name.

A separate token-boundary check

There is another boundary in the minimal example that may be worth separating:

out = model(**inputs)
kc_orig = out.past_key_values

model.generate(
    input_ids=last_token,
    past_key_values=new_cache,
)

If last_token was already included in inputs, then the cache already contains its K/V, and generation processes that token a second time.

I tested that pattern separately:

# Full prompt is already cached.
full_out = model(
    **full_inputs,
    use_cache=True,
)

full_cache = full_out.past_key_values

# This is the final token from the already processed prompt.
repeated_token = full_inputs["input_ids"][:, -1:]

Re-feeding the already cached token produced:

max absolute logit difference: 7.65625
cosine similarity:             0.8698
top-1 token:                   changed

It happened with both the original cache and the rebuilt cache. Supplying a full attention mask and explicit positions did not restore the original result.

That makes this a distinct failure mode:

cache(prompt[:-1]) + prompt[-1]
    → valid split continuation

cache(prompt) + prompt[-1]
    → prompt[-1] is processed twice

I would not infer from the shortened forum snippet that the benchmark necessarily does this. The actual evaluation code may already exclude the final token from prefill, or gen_last may mean a genuinely unprocessed token. The near-recovery with mask and positions also makes duplicate processing less likely to be the sole explanation of the reported 10.7-point loss.

Still, stating the exact boundary would make the reproduction substantially easier to interpret.

What this may imply for the proposed API

I think a from_kv_tensors()-style API could still be useful, especially for:

  • validating the number of layers;
  • validating K/V shapes, dtype, device, and batch size;
  • applying the model config so the correct cache layer types are created;
  • providing an official reconstruction route;
  • adding identity-generation regression tests;
  • making failure messages clearer than silent generation drift.

There has already been at least one constructor-specific cache issue—issue #39668—so reconstruction tests would have independent value.

However, the directly observed failure boundary seems broader than a K/V constructor.

A complete resume contract may need to distinguish:

State Meaning
Physical K/V length Number of key/value slots currently stored
Total processed length Number of tokens that have logically passed through the model
Unprocessed suffix Tokens supplied for the next forward call
Full attention mask Validity/padding state across cached prefix and current suffix
Logical next position Position assigned to the next token
Retained-token positions Original positions after eviction or selection
Per-layer cache lengths Potentially different after layer-specific compression
Model-specific metadata For example Qwen3-VL multimodal positional state

For a plain DynamicLayer, physical sequence length can be recovered from the K tensor shape. That does not necessarily recover all of the other concepts above.

This becomes more important after real compression:

  • If tokens are evicted, retained K/V slot count may differ from total processed length.
  • If different layers retain different tokens, per-layer lengths may diverge.
  • If keys already have RoPE applied, moving them to new logical positions may require rerotation or recomputation rather than merely changing a scalar sequence length.
  • Sliding or hybrid cache layers may maintain cumulative state that is not equivalent to the currently stored tensor width.
  • Multimodal models may retain additional positional metadata outside the cache.

There are nearby examples of those boundaries:

  • Issue #35168 reported that layer-dependent cache lengths behaved differently between FlashAttention 2 and SDPA/eager paths.
  • PR #39843 added support for continuing generation from cache in a FlashAttention 2 path and used backend/first-step comparisons.
  • Issue #40833 showed a different silent-output problem where left-padded cached batches attended to padding because the cache mask was incomplete.

These are not necessarily the same cause as this report, but they suggest that generation-equivalence tests should cover more than tensor identity.

Possible API directions are therefore not mutually exclusive:

  1. An ergonomic K/V constructor

    cache = DynamicCache.from_kv_tensors(
        kv_pairs,
        config=model.config.text_config,
    )
    
  2. Validation in generate()

    For example, detect a non-empty cache combined with a short suffix and an attention mask that covers only that suffix, then raise a useful error or warning instead of silently slicing to an empty input.

  3. A documented resume-state object

    Something conceptually like:

    resume_state = {
        "past_key_values": cache,
        "attention_mask": full_attention_mask,
        "processed_length": processed_length,
        "model_metadata": model_specific_state,
    }
    
  4. Model-specific export/restore hooks

    Particularly for multimodal positional state.

  5. Generation-equivalence regression tests

    At minimum:

    • original cache versus identity-rebuilt cache;
    • mask_only, position_only, both, and neither;
    • original model instance versus fresh model instance;
    • text-only and multimodal models;
    • eager, SDPA, and FlashAttention where applicable;
    • padded batches;
    • sliding/hybrid caches;
    • layer-variable compressed caches.

In the tested Transformers 5.14.1 path, the existing constructor already accepted K/V data plus a config:

DynamicCache(
    ddp_cache_data=kv_pairs,
    config=model.config.text_config,
)

and it matched both the original cache and the manual update() reconstruction in this identity test.

That does not make a clearer public factory unnecessary. It does mean the proposed factory’s responsibility may need to be specified more precisely: is it only constructing validated K/V layers, or is it also promising that generate() can resume without separately supplied mask, token-boundary, positional, and model-specific state?

Compact decision tree
1. Compare original and rebuilt caches under identical arguments
   |
   +-- Both behave identically
   |     |
   |     +-- Both fail without full mask, both work with it
   |     |      → generation-resume mask/input ambiguity
   |     |
   |     +-- Same instance works, fresh instance fails
   |     |      → model-side state outside K/V
   |     |
   |     +-- Identity cache works, compressed cache fails
   |            → logical positions, token mapping,
   |              per-layer lengths, or backend path
   |
   +-- Only rebuilt cache fails
         |
         +-- K/V tensors differ
         |      → reconstruction/copy path
         |
         +-- K/V identical, layer classes or lengths differ
         |      → config/cache-layer metadata
         |
         +-- K/V and visible cache metadata match
                → model-specific hidden state or
                  a GenerationMixin branch that distinguishes the objects

The smallest control matrix that seems to settle most of this is:

cache_variants = {
    "original": original_cache,
    "rebuilt_update": rebuilt_update,
    "rebuilt_config": rebuilt_config,
    "rebuilt_ctor": rebuilt_ctor,
}

generation_variants = {
    "neither": {},
    "mask_only": {
        "attention_mask": full_attention_mask,
    },
    "position_only": {
        "position_ids": current_position_ids,
    },
    "both": {
        "attention_mask": full_attention_mask,
        "position_ids": current_position_ids,
    },
}

Then compare the first-step logits rather than only complete decoded answers:

result = model.generate(
    input_ids=current_ids,
    past_key_values=copy.deepcopy(cache),
    max_new_tokens=1,
    do_sample=False,
    return_dict_in_generate=True,
    output_logits=True,
    **generation_kwargs,
)

first_step_logits = result.logits[0]

If the original and rebuilt cache produce exactly the same first-step logits under the same arguments, the reconstruction step is unlikely to be the branch responsible for a later benchmark difference.

Limits of this control

This control does not establish that:

  • the reported Qwen3-VL-8B MMBench-Video result is incorrect;
  • Qwen3-VL-8B behaves exactly like the 2B model;
  • video prompts behave like a single synthetic image;
  • all Transformers versions use the same generation path;
  • all attention backends behave identically;
  • actual compressed caches behave like identity-rebuilt caches;
  • all K/V compression methods are unaffected;
  • explicit position IDs are never needed;
  • rope_deltas is the cause of the reported 10.7-point loss;
  • a from_kv_tensors() API would have no value.

The control was intentionally narrow:

  • 2B model;
  • one image;
  • batch size 1;
  • one-token continuation;
  • fp16;
  • SDPA;
  • identity reconstruction;
  • Transformers 5.14.1.

Its useful result is only that the problem separated cleanly into different branches:

  1. K/V object reconstruction;
  2. generation resume inputs and full mask;
  3. already-processed versus unprocessed token boundary;
  4. Qwen3-VL model-side multimodal position state;
  5. future compression-specific position and length semantics.

Under that separation, the original and externally rebuilt caches behaved identically. The full attention mask determined whether the suffix reached the model, and Qwen3-VL’s model-side positional state determined whether the same K/V tensors retained their multimodal meaning.

So my tentative conclusion would be:

  • the benchmark observation is worth preserving;
  • the mask_only and position_only ablations would clarify which part produced the recovery;
  • the original cache should be run through exactly the same continuation arguments as the rebuilt cache;
  • gen_last should be identified explicitly as processed or unprocessed;
  • same-instance and fresh-instance controls would separate K/V reconstruction from Qwen3-VL model-side state;
  • a constructor may be useful, but validation or a broader resume-state contract may address the observed failure more directly than a K/V-only factory.

The claim that every “extract → compress → rebuild → generate” method is affected may also be safer if scoped to methods that resume through this generate() path without carrying the required mask, token boundary, logical positions, and model-specific metadata. In-forward methods, patched attention implementations, and inference engines that own their resume state may cross a different boundary.