perf: batch draft traversal and compact grammar caches

[benmyles]

Important

Superseded by upstream PR mlc-ai/xgrammar#740. Review and discussion have moved there.

Summary

This packages three related serving-path improvements: a native batched speculative-decoding traversal API, reusable bounded worker pools for matching and compilation, and compact compiled-grammar serialization. It also adds fail-closed validation, regression coverage, integration documentation, and a reproducible synthetic benchmark.

Guided review

Tip

Short on review time? Follow the numbered path below. It moves from the public contract, through the two concurrency-sensitive implementations, into the cache-format change, and ends at the tests that lock each behavior down. The source links are pinned to the reviewed commit, so they will not drift as the branch evolves.

This change is easiest to understand as one request moving through the system: Python describes a speculative draft tree, the FFI hands tensor views to C++, native workers traverse that tree and write token masks, and the compiler/cache layer makes the grammars feeding those matchers cheaper to build and store.

Review map

Step	Reviewer question	Primary code
1	What is the new public contract?	C++ matcher API, batch C++ API, Python bindings
2	Does one draft-tree traversal remain correct?	Recursive traversal, offset-aware mask view, scalar validation
3	How does batching remove dispatch overhead safely?	Batch validation and views, chunked execution
4	Can a persistent pool survive worker failures?	ThreadPool implementation, native regressions
5	How is grammar compilation bounded?	Compiler scheduling, concurrent stress test
6	Why is the v16 cache much smaller, and how is corruption rejected?	Compact serialization, nested-error propagation
7	Which tests prove the end-to-end behavior?	Speculative-decoding tests, serialization tests

flowchart LR
    A["Python engine integration"] -->|"tokens, tree, roots, output"| B["TVM FFI"]
    B --> C["BatchGrammarMatcher validation"]
    C -->|"tensor views only after all checks pass"| D["Persistent native worker pool"]
    D --> E["Per-request draft traversal"]
    E --> F["Flattened token-mask rows"]
    G["GrammarCompiler"] -->|"compiled grammar"| E
    G --> H["v16: one complete FSM + compact rule views"]
    H -->|"deserialize and validate"| E

Loading

---

1. Begin with the serving-engine contract

The public API change is deliberately small. The scalar matcher gains a root_position, while the batch matcher gains one operation that accepts all active requests at once. Start with the C++ declarations and the batch method contract, then read the corresponding scalar Python wrapper and batch Python wrapper.

In plain language, a caller can now provide:

• one shared one-dimensional tree for every request, or a separate tree in each row of a two-dimensional tensor;
• a two-dimensional draft-token tensor;
• optional sparse request-row indices;
• one root position per selected matcher; and
• one preallocated flattened output buffer.

The result is one completion flag per matcher. The output layout stays flat so C++ can write directly into the engine-owned bitmask without assembling temporary Python objects:

request i, draft node j
        │
        └── output row = i * num_nodes + j

The TVM FFI binding is intentionally mechanical: it forwards already-described tensors and options into the native matcher. That keeps policy and validation in one C++ implementation rather than duplicating it across language boundaries.

Important

The batch API is additive. Existing scalar callers keep their current path, and root_position=0 preserves the previous behavior.

Reviewer lens: public API

• Confirm shared-tree and per-request-tree shapes are unambiguous.
• Confirm sparse indices select matcher rows without changing output-row ownership.
• Confirm completion flags line up with the selected matchers.
• Confirm the default scalar root remains backward compatible.

---

2. Establish correctness in the scalar primitive first

The batch implementation is built on the same scalar traversal primitive, so the next stop is TraverseDraftTreeRecursive. It walks the draft tree depth-first, accepts each edge token against the grammar state, emits the allowed-token mask at that node, and rolls the matcher state back before visiting the next sibling. That rollback is what prevents one draft branch from contaminating another.

Two correctness details were made explicit before batching:

