[Lewis300]

Type of Change

• Bug fix
• Small feature (agreed in an issue)
• New model (export recipe only)
• New model (export recipe and new runtime support)
• Documentation update

Motivation and Context

Closes #373.

llm-runner --raw-tokens CTX --save-logits out.json (or --print-logits) does not generate. It splits the ids into a one-token context (tokens[0]) and a forced continuation (tokens[1...]) and scores that continuation instead. So --max-tokens has no effect, and the tokens document (token_id, incremental_text, top_logits) is never written. There was no way to generate from pre-tokenized input and save the logits of what was generated.

This PR makes the raw-token mode mirror the text mode, where --prompt is the context and --continuation is what gets scored:

Mode	Flags	Context	Writes
Generation with logits (default)	--raw-tokens CTX + --save-logits or --print-logits	all ids in CTX	the tokens document, one entry per generated position, up to --max-tokens
Forced-continuation scoring	add --raw-continuation CONT	all ids in CTX	the scoring document: context_tokens, continuation_tokens, positions, total_log_probability, perplexity, … with one position per id in CONT

CONT uses the same {"tokens": [...]} format as --raw-tokens, and the context can be any length.

The run refuses, before it loads the model, when --raw-continuation is given:

• without --raw-tokens;
• without --print-logits or --save-logits;
• together with the text --continuation, --batch-prototype, --image or --video (each of those takes its own path and would skip the scoring);
• with an empty context or continuation, or with either file unreadable or not decodable;
• with an id outside the model's vocabulary, or with more ids in total than the model's context length.

With --json-schema, scoring uses the grammar constrained generation uses, with the same stop tokens and whitespace limit: each continuation position's logits are masked by the schema's grammar before they are scored. Once the JSON is complete, only a stop token is allowed: the tokenizer's EOS, the additional ones from tokenizer_config.json, or a single-token --stop-tokens. A continuation id the schema does not allow, or one after the stop token, is refused with its position.

Behaviour change. If you relied on --raw-tokens + --save-logits to get the scoring document, split the input: [t0, t1, …, tn] becomes CTX = [t0] and CONT = [t1, …, tn], passed as --raw-tokens CTX --raw-continuation CONT. You can now also give the scored continuation a context longer than one token.

Changes

• llm-runner (RawTokenScoring.swift, LLMRunnerMain.swift). New --raw-continuation option. RawTokenScoring decides between generating and scoring, makes every refusal above before the model loads, and masks the scores under --json-schema. Every error that names a command-line flag lives beside it as ContinuationOptionError, with the existing messages unchanged.
• TextGenerator.swift. New evaluateRawTokens(context:continuation:), which takes a context of any length. It and the text evaluateContinuation share one private feed (scoreForcedContinuation), so both scoring modes feed the model the same way. The old evaluateRawTokens(_:) stays, deprecated, with its old split ([tokens[0]] as the context, tokens[1...] scored).
• ContinuationEvaluation.swift. The ContinuationEvaluationError cases that name llm-runner flags stay, deprecated; llm-runner no longer throws them. ContinuationEvaluationResult gains a public memberwise init. No public API is removed.
• Stop tokens in schema-constrained grammars (xgrammar_c_bridge.h/.cpp, TokenizerInfo.swift, ConstrainedGenerationSession.swift). The C bridge passed no stop token ids to xgrammar, so ConstrainedGenerationSession dropped the stopTokenIds it was given. Every schema-constrained session, generation included, used the stop tokens xgrammar detects from the vocabulary. With Qwen3 that is <|endoftext|> alone, so after the JSON the model could never choose <|im_end|>, its EOS. The bridge and TokenizerInfo now take the ids, and both decoding strategies and scoring share StopSequences.singleTokenStopIds.
• Package.swift. Adds a LLMRunnerTests test target. It adds no product.

No inference engine changes. Constrained generation's visible output is unchanged where xgrammar's detected stop token was also in the stop list (as with Qwen3); what changes is which stop token the model may choose.

Testing

• Tested locally

• Added/updated tests

• Added/updated docs

• swift build --build-tests and the full swift test suite pass on macOS.

• New tests:

  • LLMRunnerTests: argument parsing; each refusal above, including --batch-prototype, --image, --video, ids outside the vocabulary (negative and too large, in either file) and more ids than the context length; scoring under a JSON schema (an allowed continuation is scored over the allowed ids; once the JSON is complete only a stop token is allowed and is scored; a disallowed id and an id after the stop token are refused).
  • StopTokenTests: the stop tokens a grammar is given are the ones allowed once the JSON is complete, are refused before it, and end the grammar when accepted, through both ConstrainedGenerationSession and TokenizerInfo. Two of them fail without the bridge fix.
  • ContinuationEvaluationTests: generation from raw tokens writes the tokens document; scoring feeds the whole multi-token context and scores each continuation id; empty context or continuation is refused; the deprecated evaluateRawTokens(_:) keeps its split; text --continuation is unchanged.

