From 187b700406c4734e42c7726c89f9c07d5b1cd906 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 21 Jun 2026 05:16:21 +0800 Subject: [PATCH] Keep speculative grammar traversal scalar-safe EAGLE verification can pass torch scalar tensors through the speculative tree traversal before calling grammar backends. xgrammar/tvm_ffi requires Python int token ids, so normalize traversal indices and draft token ids at the boundary while leaving tensor storage unchanged. Constraint: xgrammar/tvm_ffi rejects torch scalar tensors for GrammarMatcher.accept_token Rejected: Coerce tokens inside every grammar backend | the invalid value originates in speculative traversal and should be fixed before backend dispatch Confidence: high Scope-risk: narrow Directive: Keep grammar backend calls scalar-Python typed; do not pass torch scalar tensors through accept_token Tested: remote g0034 cjy-glm5-new PYTHONPATH=python python -m pytest -q test/registered/unit/speculative/test_spec_utils.py test/registered/unit/configs/test_nsa_index_layers.py test/registered/unit/models/test_deepseek_index_skip_weight_loading.py -> 19 passed Tested: remote g0034 cjy-glm5-new py_compile for modified runtime files Not-tested: live decode replay after this commit --- python/sglang/srt/speculative/spec_utils.py | 15 ++++--- .../unit/speculative/test_spec_utils.py | 45 +++++++++++++++++++ 2 files changed, 54 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/speculative/test_spec_utils.py diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 125fd3479..ed921dafe 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -587,13 +587,14 @@ def traverse_tree( retrieve_next_sibling: torch.Tensor, parent_pos: int, ): + curr = int(curr) if curr == 0: # the first token generated by the target model, and thus it is always # accepted from the previous iteration accepted = True else: parent_bitmask = allocate_token_bitmask[parent_pos] - curr_token_id = draft_tokens[curr] + curr_token_id = int(draft_tokens[curr]) if vocab_size and curr_token_id >= vocab_size: accepted = False else: @@ -605,14 +606,15 @@ def traverse_tree( if accepted: if curr != 0: # Accept the current token - grammar.accept_token(int(draft_tokens[curr])) + grammar.accept_token(curr_token_id) 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: + next_token = int(retrieve_next_token[curr]) + if next_token != -1: # Visit the child node dfs( - int(retrieve_next_token[curr]), + next_token, retrieve_next_token, retrieve_next_sibling, curr, @@ -622,10 +624,11 @@ def traverse_tree( # Rollback the current token grammar.rollback(1) - if retrieve_next_sibling[curr] != -1: + next_sibling = int(retrieve_next_sibling[curr]) + if next_sibling != -1: # Visit the sibling node dfs( - int(retrieve_next_sibling[curr]), + next_sibling, retrieve_next_token, retrieve_next_sibling, parent_pos, diff --git a/test/registered/unit/speculative/test_spec_utils.py b/test/registered/unit/speculative/test_spec_utils.py new file mode 100644 index 000000000..6846f0130 --- /dev/null +++ b/test/registered/unit/speculative/test_spec_utils.py @@ -0,0 +1,45 @@ +import torch + +from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject +from sglang.srt.speculative.spec_utils import traverse_tree + + +class _TypeCheckingGrammar(BaseGrammarObject): + def __init__(self): + super().__init__() + self.accepted = [] + self.rollback_calls = [] + + def accept_token(self, token: int) -> None: + if not isinstance(token, int): + raise TypeError(f"expected Python int token, got {type(token)!r}") + self.accepted.append(token) + + def rollback(self, k: int): + self.rollback_calls.append(k) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + vocab_mask[idx].fill_(-1) + + def allocate_vocab_mask(self, vocab_size: int, batch_size: int, device): + return torch.zeros((batch_size, (vocab_size + 31) // 32), dtype=torch.int32) + + +def test_traverse_tree_passes_python_int_tokens_to_grammar(): + grammar = _TypeCheckingGrammar() + retrieve_next_token = torch.tensor([1, -1], dtype=torch.int32) + retrieve_next_sibling = torch.tensor([-1, -1], dtype=torch.int32) + draft_tokens = torch.tensor([0, 5], dtype=torch.int64) + allocate_token_bitmask = torch.zeros((2, 2), dtype=torch.int32) + + traverse_tree( + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + grammar, + allocate_token_bitmask, + vocab_size=64, + ) + + assert grammar.accepted == [5] + assert grammar.rollback_calls == [1]