Conversation
4 tasks
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
force-pushed
the
reasoning-generation-events
branch
from
September 14, 2026 20:56
fa17055 to
c952ef7
Compare
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)
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.
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:
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
xin the boxes that apply.pre-commit run --all-filesto format my code / installed pre-commit prior to committing changes