A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 6, 02:00 UTC
[Bugfix] Fix structured output crash on CPU due to pin_memory=True (#37706)
Signed-off-by: Willy Hardy <whardy@redhat.com> Signed-off-by: Will Hardy <whardy@redhat.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Willy Hardy committed
Mar 24, 2026 at 13:44 UTC
057fc94cbdcfbb484d4aac5f244aee1eefe17f0e
1 file changed
+22
-16
vllm/v1/structured_output/utils.py
+22
-16
@@ -17,6 +17,7 @@ from diskcache import Cache
17
import vllm.envs as envs
18
from vllm.logger import init_logger
19
from vllm.utils.import_utils import LazyLoader
20
+from vllm.utils.platform_utils import is_pin_memory_available
21
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
22
23
if TYPE_CHECKING:
@@ -106,28 +107,33 @@ def apply_grammar_bitmask(
107
# since the bitmask is already aligned with the logits.
108
skip_out_indices = len(out_indices) == logits.shape[0]
109
109
- index_tensor = None
110
- if not skip_out_indices:
111
- # xgrammar expects a python list of indices but it will actually work with
112
- # a tensor. If we copy the tensor ourselves here we can do it in a non_blocking
113
- # manner and there should be no cpu sync within xgrammar.
114
- index_tensor = torch.tensor(
115
- out_indices, dtype=torch.int32, device="cpu", pin_memory=True
116
- )
117
- index_tensor = index_tensor.to(logits.device, non_blocking=True)
110
+ if not logits.is_cpu:
111
+ index_tensor = None
112
+ if not skip_out_indices:
113
+ # xgrammar expects a python list of indices but it will actually work with
114
+ # a tensor. If we copy the tensor ourselves here we can do it in a
115
+ # non_blocking manner and there should be no cpu sync within xgrammar.
116
+ pin_memory = is_pin_memory_available()
117
+ index_tensor = torch.tensor(
118
+ out_indices, dtype=torch.int32, device="cpu", pin_memory=pin_memory
119
+ )
120
+ index_tensor = index_tensor.to(logits.device, non_blocking=True)
121
+
122
+ xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)
123
+ return
124
125
+ # CPU case, use list for indices.
126
+ indices = None if skip_out_indices else out_indices
127
# Handle dtype conversion for CPU (older xgrammar CPU kernels require float32)
128
# See: https://github.com/vllm-project/vllm/issues/31901
121
- if logits.device.type == "cpu" and logits.dtype != torch.float32:
129
+ if logits.dtype != torch.float32:
130
# Convert to float32, apply bitmask, then convert back
123
- logits_float32 = logits.to(torch.float32)
124
- xgr.apply_token_bitmask_inplace(
125
- logits_float32, grammar_bitmask, indices=index_tensor
126
- )
131
+ logits_fp32 = logits.to(torch.float32)
132
+ xgr.apply_token_bitmask_inplace(logits_fp32, grammar_bitmask, indices=indices)
133
# Copy the modified values back to the original tensor
128
- logits.copy_(logits_float32.to(logits.dtype))
134
+ logits.copy_(logits_fp32.to(logits.dtype))
135
else:
130
- xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)
136
+ xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=indices)
137
138
139
class OutlinesVocabulary: