Motivation

This PR moves the sampler out of the CUDA graph and removes the synchronization in the cuda graph replay.

• Vocab mask, penalties calculation are done between model forward and sampling
• Make sampling a part of model_runner.

Modifications

Checklist

• Format your code according to the Contributor Guide.
• Add unit tests as outlined in the Contributor Guide.
• Update documentation as needed, including docstrings or example tutorials.