Skip to content

Support Gemma 4's labeled reasoning channels Part 1 & 2 - #563

Open
abra-code wants to merge 4 commits into
ml-explore:mainfrom
abra-code:reasoning-generation-events
Open

abra-code wants to merge 4 commits into
ml-explore:mainfrom
abra-code:reasoning-generation-events

Conversation

@abra-code

Copy link
Copy Markdown
Contributor

Proposed changes

Part 1: Support Gemma 4's labeled reasoning channels
ReasoningConfig describes a reasoning protocol as a delimiter pair, which Gemma 4 does not use. Its delimiters are special tokens (soc_token id 100 <|channel>, eoc_token id 101 <channel|>) and the opening token is followed by a role label, terminated by a newline, that names the span:

<|channel>thought\n ... <channel|>the answer

A pair cannot carry that. Anchored on "<|channel>" it emits the label as the first word of the thinking; anchored on "<|channel>thought\n" it stops matching the moment the label is anything else and leaks the raw opener into the answer.

Add ReasoningChannel, carried as an optional ReasoningConfig.channel, and give ReasoningEventEmitter a label state. nil keeps the pair protocol exactly as it was, so Qwen3 and DeepSeek-R1 are unchanged and the existing ReasoningEventEmitterTests are untouched.

The label is consumed as metadata rather than emitted; a label not listed as a response label routes to reasoning, so an unknown channel stays readable and cannot leak markers into the answer; hold-back covers the longest watched delimiter, since both Gemma 4 tokens are 10 characters and a splitter sized for tears them; an opener whose label never terminates settles into a visible body instead of buffering without bound; and a stray end delimiter with no opener is swallowed, because a special token arriving alone means the prompt opened the channel and it cannot be prose the model typed.

ReasoningChannel.gemma4 lists "content" as a response label defensively rather than from observation: no shipped Gemma 4 template emits such a channel, and Google's own x-regex response schema names the answer group "content" while leaving it outside any channel. Listing it costs nothing and means a channel-wrapped answer would not be hidden in the thought stream if one ever did appear.

The four Gemma 4 model types declare the protocol themselves through ChatConventionsProviding, as NanbeigeModel already does.

Also add a token-id overload of promptEndsInsideReasoning so the prompt-tail rule has one implementation rather than two.

Part 2: Add reasoning on the streaming API as Generation.reasoning
Nothing on the standard generation path reads ModelConfiguration.reasoningConfig.
StandardTokenStreamDecoder wraps ToolCallProcessor, which only knows tool syntax.
Raw ... therefore reached Generation.chunk with its delimiters
still attached. ChatSession then copied it into the recorded assistant message,
so thinking text was replayed into every later prompt. These families' own chat
templates strip it out. Gemma 4's <|channel>thought markers are the visible form
of the same gap.

Give StandardTokenStreamDecoder an optional ReasoningEventEmitter. It splits
reasoning out after the stop-string filter and before the tool-call parser.
Order matters in both directions. Splitting first means a model that writes
<|tool_call> in its scratchpad cannot produce a phantom tool call. finish()
drains the scanner before the processor's EOS, so an end delimiter that never
arrives as text cannot strand held-back thinking. A model with no
ReasoningConfig decodes nothing and costs nothing.

Framed token protocols decode their own reasoning. makeTokenStreamDecoder
returns the protocol decoder before it builds the standard one, so the reasoning
argument is dropped there. Muse-Glimmer declares both .atem and a
ReasoningConfig, and gets only the Onyx decoder.

Add Generation.reasoning(String) and emit it instead of dropping it. Routing
reasoning out without exposing it would delete text callers get today. The
existing drop arrived with the Onyx protocol in PR #523, "Add Muse-Glimmer, a
30B agentic multimodal model, with ATEM tool calling". It covered only framed
protocols, which had no public reasoning to lose. Adding a case to a public
non-frozen enum breaks exhaustive switches downstream. Every in-repo site is
updated, including those the package scheme never builds.

DeepSeek-R1 prefills into the prompt and generates only the closing tag.
A scanner that starts outside would send the thought to the answer, then send
the answer to reasoning. Every generate() overload computes priming from the
prompt it already holds. generateTask takes it as a defaulted parameter, for
callers that drive the loop themselves. Only the prompt tail is read.

ChatSession counts a thinking-only turn as a turn that happened. Reasoning never
enters the assistant message. It is dropped from replayed history, the way these
families' templates do it. A model that runs out of maxTokens inside its thought
block therefore leaves content empty. If that looked like an empty generation,
the rollback would remove the user's message from the transcript too.

MLXFoundationModels drops .reasoning on its unconstrained path. Anything with a
resolved config goes to runReasoning instead. Reaching that path means the
caller did not ask for .reasoning and the prompt was rendered with thinking off.
Surfacing it would be the leak the capability gate exists to stop.

Co-Authored-By: Claude

Checklist

Put an x in the boxes that apply.

  • I have read the CONTRIBUTING document
  • I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
  • I have added tests that prove my fix is effective or that my feature works
  • I have updated the necessary documentation (if needed)

@abra-code

Copy link
Copy Markdown
Contributor Author

This PR supersedes the the Part 1-only PR: #504

ReasoningConfig describes a reasoning protocol as a delimiter pair, which
Gemma 4 does not use. Its delimiters are special tokens (soc_token id 100
<|channel>, eoc_token id 101 <channel|>) and the opening token is followed
by a role label, terminated by a newline, that names the span:

    <|channel>thought\n ... <channel|>the answer

A pair cannot carry that. Anchored on "<|channel>" it emits the label as the
first word of the thinking; anchored on "<|channel>thought\n" it stops
matching the moment the label is anything else and leaks the raw opener into
the answer.

