perf: batch draft traversal and compact grammar caches

[benmyles]

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 hardens DLTensor pointer/offset handling—including sampling-temperature outputs preserved during the upstream rebase—and adds fail-closed validation, regression coverage, integration documentation, and a reproducible synthetic benchmark. The branch is merged through upstream main at 3fb48bf; the compact wire format is now v17 so it cannot collide with upstream's existing v16 temperature format.

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. Every source link is pinned to the current PR head, so the code and line anchors below match the rebased implementation exactly.

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++ scalar API, batch C++ API, scalar Python wrapper, batch Python wrapper
2	Does one draft-tree traversal remain correct?	Recursive traversal, offset/alignment helper, 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 v17 cache much smaller, and how is corruption rejected?	Compact serialization, nested-error propagation
7	Which tests prove the end-to-end behavior?	Speculative decoding, offset + temperature composition, serialization

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 preserving upstream's timeout and optional temperature-output parameters. The batch matcher gains one operation that accepts all active requests at once. Start with the C++ scalar declaration 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.

Three correctness details were made explicit before batching:

1. All tensor data access adds the DLTensor base pointer and byte_offset, checks address overflow, and validates alignment. A tensor can be contiguous while still beginning at a nonzero storage offset; ignoring that field reads or writes the wrong part of the allocation.
2. Scalar validation checks the root, shape, dtype, device, true byte range, mask output, and optional temperature output before traversal. Invalid tensor metadata is rejected rather than interpreted optimistically.
3. The upstream temperature API composes with this PR's nonzero roots and offset-backed views: mask and temperature buffers are rebased independently, initialized through their effective pointers, and covered by focused composition regressions.

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

Upstream already assigned serialization version v16 to sampling-temperature state. This PR changes the optimized grammar layout again, so the rebased branch uses v17 rather than reusing an incompatible version tag.

The previous representation serialized overlapping FSM data once per rule. Large grammars therefore repeated the same edge auxiliary data hundreds of times. Version 17 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.

upstream v16                              this PR v17
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 set to v17, and the reason for skipping the already-occupied v16 tag is documented in the serialization guide.

Caution

This is an intentional on-disk cache break. A v16 or older compiled-grammar cache must be cleared during upgrade; silently interpreting two different layouts under one version would be unsafe.

On the current rebased head, the same 256-field/259-rule schema is 31.39× smaller than current upstream main (96.81% fewer serialized bytes) because the repeated FSM structure is eliminated.

---

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
Root-position, mask-offset, and temperature-offset composition	Temperature integration 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, versioning, 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.
• v17 reconstructs rule views only after complete structural validation and cannot be confused with upstream v16.
• 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. A clearly labeled historical Cursor summary remains at the bottom as an independent map of the pre-rebase implementation.

Relationship to open upstream work

Note

This audit was refreshed on 2026-07-30 against the open mlc-ai/xgrammar pull requests. This branch is already merged through upstream main commit 3fb48bf, so PRs merged before that point are incorporated here rather than merely compared below.

Executive takeaway

No open upstream PR duplicates either the native batched draft-tree traversal API or the compact shared-FSM v17 representation. One open PR directly overlaps compiler-owned thread pooling, while several others are complementary optimizations or correctness fixes in the same compiler/matcher/FSM surfaces.

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.	grammar_compiler.cc, thread_pool.h, pool tests
#699 — Make fill_next_token_bitmask stride-aware	Partial tensor-layout overlap	Handles row strides in mask filling; this PR handles base byte_offset and alignment in draft traversal and temperature outputs. Both are needed.	grammar_matcher.cc tensor helpers
#722 — Reduce grammar compilation and token-mask work	Complementary, broad overlap	Eliminates/deduplicates work inside compilation while this PR bounds how remaining work is scheduled and changes serialized FSM layout.	Compiler, matcher, compiled grammar, grammar representation
#731 — On-demand token-mask compilation	Complementary policy change	Makes adaptive masks lazily populated; this PR retains eager semantics but bounds scheduling.	Compiler, matcher, serialization, FFI
#753 — Compile grammar objects directly on cache misses	Complementary cache-path optimization	Removes a Grammar → text → Grammar round trip without changing native adaptive-mask scheduling.	grammar_compiler.cc, compiler tests
#758 — Accelerate large structural-tag compilation	Complementary parallelism with oversubscription risk	Parallelizes per-rule FSM construction and schema conversion; it must share a coherent thread budget with this PR's persistent adaptive-mask pool.	grammar_compiler.cc and compiler concurrency policy
#755 — Merge equivalent FSM states incrementally	Adjacent FSM optimization	Changes how the complete FSM is produced, while this PR changes how that FSM is viewed and serialized.	Low textual conflict; serialization size/equality validation

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 callers from multiplying native worker pools.

