fix: Gemma4Sampler silently drops cache-full condition that Sampler always reported#715
Open
DhruvTilva wants to merge 1 commit into
Open
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Gemma4Sampler.sample()performs the identical post-sampling decodesequence but the guard was never carried over when
Gemma4Samplerwasintroduced. This is a straightforward parity gap.
Why it matters for Gemma 4 specifically
Gemma 4 multimodal inference expands each
<|image|>placeholder intohundreds of soft tokens during preprocessing. This makes hitting
cache_lengthsignificantly 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 theexisting pattern in
Sampler:No logic changes. No API changes. No new dependencies.
Testing
Added
test_gemma4_sampler_cache_full_warningto_sampler_test.py.The test mocks
state.cache_info.is_full = Trueon theSamplerLoopoutput and asserts
kd.utils.status.logis called with a messagecontaining
cache_length. This test would have caught the gap beforethis fix.
Follows the same mock-based pattern used by the existing dispatch tests
in the file.