diff --git a/gemma/gm/text/_chat_sampler.py b/gemma/gm/text/_chat_sampler.py index 887ee2ad..5bd02ebb 100644 --- a/gemma/gm/text/_chat_sampler.py +++ b/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, diff --git a/gemma/gm/text/_prefill.py b/gemma/gm/text/_prefill.py index 4c70a868..49d7a0b3 100644 --- a/gemma/gm/text/_prefill.py +++ b/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, diff --git a/gemma/gm/text/_sampler.py b/gemma/gm/text/_sampler.py index 567b598e..ebe92fa2 100644 --- a/gemma/gm/text/_sampler.py +++ b/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( diff --git a/gemma/gm/text/_sampler_loop.py b/gemma/gm/text/_sampler_loop.py index 3124aad0..a00353c4 100644 --- a/gemma/gm/text/_sampler_loop.py +++ b/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,