Different implementation. #438 introduces a TaskCounter/callback-oriented pool API and substantially reshapes the existing interface. This PR's compiler path keeps submit/future semantics, compiles workloads smaller than 64 adaptive states inline, and drains all scheduled futures before captured state can unwind. Its pool implementation also defines and tests worker-exception propagation and reuse after a failed generation.

Reviewer action. Treat these as alternate implementations of the same pool-lifetime change. Whichever abstraction is retained should preserve the inline threshold, bounded ownership, drain-before-unwind behavior, and exception/recovery tests.

#699: stride-aware token-mask filling

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

Different semantics. #699 uses strides[0] to locate rows for fill_next_token_bitmask and accepts safe row-strided layouts. This PR's effective-pointer helper applies base byte_offset, rejects address overflow/misalignment, and is used by scalar/batched draft traversal and temperature outputs. A stride describes distance between elements; byte_offset describes where the logical tensor begins. Fixing one does not fix the other.

Reviewer action. Preserve both fixes. The new batch draft API should continue to reject unsupported non-contiguous traversal tensors explicitly, while existing mask filling can retain #699's broader row-stride support.

#722: compilation and token-mask work reduction

#722 precomputes/shares Earley metadata, deduplicates adaptive-mask work and stored masks, and batch-proves token-prefix subtrees. It reduces how much work is performed and says it preserves the existing serialized format. This PR changes how eager work is scheduled and intentionally introduces the compact v17 format. The designs are conceptually composable despite a large textual conflict surface.

After combining them, rerun shared-compiler concurrency stress plus serialized byte equality, matcher parity, corruption rejection, and size measurements; #722's mask deduplication complements rather than replaces the shared-FSM layout.

#731: on-demand token-mask compilation

#731 makes adaptive masks opt-in lazy and materializes them before serialization. This PR keeps eager compilation but makes large eager work persistent-pool-backed. If combined, concurrent first-use mask population must be tested through multiple matchers, and serialization must still materialize every required mask before writing v17.

#753 and #758: newer compiler work

#753 removes a cache-miss conversion round trip and should compose with this PR with only a small grammar_compiler.cc conflict. #758 introduces additional compiler parallel regions; its work partitioning is complementary, but independent pools could reintroduce oversubscription. Integration should use one explicit concurrency budget and stress simultaneous outer compiler callers.

Already incorporated from updated main

• #727 (direct shared-FSM construction) and #729 (in-place optimizer passes) are merged and included in this head. The current 31.39× serialization comparison is measured against main after those improvements.
• #724 closed; its replacement #745 merged and is included here.
• The upstream sampling-temperature work in #730 is included. That merge claimed v16, which is why this PR now uses v17 and adds root/offset/temperature composition tests.

Other open adjacent work

• #706 caches Earley FSM state properties and #711 reduces cold compilation costs; both reduce work beneath or beside this PR's scheduling layer.
• #735 reuses converter rules and is complementary to compact serialization.
• #737 is an orthogonal matcher correctness fix with textual overlap in grammar_matcher.cc.
• #752 fixes start-state equivalence merging and should be validated together with Merge equivalent finite-state machine states incrementally #755 before using new FSM output as a serialization baseline.

What appears unique to this PR

The refreshed 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;
• effective-pointer handling for nonzero DLTensor byte_offset across scalar/batched draft traversal and temperature outputs, with overflow and alignment checks;
• v17 serialization using one complete optimized FSM plus compact per-rule views, with structural corruption checks and a current-main comparison; or
• the reusable-pool exception contract tested here: propagate worker failure, drain safely, clear the observed failure, and successfully reuse the same pool.

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

