Fix spec-decode grammar mask crash under xgrammar 0.2.1 (strict tvm-ffi binding)

EAGLE verify's traverse_tree.dfs recurses with dfs(retrieve_next_token[curr], ...),
so curr (the fill_next_token_bitmask/accept_token index) degrades from a Python int
to a 0-dim torch tensor at recursion depth >= 2. xgrammar's Python signature is
unchanged across 0.1.27 -> 0.2.1, but the binding swapped pybind11 -> apache-tvm-ffi:
the old binding silently coerced the 0-dim tensor to int, the new one rejects it with
"Expected int but got ffi.Tensor", crashing the decode scheduler on grammar-constrained
(tool-call / structured-output) requests with draft depth >= 2.

Wrap the propagated indices in int(), matching upstream (PR #21722):
  - accept_token(int(draft_tokens[curr]))
  - dfs(int(retrieve_next_token[curr]), ...)
  - dfs(int(retrieve_next_sibling[curr]), ...)
Version-agnostic: correct on both 0.1.27 and 0.2.1.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-06-18 07:16:49 +00:00
parent 40738f59cf
commit a4ec2236ef

View File

@@ -605,14 +605,14 @@ def traverse_tree(
if accepted:
if curr != 0:
# Accept the current token
grammar.accept_token(draft_tokens[curr])
grammar.accept_token(int(draft_tokens[curr]))
if not grammar.is_terminated():
# Generate the bitmask for the current token
grammar.fill_vocab_mask(allocate_token_bitmask, curr)
if retrieve_next_token[curr] != -1:
# Visit the child node
dfs(
retrieve_next_token[curr],
int(retrieve_next_token[curr]),
retrieve_next_token,
retrieve_next_sibling,
curr,
@@ -625,7 +625,7 @@ def traverse_tree(
if retrieve_next_sibling[curr] != -1:
# Visit the sibling node
dfs(
retrieve_next_sibling[curr],
int(retrieve_next_sibling[curr]),
retrieve_next_token,
retrieve_next_sibling,
parent_pos,