1. The output mask view includes the DLTensor base pointer and byte_offset. A tensor can be contiguous while still beginning at a nonzero storage offset; ignoring that field writes a valid mask to the wrong part of the allocation.
2. Scalar validation checks the root, shape, dtype, device, and true byte range before traversal. Invalid tensor metadata is rejected rather than being interpreted optimistically.

This scalar path is both the compatibility path and the correctness oracle used by the batch parity tests.

Note

“Contiguous” describes strides, not necessarily the beginning of storage. The nonzero-byte_offset regression is therefore a real tensor-layout case, not an artificial malformed input.

---

3. Lift that primitive into one native batch operation

Now read the full batch implementation. It has three phases, and their order is part of the safety argument.

3.1 Validate everything before scheduling anything

Lines 1940–2026 reject wrong devices, dtypes, dimensions, row counts, roots, duplicate indices, and output sizes. No worker is launched until every request is known to be valid.

That fail-closed ordering matters: if request 47 is malformed, requests 0–46 must not have partially advanced their matcher states or written partial masks.

3.2 Build zero-copy per-request views

Lines 2028–2073 select each request’s tree, tokens, and output region by rebasing lightweight DLTensor views. The implementation does not copy the draft tensors merely to batch them. It also preserves the original storage offset while calculating each row’s address, which is the batched counterpart of the scalar offset fix.

3.3 Schedule chunks, not individual requests

Lines 2075–2099 divide the selected requests into at most one task per worker. Each task resets its output rows and traverses its contiguous chunk. This makes two formerly separate operations—mask clearing and traversal—one worker pass and avoids queueing one tiny task per request.

The pool itself is created once in the BatchGrammarMatcher constructor. The same chunking principle is also used by batched next-token mask generation, so both hot paths avoid per-call thread construction and per-request scheduling overhead.

sequenceDiagram
    participant P as Python caller
    participant B as Batch matcher
    participant W as Persistent workers
    P->>B: tensors + indices + roots + output
    B->>B: validate the complete batch
    B->>B: create zero-copy row views
    B->>W: one chunk per active worker
    par worker chunks
        W->>W: clear owned output rows
        W->>W: traverse owned requests
    end
    W-->>B: completion flags / first exception
    B-->>P: one result per matcher

Loading

Reviewer lens: native batch path

The important invariants are:

• validation completes before any observable mutation;
• one worker owns each matcher and output-row range;
• duplicate sparse indices are rejected, so two workers cannot mutate the same matcher;
• shared and per-request tree addressing produce the same scalar semantics; and
• all scheduled work finishes before tensor views and captured state leave scope.

---

4. Verify the persistent worker pool’s failure semantics

Reusing threads is only an optimization if a failed task cannot poison later calls. The central behavior lives in ThreadPool.

Workers catch task exceptions and record the first one while still marking the task complete. Then Wait() waits for the submitted generation, rethrows that exception on the caller thread, and clears it. Clearing after observation is the key reuse property: the next independent batch can run normally instead of seeing a stale failure.

The convenience ParallelFor waits before joining, so an exception cannot be lost during pool teardown.

Warning

This code intentionally does not fall back to serial execution after a worker failure. The operation fails at its boundary, all outstanding work is drained safely, and the reusable pool remains valid for the next call.

The native pool tests exercise repeated waits, Execute, Submit, worker exceptions, subsequent reuse, and ParallelFor propagation. The matcher-native regressions additionally prove that a batch exception reaches the caller and that the same matcher pool works afterward.

---

5. Apply the same bounded-concurrency model to compilation

Grammar compilation used to risk paying thread setup repeatedly, especially when several outer callers compiled grammars concurrently. The compiler now owns one bounded pool for its lifetime, so concurrent callers share the same native capacity instead of multiplying it.

The scheduling flow is:

1. Collect adaptive states first, while the compilation structures are still under straightforward caller control.
2. Compile small workloads inline and queue only sufficiently large ones. The 64-state threshold avoids spending more on synchronization than the work itself.
3. Keep futures for queued work and drain every outstanding task if one fails. Only after captured compilation state is no longer in use does the exception unwind to the caller.

