Skip to content

Fix Gemma 4 session cache across rewritten conversation turns - #283

Open
TipKnuckle wants to merge 1 commit into
youssofal:mainfrom
TipKnuckle:fix/gemma4-session-cache
Open

Fix Gemma 4 session cache across rewritten conversation turns#283
TipKnuckle wants to merge 1 commit into
youssofal:mainfrom
TipKnuckle:fix/gemma4-session-cache

Conversation

@TipKnuckle

Copy link
Copy Markdown

Summary

  • reconstruct Gemma 4 assistant reasoning from reasoning_content using its native channel markers, so echoed assistant history retokenizes to the committed stream
  • retain the pre-decode Gemma 4 prompt boundary and store it with the backend's expected pre-norm hidden state and shared-KV metadata
  • restore the longest compatible common prefix when clients rewrite a prior volatile memory suffix, then prefill only the divergent tail
  • cover native reasoning token identity and a two-turn memory rewrite with warm/cold state parity

Verification

  • python -m pytest tests/test_no_mlx_imports.py tests/test_public_cli.py tests/test_runtime_kpis.py — 268 passed
  • focused Gemma/server/session-cache regression set — 375 passed
  • python -m build — sdist and wheel built successfully
  • scripts/fresh_venv_smoke.sh — passed
  • Ruff and git diff --check — passed
  • full python -m pytest -q — only five pre-existing exact-float failures in tests/test_graphbank_compiled_verify.py; the same five failures reproduce from clean upstream commit f52f9be

Benchmark Evidence

Measured 2026-08-17 on commit a233e9a:

  • hardware: Apple M4 Max, 128 GB, macOS 26.5.2
  • model: local Gemma 4 31B QAT target/assistant pair; target mixed 4-bit/8-bit QAT (MLP projections 8-bit), assistant 4-bit
  • runtime: resolved sustained profile, MTP generation, reasoning on, SSD session cache off, default Apple fan control
  • sampler: temperature 0, seed 1729, max 128 tokens; each compared completion used 53 tokens
  • workload: second OpenAI-format turn removes the prior user's volatile <memory_context>, echoes assistant reasoning_content plus content, and injects a new memory block on the current user turn
Run Prompt Cached New prefill Prompt eval Wall
Warm rewritten turn 1,705 1,610 95 0.519 s 2.028 s
Cold control 1,705 0 1,705 8.436 s 9.944 s
Warm exact repeat 1,705 1,705 0 0.0001 s 1.545 s

The rewritten warm turn was 4.9x faster end-to-end than cold; the exact repeat was 6.44x faster. The complete returned message, including both reasoning_content and visible content, matched the cold control exactly in both warm cases.

@TipKnuckle
TipKnuckle requested a review from youssofal as a code owner August 18, 2026 03:32
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