GitLab

**[Guided decoding performance optimization]** Sending the guided
decoding bitmask in xgrammar to the GPU
(`self.token_bitmask.to(scores.device)`) is a blocking operation that
prevents the CPU from pre-launching the sampler kernels. The CPU waits
until decode is complete, then copies the bitmask over. This PR changes
the operation to async via setting `non-blocking=True`.

(Current) The CPU is blocked on a `cudaStreamSynchronize` and only
pre-empts the sampling kernels after bitmask application. Below is the
Nsys profile for one decode phase from Llama 3.1 8B.

![image](https://github.com/user-attachments/assets/8997eae1-b822-4f52-beb8-ef19a7c6b824)

With the optimization, this is no longer the case:

![image](https://github.com/user-attachments/assets/6d5ea83f-f169-4f98-a8c1-41c719b3e1e7

)

---------
Signed-off-by: Ryan N <ryan.nguyen@centml.ai>

Signed-off-by: Jason Cheng <jasoncky96@gmail.com>

Signed-off-by: Michal Adamczyk <madamczyk@habana.ai>

Signed-off-by: mgoin <michael@neuralmagic.com>

Signed-off-by: youkaichao <youkaichao@gmail.com>

Signed-off-by: mgoin <michael@neuralmagic.com>

Signed-off-by: mgoin <michael@neuralmagic.com>

Signed-off-by: Russell Bryant <rbryant@redhat.com>

Signed-off-by: Russell Bryant <rbryant@redhat.com>

Signed-off-by: mgoin <michael@neuralmagic.com>

Signed-off-by: mgoin <michael@neuralmagic.com>

Signed-off-by: mgoin <michael@neuralmagic.com>

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
Signed-off-by: mgoin <michael@neuralmagic.com>
Co-authored-by: mgoin <michael@neuralmagic.com>

Signed-off-by: Vasily Alexeev <alvasian@yandex.ru>

Signed-off-by: Travis Johnson <tsjohnso@us.ibm.com>
Co-authored-by: pavlo-ruban <pavlo.ruban@servicenow.com>
Co-authored-by: Nick Hill <nickhill@us.ibm.com>

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>
Co-authored-by: Nick Hill <nickhill@us.ibm.com>

Signed-off-by: Aaron Pham <contact@aarnphm.xyz>
Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Co-authored-by: constellate <constellate@1-ai-appserver-staging.codereach.com>
Co-authored-by: Kyle Mistele <kyle@constellate.ai>

Co-authored-by: Cyrus Leung <cyrus.tl.leung@gmail.com>

Signed-off-by: Joe Runde <Joseph.Runde@ibm.com>
Co-authored-by: Joe Runde <Joseph.Runde@ibm.com>
Co-authored-by: Joe Runde <joe@joerun.de>
Co-authored-by: Nick Hill <nickhill@us.ibm.com>
Co-authored-by: Simon Mo <simon.mo@hey.com>

Co-authored-by: Simon Mo <simon.mo@hey.com>
Co-authored-by: Breno Faria <breno.faria@intrafind.com>

Co-authored-by: Ubuntu <ubuntu@ip-172-31-13-147.ec2.internal>

Co-authored-by: SangBin Cho <sangcho@sangcho-LT93GQWG9C.local>

Co-authored-by: Simon Mo <simon.mo@hey.com>