This produces bounded parallelism without changing the compiled result. The concurrent shared-compiler stress test compares serial and parallel output while many callers reuse one compiler.

many compile_* callers
          │
          ▼
one GrammarCompiler-owned pool
          │
          ├── small adaptive-state set ──► inline
          └── large adaptive-state set ──► bounded futures ──► drain ──► return/rethrow

---

6. Read the serialization change as normalization, not compression magic

The previous optimized representation serialized overlapping FSM data once per rule. Large grammars therefore repeated the same edge auxiliary data hundreds of times. Version 16 stores the normalized structure instead:

• CompactRuleFSMView records only a rule’s start state, end states, DFA flag, and size metadata.
• Serialization writes the complete optimized FSM once, followed by those compact rule views.
• Deserialization validates and reconstructs views over that one shared FSM.

v15                                      v16
rule 0 → FSM data copy                   complete_fsm → one data copy
rule 1 → overlapping FSM data copy              ▲
rule 2 → overlapping FSM data copy              ├── rule 0 compact view
...                                              ├── rule 1 compact view
                                                 └── rule 2 compact view

The deserializer checks view count, required complete-FSM presence, start/end bounds, sorted unique end states, and recorded sizes before constructing anything. Nested grammar and adaptive-mask failures are propagated, rather than turning malformed child data into an apparently successful outer object. The format identifier is correspondingly bumped to v16.

Caution

This is an intentional on-disk cache break. A v15 compiled-grammar cache must be cleared during upgrade; silently interpreting it as v16 would be unsafe. The migration expectation is documented in the serialization guide.

The measured 32.55× reduction in this PR comes from eliminating structural duplication; matching behavior is expected to remain identical after a round trip.

---

7. Finish with the executable contracts

The tests are organized around failure modes a serving-engine reviewer should care about:

Behavior under review	Where it is locked down
Scalar/batch parity; shared and per-request trees; sparse indices; nonzero roots; invalid inputs	Python speculative-decoding coverage
True nonzero storage offsets; worker failure propagation; pool reuse after failure	Native matcher coverage
Persistent ThreadPool generations and ParallelFor exceptions	Native pool coverage
Concurrent callers sharing one compiler	Compiler stress coverage
Compact round trips, functional parity, large schemas, and corrupt-cache rejection	Serialization coverage
Reproducible scalar-versus-batch throughput	Self-contained benchmark
Expected serving-engine call sequence	Engine integration guide

Suggested final review checklist

• Public scalar behavior remains backward compatible.
• All batch metadata is validated before worker launch.
• Tensor views honor both row offsets and DLTensor byte_offset.
• No matcher or output rows can be owned by two workers.
• Worker exceptions reach the API caller and do not poison pool reuse.
• Compiler parallelism remains bounded across concurrent outer callers.
• Every queued compiler task is drained before captured state unwinds.
• v16 reconstructs rule views only after complete structural validation.
• Tests cover parity, corruption, concurrency, failure, and recovery—not only the happy path.

Note

After this guided pass, the benchmark and validation sections below provide the quantitative results and exact commands/suite outcomes. The generated Cursor review summary remains at the bottom as an independent high-level map.

Relationship to open upstream work

Note

This is a point-in-time audit of the open mlc-ai/xgrammar pull requests performed on 2026-07-27. PR status and implementation details can change after this review.

Executive takeaway

No open upstream PR found in this audit duplicates either the native batched draft-tree traversal API or the compact shared-FSM v16 serialization format. There is, however, one direct overlap in compiler-owned thread pooling and three meaningful adjacent changes that touch the same tensor, compiler, matcher, or serialization paths.