Suggested integration checklist

• Decide whether [Misc] Use persistent thread pool #438 or this PR supplies the base pool abstraction; retain bounded ownership and failure/recovery semantics.
• Preserve both fix: make fill_next_token_bitmask stride-aware #699 row-stride handling and this PR's base-offset/alignment handling.
• Compose Reduce grammar compilation and allowed-token generation work #722 work elimination with this PR's scheduling and v17 representation, then rerun stress and size tests.
• Give Accelerate large structural tag compilation #758's parallel compiler phases and adaptive-mask scheduling one explicit shared thread budget.
• If feat: add on-demand token mask compilation #731 is included, test lazy first-use population under native batch concurrency and verify complete v17 materialization.
• Re-run scalar/batch parity, temperature composition, malformed-input, corrupt-cache, worker-failure, and shared-compiler concurrency gates after 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, misaligned effective data addresses, 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 tree, mask, and optional temperature 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 compact serialization format from upstream v16 to v17; v16 is already occupied by sampling-temperature state.
• 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, mask/temperature storage-offset composition, 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 and misaligned 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

Current-head rerun at 48a446a, using the checked-in synthetic draft-tree benchmark with 96 requests × 6 nodes, vocabulary 101, and the median of 500 iterations:

Mode	Threads	Median	Requests/s	Speedup vs scalar
Scalar Python dispatch	1	0.300 ms	319,533	1.00×
Native batch	1	0.162 ms	592,149	1.85×
Native batch	2	0.159 ms	604,911	1.89×
Native batch	4	0.148 ms	650,043	2.03×
Native batch	8	0.151 ms	634,345	1.99×

The thread-count spread is small enough to show normal microbenchmark scheduling noise; the stable result is that one native batch boundary removes roughly half of scalar Python dispatch time.

Compiled serialization for the same 256-field schema and 259 optimized rules, measured from clean Release wheels built from current upstream main and the current PR head:

Revision	Format	Serialized bytes	edge_aux_data copies
Upstream main 3fb48bf	v16	22,762,528	260
This PR 48a446a	v17	725,254	1

That is a 96.81% size reduction (31.39× smaller) with exact serialize → deserialize → serialize equality and functional matcher parity. The checked-in shared-compiler stress regression also compiles 32 distinct large schemas through one compiler with eight concurrent callers, compares serial/parallel output, deserializes every result, and verifies pool reuse after idle.

Validation

All results below are from the current source state at 48a446a:

Note

48a446a changes test annotations only. The native build, GPU, and serialization results were established on parent 3e76ba2, whose implementation sources are byte-for-byte identical; the synthetic benchmark and exact Python 3.8 CI suite were rerun on the current head.

• pre-commit run --all-files: pass
• ruff check python tests examples/benchmark/bench_batch_draft_tree.py: pass
• Clean native RelWithDebInfo build with GCC 13 and warnings as errors: pass
• C++ test suite: 67 passed, 0 failed
• Focused serialization + temperature + speculative-decoding suite: 65 passed, 1 Hugging Face-token test deselected, 0 failed
• Exact Python 3.8 CI reproduction (Python 3.8.20, PyTorch 2.4.1, Triton 3.0.0, CPU-only): 3,012 passed, 58 skipped, 339 deselected, 0 failed
• Full CPU/non-HuggingFace suite: 3,012 passed, 57 skipped, 622 deselected, 0 failed
• Full GPU-enabled/non-HuggingFace suite on 2× NVIDIA B300 (compute capability 10.3): 3,068 passed, 1 unrelated MLX skip, 622 deselected, 0 failed
  • PyTorch 2.9.1+cu128, Triton 3.5.1, driver 580.159.03.
  • Triton's bundled CUDA 12.8 assembler predates sm_103a; the test used the host CUDA 13.0 assembler via TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas.
