Optimize Engine guidance for production workloads

[Baiju Meswani (baijumeswani)]

Summary

Makes constrained decoding practical for concurrent, long-running Engine workloads by amortizing grammar setup and moving per-token guidance work off the serialized Engine step path.

This change builds on Engine guidance and continuation introduced by #2412. It adds:

• model-local immutable tokenizer assets;
• a bounded, single-flight compiled grammar cache;
• request-local mutable grammar cursors;
• low-cost grammar reset for retained-request continuation;
• asynchronous cross-request mask computation using llguidance's parallel API;
• reusable host/device mask storage;
• one batched H2D transfer and CUDA mask launch for contiguous Engine decode rows;
• cache telemetry through C, C++, and Python;
• paired guidance performance and lifecycle qualification tooling.

Motivation

The original guidance path was correct, but each request independently rebuilt tokenizer/trie state and compiled its grammar. It also computed, allocated, copied, and applied masks once per request on every decode step.

On the target Qwen3.8 27B workload, the production impact and measured improvement were:

Problem	Before	After
Repeated schema request construction	~2.6 seconds	~0.5 ms p95 cached
Guided continuation	Recompiled/reset grammar state	<0.25 ms p95
B4 guided decode regression	~11%	<1.1% upper 95% bound
B8 guided decode regression	~20%	<1.0% upper 95% bound
Unique grammar cache growth	Unbounded concern	64 retained entries after 800 misses

Coding workloads repeatedly reuse the same tool and response schemas across many requests and continuation turns, so both costs should be amortized at the model and scheduled-batch levels.

Changes

Model-local guidance cache

• Creates the llguidance tokenizer/trie once per loaded model.
• Caches compiled initial constraints by guidance type, exact guidance data, and relevant options.
• Single-flights concurrent identical misses.
• Clones an independent mutable cursor for every request.
• Uses a bounded 64-entry LRU with a 16 MiB key-byte limit.
• Keeps evicted assets alive while active requests or asynchronous work still reference them.
• Removes failed compilations from the cache so retries perform a fresh compile.
• Releases all cached assets with the model.

Low-cost continuation

CloneForNewTurn() now clones from the cached initial constraint instead of cloning tokenizer state and recompiling the grammar. Every continued turn still starts from grammar state zero.

Overlapped mask computation

• Commits selected tokens on the Engine thread and marks their cursors dirty.
• Collects dirty cursors across the scheduled request batch.
• Calls llg_par_compute_mask once for the batch.
• Overlaps CPU grammar work with the next model forward.
• Waits for the result only when logits masking needs it.

Batched CUDA mask application

For contiguous CUDA Engine decode rows:

• stores mask rows contiguously;
• reuses one scheduler-owned device workspace;
• uses one H2D transfer;
• applies all guided and pass-through rows with one CUDA kernel launch.

CPU, noncontiguous, and unsupported batching cases retain the per-request fallback. A guided request never silently becomes unconstrained when the fast path is unavailable.

Correctness and ownership

• Uses copy-on-write grammar cursors for transactional and speculative clones.
• Keeps rollback checkpoints isolated from staged attempts.
• Pins grammar and tokenizer assets until asynchronous callbacks finish.
• Destroys constraints before their tokenizer.
• Publishes pending-mask futures only after the complete batch job is prepared.
• Contains all exceptions at the C callback boundary.
• Keeps synchronous and parallel stop-mask semantics identical.
• Retains the corrected padded mask-row stride for non-aligned vocabularies.
• Continues to reject Engine guidance fast-forward tokens, whose scheduler and KV accounting are not implemented.

Cache telemetry

Adds model-level counters for:

• tokenizer initializations;
• grammar hits, misses, and waits;
• grammar compilation time;
• evictions;
• retained grammar count;
• retained cache-key bytes.

The counters are available through C, C++, and Python.

Benchmarking

Extends the guidance profiler with:

• fixed-work guided/unguided comparisons;
• B1/B2/B4/B8 concurrency;
• exact 4K–256K prompts;
• JSON, regex, coding-tool catalog, repeated, and unique grammars;
• retained continuation;
• separate grammar construction and prompt-copy timing;
• paired percentage deltas and 95% confidence intervals;
• cache telemetry in result JSON;
• source-built ORT provider support.

