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

@@ -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,

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]