Skip to content

Fix Studio reasoning channel rendering - #7121

Merged
oobabooga merged 15 commits into
unslothai:mainfrom
Lyxot:fix/studio-gemma4-channel-rendering
Jul 17, 2026
Merged

oobabooga merged 15 commits into
unslothai:mainfrom
Lyxot:fix/studio-gemma4-channel-rendering

Conversation

@Lyxot

@Lyxot Lyxot commented Jul 14, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR fixes Studio inference rendering for models that emit native reasoning-channel markers, especially Gemma-family checkpoints that can output <|channel>thought / <channel|> style control text. Studio should present that content as structured reasoning (<think>...</think>) instead of leaking raw channel markers into the final assistant message.

Problem

Some Gemma-style native chat templates encode reasoning through channel delimiters instead of the <think> format Studio already understands. When those models generated through Studio, the raw native markers could appear in chat output, for example internal channel text, thought labels, or turn delimiters. The issue was visible in Studio chat for Gemma 4 MLX checkpoints, but the underlying rendering problem sat in the shared Studio inference path rather than in a single model card.

Root Cause

Studio previously inferred whether to preserve or strip special tokens from tokenizer metadata, model-level hints, or the original tokenizer state. That was not the same as the exact chat template used to render the current request. Native template fallback could select a channel-based prompt while the downstream streamer still inspected the original tokenizer/template and missed the protocol. The opposite could also happen: token metadata could make Studio preserve special tokens for requests that were not actually using native reasoning channels.

After generation started, CUDA/safetensors and MLX streaming paths treated native channel text as ordinary assistant text. That meant the backend had no consistent place to convert native reasoning-channel output into Studio's reasoning representation or to strip response-boundary control tokens safely.

What Changed

  • Detect native reasoning-channel markers from the active chat template selected for the request, not from unrelated token metadata alone.
  • Carry chat-template render metadata into generation/streaming so the streamer uses the same template decision that built the prompt.
  • Normalize native reasoning-channel output into Studio's existing <think>...</think> reasoning format.
  • Apply the normalization across CUDA/safetensors, MLX text inference, and MLX VLM streaming snapshots.
  • Preserve normal decoding behavior for models/templates that do not use native reasoning channels.
  • Keep Harmony/gpt-oss reasoning handling on its existing path, with Harmony detection taking precedence over this native-channel normalization.
  • Preserve tool delimiters by avoiding broad special-token stripping when only response-boundary tokens should be removed.
  • Harden local generation error handling so typed stream errors surface as backend errors instead of being sliced into successful model output.

Scope And Behavior Boundaries

This is intended to affect only Studio inference output rendering for native reasoning-channel templates. It does not change training, save/export, GGUF/llama.cpp execution, remote provider APIs, model loading, or tokenizer files. Non-native templates continue through the existing normal text decoding path with skip_special_tokens=True behavior where applicable.

Models with ordinary <think> output are not converted through this native-channel normalizer. Harmony/gpt-oss streams keep their dedicated reasoning parser. Cached template checks covered Gemma, Qwen, GLM, DeepSeek, FunctionGemma, and no-template/base-style cases to reduce the chance of false-positive activation.

Reviewer Notes

The main review surface is the handoff between chat-template selection and streaming. The important invariant is that if native-channel fallback renders the prompt, the downstream stream normalizer must see the same marker metadata; if the selected template does not contain native-channel markers, the normal decoding path should remain unchanged.

The stream normalizer is intentionally protocol-based rather than model-name based. That keeps the fix applicable to Gemma-family native channel templates without special-casing a specific checkpoint name.

Validation

Focused regression suite:

python -m pytest studio/backend/tests/test_think_prefill_reemit.py studio/backend/tests/test_safetensors_reasoning_stream.py studio/backend/tests/test_mlx_inference_backend.py studio/backend/tests/test_sf_client_tools_passthrough.py -q

Result:

73 passed, 3 warnings

Expanded local validation also passed before opening the PR:

python -m pytest studio/backend/tests/test_think_prefill_reemit.py studio/backend/tests/test_safetensors_reasoning_stream.py studio/backend/tests/test_mlx_inference_backend.py studio/backend/tests/test_sf_client_tools_passthrough.py studio/backend/tests/test_safetensors_tool_loop.py studio/backend/tests/test_native_template_trust_remote_code.py studio/backend/tests/test_inference_orchestrator_crash_message.py studio/backend/tests/test_inference_dispatcher_resilience.py studio/backend/tests/test_openai_tool_passthrough.py studio/backend/tests/test_responses_tool_passthrough.py studio/backend/tests/test_offline_inference_parent.py studio/backend/tests/test_audio_token_detection.py -q