• Current-main/current-head serialization size comparison: pass; clean wheels produced 22,762,528 bytes (v16) versus 725,254 bytes (v17).
• Copilot alignment-review follow-up: fb48ca5, with scalar/batch int32/int64 misalignment regressions.
• Post-rebase version-collision fix: 3e76ba2; v17 tests and upgrade documentation included.
• Python 3.8 CI follow-up: 48a446a replaces runtime-evaluated PEP 585/collections.abc annotations in the upstream max_chars tests with Python 3.8-compatible typing forms.
• Confidentiality audit of the complete diff and PR prose: no deployment names, credentials, private endpoints, or customer/workload data.

Compatibility note

The compact representation intentionally uses serialization version v17. Upstream v16 added sampling-temperature state, and this PR changes the optimized grammar layout again; reusing v16 would allow incompatible documents to share one version identifier. Existing v16 or older on-disk compiled-grammar caches must therefore be cleared after upgrading, consistent with the documented versioning contract.

Warning

Historical automated summary. Cursor generated the block below for commit 91db6bf, before the upstream merge and the v17 collision fix. Use the guided review, overlap audit, benchmarks, and validation above for the current head 48a446a.

---

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.}

[Copilot]

Pull request overview

This PR introduces serving-path performance improvements by adding a native batched draft-tree traversal API for speculative decoding, reusing bounded native thread pools across requests/compilations, and compacting compiled-grammar serialization (v16) to avoid per-rule FSM duplication while validating corruption fail-closed.

Changes:

• Add root_position to scalar draft-tree traversal and a new BatchGrammarMatcher.batch_traverse_draft_tree API across C++/FFI/Python.
• Reuse persistent native worker pools (matcher + compiler), with exception propagation and pool reusability semantics.
• Bump serialization format to v16 with compact per-rule FSM views over one shared complete FSM, plus validation, docs, tests, and a benchmark.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/python/test_speculative_decoding.py	Adds batch traversal parity/validation/offset tests.
tests/python/test_serialization.py	Updates to v16 + corruption/roundtrip tests for compact views.
tests/python/test_grammar_compiler.py	Adds concurrent compilation stress test for persistent pool reuse.
tests/cpp/test_thread_pool.cc	Adds tests for Wait() rethrow/reuse and ParallelFor exception propagation.
tests/cpp/test_grammar_matcher.cc	Adds native regressions for byte_offset, batch failure propagation, pool reuse.
python/xgrammar/matcher.py	Exposes root_position and adds Python batch traversal wrapper/docs.
include/xgrammar/matcher.h	Extends public C++ API: scalar root_position, batch traverse method.
examples/benchmark/README.md	Documents new batch draft-tree benchmark.
examples/benchmark/bench_batch_draft_tree.py	Adds reproducible scalar-vs-batch traversal benchmark.
docs/using_xgrammar/serialization.md	Documents compact FSM view format and validation behavior.
docs/using_xgrammar/engine_integration.md	Documents batch traversal integration + root_position.
docs/start/workflow_of_xgrammar.md	Updates compilation threading/pool semantics documentation.
cpp/tvm_ffi/tvm_ffi.cc	Wires new APIs through TVM FFI bindings.
cpp/support/thread_pool.h	Adds Execute-exception capture + Wait() rethrow/clear; ParallelFor waits before join.
cpp/support/json_serializer.h	Bumps global serialization version constant to v16.
cpp/grammar.cc	Implements compact per-rule FSM view serialization/deserialization with validation.
cpp/grammar_matcher.cc	Adds root_position traversal + batch traversal implementation + persistent pool usage.
cpp/grammar_impl.h	Switches Grammar (Impl) JSON (de)serialization to custom functions.
cpp/grammar_compiler.cc	Introduces persistent compiler-owned pool and bounded scheduling/drain-on-failure.
cpp/fsm.h	Exposes GetIsDFA() accessor used by compact view serialization.
cpp/compiled_grammar.cc	Propagates nested deserialization errors instead of ignoring them.

Comments suppressed due to low confidence (2)

cpp/grammar_matcher.cc:190

• Pointer arithmetic that incorporates byte_offset is immediately reinterpret_casted to int32_t* without validating alignment. If a caller passes a contiguous DLTensor with an unaligned byte_offset (or base pointer), this becomes undefined behavior on platforms requiring alignment.

Consider validating alignment (and failing closed) before casting.

  auto* data =
      static_cast<char*>(token_bitmask.data) + static_cast<std::size_t>(token_bitmask.byte_offset);
  return reinterpret_cast<int32_t*>(data) + index * buffer_size;

cpp/grammar_matcher.cc:2219

• The byte_offset-adjusted pointers for the int64 draft-tree tensors are reinterpret_casted to int64_t* without validating alignment. If a caller passes a view with an unaligned byte_offset (or base pointer), this is undefined behavior.

Please validate that (uintptr_t(data) + byte_offset) preserves int64_t alignment before casting.

  const auto* next_token_data = reinterpret_cast<const int64_t*>(
      static_cast<const char*>(retrieve_next_token->data) + retrieve_next_token->byte_offset
  );
  const auto* next_sibling_data = reinterpret_cast<const int64_t*>(
      static_cast<const char*>(retrieve_next_sibling->data) + retrieve_next_sibling->byte_offset

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[benmyles]

CUDA validation on Blackwell

I ran the CUDA-enabled test paths against the current PR head (5891cca2cfcfc012c7f9d50e2741d5d92e286ec9) on the following local environment:

• 2× NVIDIA B300 SXM6 AC (compute capability 10.3)
• NVIDIA driver 580.159.03
• PyTorch 2.9.1+cu128 (CUDA 12.8 runtime)
• Triton 3.5.1
• CUDA 13.0 ptxas 13.0.88

Note

Triton 3.5.1 bundles CUDA 12.8 ptxas, which predates the sm_103a target used for B300. I pointed Triton at the host CUDA 13 assembler with TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas; this is a local toolchain selection, not a source change or test workaround.

Focused CUDA/Triton bitmask coverage

TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas \
  /root/venv-scratch/.venv/bin/python -m pytest \
  tests/python/test_token_bitmask_operations.py -q

Result: 101 passed, 1 skipped, 0 failed in 134.84s. The single skip is the unrelated MLX backend because MLX is not installed. This exercised the native CUDA, Triton, and torch.compile implementations, including batched/strided indices, large vocabularies, non-power-of-two vocabulary sizes, and fp32/fp16/bf16 cases.

Full GPU-enabled Python suite

TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas \
  /root/venv-scratch/.venv/bin/python -m pytest \
  --ignore=3rdparty -m 'not hf_token_required' -q -o addopts=''

Result: 3,068 passed, 1 skipped, 622 deselected, 0 failed in 146.13s. The deselected tests require a Hugging Face token; the single skip is again MLX. With CUDA visible, this also exercises the GPU-backed matcher tests and the DLTensor/storage-offset paths affected by this PR.

No CUDA-specific correctness issues were found.

[benmyles]

CI follow-up: Python 3.8 matrix fixed

The rerun exposed three failed matrix cells that GitHub's top-level checks view collapsed under duplicate job names:

• Ubuntu x86_64 / Python 3.8
• Ubuntu ARM64 / Python 3.8
• Windows / Python 3.8

All three had the same root cause in tests/python/test_max_chars.py: runtime evaluation of list[int], followed by collections.abc.Sequence[int], is not supported on Python 3.8. The native build and all 67 C++ tests passed in each affected job before Python test collection stopped.

Commit 48a446a uses the Python 3.8-compatible typing.List and typing.Sequence forms.

I reproduced the affected CI configuration locally with Python 3.8.20, PyTorch 2.4.1, Triton 3.0.0, no Hugging Face token, and CUDA hidden to match the CPU-only Actions runners:

3012 passed, 58 skipped, 339 deselected, 0 failed

Additional gates on the current head:

• pre-commit run --all-files: pass
• Ruff: pass
• git diff --check: pass
• Python 3.8 collection: 3,069 selected tests collected with no errors

The PR body and all guided-review source links have also been refreshed to the current head. The newly triggered upstream workflows are awaiting external-fork approval (action_required).

[benmyles]

@Seven-Streams pushed up some test fixes for the few that were failing. lmk if you need anything else from me