Fix: Resolve KV Cache OOM & Logit Extraction Memory Spikes

gemma/gm/text/_chat_sampler.py

@@ -122,9 +122,9 @@ class ChatSampler:
   )
   forbidden_tokens: Sequence[str | int] | None = None
   stop_tokens: Sequence[str | int] | None = None
-  # TODO(epot): Support and test rolling cache.
   # TODO(epot): Add a property to show how much of the cache is used.
   cache_length: int | None = 4096
+  rolling_cache_threshold: int | None = None
   max_out_length: int = 2048
 
   # Gemma 4-specific fields (ignored for non-Gemma4 models).
@@ -377,6 +377,24 @@ def chat(
         add_tool_response_tag_after_call=not is_legacy_tool_answer,
     )
 
+    # Rolling cache: if the cache is getting full, evict the oldest turn by
+    # discarding `last_state` and re-prefilling the truncated conversation.
+    if self.rolling_cache_threshold and last_state is not None:
+      # We check the cache info dynamically using a simple threshold.
+      # Because `used_cache_length` is a JAX array, we extract its value.
+      used_length = last_state.used_cache_length
+      if hasattr(used_length, 'item'):
+        used_length = used_length.item()
+      if used_length > self.rolling_cache_threshold:
+        while len(self.turns) >= 2 and used_length > self.rolling_cache_threshold:
+          self.turns = self.turns[2:]  # Evict oldest turn
+          # Break out of loop since we'll just rebuild state entirely.
+          break
+        last_state = None
+
+    if last_state is None and self.multi_turn and self.turns:
+      prompt_text = ''.join(t.text for t in self.turns) + prompt_text
+
     # --- Dispatch to the correct sampler ---
     out = self._sample(
         prompt_text,

gemma/gm/text/_prefill.py

@@ -59,6 +59,7 @@ def prefill(
     pad_length: None | int | tuple[int, ...] = None,
     rng: PRNGKey,
     sharding: kd.sharding.ShardingTree | None,  # pyrefly: ignore[not-a-type]
+    top_k_logits: int = 0,
     vision_input=None,
     audio=None,
     audio_lengths=None,
@@ -236,6 +237,7 @@ def prefill(
       prev_turns=prev_turns,
       cache=cache,
       rng=rng,
+      top_k_logits=top_k_logits,
   )
 
 
@@ -247,6 +249,7 @@ def _make_init_state(
     prev_turns: _turn_utils.PrevTurns,
     cache: _cache_helper.Cache,
     rng: PRNGKey,
+    top_k_logits: int = 0,
 ) -> _sampler_loop.SamplingState:
   """Initial state for the sampling loop."""
 
@@ -272,10 +275,12 @@ def _make_init_state(
       predicted_tokens=jnp.zeros(
           (input.batch_size, max_out_length), dtype=jnp.int32
       ),
-      # predicted_logits=jnp.zeros(
-      #     (batch_size, self.max_out_length, out.logits.shape[-1]),
-      #     dtype=jnp.float32,
-      # ),
+      predicted_top_logits=jnp.zeros(
+          (input.batch_size, max_out_length, top_k_logits), dtype=jnp.float32
+      ) if top_k_logits > 0 else None,
+      predicted_top_indices=jnp.zeros(
+          (input.batch_size, max_out_length, top_k_logits), dtype=jnp.int32
+      ) if top_k_logits > 0 else None,
       cache=cache.cache,
       rng=rng,
       full_attention_mask=full_attention_mask,

gemma/gm/text/_sampler.py

@@ -122,6 +122,7 @@ class Sampler:
       you have a task where the model generates really long outputs.
     pad_length: If provided, pad the prompt to this length. This ensure the
       prompt is always the same length, to avoid jit re-compilation.
+    top_k_logits: Number of top logits to retain and extract. Defaults to 0 (disabled).
   """
   # pylint: enable=g-docstring-quotes
 
@@ -137,6 +138,7 @@ class Sampler:
   cache_length: int = 4096
   max_out_length: int = 2048
   pad_length: None | int | tuple[int, ...] = (256, 512, 1024)
+  top_k_logits: int = 0
 
   def __post_init__(self):
     # If not provided, initialize the tokenizer.
@@ -328,6 +330,7 @@ def sample(
         # the output buffer. However in the sampling loop, users can choose
         # to only decode a subset by setting a smaller `max_new_tokens`.
         max_out_length=self.max_out_length,
+        top_k_logits=self.top_k_logits,
     )
 
     # Max out length is static, while max_new_tokens is dynamic.
@@ -385,6 +388,7 @@ def _initialize_sampler_loop(self, sampling) -> _sampler_loop.SamplerLoop:
         sampling=sampling,
         cache_length=self.cache_length,
         special_tokens=self.tokenizer.special_tokens,
+        top_k_logits=self.top_k_logits,
     )
 
   def _get_inputs(

gemma/gm/text/_sampler_loop.py

@@ -30,7 +30,7 @@
 from gemma.gm.utils import _cache_helper
 import jax
 import jax.numpy as jnp
-from kauldron.ktyping import Bool, Int, PRNGKey, check_type, typechecked  # pylint: disable=g-multiple-import,g-importing-member
+from kauldron.ktyping import Bool, Float, Int, PRNGKey, check_type, typechecked  # pylint: disable=g-multiple-import,g-importing-member
 
 
 @flax.struct.dataclass(kw_only=True)
@@ -59,6 +59,8 @@ class SamplingState:
   last_token: Int['B']  # pyrefly: ignore[not-a-type, unknown-name]
   last_token_pos: Int['B']  # pyrefly: ignore[not-a-type, unknown-name]
   predicted_tokens: Int['B max_out_length']  # pyrefly: ignore[not-a-type]
+  predicted_top_logits: Float['B max_out_length top_k'] | None = None
+  predicted_top_indices: Int['B max_out_length top_k'] | None = None
   # TODO(epot): Better way to extract logits. This takes a lot of memory.
   # TODO(epot): Only keep the top-k logits instead? But sorting might increase
   # computation.
@@ -117,6 +119,7 @@ class SamplerLoop:
   sampling: _sampling.SamplingMethod
   cache_length: int
   special_tokens: type[_tokenizer.SpecialTokens]
+  top_k_logits: int = 0
 
   # @functools.partial(
   #     jax.jit,
@@ -248,7 +251,14 @@ def _sample_step(
 
     # Update the buffers to save the outputs.
     predicted_tokens = state.predicted_tokens.at[:, state.step].set(next_token)
-    # predicted_logits = state.predicted_logits.at[:, state.step].set(logits)
+
+    if self.top_k_logits > 0:
+      top_logits, top_indices = jax.lax.top_k(logits, k=self.top_k_logits)
+      predicted_top_logits = state.predicted_top_logits.at[:, state.step].set(top_logits)
+      predicted_top_indices = state.predicted_top_indices.at[:, state.step].set(top_indices)
+    else:
+      predicted_top_logits = state.predicted_top_logits
+      predicted_top_indices = state.predicted_top_indices
 
     # Check whether we have reached an end token.
     done = state.done | jnp.isin(next_token, jnp.asarray(self.end_tokens))
@@ -261,6 +271,8 @@ def _sample_step(
         # is still incremented, so we use previous `state.done`.
         last_token_pos=state.last_token_pos + ~state.done,
         predicted_tokens=predicted_tokens,
+        predicted_top_logits=predicted_top_logits,
+        predicted_top_indices=predicted_top_indices,
         # predicted_logits=predicted_logits,
         cache=out.cache,  # pyrefly: ignore[missing-attribute]
         rng=next_rng,