I implemented this for PyTorch too at https://github.com/Shopify/torch-grammar. I have a hacked version of text-generation-inference that uses it—happy to share that if it’s useful to anyone.
I've been meaning to play around with dumping the token probability vectors inside one of the LLM UIs. Having a diff starting point would help a bunch.