From b71db0f335b52c3becf4c60a19e8553dfb4119f2 Mon Sep 17 00:00:00 2001 From: onepunchmonk Date: Fri, 28 Aug 2026 13:54:30 +0530 Subject: [PATCH] fix: raise a clear error on stale KV cache, document cache lifecycle (#2190) --- litgpt/chat/base.py | 5 +++++ litgpt/model.py | 42 +++++++++++++++++++++++++++++----- tests/test_chat.py | 55 +++++++++++++++++++++++++++++++++++++++++++++ tests/test_model.py | 30 +++++++++++++++++++++++++ 4 files changed, 126 insertions(+), 6 deletions(-) diff --git a/litgpt/chat/base.py b/litgpt/chat/base.py index e3d22bf409..bc271ef364 100644 --- a/litgpt/chat/base.py +++ b/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( diff --git a/litgpt/model.py b/litgpt/model.py index 541860ab5b..58d49f0e27 100644 --- a/litgpt/model.py +++ b/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 diff --git a/tests/test_chat.py b/tests/test_chat.py index 02f211ee1a..415ab4935b 100644 --- a/tests/test_chat.py +++ b/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) diff --git a/tests/test_model.py b/tests/test_model.py index 8d0cf21d5e..78f27cbe90 100644 --- a/tests/test_model.py +++ b/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)