Add ReasoningChannel, carried as an optional ReasoningConfig.channel, and
give ReasoningEventEmitter a label state. nil keeps the pair protocol
exactly as it was, so Qwen3 and DeepSeek-R1 are unchanged and the existing
ReasoningEventEmitterTests are untouched.

The label is consumed as metadata rather than emitted; a label not listed as
a response label routes to reasoning, so an unknown channel stays readable
and cannot leak markers into the answer; hold-back covers the longest
watched delimiter, since both Gemma 4 tokens are 10 characters and a
splitter sized for </think> tears them; an opener whose label never
terminates settles into a visible body instead of buffering without bound;
and a stray end delimiter with no opener is swallowed, because a special
token arriving alone means the prompt opened the channel and it cannot be
prose the model typed.

ReasoningChannel.gemma4 lists "content" as a response label defensively
rather than from observation: no shipped Gemma 4 template emits such a
channel, and Google's own x-regex response schema names the answer group
"content" while leaving it outside any channel. Listing it costs nothing and
means a channel-wrapped answer would not be hidden in the thought stream if
one ever did appear.

The four Gemma 4 model types declare the protocol themselves through
ChatConventionsProviding, the seam every model now uses after PR ml-explore#502,
"Migrate all models to ChatConventionsProviding; delete the infer chains".
The shared ReasoningConfig.gemma4 preset joins the two think-tag presets that
migration left in its place.

Also add a token-id overload of promptEndsInsideReasoning so the
prompt-tail rule has one implementation rather than two.

Co-Authored-By: Claude
The constrained-generation path decided whether to run a think-then-call
phase by reading the reasoning level as `thinkingEnabled(for:) != false`,
i.e. treating an unspecified level as "think". A few lines below, the same
path rendered the prompt with `thinkingEnabled(for:) ?? defaultOn`.

Those two readings agree for every family whose template defaults thinking
on, so the split was invisible while `.templateFlag` meant Qwen3. Gemma 4 is
the first `defaultOn: false` family, and there they disagree: the gate opens
a thinking phase while the prompt tells the template not to think. Gemma 4's
31B template answers that by prefilling a closed, empty
`<|channel>thought\n<channel|>`, so no channel ever opens, the phase spends
its whole `maxTokens` budget, and a `.required` tool call never arrives.

Add a single `thinkingEnabled(for:defaultOn:)` overload and route both the
gate and the prompt through it, so the two cannot drift apart again. Behavior
is unchanged for every `defaultOn: true` model.

The accompanying tests pin the resolution against what `ReasoningConfig`
itself injects into the template. Note their header: `MLXLanguageModel` is
`@available(macOS 27.0, ...)`, so they need a macOS 27 host to execute -
building against the macOS 27 SDK alone lets the `#available` guards return
early and the suite reports green without asserting anything.
Nothing on the standard generation path reads ModelConfiguration.reasoningConfig.
StandardTokenStreamDecoder wraps ToolCallProcessor, which only knows tool syntax.
Raw <think>...</think> therefore reached Generation.chunk with its delimiters
still attached. ChatSession then copied it into the recorded assistant message,
so thinking text was replayed into every later prompt. These families' own chat
templates strip it out. Gemma 4's <|channel>thought markers are the visible form
of the same gap.

Give StandardTokenStreamDecoder an optional ReasoningEventEmitter. It splits
reasoning out after the stop-string filter and before the tool-call parser.
Order matters in both directions. Splitting first means a model that writes
<|tool_call> in its scratchpad cannot produce a phantom tool call. finish()
drains the scanner before the processor's EOS, so an end delimiter that never
arrives as text cannot strand held-back thinking. A model with no
ReasoningConfig decodes nothing and costs nothing.

Framed token protocols decode their own reasoning. makeTokenStreamDecoder
returns the protocol decoder before it builds the standard one, so the reasoning
argument is dropped there. Muse-Glimmer declares both .atem and a
ReasoningConfig, and gets only the Onyx decoder.

Add Generation.reasoning(String) and emit it instead of dropping it. Routing
reasoning out without exposing it would delete text callers get today. The
existing drop arrived with the Onyx protocol in PR ml-explore#523, "Add Muse-Glimmer, a
30B agentic multimodal model, with ATEM tool calling". It covered only framed
protocols, which had no public reasoning to lose. Adding a case to a public
non-frozen enum breaks exhaustive switches downstream. Every in-repo site is
updated, including those the package scheme never builds.

DeepSeek-R1 prefills <think> into the prompt and generates only the closing tag.
A scanner that starts outside would send the thought to the answer, then send
the answer to reasoning. Every generate() overload computes priming from the
prompt it already holds. generateTask takes it as a defaulted parameter, for
callers that drive the loop themselves. Only the prompt tail is read.

ChatSession counts a thinking-only turn as a turn that happened. Reasoning never
enters the assistant message. It is dropped from replayed history, the way these
families' templates do it. A model that runs out of maxTokens inside its thought
block therefore leaves content empty. If that looked like an empty generation,
the rollback would remove the user's message from the transcript too.

MLXFoundationModels drops .reasoning on its unconstrained path. Anything with a
resolved config goes to runReasoning instead. Reaching that path means the
caller did not ask for .reasoning and the prompt was rendered with thinking off.
Surfacing it would be the leak the capability gate exists to stop.
@davidkoski
davidkoski force-pushed the reasoning-generation-events branch from fa17055 to c952ef7 Compare September 14, 2026 20:56
@davidkoski

Copy link
Copy Markdown
Member

rebased to main, mostly to pick up #548 which had some changes in the same area

Comments shortened to comply with the new AI agent policy.
Added new tests for prompt priming and channel label too long (1 char over maxLabelLength=32 limit)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants