Yes! This is closer to the approach I took in my port of llama.cpp's grammar support to PyTorch: https://github.com/Shopify/torch-grammar/blob/main/torch_gramma

Yes! This is closer to the approach I took in my port of llama.cpp's grammar support to PyTorch: https://github.com/Shopify/torch-grammar/blob/main/torch_grammar/grammar_sampler.py#L203-L212 https://github.com/Shopify/torch-grammar/blob/main/torch_gra... ... it generates a tensor mapping each PDA stack to a map of which tokens are acceptable from that state. It seems like a much better way to do it than looping over the sampled tokens on each turn.