Upstream PR	Relationship	Short version	Likely conflict surface
#438 — Use persistent thread pool	Direct overlap	Solves the same compiler-pool lifetime/oversubscription problem with a different pool abstraction.	cpp/grammar_compiler.cc, cpp/support/thread_pool.h, pool tests
#699 — Make fill_next_token_bitmask stride-aware	Partial tensor-layout overlap	Fixes row-stride handling for mask filling; this draft fixes byte_offset handling for draft traversal. Both are needed.	cpp/grammar_matcher.cc tensor helpers
#722 — Reduce grammar compilation and token-mask work	Complementary, broad file overlap	Reduces and deduplicates the work performed inside compilation while preserving the current serialized format.	Compiler, matcher, compiled grammar, grammar representation
#731 — On-demand token-mask compilation	Complementary policy change	Makes adaptive masks opt-in lazy and materializes them before serialization; this draft keeps eager semantics but bounds their scheduling.	Compiler, matcher, compiled grammar, FFI

Detailed compare/contrast

#438: persistent thread pool

Shared goal. Both changes replace per-compilation pool creation with one pool owned by each GrammarCompiler, preventing concurrent outer callers from multiplying native worker pools.

Different implementation. mlc-ai#438 introduces a TaskCounter/callback-oriented pool API and substantially reshapes the existing ThreadPool interface. This draft's compiler path keeps the existing submit/future model, compiles workloads smaller than 64 adaptive states inline, and drains all scheduled futures before captured compiler state can unwind. Its pool implementation additionally defines and tests worker-exception propagation and reuse after a failed generation; those failure/recovery semantics are not described in mlc-ai#438.

Reviewer action. Treat these as alternate implementations of the same pool-lifetime change, not two independent features. Pick or consolidate one pool abstraction, then preserve the inline threshold, bounded concurrency, drain-before-unwind behavior, and exception/recovery tests from this draft whichever implementation becomes the base.

#699: stride-aware token-mask filling

Shared area. Both changes harden DLTensor address calculation near token-mask generation.

Different semantics. mlc-ai#699 uses strides[0] when locating rows for fill_next_token_bitmask and accepts safe row-strided layouts. This draft's scalar traversal helper and validation correctly apply the tensor's base byte_offset for draft-tree traversal, while the new batch traversal deliberately requires contiguous inputs.

A stride describes the distance between elements or rows; byte_offset describes where the logical tensor begins in its backing allocation. Fixing one does not fix the other.

Reviewer action. Preserve both fixes when rebasing: retain mlc-ai#699's row-stride semantics for mask filling and this draft's base-offset correctness for scalar and batched draft traversal. The batch API should continue to reject unsupported non-contiguous traversal tensors explicitly rather than silently interpreting them as contiguous.

#722: compilation and token-mask work reduction

Shared area. mlc-ai#722 and this draft both touch grammar_compiler.cc, grammar_matcher.cc, compiled_grammar.cc, grammar.cc, and grammar representation internals, so textual conflicts are likely.

Different layer of optimization. mlc-ai#722 precomputes/shares Earley metadata, deduplicates equivalent adaptive-mask work and stored masks, and batch-proves token-prefix subtrees. It reduces how much work compilation performs and explicitly preserves the existing serialized format. This draft changes how that work is scheduled—one bounded compiler-owned pool, an inline cutoff for small jobs, futures, and safe draining—and its v16 representation intentionally changes the format so one complete optimized FSM is stored with compact per-rule views.

Reviewer action. These optimizations are conceptually composable despite the large merge-conflict surface. After combining them, rerun the 128-schema shared-compiler concurrency stress test plus serialization byte-equality, matcher-parity, corruption-rejection, and size benchmarks. mlc-ai#722's mask deduplication should be measured together with—not used as a substitute for—the shared-FSM serialization layout.

#731: on-demand token-mask compilation

Shared area. Both changes affect adaptive-mask compilation and serialization boundaries.

Different policy. mlc-ai#731 adds opt-in lazy mask compilation, populating an adaptive mask on first state visit and materializing the complete cache before serialization to preserve the existing format. This draft retains eager compilation semantics but makes large eager work bounded and persistent-pool-backed; its cache-format change concerns the optimized FSM structure rather than when adaptive masks are created.

