Add generation loop/repetition detection to break Gemma 4 token-repetition collapse under tool/grammar-constrained decoding #432

Summary

Gemma 4 models (31B Dense and 26B-A4B MoE, including the QAT 4-bit checkpoints) have an upstream, weights-level token-repetition collapse: during generation the output degenerates into a single repeated token (or short fragment) that fills the token budget. This is documented upstream in google-deepmind/gemma#622 and reproduced across Ollama, LMStudio, vLLM, and Cloudflare Workers AI. mlxcel 0.3.2 is also affected.

The collapse is strongly amplified when the prompt carries tool declarations and/or when decoding is grammar-constrained (structured output / tool-call JSON). mlxcel currently has no countermeasure that reliably works for this failure mode, so we request an engine-side repetition/loop detection mechanism.

Environment

• mlxcel-server 0.3.2 (package 26.6.21), macOS arm64, Metal
• Model: mlx-community/gemma-4-26b-a4b-it-qat-4bit
• Launch flags: --ctx-size 262144 --jinja --parallel 2 --cache-type-k int8 --cache-type-v int8 --prefill-chunk-size 1024 (no --repeat-penalty)

Reproduction

Direct OpenAI-compatible requests to the running mlxcel socket, single user turn.

• No tools in the request: always clean.
• System prompt + 10 tool declarations, chat_template_kwargs.enable_thinking=true, fresh first turn: the reasoning/thought channel collapses into a repeated garbage token (observed: CJK such as 様様様 / 婕婕婕, box-drawing ◸◸◸ / ╿╽╿╽, and control characters) up to the thinking budget, after which a tool call still emerges. prompt_tokens is about 2150 against a 262144 context, so this is not context overflow.

Determinism and seed dependence (this proves the prompt is well-formed and the failure is a sampling trajectory, not a truncated/malformed prompt):

• Fixed seed=42, same payload, 3 runs: byte-identical output every time.
• Same payload, varying seed: clean at some seeds, degenerate at others.

Sampling comparison on the worst case (10 tools, enable_thinking=true, first turn):

Why existing mitigations are insufficient

• The Gemma-4 channel priming in server/chat_template.rs (patch_gemma4_generation_prompt) only applies when enable_thinking=true AND the rendered prompt ends exactly in <|turn>model\n. Under heavier real prompts it does not reliably fire, and even when it does the collapse still occurs (see the 3-4/5 row).
• Flat repeat_penalty and the DRY sampler shape the probability distribution. Once the logits collapse, the top-k candidates are themselves garbage, so distribution shaping cannot recover. Upstream [BUG] Gemma 4 token repetition collapse during long generation — affects both 31b Dense and 26b MoE google-deepmind/gemma#622 confirms repeat_penalty 1.0 to 1.5 has no reliable effect.

Proposed mitigation

Add N-gram repetition detection that ends a generation early when the output stream contains a short pattern that repeats. This is exactly the mechanism vLLM exposes in SamplingParams:

• max_pattern_size (int, default 0 = disabled): maximum N-gram pattern size to scan for. Must be set together with min_count.
• min_pattern_size (int, default 0, treated as 1): minimum N-gram pattern size to check. Must be <= max_pattern_size.
• min_count (int, default 0, must be >= 2): minimum number of repeats of the pattern that triggers detection.
• Action: when a pattern whose size is in [min_pattern_size, max_pattern_size] repeats at least min_count times, generation is ended early. The vLLM docs state: "If such repetition is detected, generation will be ended early." With min_pattern_size=1 this catches single-token loops such as the observed 様様様 collapse, and larger sizes catch short multi-token loops such as the the the or abcdabcd....

In vLLM these fields are opt-in (all default 0). For mlxcel we suggest either defaulting them on for the Gemma 4 family, or auto-enabling them when tools or response_format: json_schema is present, with a per-request override. A starting point that covers the loops we observed is roughly min_pattern_size=1, max_pattern_size=20, min_count=4.

Important: the detector must run on the raw generated token stream (around the reasoning-channel and tool-call parsing), so it also catches loops inside the <|channel>thought span, not just the final answer. In our reproduction the collapse most often filled the thought channel up to the thinking budget.

Preferred activation surface (for downstream serving apps)

mlxcel is bundled as a serving engine inside a desktop app, and many users download Gemma 4 and derivative checkpoints independently. To keep behavior consistent without per-user configuration or per-model file edits, we would prefer the detection to be activated at the engine level rather than only as a per-request field:

• Best: auto-enable the detection for the Gemma 4 model family by default, with sane built-in parameters, so a downstream serving app needs no configuration and end users see no toggle. This is a safety mechanism that ends generation only when a real repetition loop is detected, so a conservative default-on for this family is low risk.
• Acceptable: expose it as server-level defaults via CLI flags (the same way --repeat-penalty is exposed today) and/or auto-enable it whenever tools or response_format: json_schema is present, with a per-request override. A single launch default per family is easy for an integrator to set.
• Please avoid making a per-request SamplingParams field the only surface. That forces every integrating client to add request plumbing, and settings UI if it is user-tunable, for what is really an engine-level safety default.

Rationale: an engine-side default keeps the fix consistent across every integrating app and every independently downloaded checkpoint, with no per-user options and no per-model template edits.

References

• [BUG] Gemma 4 token repetition collapse during long generation — affects both 31b Dense and 26b MoE google-deepmind/gemma#622 (root cause: Gemma 4 token repetition collapse; reports flat repeat_penalty 1.0 to 1.5 as ineffective)
• [Bug]: Gemma 4 (31B / 26B-A4B) generates infinite repetition loops, especially with structured output (JSON schema) vllm-project/vllm#40080 (Gemma 4 infinite repetition with structured output)
• vLLM SamplingParams N-gram repetition detection fields max_pattern_size / min_pattern_size / min_count (docs.vllm.ai/en/latest/api/vllm/sampling_params/)
• gemma4:31b repetition loop during constrained JSON generation with free-text string fields ollama/ollama#15502, Feature request: repetitionPenalty on SamplerConfig + per-segment sampler override - reproducible Gemma 4 tool-arg corruption catalog inside google-ai-edge/LiteRT-LM#2249
• huggingface google/gemma-4-12B-it discussion test(vision): add qwen2.5-vl-3b-4bit warmup regression guard #38 (chat-template prior-turn reasoning re-injection amplifies the loop on multi-turn tool use)