fix: raise a clear error on stale KV cache, avoid double allocation in chat (#2190) [1/3]

litgpt/chat/base.py

@@ -88,6 +88,11 @@ def process_prompt(
         max_returned_tokens = encoded_prompt.size(0) + max_new_tokens
         if first_turn or max_returned_tokens > model.max_seq_length:
             model.max_seq_length = max_returned_tokens
+            if not first_turn:
+                # explicit destroy before reallocate: growing the cache in place would otherwise
+                # hold the old (too-small) and new kv caches in memory at the same time, needlessly
+                # doubling peak memory for a REPL turn that runs for as long as the chat session does
+                model.clear_kv_cache()
             model.set_kv_cache(batch_size=1, device=fabric.device)
 
     y: Iterator[torch.Tensor] = generate(

litgpt/model.py

@@ -61,11 +61,9 @@ def max_seq_length(self, value: int) -> None:
         elif value != self.cos.size(0):
             self.cos, self.sin = self.rope_cache(device=self.cos.device)
         # the mask and kv cache size will get updated on `set_kv_cache`. we cannot update it here because we don't know
-        # if the kv cache is expected
-        if self.mask_cache is not None and self.mask_cache.shape[-1] < value:
-            print(
-                f"Warning: KV cache has length {self.mask_cache.shape[-1]} < {value} = max_seq_length. Call 'set_kv_cache' before doing any forwards!"
-            )
+        # if the kv cache is expected. we also don't raise if the existing cache is now too small: callers (e.g.
+        # the chat loop) commonly grow `max_seq_length` and then immediately call `set_kv_cache` to resize it.
+        # `forward` raises a clear error if a stale, too-small cache actually ends up being used.
 
     def reset_parameters(self) -> None:
         # Trigger resetting the rope-cache
@@ -134,7 +132,15 @@ def forward(
                 sin = sin.unsqueeze(0)
             if self.mask_cache is None:
                 raise TypeError("You need to call `gpt.set_kv_cache()`")
-            mask = batched_index_select(self.mask_cache, 2, input_pos)
+            try:
+                mask = batched_index_select(self.mask_cache, 2, input_pos)
+            except IndexError as ex:
+                raise RuntimeError(
+                    f"KV cache has length {self.mask_cache.shape[-1]}, which is too small for the requested "
+                    f"`input_pos` (max index {int(input_pos.max())}). Call `gpt.set_kv_cache(...)` again with a "
+                    "large enough `max_seq_length` before forwarding, otherwise the stale cache would silently "
+                    "produce incorrect attention results or crash with a cryptic index error."
+                ) from ex
             if mask.dim() > 4:
                 # the mask cache has a batch dim of 1 in addition to the one
                 # we get if input_pos has a batch dimension
@@ -279,6 +285,23 @@ def set_kv_cache(
         device: torch.device | None = None,
         dtype: torch.dtype | None = None,
     ) -> None:
+        """Allocates (or reallocates) the key-value cache and attention mask cache for inference.
+
+        This is the "init" step of the KV cache lifecycle: call it once before autoregressive
+        generation (with `input_pos` passed to `forward`), and call it again with a larger
+        `max_seq_length` to grow the cache. Existing cache tensors are dropped and replaced, not
+        resized in place. Pair with `clear_kv_cache()` to explicitly release the cache once
+        generation is done and the memory is needed elsewhere.
+
+        Args:
+            batch_size: The batch size the cache should be allocated for.
+            max_seq_length: Maximum sequence length the cache should support. Defaults to
+                `self.max_seq_length`.
+            rope_cache_length: Length of the rotary position embedding cache. Defaults to
+                `self.rope_cache_length()`.
+            device: Device to allocate the cache tensors on.
+            dtype: Dtype of the cache tensors.
+        """
         if rope_cache_length is None:
             rope_cache_length = self.rope_cache_length()
 
@@ -301,6 +324,13 @@ def set_kv_cache(
             self.mask_cache = build_mask_cache(max_seq_length, device)
 
     def clear_kv_cache(self) -> None:
+        """Releases the key-value cache and attention mask cache, allowing the underlying
+        tensors to be garbage collected.
+
+        This is the explicit "destroy" step of the KV cache lifecycle. It is safe to call
+        `set_kv_cache(...)` again afterwards to reallocate. Calling `forward(..., input_pos=...)`
+        after `clear_kv_cache()` without calling `set_kv_cache()` again raises `TypeError`.
+        """
         self.mask_cache = None
         for block in self.transformer.h:
             block.attn.kv_cache = None

tests/test_chat.py

@@ -66,6 +66,61 @@ def multinomial(*_, **__):
         assert actual_list == expected, (actual_list, expected)
 
 
+def test_process_prompt_clears_kv_cache_before_growing(monkeypatch):
+    # process_prompt grows the kv cache mid-session whenever a turn's max_returned_tokens exceeds
+    # the current model.max_seq_length. It must call clear_kv_cache() before set_kv_cache() on that
+    # path (matching the LLM.generate() dynamic-growth path in api.py) so the old, too-small cache
+    # is dropped instead of coexisting with the newly allocated one.
+    monkeypatch.setattr(chat, "generate", lambda *a, **k: iter([]))
+
+    model = MagicMock()
+    model.max_seq_length = 10
+    model.mask_cache = None  # first turn: no cache yet
+    tokenizer = MagicMock()
+    tokenizer.encode.return_value = torch.zeros(3, dtype=torch.long)
+    tokenizer.decode_stream.side_effect = lambda *a, **k: iter(["x"])
+    prompt_style = MagicMock()
+    prompt_style.apply.return_value = "prompt"
+    fabric = MagicMock()
+
+    common_kwargs = dict(
+        model=model,
+        tokenizer=tokenizer,
+        prompt_style=prompt_style,
+        fabric=fabric,
+        temperature=1.0,
+        max_new_tokens=5,
+        top_k=None,
+        top_p=1.0,
+        stop_tokens=(),
+    )
+
+    # turn 1 (first turn): allocates the cache, nothing to clear yet
+    chat.process_prompt("hi", **common_kwargs)
+    model.set_kv_cache.assert_called_once()
+    model.clear_kv_cache.assert_not_called()
+
+    # turn 2: still fits in the (now 8-token) cache, no reallocation at all
+    model.mask_cache = MagicMock()
+    tokenizer.encode.return_value = torch.zeros(2, dtype=torch.long)
+    common_kwargs["max_new_tokens"] = 2
+    chat.process_prompt("hi again", **common_kwargs)
+    model.set_kv_cache.assert_called_once()
+    model.clear_kv_cache.assert_not_called()
+
+    # turn 3: conversation has grown past max_seq_length, cache must grow -> clear before set
+    tokenizer.encode.return_value = torch.zeros(20, dtype=torch.long)
+    common_kwargs["max_new_tokens"] = 20
+    chat.process_prompt("a much longer prompt", **common_kwargs)
+    assert model.clear_kv_cache.call_count == 1
+    assert model.set_kv_cache.call_count == 2
+    # clear must happen strictly before the second (growing) set_kv_cache call
+    call_names = [str(c) for c in model.mock_calls]
+    clear_index = next(i for i, c in enumerate(call_names) if c.startswith("call.clear_kv_cache"))
+    set_indices = [i for i, c in enumerate(call_names) if c.startswith("call.set_kv_cache")]
+    assert clear_index < set_indices[-1]
+
+
 def test_decode():
     checkpoint_dir = auto_download_checkpoint("EleutherAI/pythia-14m")
     tokenizer = Tokenizer(checkpoint_dir)

tests/test_model.py

@@ -1454,6 +1454,36 @@ def generate(logits):
         input_pos = input_pos[-1:] + 1
 
 
+def test_stale_kv_cache_raises_clear_error():
+    config = Config(block_size=25, padded_vocab_size=5, n_layer=2, n_head=2, n_embd=8)
+    model = GPT(config)
+    model.max_seq_length = 10
+    model.set_kv_cache(1)
+    model.max_seq_length = 25  # grow without resizing the kv cache
+
+    idx = torch.randint(0, config.padded_vocab_size, (1, 1))
+    input_pos = torch.tensor([15])  # beyond the cache built for max_seq_length=10
+    with pytest.raises(RuntimeError, match="Call `gpt.set_kv_cache"):
+        model(idx, input_pos)
+
+
+@torch.inference_mode()
+def test_kv_cache_full_context_length():
+    # exercises `set_kv_cache`/`forward` at a real model's full block_size and a realistic batch
+    # size (issue #2190, "test with full context lengths and realistic batch sizes"), rather than
+    # only the artificial block_size=25 tiny configs used elsewhere in this file.
+    config = Config.from_name("pythia-14m")
+    model = GPT(config)
+    batch_size = 4
+    model.set_kv_cache(batch_size)
+
+    idx = torch.randint(0, config.padded_vocab_size, (batch_size, config.block_size))
+    input_pos = torch.arange(config.block_size)
+    logits = model(idx, input_pos)
+
+    assert logits.shape == (batch_size, config.block_size, config.padded_vocab_size)
+
+
 @torch.inference_mode()
 def test_model_kv_cache_amp():
     config = Config.from_name("pythia-14m", n_layer=2)