Reviewer action. If the two are combined, test concurrent first-use mask population through multiple matchers and verify that serialize/deserialize still materializes and validates every required mask under v16. The native batch draft path should also be exercised against lazily populated masks because it increases concurrent matcher pressure.

What appears unique to this draft

The open-PR audit did not find another implementation of:

• one native BatchGrammarMatcher.batch_traverse_draft_tree operation with shared/per-request trees, sparse indices, per-matcher roots, flattened output ownership, and validation-before-scheduling;
• correct nonzero DLTensor byte_offset handling specifically for scalar and batched draft traversal;
• v16 serialization using one complete optimized FSM plus compact per-rule views, with structural corruption checks and measured 32.55× size reduction; or
• the exact reusable-pool exception contract tested here: propagate a worker failure to the caller, drain safely, clear the observed failure, and successfully reuse the same pool.

The batch API builds on the scalar traversal foundations already merged in #490 and #613; it lifts that work into a serving-engine-oriented native batch boundary rather than replacing it.

Other adjacent optimization PRs

The current cold-compilation/converter optimization stack—#711, #724, #727, #729, and #735—reduces parsing, converter, FSM construction, and rule-reuse costs. Those PRs are complementary rather than duplicates of the batch traversal, persistent scheduling semantics, or compact serialization introduced here, although rebases may conflict in compiler and grammar-representation files.

Suggested integration checklist

• Decide whether [Misc] Use persistent thread pool mlc-ai/xgrammar#438 or this draft supplies the base pool abstraction; retain bounded compiler ownership and the failure/recovery contract.
• Rebase the matcher helpers with fix: make fill_next_token_bitmask stride-aware mlc-ai/xgrammar#699 and preserve both row-stride and base-byte_offset correctness.
• Compose Reduce grammar compilation and allowed-token generation work mlc-ai/xgrammar#722's work elimination/deduplication with this draft's scheduling and serialization changes, then rerun combined stress and size measurements.
• If feat: add on-demand token mask compilation mlc-ai/xgrammar#731 is included, test lazy first-use population under native batch concurrency and verify v16 materialization/round trips.
• Re-run all scalar/batch parity, malformed-input, corrupt-cache, worker-failure, and shared-compiler concurrency gates after the final conflict resolution.

Changes

Batched speculative decoding

• Adds BatchGrammarMatcher.batch_traverse_draft_tree across the C++, TVM FFI, and Python APIs.
• Supports one shared 1D draft tree or per-request 2D trees, a 2D draft-token batch, sparse request-row indices, and per-matcher root_positions.
• Returns a completion flag per matcher and fills a flattened (batch_size * num_nodes, mask_words) bitmask.
• Rejects non-contiguous/wrong-device/wrong-dtype inputs, invalid roots, shape mismatches, out-of-range rows, and duplicate indices before launching workers.
• Extends scalar traverse_draft_tree with root_position and correctly honors DLTensor byte_offset for offset-backed contiguous tensors.
• Keeps one worker pool per BatchGrammarMatcher instead of recreating and joining threads on every call.
• Fuses bitmask reset with traversal and chunks work into at most one task per worker, eliminating a barrier and per-request queue overhead. The same chunking optimization is applied to batched next-token mask generation.

Bounded persistent grammar compilation

• Keeps one native worker pool per GrammarCompiler, shared by concurrent compile_* calls.
• Collects adaptive states before scheduling and compiles small grammars inline when queueing would cost more than it saves.
• Uses futures so native worker failures propagate to the caller, and drains outstanding work before unwinding captured compilation state.
• Prevents concurrent outer callers from multiplying native worker pools while retaining parallel compilation for large grammars.

Compact, validated serialization

