feat(engine): schema/grammar-constrained tool-call argument decoding (#374)

[pekkah]

Closes #374.

What

Opt-in schema/grammar-constrained decoding that forces a tool call's arguments to satisfy the supplied JSON Schema, expressed in each architecture's native call syntax. Default off → byte-identical to today. Enabled via SHARPI_TOOL_GRAMMAR=1 or SharpInferenceServerOptions.ToolGrammar. Scoped to Gemma 4 end-to-end; the adapter hook generalizes to other families (follow-up for Qwen/Llama JSON).

Why

On gemma-4-12b-it-qat-q4_0 (temp 0) the call envelope is always correct but arguments are wrong deterministically — a model prior, not an engine bug:

Prompt	before	after
"weather in Berlin" (location required)	{}	{"location":"Berlin"}
"population of France" (query: string)	{"queries":["…"]}	{"query":"…"}
enum unit∈{celsius,fahrenheit}	{}	{"location":"Berlin","unit":"celsius"}
tags: array<string>	—	{"tags":["urgent","bug"]}

Constrained decoding makes the bad token sequences unreachable instead of hoping the model conforms.

Approach

Per-adapter grammar. IToolCallAdapter.BuildArgumentConstraint(tools, vocab) (default null; Gemma4 override) so each family owns its wire format.

The hard part — terminal→token mapping. A byte-level pushdown automaton over Gemma's <|tool_call>call:NAME{key:<|"|>val<|"|>,n:3}<tool_call|> syntax. A candidate token is allowed iff replaying its bytes keeps the FSM alive. Matching is byte-level for the structural skeleton (so multi-token keys and merged delimiters like :[ / ]} work) and token-level for the <|"|> quote special token and free string content. The grammar enforces: only declared keys, every required key exactly once, value shape per declared type (strings quoted with free content; numbers/bool/null bare; arrays […]; nested objects {…}), and enum restriction. It engages the instant the tool name completes — before { — so Gemma's merged {} empty-args token can't slip required arguments past the mask. Allocation-free hot path (reused mask/scratch buffers + first-byte pruning over the 262k vocab; warm constrained request ~0.33s, overhead negligible).

Engine. SamplingParams.Constraint; the InferenceEngine decode loop masks logits while constraining and advances the constraint per token (disables the device-side argmax fast path + MTP for constrained requests). ContinuousBatchingEngine warns once and ignores it (not yet wired — noted as follow-up).

Server. ToolGrammar option + SHARPI_TOOL_GRAMMAR env; OpenAI/Anthropic endpoints build the constraint from req.Tools on tool-active requests (OpenAI honors tool_choice:"none").

Verification (CUDA, gemma-4-12b-it-qat-q4_0, temp 0)

All failure cases above flip to correct; plain text, already-valid calls, and tool_choice:"none" are byte-identical to default-off; streaming honors the constraint.

Tests

• ToolSchemaParseTests — schema derivation (types/required/enum/array/nested).
• ToolGrammarMockTests — decode-time conformance via a fake Gemma tokenizer (CI-runnable, no GGUF): required-key, foreign-key rejection, string/enum/number/array/nested-object value shapes, early-engage, reset.
• ToolGrammarConstraintTests — model-gated conformance against the real GGUF.

190 Tests.Core + 125 Tests.Server pass; full Release build clean under TreatWarningsAsErrors (trim/AOT analyzers + InvariantGlobalization on).

Follow-ups (filed)

• Tool-arg grammar: Qwen/Llama JSON-syntax constraints (#374 follow-up) #376 � Qwen/Llama JSON-syntax argument grammars via the same adapter hook.
• Tool-arg grammar: wire constraint into ContinuousBatchingEngine (#374 follow-up) #377 � wire the constraint into ContinuousBatchingEngine (today it warns + ignores).
• Tool-arg grammar: constrain partially-typed tools (Any/open-object values) (#374 follow-up) #378 � constrain partially-typed tools (Any-typed value / open object currently left unconstrained).

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces schema- and grammar-constrained decoding for tool-call arguments, specifically supporting Gemma 4's native wire format. It adds a state machine constraint (GemmaToolArgumentConstraint), JSON-Schema parsing to ToolSchema, and integrates these constraints into the single-user InferenceEngine and server endpoints, while providing a warning in the ContinuousBatchingEngine. The review feedback highlights several critical performance optimizations on the hot path of token generation—such as optimizing large vocabulary loops, replacing Array.Copy with manual loops for shallow stacks, and avoiding HashSet.Contains lookups—as well as preventing potential stack overflows during recursive schema compilation by enforcing a depth limit.

Important

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

[gemini-code-assist]

To prevent potential StackOverflowException when compiling maliciously crafted or extremely deeply nested schemas, we should enforce a compilation depth limit. Since MaxDepth is already defined as 32 for the runtime frame stack, we can pass a depth parameter through the compilation methods and return null if it is exceeded.

    private static CompiledObject? TryCompileObject(ToolSchemaObject obj, int depth = 0)
    {
        if (depth >= MaxDepth) return null;
        if (obj.Open || obj.Properties.Count is 0 or > 64) return null;
        var keys = new byte[obj.Properties.Count][];
        var values = new CompiledNode[obj.Properties.Count];
        ulong reqMask = 0;
        for (int i = 0; i < obj.Properties.Count; i++)
        {
            var p = obj.Properties[i];
            if (p.Name.Length == 0) return null;
            keys[i] = ToolSchema.Utf8(p.Name);
            var node = TryCompileNode(p.Value, depth + 1);
            if (node is null) return null;                  // a non-constrainable value → skip the tool
            values[i] = node;
            if (p.Required) reqMask |= 1UL << i;
        }
        return new CompiledObject { KeyBytes = keys, Values = values, RequiredMask = reqMask };
    }

[gemini-code-assist]

Enforce the compilation depth limit in TryCompileNode as well to prevent stack overflow during recursive schema compilation.

    private static CompiledNode? TryCompileNode(ToolSchemaNode node, int depth = 0)
    {
        if (depth >= MaxDepth) return null;
        switch (node.Kind)
        {
            case JsonSchemaKind.String:
                byte[][]? strEnum = node.EnumValues is { } se ? Encode(se) : null;
                return new CompiledNode { Kind = JsonSchemaKind.String, Literals = strEnum, QuotedLiterals = strEnum is not null };

            case JsonSchemaKind.Number:
            case JsonSchemaKind.Integer:
                byte[][]? numEnum = node.EnumValues is { } ne ? Encode(ne) : null;
                return new CompiledNode
                {
                    Kind = node.Kind,
                    Literals = numEnum,
                    QuotedLiterals = false,
                    IntegerOnly = node.Kind == JsonSchemaKind.Integer,
                };

            case JsonSchemaKind.Boolean:
                return new CompiledNode { Kind = JsonSchemaKind.Boolean, Literals = s_boolLiterals };

            case JsonSchemaKind.Null:
                return new CompiledNode { Kind = JsonSchemaKind.Null, Literals = s_nullLiterals };

            case JsonSchemaKind.Array:
                // An untyped array (no items) isn't fully constrainable — skip the tool.
                if (node.Items is null) return null;
                var item = TryCompileNode(node.Items, depth + 1);
                return item is null ? null : new CompiledNode { Kind = JsonSchemaKind.Array, Items = item };

            case JsonSchemaKind.Object:
                if (node.Object is null) return null;
                var obj = TryCompileObject(node.Object, depth + 1);
                return obj is null ? null : new CompiledNode { Kind = JsonSchemaKind.Object, Object = obj };

            default:
                return null;   // Any / unknown — not constrainable
        }
    }

[gemini-code-assist]

In the performance-critical hot path of token generation, calling _forbidden.Contains(id) inside a loop of size 262,144 (the vocabulary size) for every generated token introduces significant overhead. Since _forbidden typically contains only a few EOG token IDs, we can set those specific IDs to -inf first, and then use a simple, extremely fast float comparison loop to count kept tokens, completely avoiding HashSet.Contains overhead inside the 262k loop.

        if (top.Kind == FK.Str && top.State == SContent)
        {
            // Any non-forbidden token stays in content; the quote token closes. Forbid only EOG.
            foreach (int id in _forbidden)
            {
                if (id >= 0 && id < buf.Length)
                {
                    buf[id] = float.NegativeInfinity;
                }
            }
            for (int id = 0; id < buf.Length; id++)
            {
                if (!float.IsNegativeInfinity(buf[id]))
                    kept++;
            }
            return kept;
        }

[gemini-code-assist]

Instead of looping 262,144 times with a conditional branch to set everything except _quoteId to -inf, we can use Array.Fill to set the entire buffer to -inf and then restore the original logit of _quoteId. This is orders of magnitude faster and avoids the overhead of a large loop with branches.

        if ((top.Kind == FK.Str && top.State == SExpectOpen)
            || (top.Kind == FK.StrEnum && top.State == SeExpectOpen))
        {
            float originalQuoteLogit = buf[_quoteId];
            Array.Fill(buf, float.NegativeInfinity);
            if (!float.IsNegativeInfinity(originalQuoteLogit))
            {
                buf[_quoteId] = originalQuoteLogit;
                return 1;
            }
            return 0;
        }

[gemini-code-assist]

Since SimulateToken is called for every candidate token in the vocabulary (which can be thousands of candidates per generated token), using Array.Copy twice to save and restore the stack introduces noticeable overhead. Since _depth is typically very small (usually 1 to 3), a manual loop to copy the active frames is significantly faster because it avoids native transition, runtime type checking, and bounds checking overhead.

    private bool SimulateToken(int token)
    {
        int savedDepth = _depth;
        for (int i = 0; i < savedDepth; i++) _scratch[i] = _stack[i];
        bool ok = RunToken(token);
        for (int i = 0; i < savedDepth; i++) _stack[i] = _scratch[i];
        _depth = savedDepth;
        return ok;
    }

[pekkah]

Thanks @gemini-code-assist — addressed in ca28054:

• TryCompileObject / TryCompileNode depth limit (both high): added a MaxDepth cap on compilation. The schema parser was already depth-capped in a prior commit, but the ToolSchema records are public, so an in-memory deeply-nested schema would bypass that — good catch on the defense-in-depth.
• String-open mask (medium): now Array.Fill(-inf) + restore the quote logit, no branchy 262k loop.
• SimulateToken save/restore (medium): manual frame-copy loop instead of Array.Copy (depth is 1–3), avoiding per-call overhead on the masking hot path.
• Free-string-content mask (medium): already EOG-direct as of the prior review commit (e408688) — masks the handful of EOG ids rather than testing all 262k against the set.

22 grammar/schema tests pass; build clean under TreatWarningsAsErrors.

[gemini-code-assist]

Thanks for the update, @pekkah. The depth-capping in TryCompileObject and TryCompileNode is a solid defensive measure against malicious or malformed schemas, and the manual frame-copy loop is a great optimization for the hot path. The move to Array.Fill(-inf) for the string-open mask is also much cleaner than the previous branchy loop. The implementation looks robust and well-verified.