Skip to content

fix: Gemma4Sampler silently drops cache-full condition that Sampler always reported#715

Open
DhruvTilva wants to merge 1 commit into
google-deepmind:mainfrom
DhruvTilva:fix/gemma4-sampler-cache-full-warning
Open

fix: Gemma4Sampler silently drops cache-full condition that Sampler always reported#715
DhruvTilva wants to merge 1 commit into
google-deepmind:mainfrom
DhruvTilva:fix/gemma4-sampler-cache-full-warning

Conversation

@DhruvTilva

Copy link
Copy Markdown

Description

## What fixes

When a Gemma 4 prompt + generated output exceeds `cache_length`, the
`SamplerLoop` exits silently via its `~state.cache_info.is_full` stopping
condition. The user gets a truncated (or empty) response with no indication
of why generation stopped.

`Sampler._decode_state()` has always guarded against this with a
`kd.utils.status.log()` call after sampling completes:

```python
# _sampler.py – present since before Gemma 4 was introduced
if state.cache_info.is_full.item():
    kd.utils.status.log(
        'Cache is full! Try increasing `Sampler.cache_length`. ...'
    )

Gemma4Sampler.sample() performs the identical post-sampling decode
sequence but the guard was never carried over when Gemma4Sampler was
introduced. This is a straightforward parity gap.

Why it matters for Gemma 4 specifically

Gemma 4 multimodal inference expands each <|image|> placeholder into
hundreds of soft tokens during preprocessing. This makes hitting
cache_length significantly more likely than in text-only scenarios —
and more surprising, since the user's text prompt may look well within
budget.

Change

Added the missing check in Gemma4Sampler.sample(), mirroring the
existing pattern in Sampler:

if state.cache_info.is_full.item():
    kd.utils.status.log(
        'Cache is full! Try increasing `Gemma4Sampler.cache_length`. Current:'
        f' {self.cache_length}'
    )

No logic changes. No API changes. No new dependencies.

Testing

Added test_gemma4_sampler_cache_full_warning to _sampler_test.py.
The test mocks state.cache_info.is_full = True on the SamplerLoop
output and asserts kd.utils.status.log is called with a message
containing cache_length. This test would have caught the gap before
this fix.

Follows the same mock-based pattern used by the existing dispatch tests
in the file.

Sampler._decode_state() checks state.cache_info.is_full and emits a
kd.utils.status.log() diagnostic when the KV cache is exhausted during
generation. Gemma4Sampler.sample() performs the same post-sampling
decode sequence but omits this guard entirely, causing silently
truncated responses when cache_length is exceeded.

This is especially impactful for multimodal Gemma 4 inference where
image tokens significantly expand the effective context length.

Add the equivalent check to Gemma4Sampler to match Sampler behaviour,
and add a regression test that would have caught this gap.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant