diff --git a/CHANGELOG.md b/CHANGELOG.md index 1465b79..bdb0cab 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- `CompletionOptions.thinkingBudget.pauseGrammar: false` keeps a grammar applied inside a reasoning block, so a grammar that takes the block and its end tag itself, such as one for a prompt that ends inside the block, sees the budget's forced ``. It needs the new `thinkingBudgetPauseGrammar` capability; omitted, the grammar pauses as before. - WebGPU text-to-speech applies the GGUF's suppress tokens in the talker sampler, as upstream does. Qwen3-TTS no longer fails with `code out of range` when a codec control token is sampled. - `createCompletion` can explicitly count prompt tokens in penalties with a grammar using `penaltyPromptTokensWithGrammar`; a matching loaded-core capability is required, and existing grammar defaults stay unchanged. diff --git a/CMakeLists.txt b/CMakeLists.txt index d040ffb..b7db2ec 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -235,7 +235,7 @@ set(LLAMADART_WEBGPU_LINK_OPTIONS "-sEXPORT_NAME=createLlamaWebGpuCoreModule" "-sENVIRONMENT=web,worker" "-sEXPORTED_RUNTIME_METHODS=['FS','ccall','UTF8ToString']" - "-sEXPORTED_FUNCTIONS=['_main','_llamadart_webgpu_probe','_llamadart_webgpu_supports_pthreads','_llamadart_webgpu_pthread_pool_size','_llamadart_webgpu_backends_json','_llamadart_webgpu_last_error','_llamadart_webgpu_set_log_level','_llamadart_webgpu_load_model','_llamadart_webgpu_load_model_from_url','_llamadart_webgpu_free_model','_llamadart_webgpu_mmproj_load','_llamadart_webgpu_mmproj_free','_llamadart_webgpu_mmproj_supports_vision','_llamadart_webgpu_mmproj_supports_audio','_llamadart_webgpu_media_clear_pending','_llamadart_webgpu_media_add_file','_llamadart_webgpu_media_add_encoded','_llamadart_webgpu_media_add_rgb','_llamadart_webgpu_media_add_audio_f32','_llamadart_webgpu_tokenize_to_json','_llamadart_webgpu_tokenize_bytes_to_json','_llamadart_webgpu_last_tokens_json','_llamadart_webgpu_state_save_file','_llamadart_webgpu_state_load_file','_llamadart_webgpu_test_state_load_failures','_llamadart_webgpu_detokenize_from_json','_llamadart_webgpu_last_detokenized','_llamadart_webgpu_last_detokenized_bytes_json','_llamadart_webgpu_embed_to_json','_llamadart_webgpu_embed_bytes_to_json','_llamadart_webgpu_embed_parts_to_json','_llamadart_webgpu_embedding_capabilities_json','_llamadart_webgpu_last_embedding_json','_llamadart_webgpu_score_next_token_to_json','_llamadart_webgpu_score_next_token_bytes_to_json','_llamadart_webgpu_last_next_token_scores_json','_llamadart_webgpu_generate','_llamadart_webgpu_begin_generation','_llamadart_webgpu_begin_generation_bytes','_llamadart_webgpu_next_token','_llamadart_webgpu_last_piece','_llamadart_webgpu_end_generation','_llamadart_webgpu_request_cancel','_llamadart_webgpu_last_output','_llamadart_webgpu_last_output_bytes_json','_llamadart_webgpu_last_generation_usage_json','_llamadart_webgpu_completion_capabilities_json','_llamadart_webgpu_set_next_speculative','_llamadart_webgpu_set_next_grammar','_llamadart_webgpu_add_next_grammar_trigger','_llamadart_webgpu_draft_model_load','_llamadart_webgpu_draft_model_free','_llamadart_webgpu_draft_model_info_json','_llamadart_webgpu_heap_headroom_bytes','_llamadart_webgpu_get_context_size','_llamadart_webgpu_model_meta_json','_llamadart_webgpu_tts_api_version','_llamadart_webgpu_tts_info_json','_llamadart_webgpu_tts_start','_llamadart_webgpu_tts_step','_llamadart_webgpu_tts_progress_json','_llamadart_webgpu_tts_write_pcm','_llamadart_webgpu_tts_reset','_llamadart_webgpu_decision_api_version','_llamadart_webgpu_decision_capabilities_json','_llamadart_webgpu_decision_load','_llamadart_webgpu_decision_head_info_json','_llamadart_webgpu_decision_run','_llamadart_webgpu_decision_free','_llamadart_webgpu_lora_api_version','_llamadart_webgpu_lora_load','_llamadart_webgpu_lora_set','_llamadart_webgpu_lora_remove','_llamadart_webgpu_lora_clear','_llamadart_webgpu_shutdown']" + "-sEXPORTED_FUNCTIONS=['_main','_llamadart_webgpu_probe','_llamadart_webgpu_supports_pthreads','_llamadart_webgpu_pthread_pool_size','_llamadart_webgpu_backends_json','_llamadart_webgpu_last_error','_llamadart_webgpu_set_log_level','_llamadart_webgpu_load_model','_llamadart_webgpu_load_model_from_url','_llamadart_webgpu_free_model','_llamadart_webgpu_mmproj_load','_llamadart_webgpu_mmproj_free','_llamadart_webgpu_mmproj_supports_vision','_llamadart_webgpu_mmproj_supports_audio','_llamadart_webgpu_media_clear_pending','_llamadart_webgpu_media_add_file','_llamadart_webgpu_media_add_encoded','_llamadart_webgpu_media_add_rgb','_llamadart_webgpu_media_add_audio_f32','_llamadart_webgpu_tokenize_to_json','_llamadart_webgpu_tokenize_bytes_to_json','_llamadart_webgpu_last_tokens_json','_llamadart_webgpu_state_save_file','_llamadart_webgpu_state_load_file','_llamadart_webgpu_test_state_load_failures','_llamadart_webgpu_detokenize_from_json','_llamadart_webgpu_last_detokenized','_llamadart_webgpu_last_detokenized_bytes_json','_llamadart_webgpu_embed_to_json','_llamadart_webgpu_embed_bytes_to_json','_llamadart_webgpu_embed_parts_to_json','_llamadart_webgpu_embedding_capabilities_json','_llamadart_webgpu_last_embedding_json','_llamadart_webgpu_score_next_token_to_json','_llamadart_webgpu_score_next_token_bytes_to_json','_llamadart_webgpu_last_next_token_scores_json','_llamadart_webgpu_generate','_llamadart_webgpu_begin_generation','_llamadart_webgpu_begin_generation_bytes','_llamadart_webgpu_next_token','_llamadart_webgpu_last_piece','_llamadart_webgpu_end_generation','_llamadart_webgpu_request_cancel','_llamadart_webgpu_last_output','_llamadart_webgpu_last_output_bytes_json','_llamadart_webgpu_last_generation_usage_json','_llamadart_webgpu_completion_capabilities_json','_llamadart_webgpu_set_next_speculative','_llamadart_webgpu_set_next_grammar','_llamadart_webgpu_add_next_grammar_trigger','_llamadart_webgpu_set_next_grammar_reasoning_pause','_llamadart_webgpu_draft_model_load','_llamadart_webgpu_draft_model_free','_llamadart_webgpu_draft_model_info_json','_llamadart_webgpu_heap_headroom_bytes','_llamadart_webgpu_get_context_size','_llamadart_webgpu_model_meta_json','_llamadart_webgpu_tts_api_version','_llamadart_webgpu_tts_info_json','_llamadart_webgpu_tts_start','_llamadart_webgpu_tts_step','_llamadart_webgpu_tts_progress_json','_llamadart_webgpu_tts_write_pcm','_llamadart_webgpu_tts_reset','_llamadart_webgpu_decision_api_version','_llamadart_webgpu_decision_capabilities_json','_llamadart_webgpu_decision_load','_llamadart_webgpu_decision_head_info_json','_llamadart_webgpu_decision_run','_llamadart_webgpu_decision_free','_llamadart_webgpu_lora_api_version','_llamadart_webgpu_lora_load','_llamadart_webgpu_lora_set','_llamadart_webgpu_lora_remove','_llamadart_webgpu_lora_clear','_llamadart_webgpu_shutdown']" "-lwasmfs_fetch.js" "${LLAMADART_WEBGPU_EXCEPTION_CATCHING_ALLOWED}" "--post-js=${LLAMADART_WEBGPU_CORE_POST_JS}" diff --git a/docs/api.md b/docs/api.md index 6c9b6e2..7182c51 100644 --- a/docs/api.md +++ b/docs/api.md @@ -326,7 +326,7 @@ Common `options` keys: | `grammarLazy` | `true` applies `grammar` only after a `grammarTriggers` entry fires, as llama.cpp's lazy grammar (`llama-server`'s `grammar_lazy`) does; see [Lazy grammars](#lazy-grammars). Needs a non-empty `grammar` and at least one trigger. Omitted or `null` means `false`. | | `grammarTriggers` | Array of `{ type: 'word' \| 'pattern' \| 'pattern-full', value }` or `{ type: 'token', token }` that fire a lazy grammar. Always validated; applied only when `grammarLazy` is `true`. | | `seed` | Integer seed; random when omitted. | -| `thinkingBudget` | `{ maxTokens, startTag, endTag, forcedMessage? }`: llama.cpp's reasoning budget. After `maxTokens` tokens inside a block opened by `startTag`, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` (default empty) and `endTag`. A block the prompt leaves open counts from the start of generation, so `maxTokens: 0` forces the end at once. Each new `startTag` gets a fresh budget. A grammar constrains `startTag` like any other text, then pauses until `endTag` completes. `maxTokens` is an integer from `0` to `2147483647`; the tags must be non-empty. Each tag is tokenized on its own, as native llamadart does: with a vocabulary that adds a leading space to text, such as Llama 2's SentencePiece, a tag in the prompt matches only after a space, and the forced end tag starts with one. Text-only prompts: with `parts` it rejects. | +| `thinkingBudget` | `{ maxTokens, startTag, endTag, forcedMessage?, pauseGrammar? }`: llama.cpp's reasoning budget. After `maxTokens` tokens inside a block opened by `startTag`, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` (default empty) and `endTag`. A block the prompt leaves open counts from the start of generation, so `maxTokens: 0` forces the end at once. Each new `startTag` gets a fresh budget. A grammar constrains `startTag` like any other text, then pauses until `endTag` completes, unless `pauseGrammar` is `false`: that keeps the grammar applied inside the block, for one that takes the block and `endTag` itself (such as one written for a prompt that ends inside the block; paused, it resumes at its start and takes the reply as reasoning) and so accepts `forcedMessage` and `endTag` wherever the budget can run out. Omitted, `null` or `true` pauses it; `false` needs the `thinkingBudgetPauseGrammar` capability. `maxTokens` is an integer from `0` to `2147483647`; the tags must be non-empty. Each tag is tokenized on its own, as native llamadart does: with a vocabulary that adds a leading space to text, such as Llama 2's SentencePiece, a tag in the prompt matches only after a space, and the forced end tag starts with one. Text-only prompts: with `parts` it rejects. | | `onUsage(usage)` | Called at most once, with the `CompletionUsage` of the last attempt whose generation returned, before the promise resolves or rejects with an `AbortError`. Not called when no generation returned: other rejections, a skipped multimodal warmup, or a worker failure after a cancel. | | `onToken(piece, currentText)` | Token callback. By default `piece` is a `Uint8Array` containing complete UTF-8 characters; see [Streamed text](#streamed-text). `currentText` is the text streamed so far. Direct runtime mode provides it by default; worker mode provides `''` unless `emitCurrentTextOnToken: true` is set. | | `signal` | `AbortSignal`; aborting cancels this operation and rejects with `AbortError`. | @@ -365,7 +365,7 @@ An invalid `minP`, a `presencePenalty` that is not finite as a 32-bit float, an invalid `thinkingBudget`, or a malformed `grammarLazy`, `grammarTriggers` or `grammarRoot` rejects with a `RangeError` or `TypeError` before generation starts, and `grammarLazy: true` or another `grammarRoot` without a non-empty -`grammar`, or a lazy grammar without triggers, with an `Error`. A nonzero `minP` or `presencePenalty`, any `thinkingBudget`, and +`grammar`, or a lazy grammar without triggers, with an `Error`. A nonzero `minP` or `presencePenalty`, any `thinkingBudget`, `thinkingBudget.pauseGrammar: false`, and `grammarLazy: true` or a `grammarRoot` other than `root` reject when the loaded core does not apply that option. @@ -400,7 +400,7 @@ completion rejects with `The lazy grammar rejected the text that triggered it`. The core stays usable after either, and worker mode rethrows both without falling back to the main thread; like an invalid grammar, each rejection leaks a few KiB at most. Triggers match only generated text, not the -prompt, and a `thinkingBudget` block pauses a lazy grammar, so text inside the +prompt, and a `thinkingBudget` block pauses a lazy grammar (unless `pauseGrammar` is `false`), so text inside the block fires no trigger. The bridge streams every special token's text, so llama.cpp's `preserved_tokens` has no equivalent here. @@ -430,7 +430,7 @@ cancelled generation once it returns. getCompletionCapabilities(): Promise ``` -Returns `{ presencePenalty, minP, thinkingBudget, penaltyPromptTokens, penaltyPromptTokensWithGrammar, lazyGrammar, speculativeDecoding }`: +Returns `{ presencePenalty, minP, thinkingBudget, thinkingBudgetPauseGrammar, penaltyPromptTokens, penaltyPromptTokensWithGrammar, lazyGrammar, speculativeDecoding }`: whether the loaded core applies each `createCompletion` option. `lazyGrammar` covers `grammarLazy`, `grammarTriggers` and `grammarRoot`. `penaltyPromptTokensWithGrammar` separately permits the explicit grammar opt-in; diff --git a/js/llama_webgpu_bridge.d.ts b/js/llama_webgpu_bridge.d.ts index 7133e51..7b269ae 100644 --- a/js/llama_webgpu_bridge.d.ts +++ b/js/llama_webgpu_bridge.d.ts @@ -241,7 +241,7 @@ export type GrammarTrigger = | { type: 'word' | 'pattern' | 'pattern-full'; value: string } | { type: 'token'; token: number }; -/** llama.cpp's reasoning budget: after `maxTokens` tokens inside a reasoning block, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` and `endTag`. A grammar pauses after `startTag` until `endTag` completes. */ +/** llama.cpp's reasoning budget: after `maxTokens` tokens inside a reasoning block, plus any that finish a UTF-8 character, the sampler forces `forcedMessage` and `endTag`. A grammar pauses after `startTag` until `endTag` completes, unless `pauseGrammar` is false. */ export interface ThinkingBudgetOptions { /** Integer from 0 to 2147483647. 0 forces the end as soon as a block opens, or at once when the prompt ends inside one. */ maxTokens: number; @@ -251,6 +251,8 @@ export interface ThinkingBudgetOptions { endTag: string; /** Text forced before `endTag` when the budget runs out; empty by default. */ forcedMessage?: string; + /** `false` keeps `grammar` applied inside a reasoning block, for a grammar that takes the block and `endTag` itself, such as one written for a prompt that ends inside the block; paused, it would resume at its start and take the reply as reasoning. It must accept `forcedMessage` and `endTag` wherever the budget can run out. Omitted, null or true pauses the grammar. Needs the `thinkingBudgetPauseGrammar` capability. */ + pauseGrammar?: boolean | null; } /** Completion options the loaded core applies. Every flag is false until a model load initializes the core. */ @@ -261,6 +263,8 @@ export interface CompletionCapabilities { minP: boolean; /** `CompletionOptions.thinkingBudget` is applied. */ thinkingBudget: boolean; + /** `CompletionOptions.thinkingBudget.pauseGrammar` is applied. A core without it pauses every grammar inside a reasoning block. */ + thinkingBudgetPauseGrammar: boolean; /** The penalties count the prompt, and `CompletionOptions.penaltyPromptTokens` is applied. A core without it never counts the prompt, so `penaltyPromptTokens: false` holds there too. */ penaltyPromptTokens: boolean; /** The loaded core applies CompletionOptions.penaltyPromptTokensWithGrammar. Only literal true permits the opt-in. */ diff --git a/js/llama_webgpu_bridge.js b/js/llama_webgpu_bridge.js index 0c6bedc..102b988 100644 --- a/js/llama_webgpu_bridge.js +++ b/js/llama_webgpu_bridge.js @@ -239,6 +239,7 @@ var NO_COMPLETION_CAPABILITIES = Object.freeze({ presencePenalty: false, minP: false, thinkingBudget: false, + thinkingBudgetPauseGrammar: false, penaltyPromptTokens: false, penaltyPromptTokensWithGrammar: false, lazyGrammar: false, @@ -284,10 +285,16 @@ function thinkingBudget(value, hasMediaParts) { throw new TypeError("CompletionOptions.thinkingBudget.forcedMessage must be a string."); } withoutNul(forcedMessage, "CompletionOptions.thinkingBudget.forcedMessage"); + const pauseGrammar = budget.pauseGrammar ?? true; + if (typeof pauseGrammar !== "boolean") { + throw new TypeError( + `CompletionOptions.thinkingBudget.pauseGrammar must be a boolean; got ${String(pauseGrammar)}.` + ); + } if (hasMediaParts) { throw new Error("CompletionOptions.thinkingBudget supports text-only prompts; remove the media parts."); } - return { maxTokens, startTag, endTag, forcedMessage }; + return { maxTokens, startTag, endTag, forcedMessage, pauseGrammar }; } function speculativeInteger(config, name, min, max) { const value = config[name]; @@ -545,6 +552,7 @@ function completionCapabilitiesFrom(raw) { presencePenalty: parsed?.presencePenalty === true, minP: parsed?.minP === true, thinkingBudget: parsed?.thinkingBudget === true, + thinkingBudgetPauseGrammar: parsed?.thinkingBudgetPauseGrammar === true, penaltyPromptTokens: parsed?.penaltyPromptTokens === true, penaltyPromptTokensWithGrammar: parsed?.penaltyPromptTokensWithGrammar === true, lazyGrammar: parsed?.lazyGrammar === true, @@ -557,6 +565,7 @@ function requireCompletionCapabilities(sampling, capabilities) { sampling.minP !== 0 && !capabilities.minP ? "minP" : null, sampling.presencePenalty !== 0 && !capabilities.presencePenalty ? "presencePenalty" : null, sampling.thinkingBudget != null && !capabilities.thinkingBudget ? "thinkingBudget" : null, + sampling.thinkingBudget?.pauseGrammar === false && !capabilities.thinkingBudgetPauseGrammar ? "thinkingBudget.pauseGrammar" : null, sampling.grammarOptions?.lazy === true && !capabilities.lazyGrammar ? "grammarLazy" : null, sampling.grammarOptions != null && sampling.grammarOptions.root !== "root" && !capabilities.lazyGrammar ? "grammarRoot" : null, sampling.speculative != null && !Object.values(capabilities.speculativeDecoding).includes(true) ? "speculativeDecoding" : null @@ -4570,6 +4579,9 @@ var LlamaWebGpuBridgeRuntime = class { if (sampling.grammarOptions) { this._setNextGrammar(sampling.grammarOptions); } + if (sampling.thinkingBudget?.pauseGrammar === false) { + this._core.ccall("llamadart_webgpu_set_next_grammar_reasoning_pause", null, ["number"], [0]); + } const beginRc = Number( await this._core.ccall( input.exportName, @@ -6885,7 +6897,7 @@ var LlamaWebGpuBridge = class { } async _createCompletionUnlocked(prompt, options = {}) { const sampling = resolveCompletionSamplingOptions(options); - if (sampling.penaltyPromptTokensWithGrammar) { + if (sampling.penaltyPromptTokensWithGrammar || sampling.thinkingBudget?.pauseGrammar === false) { requireCompletionCapabilities(sampling, await this._getCompletionCapabilitiesUnlocked()); } const isWarmup = options?.warmup === true; diff --git a/js/src/bridge.ts b/js/src/bridge.ts index 2429f96..83a763e 100644 --- a/js/src/bridge.ts +++ b/js/src/bridge.ts @@ -2006,7 +2006,7 @@ export class LlamaWebGpuBridge implements PublicLlamaWebGpuBridge { options: BridgeCompletionOptions = {}, ): Promise { const sampling = resolveCompletionSamplingOptions(options); - if (sampling.penaltyPromptTokensWithGrammar) { + if (sampling.penaltyPromptTokensWithGrammar || sampling.thinkingBudget?.pauseGrammar === false) { // Probe the actual worker too: an older worker wrapper may ignore options. requireCompletionCapabilities(sampling, await this._getCompletionCapabilitiesUnlocked()); } diff --git a/js/src/internal/completion_options.ts b/js/src/internal/completion_options.ts index 3a2e9ca..b318ad7 100644 --- a/js/src/internal/completion_options.ts +++ b/js/src/internal/completion_options.ts @@ -12,13 +12,21 @@ import { withoutNul } from './text.ts'; export interface CompletionSamplingOptions { minP: number; presencePenalty: number; - thinkingBudget: Required | null; + thinkingBudget: ResolvedThinkingBudget | null; penaltyPromptTokens: boolean; penaltyPromptTokensWithGrammar: boolean; speculative: ResolvedSpeculativeOptions | null; grammarOptions: ResolvedGrammarOptions | null; } +export interface ResolvedThinkingBudget { + maxTokens: number; + startTag: string; + endTag: string; + forcedMessage: string; + pauseGrammar: boolean; +} + /** * A lazy grammar or a non-default grammar root, for the core's * set_next_grammar. Null when the grammar is strict and starts at `root`. @@ -80,6 +88,7 @@ export const NO_COMPLETION_CAPABILITIES: Readonly = Obje presencePenalty: false, minP: false, thinkingBudget: false, + thinkingBudgetPauseGrammar: false, penaltyPromptTokens: false, penaltyPromptTokensWithGrammar: false, lazyGrammar: false, @@ -109,7 +118,7 @@ function tag(budget: Record, name: 'startTag' | 'endTag'): stri return withoutNul(value, `CompletionOptions.thinkingBudget.${name}`); } -function thinkingBudget(value: unknown, hasMediaParts: boolean): Required | null { +function thinkingBudget(value: unknown, hasMediaParts: boolean): ResolvedThinkingBudget | null { if (value == null) { return null; } @@ -131,10 +140,16 @@ function thinkingBudget(value: unknown, hasMediaParts: boolean): Required', endTag: '' }; const THINK_BUDGET_TOKENS = 8; const FORCED_MESSAGE = 'Done.'; +// Takes the reasoning block and its end tag itself, as one written for a prompt +// that ends inside the block does. +const THINK_BODY_GRAMMAR = String.raw`root ::= think-body "" ("yes" | "no") +think-body ::= ( [^<] | "<" [^/] )*`; // Each run's prompt and options; each "-first" run warms the prompt cache. export const THINKING_RUNS = Object.freeze([ ['open-first', THINK_PROMPT, { ...THINK_GREEDY, nPredict: 1 }], @@ -215,6 +225,24 @@ export const THINKING_RUNS = Object.freeze([ grammar: YES_NO_GRAMMAR, thinkingBudget: { ...THINK_TAGS, maxTokens: THINK_BUDGET_TOKENS }, }], + ['open-budget-taking-grammar', THINK_PROMPT, { + ...THINK_GREEDY, + nPredict: 24, + grammar: THINK_BODY_GRAMMAR, + thinkingBudget: { ...THINK_TAGS, maxTokens: THINK_BUDGET_TOKENS, pauseGrammar: false }, + }], + ['open-budget-rejecting-grammar', THINK_PROMPT, { + ...THINK_GREEDY, + nPredict: 24, + grammar: String.raw`root ::= [^<]*`, + thinkingBudget: { ...THINK_TAGS, maxTokens: THINK_BUDGET_TOKENS, pauseGrammar: false }, + }], + ['open-budget-after-rejecting', THINK_PROMPT, { + ...THINK_GREEDY, + nPredict: 16, + grammar: YES_NO_GRAMMAR, + thinkingBudget: { ...THINK_TAGS, maxTokens: THINK_BUDGET_TOKENS }, + }], ['closed-first', CLOSED_THINK_PROMPT, { ...THINK_GREEDY, nPredict: 1 }], ['closed-plain', CLOSED_THINK_PROMPT, { ...THINK_GREEDY, nPredict: 4 }], ['closed-budget-zero', CLOSED_THINK_PROMPT, { @@ -555,10 +583,15 @@ export function renderHarness(nCtx, memoryModes) { const thinkingRuns = {}; for (const [name, prompt, options] of config.thinkingRuns) { - thinkingRuns[name] = await bridge.createCompletion(prompt, { - ...options, - tokenEventEncoding: 'text', - }); + try { + thinkingRuns[name] = await bridge.createCompletion(prompt, { + ...options, + tokenEventEncoding: 'text', + }); + } catch (caught) { + if (name !== 'open-budget-rejecting-grammar') throw caught; + thinkingRuns[name] = \`ERR \${caught && caught.name}: \${errorText(caught)}\`; + } } let invalidThinkingBudgetError = null; try { @@ -801,6 +834,18 @@ export function validateThinkingBudget(entry) { [`${endTag}yes`, `${endTag}no`].includes(afterPrefix(runs['open-budget-grammar'])), 'the grammar did not pause inside the block and constrain the text after it', ); + expect( + [`${endTag}yes`, `${endTag}no`].includes(afterPrefix(runs['open-budget-taking-grammar'])), + 'the grammar did not stay applied inside the block through the forced end tag', + ); + expect( + /^ERR .*Grammar rejected every candidate token/.test(runs['open-budget-rejecting-grammar']), + 'a grammar rejecting the forced end tag did not fail with the grammar-rejection error', + ); + expect( + [`${endTag}yes`, `${endTag}no`].includes(afterPrefix(runs['open-budget-after-rejecting'])), + 'the default pausing grammar changed after a rejected forced end tag', + ); expect(runs['closed-budget-zero'] === runs['closed-plain'], 'a budget changed output after the prompt closed the block'); const error = pyGet(entry, 'invalidThinkingBudgetError'); if (typeof error !== 'string' || !error.startsWith('RangeError: CompletionOptions.thinkingBudget.maxTokens')) { diff --git a/src/core/exports_generation.inc b/src/core/exports_generation.inc index 895fef6..352cb8b 100644 --- a/src/core/exports_generation.inc +++ b/src/core/exports_generation.inc @@ -172,6 +172,14 @@ EMSCRIPTEN_KEEPALIVE int32_t llamadart_webgpu_add_next_grammar_trigger( return 0; } +// Sets whether the grammar of the next begin_generation or generate call +// pauses inside a thinking-budget reasoning block. Call it after +// llamadart_webgpu_set_next_grammar, which restores the default, pausing. +EMSCRIPTEN_KEEPALIVE void llamadart_webgpu_set_next_grammar_reasoning_pause( + int32_t pause) { + g_next_grammar.pause_while_reasoning = pause != 0; +} + // Strategy flags report what the loaded models can run: draft-mtp needs the // target's MTP head, and draft-model strategies a matching draft model. EMSCRIPTEN_KEEPALIVE const char * llamadart_webgpu_completion_capabilities_json() { @@ -196,6 +204,7 @@ EMSCRIPTEN_KEEPALIVE const char * llamadart_webgpu_completion_capabilities_json( }; g_completion_capabilities_json = "{\"presencePenalty\":true,\"minP\":true,\"thinkingBudget\":true," + "\"thinkingBudgetPauseGrammar\":true," "\"penaltyPromptTokens\":true,\"penaltyPromptTokensWithGrammar\":true,\"lazyGrammar\":true," "\"speculativeDecoding\":{"; bool first = true; diff --git a/src/core/generation.inc b/src/core/generation.inc index 27f4cdc..c48f45a 100644 --- a/src/core/generation.inc +++ b/src/core/generation.inc @@ -54,12 +54,14 @@ llama_sampler * create_sampler( } // Mirrors llamadart-native's reasoning-budget sampler: the grammar pauses -// while a reasoning block is open or its end is being forced. +// while a reasoning block is open or its end is being forced, unless the +// request keeps it applied (pause_grammar_while_reasoning there). bool grammar_is_active() { if (g_active_grammar == nullptr) { return false; } - if (g_active_reasoning_budget == nullptr) { + if (g_active_reasoning_budget == nullptr || + !g_active_grammar_pauses_while_reasoning) { return true; } const common_reasoning_budget_state state = @@ -132,8 +134,11 @@ llama_token sample_next_token(const int32_t batch_index = -1) { } token = cur_p.data[cur_p.selected].id; // Every candidate is rejected only when the grammar can't be continued - // with this vocabulary; accepting the token would abort. - if (!grammar_accepts_token(token)) { + // with this vocabulary; accepting the token would abort. A -inf pick means + // the grammar and the reasoning budget's forced end tag excluded each other, + // and the budget would otherwise finish the block without its end tag. + if (cur_p.data[cur_p.selected].logit == -INFINITY || + !grammar_accepts_token(token)) { set_error("Grammar rejected every candidate token"); return LLAMA_TOKEN_NULL; } @@ -576,6 +581,7 @@ int32_t begin_generation_impl( g_active_reasoning_budget = reasoning_budget_ptr; g_active_grammar = grammar_sampler.release(); g_active_grammar_lazy = g_active_grammar != nullptr && grammar_options.lazy; + g_active_grammar_pauses_while_reasoning = grammar_options.pause_while_reasoning; g_generation_active = true; return 0; } diff --git a/src/core/state.inc b/src/core/state.inc index 6b5adc1..a85229c 100644 --- a/src/core/state.inc +++ b/src/core/state.inc @@ -28,6 +28,9 @@ llama_sampler * g_active_grammar = nullptr; // A lazy grammar can reject its trigger text when the trigger fires inside // llama_sampler_accept, which throws. bool g_active_grammar_lazy = false; +// False for a grammar that takes the reasoning block itself, so it stays +// applied while g_active_reasoning_budget counts or forces the block's end. +bool g_active_grammar_pauses_while_reasoning = true; // Borrowed from g_active_sampler's chain, which owns and frees it. llama_sampler * g_active_reasoning_budget = nullptr; std::vector g_sampler_candidates; @@ -120,6 +123,7 @@ speculative_request g_next_speculative; struct grammar_request { std::string root = "root"; bool lazy = false; + bool pause_while_reasoning = true; std::vector triggers; }; grammar_request g_next_grammar; diff --git a/src/core/support.inc b/src/core/support.inc index 24286d0..d9c3d57 100644 --- a/src/core/support.inc +++ b/src/core/support.inc @@ -346,6 +346,7 @@ void end_generation_state() { g_active_grammar = nullptr; } g_active_grammar_lazy = false; + g_active_grammar_pauses_while_reasoning = true; g_generation_active = false; g_last_piece.clear(); g_cancel_requested = false; diff --git a/tests/bridge/completion_sampling_options_test.mjs b/tests/bridge/completion_sampling_options_test.mjs index fdcef42..d5df9a6 100644 --- a/tests/bridge/completion_sampling_options_test.mjs +++ b/tests/bridge/completion_sampling_options_test.mjs @@ -15,6 +15,7 @@ const SUPPORTED = { presencePenalty: true, minP: true, thinkingBudget: true, + thinkingBudgetPauseGrammar: true, penaltyPromptTokens: true, penaltyPromptTokensWithGrammar: true, lazyGrammar: true, @@ -24,6 +25,7 @@ const NONE = { presencePenalty: false, minP: false, thinkingBudget: false, + thinkingBudgetPauseGrammar: false, penaltyPromptTokens: false, penaltyPromptTokensWithGrammar: false, lazyGrammar: false, @@ -58,7 +60,7 @@ function nativeParameterNames(name) { function coreCapabilityJson() { const body = functionBody('llamadart_webgpu_completion_capabilities_json()'); const literal = body.match( - /"\{\\"presencePenalty\\":true,\\"minP\\":true,\\"thinkingBudget\\":true,"\s*"\\"penaltyPromptTokens\\":true,\\"penaltyPromptTokensWithGrammar\\":true,\\"lazyGrammar\\":true,"/, + /"\{\\"presencePenalty\\":true,\\"minP\\":true,\\"thinkingBudget\\":true,"\s*"\\"thinkingBudgetPauseGrammar\\":true,"\s*"\\"penaltyPromptTokens\\":true,\\"penaltyPromptTokensWithGrammar\\":true,\\"lazyGrammar\\":true,"/, ); assert.ok(literal, 'the capability export must start with its fixed flags'); const strategies = [...body.matchAll(/\{"([a-z0-9-]+)",\s*([^}]*)\}/g)]; @@ -67,6 +69,7 @@ function coreCapabilityJson() { presencePenalty: true, minP: true, thinkingBudget: true, + thinkingBudgetPauseGrammar: true, penaltyPromptTokens: true, penaltyPromptTokensWithGrammar: true, lazyGrammar: true, @@ -240,6 +243,89 @@ const CASES = [ assert.deepEqual(structuredClone(sent).thinkingBudget, BUDGET); }], + ['thinkingBudget.pauseGrammar false is sent to the core after the grammar and before begin', async () => { + const pauseBudget = (pauseGrammar) => ({ ...BUDGET, pauseGrammar }); + const grammar = 'root ::= think-body "" "ok"\nthink-body ::= [^<]*\ncall ::= "ok"'; + const caps = '{"thinkingBudget":true,"thinkingBudgetPauseGrammar":true,"lazyGrammar":true}'; + const trace = (capabilityJson) => { + const { bridge, core } = directBridge(capabilityJson); + const calls = []; + const ccall = core.ccall.bind(core); + core.ccall = (name, returnType, argTypes, args) => { + if (name === 'llamadart_webgpu_set_next_grammar') { + calls.push([name, args]); + return 0; + } + if (name === 'llamadart_webgpu_set_next_grammar_reasoning_pause') { + calls.push([name, args]); + return undefined; + } + if (name === 'llamadart_webgpu_begin_generation_bytes') { + calls.push([name]); + } + return ccall(name, returnType, argTypes, args); + }; + return { bridge, calls }; + }; + { + const { bridge, calls } = trace(caps); + await bridge.createCompletion('hello', { grammar, grammarRoot: 'call', thinkingBudget: pauseBudget(false) }); + assert.deepEqual(calls, [ + ['llamadart_webgpu_set_next_grammar', ['call', 0]], + ['llamadart_webgpu_set_next_grammar_reasoning_pause', [0]], + ['llamadart_webgpu_begin_generation_bytes'], + ]); + } + for (const pauseGrammar of [undefined, null, true]) { + const { bridge, calls } = trace(caps); + await bridge.createCompletion('hello', { grammar, thinkingBudget: pauseBudget(pauseGrammar) }); + assert.deepEqual(calls, [['llamadart_webgpu_begin_generation_bytes']]); + } + const { bridge, calls } = trace('{"thinkingBudget":true}'); + await assert.rejects( + bridge.createCompletion('hello', { grammar, thinkingBudget: pauseBudget(false) }), + /^Error: The loaded WebGPU core does not support CompletionOptions\.thinkingBudget\.pauseGrammar\.$/, + ); + await bridge.createCompletion('hello', { grammar, thinkingBudget: pauseBudget(true) }); + assert.deepEqual(calls, [['llamadart_webgpu_begin_generation_bytes']]); + await assert.rejects( + bridge.createCompletion('hello', { thinkingBudget: pauseBudget(0) }), + /^TypeError: CompletionOptions\.thinkingBudget\.pauseGrammar must be a boolean; got 0\.$/, + ); + assert.equal((await bridge.getCompletionCapabilities()).thinkingBudgetPauseGrammar, false); + }], + + ['an older worker wrapper cannot silently ignore thinkingBudget.pauseGrammar false', async () => { + const calls = []; + const bridge = createWorkerBridge({ + _callWorker: async (method) => { + calls.push(method); + return method === 'getCompletionCapabilities' ? { thinkingBudget: true } : 'ignored'; + }, + }); + calls.length = 0; + await assert.rejects( + bridge.createCompletion('hello', { thinkingBudget: { ...BUDGET, pauseGrammar: false } }), + /CompletionOptions\.thinkingBudget\.pauseGrammar\.$/, + ); + assert.deepEqual(calls, ['getCompletionCapabilities']); + }], + + ['a worker completion forwards thinkingBudget.pauseGrammar', async () => { + let sent = null; + const bridge = createWorkerBridge({ + _callWorker: async (method, args) => { + if (method === 'getCompletionCapabilities') { + return { thinkingBudget: true, thinkingBudgetPauseGrammar: true }; + } + sent = args[1]; + return 'x'; + }, + }); + await bridge.createCompletion('hello', { thinkingBudget: { ...BUDGET, pauseGrammar: false } }); + assert.deepEqual(structuredClone(sent).thinkingBudget, { ...BUDGET, pauseGrammar: false }); + }], + ['a core without the probe reports nothing and rejects a nonzero option', async () => { const { bridge, begins } = directBridge(undefined); assert.deepEqual(await bridge.getCompletionCapabilities(), NONE); @@ -430,6 +516,7 @@ const CASES = [ ['the native grammar pauses while the reasoning budget is inside a block', async () => { const gate = functionBody('bool grammar_is_active()'); + assert.match(gate, /!g_active_grammar_pauses_while_reasoning/); assert.match(gate, /return state == REASONING_BUDGET_IDLE \|\| state == REASONING_BUDGET_DONE;/); const sample = functionBody('llama_token sample_next_token('); const gateIndex = sample.indexOf('const bool use_grammar = grammar_is_active();'); diff --git a/tests/bridge/thinking_budget_grammar_native_test.mjs b/tests/bridge/thinking_budget_grammar_native_test.mjs new file mode 100644 index 0000000..4231796 --- /dev/null +++ b/tests/bridge/thinking_budget_grammar_native_test.mjs @@ -0,0 +1,93 @@ +import assert from 'node:assert/strict'; +import { spawnSync } from 'node:child_process'; +import { mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import path from 'node:path'; +import { readNativeCoreSource } from './native_core_source.mjs'; + +const source = readNativeCoreSource(); +const body = (signature) => { + const start = source.indexOf(signature); + assert.ok(start >= 0, `${signature} not found in the native core`); + return source.slice(start, source.indexOf('\n}\n', start) + 2); +}; +const isActive = body('bool grammar_is_active() {'); +const begin = body('int32_t begin_generation_impl('); +const end = body('void end_generation_state() {'); + +// The production predicate, compiled against llama.cpp's reasoning-budget +// states, decides per state whether the grammar is applied. +const fixture = ` +#include +#include +enum common_reasoning_budget_state { + REASONING_BUDGET_IDLE, + REASONING_BUDGET_COUNTING, + REASONING_BUDGET_FORCING, + REASONING_BUDGET_WAITING_UTF8, + REASONING_BUDGET_DONE, +}; +struct llama_sampler { common_reasoning_budget_state state; }; +common_reasoning_budget_state common_reasoning_budget_get_state(const llama_sampler * sampler) { + return sampler->state; +} +llama_sampler * g_active_grammar = nullptr; +llama_sampler * g_active_reasoning_budget = nullptr; +bool g_active_grammar_pauses_while_reasoning = true; +${isActive} +int main() { + llama_sampler grammar{REASONING_BUDGET_IDLE}; + llama_sampler budget{REASONING_BUDGET_IDLE}; + const common_reasoning_budget_state states[] = { + REASONING_BUDGET_IDLE, REASONING_BUDGET_COUNTING, REASONING_BUDGET_FORCING, + REASONING_BUDGET_WAITING_UTF8, REASONING_BUDGET_DONE}; + for (const bool pauses : {true, false}) { + g_active_grammar_pauses_while_reasoning = pauses; + g_active_grammar = nullptr; + g_active_reasoning_budget = &budget; + assert(!grammar_is_active()); + g_active_grammar = &grammar; + g_active_reasoning_budget = nullptr; + assert(grammar_is_active()); + g_active_reasoning_budget = &budget; + for (const common_reasoning_budget_state state : states) { + budget.state = state; + const bool outside = state == REASONING_BUDGET_IDLE || state == REASONING_BUDGET_DONE; + assert(grammar_is_active() == (outside || !pauses)); + } + } +} +`; +const directory = mkdtempSync(path.join(tmpdir(), 'bridge-thinking-budget-grammar-')); +try { + const cpp = path.join(directory, 'contract.cpp'); + const executable = path.join(directory, 'contract'); + writeFileSync(cpp, fixture); + const compiled = spawnSync(process.env.CXX || 'c++', ['-std=c++17', cpp, '-o', executable], { encoding: 'utf8' }); + assert.equal(compiled.status, 0, compiled.stderr); + const ran = spawnSync(executable, [], { encoding: 'utf8' }); + assert.equal(ran.status, 0, `${ran.signal}: ${ran.stderr}`); +} finally { + rmSync(directory, { recursive: true, force: true }); +} + +// A request sets the flag for its own generation only. +assert.match(begin, /g_active_grammar_pauses_while_reasoning = grammar_options\.pause_while_reasoning;/); +assert.match(end, /g_active_grammar_pauses_while_reasoning = true;/); + +const cmake = readFileSync(new URL('../../CMakeLists.txt', import.meta.url), 'utf8'); +const exported = cmake.match(/-sEXPORTED_FUNCTIONS=\[([^\]]*)\]/)[1]; +assert.ok( + exported.includes("'_llamadart_webgpu_set_next_grammar_reasoning_pause'"), + 'llamadart_webgpu_set_next_grammar_reasoning_pause is not in EXPORTED_FUNCTIONS', +); + +// A grammar that rejects the budget's forced token leaves every logit at -inf +// after the grammar-first resample; the sampler must fail instead of letting +// the budget finish the block without its end tag. +const sample = body('llama_token sample_next_token('); +assert.match(sample, /cur_p\.data\[cur_p\.selected\]\.logit == -INFINITY/); +assert.ok( + sample.indexOf('.logit == -INFINITY') < sample.indexOf('Grammar rejected every candidate token'), + 'the -inf pick is rejected with the grammar-rejection error', +); diff --git a/tests/smoke/browser_smoke_support_test.mjs b/tests/smoke/browser_smoke_support_test.mjs index 475efae..44a01f2 100644 --- a/tests/smoke/browser_smoke_support_test.mjs +++ b/tests/smoke/browser_smoke_support_test.mjs @@ -292,6 +292,9 @@ assert.deepEqual(webGpuLaunchArgs('linux'), [ 'open-budget': 'T x', 'open-budget-zero': ' Done.x', 'open-budget-grammar': 'Tyes', + 'open-budget-taking-grammar': 'Tno', + 'open-budget-rejecting-grammar': 'ERR Error: Generation step failed: Grammar rejected every candidate token', + 'open-budget-after-rejecting': 'Tyes', 'closed-first': 'c', 'closed-plain': 'c', 'closed-budget-zero': 'c', @@ -430,12 +433,15 @@ assert.deepEqual(webGpuLaunchArgs('linux'), [ }; return [ runsCase( - { 'open-prefix': 'T', 'open-budget': 'T', 'open-budget-grammar': 'Tyes' }, + { 'open-prefix': 'T', 'open-budget': 'T', 'open-budget-grammar': 'Tyes', 'open-budget-taking-grammar': 'Tno', 'open-budget-after-rejecting': 'Tyes' }, 'the model closed the block itself, so the budget proves nothing', ), runsCase({ 'open-budget': 'T more' }, "maxTokens 8 did not force '' after the unbudgeted prefix"), runsCase({ 'open-budget-zero': 'Done.' }, "maxTokens 0 did not start with the forced 'Done.'"), runsCase({ 'open-budget-grammar': 'Tmaybe' }, 'the grammar did not pause inside the block and constrain the text after it'), + runsCase({ 'open-budget-taking-grammar': 'Tmaybe' }, 'the grammar did not stay applied inside the block through the forced end tag'), + runsCase({ 'open-budget-rejecting-grammar': 'T' }, 'a grammar rejecting the forced end tag did not fail with the grammar-rejection error'), + runsCase({ 'open-budget-after-rejecting': 'Tmaybe' }, 'the default pausing grammar changed after a rejected forced end tag'), runsCase({ 'closed-budget-zero': 'd' }, 'a budget changed output after the prompt closed the block'), [mutate(grammarPayload(), (p) => { p.modeResults[0].invalidThinkingBudgetError = null; }), both, [ 'wasm32 direct: invalid thinking budget: expected a RangeError, got None', @@ -481,7 +487,7 @@ assert.deepEqual(webGpuLaunchArgs('linux'), [ p.modeResults[3].capabilitiesBeforeLoad.minP = true; p.modeResults[3].capabilitiesAfterLoad = null; }), both, [ - "wasm64 worker: capabilities before the load: expected none, got {'presencePenalty': False, 'minP': True, 'thinkingBudget': False, " + "wasm64 worker: capabilities before the load: expected none, got {'presencePenalty': False, 'minP': True, 'thinkingBudget': False, 'thinkingBudgetPauseGrammar': False, " + "'penaltyPromptTokens': False, 'penaltyPromptTokensWithGrammar': False, 'lazyGrammar': False}", 'wasm64 worker: capabilities after the load: expected all, got None', `wasm64 worker: sampler outputs: {${Object.entries(samplerRuns()).map(([n, t]) => `'${n}': '${t}'`).join(', ')}}`,