Result:

756 passed, 4 warnings

Additional checks:

  • Cached-template scan detected native markers for Gemma 4 instruction and DiffusionGemma instruction templates.
  • Cached-template scan did not activate native-channel handling for Qwen3, Qwen3.5, Gemma3, GLM-OCR, DeepSeek, FunctionGemma, or Gemma 4 base/no-template cases.
  • Real local MLX smoke check with unsloth/gemma-4-E4B-it-UD-MLX-4bit returned clean output for both thinking enabled and disabled, without raw channel marker leakage.
  • Synthetic CUDA text-stream validation converted native-channel output to clean assistant text/reasoning output.
  • Synthetic CUDA VLM validation covered natural completion, generation failure propagation, cancellation behavior, EOS cleanup, and lock release.

@Lyxot
Lyxot requested a review from danielhanchen as a code owner July 14, 2026 12:45
Copilot AI review requested due to automatic review settings July 14, 2026 12:45
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Repo admins can enable using credits for code reviews in their settings.

@Lyxot

Lyxot commented Jul 14, 2026

Copy link
Copy Markdown
Collaborator Author

@codex review

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for native reasoning channels (such as Gemma's native thought-channel protocol) and normalizes them into canonical tags during streaming across both Transformers and MLX backends. It also improves streaming error handling and client-safe error propagation. The review feedback highlights a few issues: a fallback loop in _selected_chat_template_strings that breaks prematurely on exceptions instead of continuing, potential infinite loop risks in the custom ReasoningChannelNormalizer parser, and unreachable dead code in mlx_inference.py where the generation_failed flag is set right before an exception is re-raised.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread studio/backend/core/inference/chat_template_helpers.py Outdated
Comment thread studio/backend/core/inference/chat_template_helpers.py
Comment thread studio/backend/core/inference/mlx_inference.py

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5d6a4c4410

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread studio/backend/core/inference/chat_template_helpers.py
Comment thread studio/backend/core/inference/chat_template_helpers.py Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Note

Copilot couldn't run its full agentic review because no GitHub Actions runner was available. Make sure your repository has a runner available to run Copilot's review, or add a copilot-setup-steps.yml file specifying one with the runs-on attribute. See the docs for more details.

Fixes Studio’s streaming/rendering of native reasoning-channel output (notably Gemma-style <|channel>thought ... <channel|>) by detecting markers from the active chat template, normalizing native channel text into <think>...</think>, and ensuring streamed backend errors are surfaced as errors (not assistant text).

Changes:

  • Detect reasoning-channel markers from the request’s selected chat template and propagate that metadata into streaming backends.
  • Normalize native reasoning-channel output into Studio’s canonical <think>...</think> format across safetensors/CUDA, MLX text, and MLX VLM streaming.
  • Harden streaming error propagation (typed streamed errors + thread failures) so errors don’t leak as model text.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 2 comments.

Show a summary per file
File Description
studio/backend/core/inference/chat_template_helpers.py Adds active-template-based detection of native reasoning markers plus stream normalizers and render metadata plumbing.
studio/backend/core/inference/inference.py Updates Transformers streaming to use a protocol-aware streamer and propagate typed generation/thread errors.
studio/backend/core/inference/mlx_inference.py Carries template metadata into MLX generation and normalizes native reasoning channels for text/VLM.
studio/backend/core/inference/orchestrator.py Adds public error support and an exception form to propagate streamed errors across generator boundaries.
studio/backend/routes/inference.py Treats GenStreamError as an error event (not content) across multiple streaming routes and improves error propagation.
studio/backend/tests/test_think_prefill_reemit.py Adds unit coverage for template-based marker detection + normalizer behavior + native-template fallback metadata.
studio/backend/tests/test_*.py Adds streaming regression tests for native reasoning normalization and error handling across backends.
Comments suppressed due to low confidence (2)

studio/backend/core/inference/chat_template_helpers.py:1

  • This uses truthiness (if native_prompt:) to decide whether the native render succeeded. If the native template render legitimately returns an empty prompt (e.g., edge-case message set / template behavior), \"\" is falsy and the function will incorrectly fall back to formatted_prompt. Prefer checking explicitly for None (the documented failure sentinel), e.g. if native_prompt is not None: (and for ChatTemplateRenderResult, check native_prompt is not None rather than .prompt truthiness).
# SPDX-License-Identifier: AGPL-3.0-only

studio/backend/routes/inference.py:1

  • There’s now an inconsistency in the generic internal-error string: _friendly_error returns \"An internal error occurred\" while _friendly_gen_stream_error falls back to \"An internal error occurred.\" (with a period). This can create inconsistent UI output and brittle tests/string matching. Consider standardizing on one exact message (including punctuation) across both helpers.
# SPDX-License-Identifier: AGPL-3.0-only

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread studio/backend/core/inference/orchestrator.py Outdated
Comment thread studio/backend/routes/inference.py
@oobabooga

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Hooray!

Reviewed commit: d14e9ad91c

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

@oobabooga

Copy link
Copy Markdown
Member

Verified live on CUDA with unsloth/gemma-4-e2b-it-unsloth-bnb-4bit.

Tools + thinking path: fixed.

  • main: reasoning leaks into content as thought\n..., reasoning_content empty.
  • this PR: clean reasoning_content, no marker leakage.

Error-propagation hardening: confirmed end to end.

Two blockers.

1. CI red on all four backend "(Python 3.x)" jobs

test_native_reasoning_streamer_selected_and_errors_raise in test_safetensors_reasoning_stream.py imports core.inference.inference, which pulls in unsloth/peft/torch. Those jobs are hermetic and lack them, so the test fails at collection:

ModuleNotFoundError: No module named 'peft'

Sibling tests guard this with try/except ImportError. This test executes the backend and references _GenerationThreadError, so a None-fallback does not apply. Use a skip:

inf = pytest.importorskip("core.inference.inference")

2. Plain (no-tools) thinking chat unchanged for gemma-4 on CUDA/safetensors

Output is byte-identical to main and still leaks thought\n... into visible content.

Cause:

  • The CUDA text path installs Unsloth's registry gemma-4 template via get_chat_template() before template-based detection runs.
  • That template has no <|channel>thought marker; only gemma-4-thinking does.
  • detect_reasoning_channel_markers() therefore returns None, and the streamer keeps skip_special_tokens=True.

The model emits the channel protocol in plain thinking mode, so the normalizer would handle it if detection also consulted the model's native or thinking template on this path.

The PR description reads as a general Studio-chat fix. As written it does not hold for default CUDA chat.

@Lyxot

Lyxot commented Jul 17, 2026

Copy link
Copy Markdown
Collaborator Author

Thanks, both blockers should be addressed in b4db53487.

  • CI failure: changed the backend-heavy streamer test to use pytest.importorskip("core.inference.inference"), so environments without optional full inference deps like peft skip that specific test instead of failing import collection.
  • Gemma 4 no-tools thinking on CUDA/safetensors: marker detection now falls back to cached native chat_template_info["template"] metadata when the active Unsloth override template has no native thought-channel markers. The rendered prompt is unchanged; this only selects the reasoning-aware streamer. Tested with unsloth/gemma-4-E2B-it

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: b4db534874

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread studio/backend/core/inference/chat_template_helpers.py
@oobabooga

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 6488f2c69f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread studio/backend/core/inference/mlx_inference.py
@oobabooga
oobabooga force-pushed the fix/studio-gemma4-channel-rendering branch from fdef271 to 4708c64 Compare July 17, 2026 19:17
@oobabooga

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 4708c64a8e

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread studio/backend/routes/inference.py
Comment thread studio/backend/core/inference/inference.py
Comment thread studio/backend/core/inference/chat_template_helpers.py Outdated
@oobabooga

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9824d2f6cc

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread studio/backend/routes/inference.py Outdated
Comment thread studio/backend/core/inference/inference.py
Comment thread studio/backend/core/inference/chat_template_helpers.py
@oobabooga

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 5ac5f8be8d

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

Comment thread studio/backend/routes/inference.py Outdated
@oobabooga

Copy link
Copy Markdown
Member

@codex review

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review: Didn't find any major issues. Breezy!

Reviewed commit: cd942d021f

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

@oobabooga

Copy link
Copy Markdown
Member

Thanks @Lyxot, here is a final test.

I verified your fix end to end on CUDA with unsloth/gemma-4-e2b-it-unsloth-bnb-4bit, comparing main against this branch merged with current main. On main, a thinking request returns content that begins with a literal "thought" line and an empty reasoning_content, and with tools attached the call never parses (tool_calls null). On your branch the same requests return clean content, populated reasoning_content, and a parsed get_weather call, streaming and non-streaming. enable_thinking=false output is byte-identical. The model is vision-capable, so the text-only VLM gating path ran too.

My follow-up commits narrow marker detection to the template each request actually renders, re-gate think parsing on reasoning capability, and pass tool context into detection. All three are no-ops when the template lacks the channel markers; scanning the 24 model templates in my HF cache trips detection only for the four gemma-4 variants.

Merging now.

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.

3 participants