• Bumps the serialization format from v15 to v16.
• Serializes the complete optimized FSM once, with each rule represented by a compact [start, ends, is_dfa, edge_num, node_num] view.
• Reconstructs per-rule views over the shared complete FSM during deserialization.
• Validates view counts, complete-FSM presence, start/end bounds, sorted unique ends, and size metadata before constructing views.
• Propagates nested grammar and adaptive-mask deserialization failures instead of silently accepting malformed data.

Tests, docs, and benchmark

• Adds scalar/batch parity, shared/per-request tree, sparse-index, nonzero-root, storage-offset, persistent-pool reuse, and invalid-input tests.
• Adds concurrent compiler stress and serial/parallel equality tests.
• Adds compact round-trip, large-schema, functional, and corruption-rejection tests.
• Adds native regressions for nonzero DLTensor offsets, worker-exception propagation, persistent-pool reuse after failure, and ParallelFor propagation.
• Documents engine integration, worker-pool behavior, compact serialization, cache-version implications, and the new API.
• Adds examples/benchmark/bench_batch_draft_tree.py, which is self-contained and downloads no model.

Performance results

Synthetic draft-tree benchmark, 96 requests × 6 nodes, vocabulary 101, median of 500 iterations in the same container:

Mode	Threads	Median	Requests/s	Speedup vs scalar
Scalar Python dispatch	1	0.279 ms	344,606	1.00×
Native batch	1	0.142 ms	675,048	1.96×
Native batch	2	0.132 ms	726,136	2.11×
Native batch	4	0.128 ms	748,089	2.17×
Native batch	8	0.133 ms	719,708	2.09×

Compiled serialization for the same 256-field schema and 259 optimized rules:

Revision	Serialized bytes	edge_aux_data copies
Fork main (v15)	22,735,820	260
This branch (v16)	698,546	1

That is a 96.93% size reduction (32.55× smaller) with exact serialize → deserialize → serialize equality and functional matcher parity.

An additional stress gate compiled and deserialized 128 distinct large schemas through one shared compiler with 16 concurrent callers and max_threads=12; serial/parallel output parity and all round trips passed.

Validation

• pre-commit run --all-files: pass
• ruff check python tests examples/benchmark/bench_batch_draft_tree.py: pass
• Native Release/RelWithDebInfo build with GCC 13: pass
• C++ test suite: 66 passed, 0 failed
• Speculative-decoding suite: 13 passed
• New serialization/compiler stress tests: pass
• Cursor Bugbot review after follow-up fixes: pass
• Full CPU/non-HuggingFace suite: 2,905 passed, 57 skipped, 6 failed, 622 deselected
  • The same six expected-output/regex-macro failures reproduce 6/6 on untouched shiftup-ai/xgrammar:main in the identical disposable environment; none touch code changed here.
• Confidentiality audit of the complete diff: no deployment names, credentials, private paths, endpoints, or customer/workload data.

Compatibility note

The compact representation intentionally changes the serialization version to v16. Existing v15 on-disk compiled-grammar caches must be cleared after upgrading, consistent with the documented versioning contract.

---

Note

High Risk
Touches core matcher/compiler concurrency, serialization format bump to v16 (cache break), and speculative decoding correctness; regressions would affect token masks or compiled grammar loads at scale.

Overview

Title

Batch draft traversal, persistent worker pools, and compact grammar caches (serialization v16)

Intent

Serving workloads need speculative draft-tree mask generation and grammar compilation to scale without per-call thread churn or multi-megabyte serialized caches. This PR packages those paths behind native batch APIs, shared bounded pools, and a compact FSM representation while failing closed on corrupt serialized data.

Details

