xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-01-01 17:58:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BUGFIX] Move scores to float32 in case of running xgrammar on cpu (#12152)

Browse Source 
Signed-off-by: Michal Adamczyk <madamczyk@habana.ai>
This commit is contained in:
Michal Adamczyk 2025-01-19 04:12:05 +01:00 committed by GitHub
parent 7a8a48d51e
commit 4e94951bb1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
vllm/model_executor/guided_decoding/xgrammar_decoding.py
View File

@ -298,8 +298,11 @@ class XGrammarLogitsProcessor:
# token_bitmask is a CPU tensor for use with accept_token and
# fill_next_token_bitmask so we move it to the device of scores
device_type = scores.device.type
dtype = scores.dtype
if device_type != "cuda":
scores = scores.to("cpu").unsqueeze(0)
# xgrammar on cpu only supports float32 scores
# see: https://github.com/mlc-ai/xgrammar/blob/c1b64920cad24f44f235778c1c00bb52d57da01a/python/xgrammar/kernels/apply_token_bitmask_inplace_cpu.py#L22
scores = scores.to("cpu").float().unsqueeze(0)
# Note: In this method, if the tensors have different dimensions
# on CPU device fails, but on GPU it runs without error. Hence the
@ -307,7 +310,7 @@ class XGrammarLogitsProcessor:
xgr.apply_token_bitmask_inplace(scores,
self.token_bitmask.to(scores.device))
if device_type != "cuda":
scores = scores.to(device_type).squeeze()
scores = scores.to(dtype).to(device_type).squeeze()
return scores