• End to end on Qwen3-0.6B with --inference-engine-variant coreai-sequential:

  • --raw-tokens ctx.json --save-logits gen.json --max-tokens 4 generates and writes the tokens document.
  • --raw-tokens ctx.json --raw-continuation cont.json --save-logits out.json scores every continuation id against the whole context.
  • Adding --json-schema with {"a": integer} and a continuation of {"a": 1}: every position's probability is at least its unmasked value (total log-probability −23.3 unmasked, −6.1 masked); {"b": 1} is refused at position 1.
  • {"a": 1} followed by an end token: <|im_end|> is scored at −7.454 in total, the same as without it, and <|endoftext|> at −22.235. Before the stop-token fix, <|im_end|> was refused and <|endoftext|> scored as certain (−7.454).
  • --image, --video, --batch-prototype, an id of 10,000,000 and a 50,000-id context are each refused before the model loads.

[stikves]

This seems to add a new runtime path, this is not entirely tested (for example zero input, or context length mismatches)

Are we treating the symptom, instead of the root cause here?

[carinapeng]

This change makes sense and thanks for fixing for the eval correctness

Pls address the opens before merging, this PR looks fine in making raw tokens the context. +1 on the question of whether we are treating the root cause of the issue

[Lewis300]

This seems to add a new runtime path, this is not entirely tested (for example zero input, or context length mismatches)

Are we treating the symptom, instead of the root cause here?

Thanks. On testing: the new path is --raw-continuation, and every refusal it makes happens before the model loads, each with its own test in RawTokenScoringTests. That covers an empty context or continuation, an unreadable or malformed file, a negative or out-of-vocabulary id in either file, and context + continuation longer than the model's context length. I've also added the zero-input case for generation: an empty --raw-tokens used to reach the engine and is now refused up front.

On root cause: the bug wasn't the split itself. llm-runner guessed "score" from --raw-tokens plus a logits flag, then hard-coded tokens[0] as the context. That made generating from raw tokens with logits impossible, and it gave scoring only a one-token context. Fixing just the split (say, a --context-length N) would have kept the guess, so --raw-tokens --save-logits still couldn't generate. This PR removes the guess. --raw-tokens is always the context, and scoring happens only when you name the ids with --raw-continuation, the same way text mode has --prompt and --continuation. The library side follows suit: evaluateRawTokens(context:continuation:) takes an explicit context, the old overload is deprecated with its behaviour pinned by a test, and both scoring modes share one feed.

[stikves]

is constrained the same way generation is.

Is a strong statement. Do we check the paths for EOS handling (getting the additional EOS tokens from tokenizer)?

[stikves]

Mostly good, please check the claim for EOS handling, and update the comments.

[Lewis300]

is constrained the same way generation is.

Is a strong statement. Do we check the paths for EOS handling (getting the additional EOS tokens from tokenizer)?

Good catch. The xgrammar C bridge passed no stop token ids to xgrammar's TokenizerInfo (std::nullopt), so ConstrainedGenerationSession dropped the stopTokenIds it was given. Every schema-constrained session, generation included, used the stop tokens xgrammar detects from the vocabulary instead of ours. The "Constrained session created (… stopTokenIds=[…])" log line showed ids that never reached the grammar.

With Qwen3-0.6B, scoring {"a": 1} under {"a": integer}:

Continuation ends with	Before	After
<|im_end|> (its EOS)	refused: not allowed by --json-schema	scored, −7.454 in total (the model puts all its probability there)
<|endoftext|>	scored as certain, −7.454	scored by the model's own logits, −22.235

Fixed in 19b92fa:

• The bridge and TokenizerInfo take the stop token ids, and ConstrainedGenerationSession passes them on.
• Scoring builds its grammar with the stop tokens generation uses (the tokenizer's EOS, the additional ones from tokenizer_config.json, and single-token --stop-tokens) and the same --json-max-whitespace limit. Both decoding strategies and scoring now share StopSequences.singleTokenStopIds.
• New StopTokenTests check that the given stop tokens are the ones allowed once the JSON is complete, and only then; two of them fail without the fix. RawTokenScoringTests now cover a continuation that ends in a stop token, and one that runs past it.

This changes constrained generation too: once the JSON is complete, the model can now choose any of its stop tokens. Visible output is unchanged with Qwen3, because both of its end tokens were already in the stop list.