Skip to content

fix: correct RecurrentGemma prefill state contract - #2300

Open
inureyes wants to merge 4 commits into
mainfrom
fix/issue-2288-recurrent-gemma-chunk-contract
Open

inureyes wants to merge 4 commits into
mainfrom
fix/issue-2288-recurrent-gemma-chunk-contract

Conversation

@inureyes

@inureyes inureyes commented Oct 13, 2026 •

Copy link
Copy Markdown
Member

Summary

  • Disable cache-level chunked prefill for RecurrentGemma because multi-token no-ID forwards intentionally reset model-owned RGLRU state.
  • Add a loaded one-layer regression proving independent no-ID prompts reset and stable sequence-ID chunks match one-pass recurrent execution.

State and memory contract

Current DirectEngine prefill carries one stable sequence ID, so its state handling does not enter the no-ID reset. The capability is global, however, so returning false makes every honoring caller, including DirectEngine, use a whole-prompt graph. Long prompts can therefore use more peak transient memory than chunk-bounded prefill.

Test plan

  • Root and surgery CUDA test-fast build graph: passed in 14 minutes 25 seconds
  • GPU-locked recurrent_gemma_ focused suite: 5 passed in 0.84 seconds, including the loaded-model contract and existing model-owned state tests
  • git diff --check
  • Repository integration contracts: all stages passed, including 123 Python tests and 49 WebUI fixtures

Closes #2288

Disable chunked prefill because no-ID multi-token forwards reset model-owned RGLRU state while preserving stable-ID continuation.

Add a loaded one-layer regression that compares reset prompts with a fresh model and stable-ID split execution with one-pass logits. Focused CUDA validation is running under the orchestrator.

Refs #2288
@inureyes inureyes added status:review Under review type:bug Bug fixes, error corrections, or issue resolutions priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) labels Oct 13, 2026
Preserve the newly merged MLX pin and overlay baseline declarations before validating the model prefill correction against current main. The model-specific implementation and synthetic regression are unchanged.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:inference Generation, sampling, decoding (incl. speculative, DRY) area:models Model architectures, weights, loading, metadata priority:medium Medium priority status:review Under review type:bug Bug fixes, error corrections, or issue resolutions

Projects

None yet

Development

Successfully merging this pull request may close these issues.

fix(recurrent_gemma): correct the no-ID chunked prefill contract

1 participant