• Batched speculative decoding: BatchGrammarMatcher.batch_traverse_draft_tree (C++, FFI, Python) traverses one tree per matcher, supports shared or per-request 2D trees, sparse indices, per-matcher root_positions, and flattened (batch_size × num_nodes) mask rows; scalar traverse_draft_tree gains root_position and honors DLTensor byte_offset.
• Matcher batching: BatchFillNextTokenBitmask reuses one pool per BatchGrammarMatcher, chunks work per worker, and propagates worker exceptions; new BatchTraverseDraftTree validates contiguity, dtypes, shapes, unique indices, and roots before scheduling.
• Compilation: GrammarCompiler keeps one persistent ThreadPool; small grammars inline adaptive mask work below a state threshold; large work uses Submit with drain-on-exception semantics.
• Serialization v16: Optimized grammars store one complete_fsm plus per-rule compact views [start, ends, is_dfa, edge_num, node_num]; deserialization validates counts, bounds, sorted unique ends, and nested grammar/mask errors; CompiledGrammar deserialization checks nested failures.
• ThreadPool: Execute tasks record first exception; Wait() rethrows and clears; ParallelFor calls Wait() before Join().
• Docs, tests, benchmark: Engine integration and serialization docs; C++/Python parity and corruption tests; bench_batch_draft_tree.py.

Roadmap for Reviewers

• Start with include/xgrammar/matcher.h and python/xgrammar/matcher.py for the public API surface (root_position, batch_traverse_draft_tree).
• cpp/grammar_matcher.cc — core traversal, bitmask offsets, batch validation, and pool chunking (large file; focus on TraverseDraftTreeRecursive, BatchTraverseDraftTree, BatchFillNextTokenBitmask).
• cpp/grammar_compiler.cc — persistent pool, kMinParallelStates, and exception-safe scheduling.
• cpp/grammar.cc + cpp/grammar_impl.h — compact serialize/deserialize and view reconstruction over complete_fsm.
• cpp/support/thread_pool.h — exception propagation and pool reuse (small, high leverage).
• cpp/compiled_grammar.cc — propagated deserialize errors (quick skim).
• cpp/tvm_ffi/tvm_ffi.cc — FFI bindings (mostly mechanical).
• Tests: tests/python/test_speculative_decoding.py, tests/python/test_serialization.py, tests/cpp/test_grammar_matcher.cc — skim for coverage intent.
• Docs / examples — skim unless reviewing user-facing contracts.

Diagrams

CompiledGrammar v16 (optimized)
┌─────────────────────────────────────┐
│ complete_fsm (once)                 │
│  edges + edge_aux_data              │
└─────────────────────────────────────┘
         ▲ views reference shared states
┌────────┴────────┬────────┬───────────┐
│ rule 0 view     │ rule 1 │ ...       │
│ start, ends,    │ ...    │           │
│ is_dfa, sizes   │        │           │
└─────────────────┴────────┴───────────┘

Batch draft traverse (one BatchGrammarMatcher pool)
matchers[0..N] ──► chunk rows ──► worker pool ──► mask rows [i*nodes:(i+1)*nodes]

^{Reviewed by Cursor Bugbot for commit 91db6bf. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit a332f5b. Configure here.}

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 91db6bf843

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Validate draft-tree link indices before dispatch

When either tree tensor contains a child or sibling value other than -1 or a valid node index, batch validation still succeeds and TraverseDraftTreeRecursive uses that value directly to index the token, tree, and bitmask buffers. A malformed request can therefore cause an out-of-bounds native read/write and crash the process instead of raising a validation error; validate every selected/shared tree row's link values before scheduling workers.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Reject root positions outside the int32 range

When Python supplies a root_position that fits in int64_t but not int32_t, this unchecked narrowing can turn it into an unrelated valid node—for example, 2**32 commonly becomes zero—so traversal silently uses the wrong root instead of reporting an out-of-range argument. Check the value against the int32_t bounds before casting it.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Reject aliases of the same matcher before parallel traversal

When the matcher list contains the same GrammarMatcher more than once with different row indices—for example, Python passes [matcher, matcher]—the unique-index check succeeds even though copied matchers share the same PImpl. Separate workers then concurrently call AcceptToken, Rollback, and mask generation on the same parser state and scratch buffers, causing a data race that can corrupt the matcher or produce invalid masks; reject duplicate implementation pointers before dispatch.

Useful? React with 👍 / 👎.