Performance

Qwen3.8 27B, A100 80 GB, CUDA 12.8, SM80, INT4 weights, INT8 KV, CUDA graphs, 128 fixed generated tokens:

Concurrency	Guided ITL mean change	ITL upper 95% bound	Throughput mean change
B1	+0.56%	+0.59%	-0.56%
B2	+0.50%	+0.58%	-0.50%
B4	+0.95%	+1.00%	-0.94%
B8	+0.83%	+0.92%	-0.82%

Cached JSON-schema request construction:

• p50: 0.46 ms
• p95: 0.49 ms

Cached 16-tool catalog construction:

• p50: 0.53 ms
• p95: 0.54 ms

[Copilot]

Pull request overview

This PR makes constrained decoding (“guidance”) viable for concurrent, long-running Engine workloads by caching model-scoped guidance assets (tokenizer + compiled grammars), overlapping mask computation across requests, and batching CUDA mask application to reduce per-token overhead on the Engine step path.

Changes:

• Introduces a model-local tokenizer + compiled-grammar cache with bounded LRU retention, single-flight compilation, and lifecycle-safe eviction.
• Moves guidance mask computation onto an asynchronous, batched llguidance-parallel path and adds a contiguous-row CUDA fast path for applying masks.
• Exposes guidance cache telemetry through C/C++/Python and adds targeted unit tests + benchmarking extensions.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
test/python/models/test_n_gram_decoding.py	Makes the “guidance unavailable” skip detection case-insensitive.
test/engine/request_lifecycle_tests.cpp	Adds USE_GUIDANCE-gated lifecycle and cache behavior tests (reuse, single-flight, eviction safety, continuation reset).
test/engine/guidance_request_validation_tests.cpp	Adds model-free tests for centralized guidance request validation and build-availability behavior.
test/engine/guidance_processor_tests.cpp	Adds model-light tests for Request guidance bookkeeping using a fake ConstrainedLogitsProcessor.
test/engine/engine_test_helpers.h	Centralizes deterministic “scripted logits” helpers to share across guidance tests.
src/python/python.cpp	Exposes guidance cache counters to Python via Model.get_guidance_cache_stats().
src/ort_genai.h	Adds C++ convenience wrapper OgaModel::GetGuidanceCacheCount.
src/ort_genai_c.h	Adds C API declaration for querying guidance cache counters.
src/ort_genai_c.cpp	Implements C API OgaModelGetGuidanceCacheCount using model cache stats.
src/models/model.h	Adds model-owned, lazily initialized guidance cache state with mutex protection.
src/generators.h	Adds forward declaration for GuidanceCacheState.
src/engine/scheduled_requests.h	Extends batched sampling plan with reusable guidance mask host/device storage and new guidance scheduling hooks.
src/engine/scheduled_requests.cpp	Applies batched guidance masks for contiguous CUDA rows and schedules async mask computation post-step.
src/engine/request.h	Threads “guidance already applied” flag through generation paths; exposes mask accessors for scheduler batching.
src/engine/request.cpp	Centralizes guidance request validation/creation and makes continuation reset transactional via CloneForNewTurn().
src/engine/engine.cpp	Schedules guidance mask computation after dynamic-step commit.
src/constrained_logits_processor.h	Adds GetReadyMask() / CloneForNewTurn() APIs; introduces cache stats + validation helpers.
src/constrained_logits_processor.cpp	Implements tokenizer/grammar cache, async parallel mask computation, and batched/ready-mask plumbing.
benchmark/python/benchmark_engine_guidance.py	Extends benchmark to cover concurrency, cache effects, fixed-token comparisons, continuation turns, and richer reporting.
benchmark/engine/scenario_dispatcher.cpp	Allows omitted provider plugin path (for source-built ORT legacy provider loading).
benchmark/engine/README.md	Updates documentation to mark execution_provider_library as optional with source-built ORT.

Suppressed comments (1)

src/constrained_logits_processor.cpp:543

