Skip to content

Count single-token reasoning delimiters only when generated as that token (#48) - #94

Merged
lukealonso merged 1 commit into
NVIDIA:mainfrom
mark-yong:fix/think-delimiter-tokens
Oct 26, 2026
Merged

lukealonso merged 1 commit into
NVIDIA:mainfrom
mark-yong:fix/think-delimiter-tokens

Conversation

@mark-yong

@mark-yong mark-yong commented Oct 26, 2026 •

Copy link
Copy Markdown
Contributor

Fixes #48.

Request

ThinkReasoning finds reasoning delimiters by matching decoded text. When a model writes <think> or </think> as ordinary text, for example while explaining a chat template, Helios treats the mention as a real delimiter: text moves between reasoning_content and content, and the tag disappears (#48). The real delimiter is a single added token (248068 and 248069 on Qwen 3.8); a mention arrives as ordinary pieces such as . The issue reproduces the defect in 5 of 7 live Qwen 3.8 runs and in parser tests for Qwen and GLM.
Structured-output grammars have the same blind spot: the reasoning region ends on the delimiter's text, so a spelled </think> inside reasoning ends it and the rest of the reasoning is forced into the answer grammar.

Change

  • reasoning_delimiter_tokens (helios/src/model/prompt.rs) encodes the contract's reasoning open and close with the frontend tokenizer and keeps each one that encodes to exactly one token and decodes back to the same text. Server::new derives the list once; GrammarCompiler stores it and each request passes it to prepare_output, so the parser and the grammar always use the same list.
  • TaggedOutput and ThinkReasoning (helios/src/output/tagged.rs) accept a listed delimiter only where its token was generated. Input::Token(id) arrives before its decoded text, so the parser records the byte offset where a listed token's text will start and holds text that could still become that delimiter until the text arrives. Text at any other offset, including a full spelled-out tag, is ordinary reasoning or answer text. The MiniMax leading-close check (initial_close) follows the same rule.
  • At finish, text held for a delimiter whose text never arrived is parsed as ordinary text. Cancelled and failed requests discard it, as they discard all other output.
  • OutputPipeline::gates_reasoning_tokens excludes DSML and XTML, which recognize their markup in their own state machines, and lists every tool parser kind so a new kind cannot be added without deciding.
  • GrammarCompiler (helios/src/grammar.rs) references each delimiter by token ID (<[id]>), which llguidance matches against the generated token before its bytes. Reasoning therefore ends only at the token, while the token's bytes stay ordinary text, so a structured answer may still contain it and returns it as text. The grammar does not ban the delimiter tokens from structured answers: Qwen 3.8 generates the <think> token inside JSON strings, and a ban pushes the model to substitute <tool_call> or a blank. Referencing the delimiters by ID leaves the answer grammar the same as without this change.
  • docs/api.md and helios/model-contract.md state the rule and its limitations.
    vLLM takes the same approach: its parsers declare token_id_terminals and TokenIDScanner matches them by token ID. SGLang's structured-output backend ends reasoning on think_end_ids through TokenSequenceMatcher.
    No Protobuf, metadata or package change. Existing packages pick up the fix with the new Flare build.

Scope across recipes

The rule depends on how each checkpoint's tokenizer encodes its declared delimiters, not on model identity. Delimiters resolved from real checkpoint tokenizers:

Recipe Checkpoint tokenizer Gated delimiters
qwen38 Qwen3.8-Flash-Next-FP4-Local <think> 248068, </think> 248069
glm53 GLM-5.3-Flash-NVFP4 <think> 154841, </think> 154842
mimo_v2 MiMo-V2.6-Flash-MOPD <think> 151667, </think> 151668
nemotron_h Nemotron-3.5-Super-VL-120B-A12B-NVFP4 <think> 12, </think> 13
minimax_m3 MiniMax-M3.1-Flash <mm:think> 200059, </mm:think> 200060
step5 Step-5-Preview-NVFP4 <think> 128798, </think> 128799
gemma4 gemma-4-E4B-it-NVFP4 <channel|> 101; <|channel>thought is two tokens and stays text-matched
deepseek4, deepseek41 DeepSeek-V4.1-Flash none; the official DSML parser recognizes its own markup
kimi3 not checked none; XTML already gates its control tokens

If a delimiter is not one token, it falls back to text matching, the behavior without this change, so for that recipe #48 is not fixed and nothing regresses.

Behavior changes

  • A spelled-out delimiter is now kept as text in reasoning_content or content instead of switching modes.
  • MiniMax: a leading </mm:think> in output that starts in content is removed only when it is the token. Spelled out, it now stays in content.
  • Structured output: a spelled </think> in reasoning no longer ends reasoning; only the token does.

Validation

All runs on a 32-core x86_64 build host, base 82cb15a0 with this change applied, --features server,runtime-package unless noted. Checkpoint paths, test descriptions and further live observations are in the first comment on this PR, with the live harness scripts.

  • Helios lib suite: 1220 passed, 21 ignored, 5 failed. The tree at 82cb15a0 without this change fails the same five on the same host: four engram::reader tests that cannot register staging memory (Cannot allocate memory), and package::cache::selection_tests::failed_default_clear_keeps_the_selected_entry, which cannot fail a write as root.
  • All 12 integration targets pass.
  • test_requests, the test-tools binary with an updated prepare_output call site, builds with --features server,test-tools.
  • New unit tests: output::pipeline::tests::{single_token_reasoning_delimiters_leave_spelled_tags_as_text, finish_releases_text_held_for_a_delimiter_token, channel_thought_gates_only_its_single_token_close}, model::prompt::tests::single_token_reasoning_delimiters_gate_generated_output, grammar::tests::single_token_reasoning_delimiters_end_reasoning_only_as_tokens, server::logprobs::tests::spelled_reasoning_delimiter_in_the_answer_receives_output_text_logprobs.
  • Checkpoint tests (--ignored), all passing: tests/reasoning_delimiters.rs 8 of 8 (the table above), tests/glm_frontend.rs 1 of 1, tests/deepseek_frontend.rs 2 of 2, grammar::tests::qwen_output_envelope_uses_checkpoint_token_bytes 1 of 1.
  • Formatting: rustfmt --check (edition 2024) reports no difference this change introduces.

Live Qwen 3.8 Flash, non-streaming

One package, Qwen3.8-Flash-Next (GB202), exported at bce032c0 from nvidia/Qwen3.8-Flash-Next-FP4-Local@b9ce18a6, concurrency 32, MTP draft tokens 5, served by Flare containers built from 3267c054 and from this branch on two RTX PRO 6000 Max-Q GPUs (SM120) with identical flags. Requests went to /v1/chat/completions with return_token_ids, fixed seeds, max_tokens 4096 and the model's default sampling. A response counts as correct when reasoning_content and content equal the generated token IDs split at the first </think> token (248069). Sampling was not bit-reproducible across runs, so the builds are compared by distribution rather than as paired outputs.

Requests 3267c054 correct This branch correct
#48 reproduction prompt, 8 seeds 1 of 8 8 of 8
Answer asked to write both tags, thinking on, 4 seeds 1 of 4 4 of 4
Same, thinking off, 4 seeds 3 of 4 4 of 4
Ordinary question with logprobs, 4 seeds 4 of 4 4 of 4
JSON schema asking for both tags in a string, 68 seeds 22 of 68 68 of 68
  • In all 84 branch responses with thinking on, reasoning ended at a generated </think> token; the model never wrote its real close in pieces.
  • On 3267c054, 46 of 68 JSON responses returned reasoning that differs from the tokens. In 37 the model never generated </think>: a spelled </think> in reasoning ended the grammar's reasoning region, forcing the rest of the reasoning into the answer grammar.

Live Qwen 3.8 Flash, streaming

The same package and images, with "stream": true and return_token_ids. Reasoning and content were rebuilt from the deltas and compared with the streamed token IDs split at the first </think> token. The images were built before the branch was rebased onto 82cb15a0; the commits after 3267c054 do not touch Helios output or grammar code.

Requests 3267c054 correct This branch correct
#48 reproduction prompt, 8 seeds 2 of 8 8 of 8
Answer asked to write both tags, thinking off, 4 seeds 3 of 4 4 of 4
Same, thinking on, 4 seeds 1 of 4 4 of 4
Ordinary question with logprobs, 4 seeds 4 of 4 4 of 4
JSON schema asking for both tags in a string, 8 seeds 3 of 8 8 of 8
Reasoning that writes both tags, then a get_weather tool call, 8 seeds: reasoning matches the tokens 4 of 8 7 of 8
  • On 3267c054, 9 of the 12 responses to the first and third prompts switched between reasoning and content deltas more than once. On the branch, every response streamed all reasoning before any content.
  • The branch's one tool-call reasoning mismatch is a generated <think> token inside reasoning, which is removed as a delimiter. 3267c054 removes it too, along with spelled <think> text.

Live Gemma 4 E4B

One package (unsloth/gemma-4-E4B-it-NVFP4@1c363766, GB202, concurrency 4), exported at 3267c054 with FlashInfer 344e4837, the pin from #85, served by the same images and GPUs, with identical flags and criteria as the Qwen run. Gemma closes reasoning with the single token <channel|> (101), while its <|channel>thought open stays text-matched. Each response was split at the first generated 101 token, with the <|channel>thought prefix removed from the reasoning half.

Requests 3267c054 correct This branch correct
Prompt asking the answer to write <channel|>, thinking on, 8 seeds 0 of 8 8 of 8
Same, thinking off, 4 seeds 3 of 4 4 of 4
Ordinary question with logprobs, thinking on, 8 seeds 8 of 8 8 of 8
JSON schema asking for <channel|> in a string, thinking on, 8 seeds 8 of 8 8 of 8
  • All 8 thinking-on reproduction responses on each build generated the <channel|> token and also spelled <channel|> in the text. On 3267c054, reasoning ended at the spelled marker inside the reasoning, and the rest of the reasoning was returned as content.
  • The model generated no reasoning in any JSON response, so these requests did not exercise the grammar's reasoning region on Gemma.
    Not run: live runs on any recipe other than Qwen 3.8 and Gemma 4 E4B, streaming on Gemma, and /v1/messages requests; image_admission and instruction_boundary, whose reference fixture hash does not match the available Qwen 3.8 snapshots; Heliopause, Python and CUDA tests.

Limitations

  • Delimiters that span several tokens remain text-matched: Gemma 4's <|channel>thought open, and DeepSeek's delimiters through the official parser.
  • A delimiter token the model generates as text inside reasoning still acts as a delimiter: a generated <think> is removed from reasoning, and a generated </think> ends it. Qwen 3.8 did this once in 32 thinking-on streaming responses.
  • If a model writes its real closing delimiter in ordinary pieces, the rest of the output stays in reasoning_content. Under a structured-output grammar, reasoning then cannot end, so the request runs to its token limit without an answer. Qwen 3.8 never wrote its close in pieces in this branch's live runs. On 3267c054, 37 of 68 JSON responses never generated it and produced answers only because the grammar ended reasoning at spelled text. vLLM's token-ID terminals carry the same assumption.
  • The tokenizer shows that a delimiter is one token, not that the model generates that token. The live runs confirm it for Qwen 3.8 and for Gemma 4's close; GLM, MiMo and MiniMax remain unconfirmed.
  • When a delimiter token's text never arrives before finish, the held text is emitted as text and reasoning is not closed. vLLM's flush_pending emits the deferred terminal instead. The difference only affects the final reasoning state, since no answer text follows.
    Written with AI assistance; reviewed before opening.

…oken

A model that writes <think> or </think> as ordinary text, for example
while explaining a chat template, was treated as emitting a real
reasoning delimiter: text moved between reasoning_content and content
and the tag disappeared (NVIDIA#48). Structured-output grammars ended the
reasoning region on the same text.

Derive, once per server, the contract's reasoning open and close that
the frontend tokenizer encodes as exactly one round-tripping token.
TaggedOutput and ThinkReasoning accept such a delimiter only at the
byte offset where its generated token's text starts, holding possible
delimiter text until it arrives and releasing it as text at finish.
The MiniMax leading-close check follows the same rule. The grammar
references the same delimiters by token ID, so reasoning ends only at
the token while its bytes remain ordinary answer text. DSML and XTML
keep recognizing their own markup. Multi-token delimiters, such as
Gemma 4's <|channel>thought open, stay text-matched. No package or
metadata change.

Validation: on base 82cb15a, Helios lib 1220 passed, 21 ignored, 5
failed; main fails the same five on that host (four engram::reader
staging-memory tests, one package-cache test that cannot fail a write
as root). All 12 integration targets pass; test_requests builds with
server,test-tools. Ignored checkpoint tests pass: delimiter resolution
for eight recipe tokenizers, glm_frontend, deepseek_frontend and the
Qwen 3.8 grammar envelope. Live Flare on two RTX PRO 6000 GPUs, base
3267c05 against this change: Qwen 3.8 non-streaming, streaming and
JSON-schema requests and Gemma 4 E4B requests all split reasoning at
the generated delimiter token on this change. One streamed Qwen
tool-call response generated a <think> token inside reasoning, which
is still removed as a delimiter.

Not run: live runs on other recipes, Gemma streaming, /v1/messages,
image_admission and instruction_boundary (fixture hash mismatch),
Heliopause, Python and CUDA tests. A model that writes its real close
in ordinary pieces stays in reasoning, and under a structured-output
grammar runs to its token limit; Qwen 3.8 never did so on this change.

Design-Signed-Off-By: Mark Yong via Cursor Agent (Auto); design-review=babb
@mark-yong

Copy link
Copy Markdown
Contributor Author

Detail behind the validation in the PR description: what the new tests cover, the checkpoint paths, further live observations, and the live harness for the tables. Each harness script sends requests to a running Flare server with return_token_ids, splits the returned token IDs at the generated closing delimiter token, and compares the two halves with reasoning_content and content. Expectations come from the token IDs, not from Flare's parsing.

What the new tests cover

Unit tests:

  • output::pipeline::tests::single_token_reasoning_delimiters_leave_spelled_tags_as_text: spelled tags stay text and token delimiters switch modes, with the Qwen 3.8, GLM 5.3 and MiMo contracts, plus MiniMax's leading close and quoted </mm:think>.
  • output::pipeline::tests::finish_releases_text_held_for_a_delimiter_token: held text is released for every termination cause, from both reasoning and content.
  • output::pipeline::tests::channel_thought_gates_only_its_single_token_close: Gemma's mix, where only the close is a token. A spelled close stays text in reasoning and content, the text-matched open still opens reasoning, a partial open that never completes is released, a stray close token at the start of content is removed while a spelled one stays, and a tool call still ends reasoning.
  • model::prompt::tests::single_token_reasoning_delimiters_gate_generated_output: the derivation, including a tokenizer whose single-token encoding fails the round trip and a DSML contract that yields no delimiters.
  • grammar::tests::single_token_reasoning_delimiters_end_reasoning_only_as_tokens: grammar reasoning ends only at the token, spelled tags stay reasoning, and a structured answer accepts both the tokens and spelled tags.
  • server::logprobs::tests::spelled_reasoning_delimiter_in_the_answer_receives_output_text_logprobs: a spelled delimiter's tokens keep their logprobs as answer text.

Checkpoint tests (--ignored):

  • tests/reasoning_delimiters.rs: the delimiter list resolved from each tokenizer in the description's recipe table.
  • tests/glm_frontend.rs: runs against a reference generated by tests/frontend/generate_glm_fixtures.py from the tokenizer, template and processor files of nvidia/GLM-5.3-Flash-FP4-Local@76348ddb, the recipe's pinned revision, with the repo's pinned Transformers and tokenizers. It checks that the delimiters resolve to 154841/154842 and runs requests through prepare_request.
  • tests/deepseek_frontend.rs: now asserts the DeepSeek list is empty.
  • grammar::tests::qwen_output_envelope_uses_checkpoint_token_bytes, on the Qwen 3.8 tokenizer: with and without reasoning that spells </think> before the real </think> token, a JSON answer accepts the tags as delimiter tokens, in ordinary pieces and per character.

Checkpoint paths used:

FLARE_TEST_SNAPSHOT=/models/hf-cache/flare/models--nvidia--Qwen3.8-Flash-Next-FP4-Local/snapshots/b9ce18a6887e18c089fdea8509a33e63429a3300
FLARE_GLM_TEST_SNAPSHOT=nvidia/GLM-5.3-Flash-FP4-Local@76348ddb2be2e35e39ecfc00e75dc07703f4f3d5 (non-weight files)
FLARE_STEP5_TEST_SNAPSHOT=stepfun-ai/Step-5-Preview-NVFP4@bc246cc2848e (tokenizer files)
FLARE_MINIMAX_TEST_SNAPSHOT=MiniMaxAI/MiniMax-M3.1-Flash@262ae6830934 (tokenizer files)
FLARE_MIMO_TEST_SNAPSHOT=/models/XiaomiMiMo/MiMo-V2.6-Flash-MOPD
FLARE_NEMOTRON_TEST_SNAPSHOT=/models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-NVFP4
FLARE_DEEPSEEK_TEST_SNAPSHOT=/models/safetensors/DeepSeek-V4.1-Flash-1017-1bae0e10
FLARE_GEMMA4_TEST_SNAPSHOT=unsloth/gemma-4-E4B-it-NVFP4@1c363766fcfe575ac17a6d544963201a0d9b24c3

rustfmt --check does report differences in the changed files, but it reports the same ones on main.

Further live observations

Both images were built from the repo Dockerfile with the #76 workaround (copy ignition/recipes/ into the build stage).

Qwen 3.8, non-streaming:

  • 14 of the branch's 63 completed JSON answers contained a delimiter token inside a string, returned as <think>/</think> text.
  • Where reasoning ended at the real token on both builds, the JSON literal field contained both tags in 11 of 21 answers on 3267c054 and 24 of 63 on the branch. On both builds the field is otherwise often a substitute: it contains <tool_call> or is blank in 7 of 21 and 22 of 63; this is the model's choice under an identical answer grammar.
  • Output-limit stops in the JSON requests: 10 of 68 on 3267c054, 5 of 68 on the branch, all whitespace runs.

Qwen 3.8, streaming:

  • 5 of 8 branch tool-call responses made the call, and all 5 parsed to get_weather with {"city": "Paris"}. One of them also wrote <tool_call> as text in its answer, which stayed text.
  • In five branch responses (three of them tool-call requests) and in several base responses, reasoning ended mid-sentence where the model was about to write a chat-template token. In a probed stream, the model had generated <|im_start|> and then <|im_end|>, which ends generation. This is the model's output and occurs on both builds.

Gemma 4 E4B:

  • The thinking-off miss on 3267c054 removed a spelled <channel|> from the answer.
  • All 16 JSON answers parsed, and 5 per build spelled the marker inside a string.
think48_live.py: Qwen 3.8, non-streaming
"""Live check for #48: compare Flare's reasoning/content split with the generated token IDs.

Usage: python think48_live.py BASE_URL LABEL OUT.jsonl TOKENIZER_JSON
"""

import json
import sys
import urllib.request

from tokenizers import Tokenizer

THINK, END = 248068, 248069
IM_END = None

REPRO = (
    "Qwen chat templates wrap the model's reasoning in a pair of XML-style tags named think: "
    "an opening tag and a closing tag with a slash. In two or three sentences, explain what they do, "
    "and write both tags out literally with their angle brackets in your answer."
)
ANSWER = (
    "Some chat templates mark hidden reasoning with an XML-style tag named think. Write one short "
    "paragraph that shows the opening tag and the closing tag (with the slash) literally, with their "
    "angle brackets, and then a sentence after them."
)
PLAIN = "A train leaves at 09:40 and arrives at 13:05. How long is the trip? Answer in one sentence."
JSON_PROMPT = (
    "Return JSON. In `literal`, write the opening and closing reasoning tags used by Qwen chat templates "
    "(an XML-style tag named think, and the same tag with a slash), literally with angle brackets. "
    "In `explanation`, explain them in one sentence."
)
SCHEMA = {
    "type": "json_schema",
    "json_schema": {
        "name": "tags",
        "strict": True,
        "schema": {
            "type": "object",
            "properties": {"literal": {"type": "string"}, "explanation": {"type": "string"}},
            "required": ["literal", "explanation"],
            "additionalProperties": False,
        },
    },
}


def cases():
    for seed in range(1, 9):
        yield "repro", seed, {"messages": [{"role": "user", "content": REPRO}]}
    for seed in range(1, 5):
        yield "answer_nothink", seed, {
            "messages": [{"role": "user", "content": ANSWER}],
            "chat_template_kwargs": {"enable_thinking": False},
        }
    for seed in range(1, 5):
        yield "answer_think", seed, {"messages": [{"role": "user", "content": ANSWER}]}
    for seed in range(1, 5):
        yield "plain_logprobs", seed, {
            "messages": [{"role": "user", "content": PLAIN}],
            "logprobs": True,
        }
    for seed in range(1, 5):
        yield "json_schema", seed, {
            "messages": [{"role": "user", "content": JSON_PROMPT}],
            "response_format": SCHEMA,
        }


def post(url, body):
    request = urllib.request.Request(
        url + "/v1/chat/completions",
        json.dumps(body).encode(),
        {"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(request, timeout=900) as response:
            return response.status, json.loads(response.read())
    except urllib.error.HTTPError as error:
        return error.code, {"error": error.read().decode(errors="replace")}


def expected(tokenizer, ids, thinking):
    """Split the generated tokens at the generated `</think>` token, as the model intended."""
    ids = list(ids)
    if ids and ids[-1] == IM_END:
        ids = ids[:-1]
    if not thinking:
        return "", tokenizer.decode(ids, skip_special_tokens=False)
    if END not in ids:
        return tokenizer.decode(ids, skip_special_tokens=False), ""
    at = ids.index(END)
    return (
        tokenizer.decode(ids[:at], skip_special_tokens=False),
        tokenizer.decode(ids[at + 1 :], skip_special_tokens=False),
    )


def main():
    url, label, out, tokenizer_path = sys.argv[1:5]
    tokenizer = Tokenizer.from_file(tokenizer_path)
    global IM_END
    IM_END = tokenizer.token_to_id("<|im_end|>")
    assert tokenizer.token_to_id("<think>") == THINK and tokenizer.token_to_id("</think>") == END
    with open(out, "a") as sink:
        for case, seed, body in cases():
            body = {**body, "seed": seed, "max_tokens": 4096, "return_token_ids": True}
            status, response = post(url, body)
            record = {"label": label, "case": case, "seed": seed, "status": status}
            if status != 200:
                record["error"] = response.get("error")
            else:
                choice = response["choices"][0]
                message = choice["message"]
                ids = choice.get("token_ids") or []
                thinking = case != "answer_nothink"
                want_reasoning, want_content = expected(tokenizer, ids, thinking)
                got_reasoning = message.get("reasoning_content") or ""
                got_content = message.get("content") or ""
                record.update(
                    finish=choice.get("finish_reason"),
                    tokens=len(ids),
                    close_tokens=ids.count(END),
                    open_tokens=ids.count(THINK),
                    spelled=sum(
                        text.count(tag)
                        for text in (want_reasoning, want_content)
                        for tag in ("<think>", "</think>")
                    ),
                    reasoning_ok=got_reasoning.strip() == want_reasoning.strip(),
                    content_ok=got_content.strip() == want_content.strip(),
                    reasoning=got_reasoning,
                    content=got_content,
                    want_reasoning=want_reasoning,
                    want_content=want_content,
                )
                if case == "plain_logprobs":
                    entries = (choice.get("logprobs") or {}).get("content") or []
                    record["logprob_entries"] = len(entries)
                if case == "json_schema":
                    try:
                        record["json"] = json.loads(got_content)
                    except ValueError:
                        record["json"] = None
            sink.write(json.dumps(record) + "\n")
            sink.flush()
            print(
                label, case, seed, status,
                record.get("finish"), "spelled", record.get("spelled"),
                "close", record.get("close_tokens"),
                "ok", record.get("reasoning_ok"), record.get("content_ok"),
                flush=True,
            )


if __name__ == "__main__":
    main()
think48_stream_live.py: Qwen 3.8, streaming (imports the prompts from `think48_live.py`)
"""Streaming live check for #48 on Qwen 3.8: rebuild the reasoning/content split from SSE deltas
and compare it with the streamed token IDs.

Usage: python think48_stream_live.py BASE_URL LABEL OUT.jsonl TOKENIZER_JSON
"""

import json
import sys
import urllib.request

from tokenizers import Tokenizer

from think48_live import ANSWER, END, JSON_PROMPT, PLAIN, REPRO, SCHEMA, THINK

TOOL_PROMPT = (
    "First, in your reasoning, write out literally the opening and closing think tags that Qwen chat "
    "templates use (an XML-style tag named think, and the same tag with a slash). Then call get_weather "
    "for Paris."
)
TOOLS = [{
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Current weather for a city.",
        "parameters": {
            "type": "object",
            "properties": {"city": {"type": "string"}},
            "required": ["city"],
        },
    },
}]


def cases():
    for seed in range(1, 9):
        yield "repro", seed, True, {"messages": [{"role": "user", "content": REPRO}]}
    for seed in range(1, 5):
        yield "answer_nothink", seed, False, {
            "messages": [{"role": "user", "content": ANSWER}],
            "chat_template_kwargs": {"enable_thinking": False},
        }
    for seed in range(1, 5):
        yield "answer_think", seed, True, {"messages": [{"role": "user", "content": ANSWER}]}
    for seed in range(1, 5):
        yield "plain_logprobs", seed, True, {
            "messages": [{"role": "user", "content": PLAIN}],
            "logprobs": True,
        }
    for seed in range(1, 9):
        yield "json_schema", seed, True, {
            "messages": [{"role": "user", "content": JSON_PROMPT}],
            "response_format": SCHEMA,
        }
    for seed in range(1, 9):
        yield "tool_call", seed, True, {
            "messages": [{"role": "user", "content": TOOL_PROMPT}],
            "tools": TOOLS,
        }


def stream(url, body):
    request = urllib.request.Request(
        url + "/v1/chat/completions",
        json.dumps(body).encode(),
        {"Content-Type": "application/json"},
    )
    chunks = []
    try:
        with urllib.request.urlopen(request, timeout=900) as response:
            for line in response:
                line = line.decode().strip()
                if not line.startswith("data:"):
                    continue
                data = line[5:].strip()
                if data == "[DONE]":
                    break
                chunks.append(json.loads(data))
        return 200, chunks
    except urllib.error.HTTPError as error:
        return error.code, error.read().decode(errors="replace")


def assemble(chunks):
    reasoning, content, ids, order, logprobs = [], [], [], [], 0
    calls, finish = {}, None
    for chunk in chunks:
        for choice in chunk.get("choices") or []:
            ids.extend(choice.get("token_ids") or [])
            delta = choice.get("delta") or {}
            if delta.get("reasoning_content"):
                reasoning.append(delta["reasoning_content"])
                order.append("r")
            if delta.get("content"):
                content.append(delta["content"])
                order.append("c")
            for call in delta.get("tool_calls") or []:
                entry = calls.setdefault(call.get("index", 0), {"name": "", "arguments": ""})
                function = call.get("function") or {}
                entry["name"] += function.get("name") or ""
                entry["arguments"] += function.get("arguments") or ""
                order.append("t")
            logprobs += len((choice.get("logprobs") or {}).get("content") or [])
            finish = choice.get("finish_reason") or finish
    switches = sum(1 for a, b in zip(order, order[1:]) if a != b)
    return {
        "reasoning": "".join(reasoning),
        "content": "".join(content),
        "ids": ids,
        "reasoning_deltas": len(reasoning),
        "content_deltas": len(content),
        "order": "".join(dict.fromkeys(order)),
        "switches": switches,
        "calls": list(calls.values()),
        "logprob_entries": logprobs,
        "finish": finish,
    }


def expected(tokenizer, ids, thinking, im_end):
    ids = list(ids)
    while ids and ids[-1] == im_end:
        ids.pop()
    decode = lambda part: tokenizer.decode(part, skip_special_tokens=False)
    if not thinking:
        return "", decode(ids)
    if END not in ids:
        return decode(ids), ""
    at = ids.index(END)
    return decode(ids[:at]), decode(ids[at + 1 :])


def main():
    url, label, out, tokenizer_path = sys.argv[1:5]
    tokenizer = Tokenizer.from_file(tokenizer_path)
    im_end = tokenizer.token_to_id("<|im_end|>")
    assert tokenizer.token_to_id("<think>") == THINK and tokenizer.token_to_id("</think>") == END
    with open(out, "a") as sink:
        for case, seed, thinking, body in cases():
            body = {**body, "seed": seed, "max_tokens": 4096, "stream": True, "return_token_ids": True}
            status, chunks = stream(url, body)
            record = {"label": label, "case": case, "seed": seed, "status": status}
            if status != 200:
                record["error"] = chunks
            else:
                got = assemble(chunks)
                want_reasoning, want_content = expected(tokenizer, got["ids"], thinking, im_end)
                record.update(
                    {key: value for key, value in got.items() if key != "ids"},
                    tokens=len(got["ids"]),
                    close_tokens=got["ids"].count(END),
                    reasoning_ok=got["reasoning"].strip() == want_reasoning.strip(),
                    want_reasoning=want_reasoning,
                    want_content=want_content,
                )
                if case == "tool_call":
                    # The tool call's markup is in the tokens, not in `content`; check the call itself.
                    parsed = []
                    for call in got["calls"]:
                        try:
                            parsed.append((call["name"], json.loads(call["arguments"] or "null")))
                        except ValueError:
                            parsed.append((call["name"], None))
                    record["parsed_calls"] = parsed
                    record["content_ok"] = (
                        "<tool_call>" not in got["content"]
                        and any(name == "get_weather" and isinstance(args, dict) for name, args in parsed)
                    )
                else:
                    record["content_ok"] = got["content"].strip() == want_content.strip()
                if case == "json_schema":
                    try:
                        record["json"] = json.loads(got["content"])
                    except ValueError:
                        record["json"] = None
            sink.write(json.dumps(record) + "\n")
            sink.flush()
            print(
                label, case, seed, status, record.get("finish"),
                "close", record.get("close_tokens"), "order", record.get("order"),
                "deltas", record.get("reasoning_deltas"), record.get("content_deltas"),
                "ok", record.get("reasoning_ok"), record.get("content_ok"),
                flush=True,
            )


if __name__ == "__main__":
    main()
think48_gemma_live.py: Gemma 4
"""Live check for #48 on Gemma 4: compare Flare's reasoning/content split with the generated token IDs.

Gemma opens reasoning with `<|channel>` (100) followed by `thought`, which Helios still
matches as text, and closes it with the single token `<channel|>` (101).

Usage: python think48_gemma_live.py BASE_URL LABEL OUT.jsonl TOKENIZER_JSON
"""

import json
import sys
import urllib.request

from tokenizers import Tokenizer

CHANNEL, CLOSE = 100, 101
ENDS = {1, 106, 50}

REPRO = (
    "Gemma chat templates end the model's hidden thinking with a closing channel marker: a less-than sign, "
    "the word channel, a vertical bar, and a greater-than sign. In two or three sentences, explain what it "
    "does, and write the marker out literally in your answer."
)
PLAIN = "A train leaves at 09:40 and arrives at 13:05. How long is the trip? Answer in one sentence."
JSON_PROMPT = (
    "Return JSON. In `literal`, write Gemma's closing channel marker literally: a less-than sign, the word "
    "channel, a vertical bar, and a greater-than sign. In `explanation`, explain it in one sentence."
)
SCHEMA = {
    "type": "json_schema",
    "json_schema": {
        "name": "marker",
        "strict": True,
        "schema": {
            "type": "object",
            "properties": {"literal": {"type": "string"}, "explanation": {"type": "string"}},
            "required": ["literal", "explanation"],
            "additionalProperties": False,
        },
    },
}
THINK = {"enable_thinking": True}


def cases():
    for seed in range(1, 9):
        yield "repro", seed, True, {"messages": [{"role": "user", "content": REPRO}], "chat_template_kwargs": THINK}
    for seed in range(1, 5):
        yield "repro_nothink", seed, False, {"messages": [{"role": "user", "content": REPRO}]}
    for seed in range(1, 9):
        yield "plain_logprobs", seed, True, {
            "messages": [{"role": "user", "content": PLAIN}],
            "chat_template_kwargs": THINK,
            "logprobs": True,
        }
    for seed in range(1, 9):
        yield "json_schema", seed, True, {
            "messages": [{"role": "user", "content": JSON_PROMPT}],
            "chat_template_kwargs": THINK,
            "response_format": SCHEMA,
        }


def post(url, body):
    request = urllib.request.Request(
        url + "/v1/chat/completions",
        json.dumps(body).encode(),
        {"Content-Type": "application/json"},
    )
    try:
        with urllib.request.urlopen(request, timeout=900) as response:
            return response.status, json.loads(response.read())
    except urllib.error.HTTPError as error:
        return error.code, {"error": error.read().decode(errors="replace")}


def expected(tokenizer, ids):
    """Split the generated tokens at the generated `<channel|>` token, as the model intended."""
    ids = list(ids)
    while ids and ids[-1] in ENDS:
        ids.pop()
    decode = lambda part: tokenizer.decode(part, skip_special_tokens=False)
    if CLOSE not in ids:
        return "", decode(ids)
    at = ids.index(CLOSE)
    reasoning = decode(ids[:at])
    if reasoning.startswith("<|channel>thought"):
        reasoning = reasoning[len("<|channel>thought") :]
    return reasoning, decode(ids[at + 1 :])


def main():
    url, label, out, tokenizer_path = sys.argv[1:5]
    tokenizer = Tokenizer.from_file(tokenizer_path)
    assert tokenizer.token_to_id("<|channel>") == CHANNEL and tokenizer.token_to_id("<channel|>") == CLOSE
    with open(out, "a") as sink:
        for case, seed, thinking, body in cases():
            body = {**body, "seed": seed, "max_tokens": 4096, "return_token_ids": True}
            status, response = post(url, body)
            record = {"label": label, "case": case, "seed": seed, "thinking": thinking, "status": status}
            if status != 200:
                record["error"] = response.get("error")
            else:
                choice = response["choices"][0]
                message = choice["message"]
                ids = choice.get("token_ids") or []
                want_reasoning, want_content = expected(tokenizer, ids)
                got_reasoning = message.get("reasoning_content") or ""
                got_content = message.get("content") or ""
                text = tokenizer.decode(ids, skip_special_tokens=False)
                record.update(
                    finish=choice.get("finish_reason"),
                    tokens=len(ids),
                    close_tokens=ids.count(CLOSE),
                    channel_tokens=ids.count(CHANNEL),
                    first_ids=ids[:4],
                    spelled_close=text.count("<channel|>") - ids.count(CLOSE),
                    reasoning_ok=got_reasoning.strip() == want_reasoning.strip(),
                    content_ok=got_content.strip() == want_content.strip(),
                    reasoning=got_reasoning,
                    content=got_content,
                    want_reasoning=want_reasoning,
                    want_content=want_content,
                )
                if case == "plain_logprobs":
                    entries = (choice.get("logprobs") or {}).get("content") or []
                    record["logprob_entries"] = len(entries)
                if case == "json_schema":
                    try:
                        record["json"] = json.loads(got_content)
                    except ValueError:
                        record["json"] = None
            sink.write(json.dumps(record) + "\n")
            sink.flush()
            print(
                label, case, seed, status, record.get("finish"),
                "close", record.get("close_tokens"), "spelled", record.get("spelled_close"),
                "ok", record.get("reasoning_ok"), record.get("content_ok"),
                flush=True,
            )


if __name__ == "__main__":
    main()

The non-streaming Qwen table's 68 JSON-schema requests are seeds 1 to 68. think48_live.py sends seeds 1 to 4; two wrappers, run from the directory holding think48_live.py, add the rest:

Wrappers for JSON seeds 5 to 68
# all.py: the harness cases plus JSON seeds 5 to 20
import think48_live as t
base_cases = t.cases
def cases():
    yield from base_cases()
    for seed in range(5, 21):
        yield "json_schema", seed, {"messages": [{"role": "user", "content": t.JSON_PROMPT}], "response_format": t.SCHEMA}
t.cases = cases
t.main()
# json48.py: JSON seeds 21 to 68
import think48_live as t
t.cases = lambda: (("json_more", seed, {"messages": [{"role": "user", "content": t.JSON_PROMPT}], "response_format": t.SCHEMA}) for seed in range(21, 69))
t.main()

@lukealonso
lukealonso merged commit de6f53e into NVIDIA:main Oct 26, 2026
tunamitom added a commit to tunamitom/flare-local-inference that referenced this pull request Nov 11, 2026
Flag plumbing for the transplanted upstream NVIDIA#94/NVIDIA#95 direct-table geometry:
--trellis-prefill-table 0|1|2 (default off) stamps B12X_W4A16_PREFILL_TABLE
before any kernel compiles and records the selection in the plan settings,
the same contract as --trellis-prefill-shapes (W3 R3). 1 = 48-row blocks,
2 stages, tile 128,128,64,256 for 1536..2560-row launches; 2 = additionally
64-row blocks, 3 stages, tile 64,128,64,128 for >=2560-row launches.
Decode / verify / draft launches never see the switch (kernel-side gate).

Design-Signed-Off-By: Claude (hermes-agent); design-review=a38f
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.

Literal <think>/</think> text in model output is parsed as reasoning delimiters

2 participants