Fix _remove_answer_from_cache crash with QuantizedCache(backend='hqq')

[bonginn]

PR description

Fixes #289. Following @maxjeblick's review, _remove_answer_from_cache no longer slices a QuantizedCache: it saves each layer's state before an answer and restores it afterwards, as in the prefix caching pattern of the transformers docs.

Slicing can't restore a quantized cache: QuantizedLayer flushes its full-precision residual into the quantized storage once it reaches residual_length, so the question and answer end up inside the quantized tensors. With more than one question, slicing failed on both backends:

• hqq: _quantized_keys is a (qtensor, meta) tuple (the original TypeError), and cumulative_length was never reset, so the next question hit an attention mask size mismatch.
• quanto: slicing a WeightQBitsTensor returns a plain tensor that dequantizes to float32 (Expected query, key, and value to have the same dtype).

A shallow copy of vars(layer) is enough because QuantizedLayer.update assigns new tensors instead of writing in place. The quantized branch returns before the DynamicCache slicing, which is unchanged.

Tests: added test_pipeline_with_quantized_cache_multiple_questions (quanto and hqq, residual_length=4 so the residual flushes while answering). It checks that the cache is back to the compressed context length and that each answer matches a fresh single-question call. It fails on main for both backends and passes with this change. Also fixed the README example, which was missing the required config argument.

Checklist

Before submitting a PR, please make sure:

• Tests are working (make test)
• Code is formatted correctly (make style, on errors try fix with make format)
• Copyright header is included
• All commits are signed-off using git commit -s

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[maxjeblick]

Hi @bonginn , thanks a lot for spotting the bug, and opening a PR.

I reviewed the PR; the current implementation doesn't reset the hqq cache properly.
Restoring a quantized cache after the generate step would require quite a bit of "surgery" on the cache object, and would be out of scope for this library.

Transformers library also faces this issue (reusable quantized cache), and proposes to copy over the whole cache before generation, then restoring it.

We opt to copy HF's pattern to support quantized cache, althtough it temporarily copies the kv-cache. A proper fix would require an implementation within the quantized cache class (something like cache.remove_last_tokens(N)), and is out of scope here.

We are happy if you are interested in adapting your PR, below you find a detailed assesment by an agent that may help.

---

Assessment: the PR removes the TypeError, but _remove_answer_from_cache still doesn't restore a quantized cache to the context. A QuantizedLayer keeps its newest tokens in a full-precision residual (layer.keys) and tracks its length in cumulative_length. When the residual fills up (residual_length), everything is re-quantized together and the residual becomes an empty 1-D tensor. Slicing therefore can't remove the question and answer. With this PR, hqq still has these bugs:

• Later questions on the same context see the previous questions and answers.
• If the residual was flushed during an answer, the next question fails with an attention mask size mismatch.
• If the flush happens on the last decoding step, slicing the empty residual raises IndexError, even with a single question.

The new test asks a single short question, so it covers none of these cases.

Suggested fix: create a copy of the layer state that is taken before each answer and restored afterwards. This follows the prefix caching pattern in the transformers docs, which deep-copies the prefilled cache before each generation. A shallow copy is enough here, because QuantizedLayer.update assigns new tensors to its attributes instead of writing into the existing ones:

Sketch of implementation, needs to be implemented:

# in _forward, before generate_answer
layer_states = [vars(layer).copy() for layer in cache.layers] if isinstance(cache, QuantizedCache) else None
...
self._remove_answer_from_cache(cache, cache_seq_lengths, layer_states)

# in _remove_answer_from_cache, replacing the QuantizedCache branch
if isinstance(cache, QuantizedCache):
    for layer, state in zip(cache.layers, layer_states):
        vars(layer).update(state)
    return

• Tests: keep the quanto/hqq parametrization and add a two-question test with QuantizedCache(..., residual_length=4), so that the residual is flushed during generation. The test should check that cache.get_seq_length() equals the context length after the call, and that each answer matches a fresh single-question call.

• README: the example is missing the required config argument; it should be QuantizedCache(backend="quanto", config=pipe.model.config, nbits=4).

[bonginn]

Thanks a lot for the detailed review @maxjeblick! I've updated the PR following your suggestion: the layer states are saved before each answer and restored afterwards instead of slicing the quantized cache. I added the two-question test with residual_length=4 for quanto and hqq (it fails on main for both backends and passes now) and fixed the README example. I also updated the PR description.

[maxjeblick]

LGTM, thanks a lot for this fix!

[maxjeblick]

/ok to test aef39ea