187b700406
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
46 lines
1.4 KiB
Python
46 lines
1.4 KiB
Python
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]
|