• The mask-bit semantics described here are inverted relative to the actual masking logic (and the CUDA AddLogitsMask kernel): a set bit means the token is allowed; an unset bit means it should be masked to -inf. Please update the comment to match the behavior to avoid future confusion/misuse.

      // mask is a 32-bit integer, where each bit corresponds to a token in the vocabulary.
      // If the bit is set, the corresponding token is masked (i.e., its logit is set to the lowest possible value).
      subspan[i] = mask[i / 32] & (uint32_t{1} << (i % 32)) ? subspan[i] : std::numeric_limits<float>::lowest();

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Tianlei Wu (tianleiwu)]

Review: Optimize Engine guidance for production workloads

Verdict: COMMENT. The design is sound and unusually well-defended. I did not find a correctness blocker; everything below is hardening around cache accounting, error signalling, and the new public C API shape.

Verified against the implementation

• Sync/parallel stop-mask parity is real. ComputeMasks() sets the EOS bit on LlgMaskResult::is_stop, and upstream ffi_par.rs::par_compute_mask_inner does the same (add_eos = r.is_stop() -> *mask_dest.add(eos/32) |= 1 << (eos%32), using constraint.tok_trie().eos_token(), which is the tok_eos this code passes as config.model.eos_token_id[0]). The claim in the description holds.
• Async lifetime is correct. ParallelMaskJob holds shared_ptr copies of both the cursors and the grammar assets (which own the tokenizer), so a request or model destroyed mid-flight cannot free memory a rayon worker is writing. Dropping those references before satisfying the promises is a nice touch - the woken request sees unique ownership and skips an unnecessary COW clone.
• Destruction ordering is correct in all three asset types (llg_tokenizer last in GuidanceTokenizerAsset, initial_constraint after tokenizer in GuidanceGrammarAsset, llg_constraints_ after grammar_asset_ in the processor), satisfying llguidance's "must not free the tokenizer while constraints exist" rule.
• The batched CUDA path never silently unconstrains. CollectBatchedGuidanceMasks returns FallbackRequired for any guided row it cannot fill and Ready only when every eligible guided row was masked; guidance_applied is only consulted for rows gated on the same IsChunkComplete() predicate. AddLogitsMask assigns float::lowest(), so re-application after a transaction rollback is idempotent.
• mask_dirty_ and pending_masks_.valid() are mutually exclusive on every path, including through Clone(), so the logic_error guard in ScheduleGuidanceMaskComputation is not reachable from speculative clones.
• CloneForNewTurn() reads no mutable source state, which is exactly what makes the clone-then-swap sequence in Request::Continue() retry-safe. Good that there is a test for the failing-reset case.

Cross-cutting notes (files not in this diff)

Implicit hard dependency on llguidance's rayon feature. cmake/external/onnxruntime_external_deps.cmake calls corrosion_import_crate() without FEATURES, so this relies on rayon staying in llguidance's default = ["lark", "rayon", "referencing", "ahash"]. Without it, llg_par_compute_mask calls cc.set_error(...) on every constraint, which clears the parser (self.constraint = None) permanently - so every guided request in such a build would fail unrecoverably rather than fall back to llg_compute_mask. Now that the Engine decode path depends on the parallel API, pinning FEATURES rayon explicitly would make that dependency non-silent.

Related: the "leaves failed cursors dirty so mask construction retries" comments are accurate for injected/transient failures, but a genuine llguidance mask error poisons the cursor for good (llg_get_error keeps returning the same pointer until the constraint is freed), so the retry will fail every step. That is arguably the right failure mode - failing loudly beats silently unconstraining - but the comments read as if recovery is expected.

Positives

Grammar/tokenizer caching, single-flight, COW cursors, and async mask overlap are each individually testable and each individually tested (concurrency, eviction with live processors, failed compilation not retained, pending mask outliving a removed processor, rollback isolation, Continue() transactionality). Consolidating validation into ValidateGuidanceRequest() removes the duplicated and slightly divergent checks from Request's constructor. The docs accurately match the implemented limits, including the subtlety that cached_grammars counts ready entries while cached_key_bytes also covers in-flight compilations.

Specific file/line feedback is inline.