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 15, 19:45 UTC
[V1] Prevent xgrammar from breaking TPU support (#14575)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Mar 10, 2025 at 19:06 UTC
04421dff8a05a81bbd5ed19545175066163a4082
2 files changed
+11
-2
vllm/v1/engine/processor.py
+4
@@ -4,6 +4,7 @@ import time
4
from collections.abc import Mapping
5
from typing import Optional, Union
6
7
+import vllm.platforms
8
from vllm.config import VllmConfig
9
from vllm.inputs import (INPUT_REGISTRY, InputRegistry, ProcessorInputs,
10
PromptType, SingletonInputsAdapter)
@@ -133,6 +134,9 @@ class Processor:
134
if self.vllm_config.speculative_config:
135
raise ValueError("Structured output is not supported with "
136
"speculative decoding.")
137
+ if vllm.platforms.current_platform.is_tpu():
138
+ raise ValueError("Structured output is not supported on TPU.")
139
+
140
validate_structured_output_request(params)
141
142
def process_inputs(
vllm/v1/structured_output/__init__.py
+7
-2
@@ -17,6 +17,7 @@ from vllm.v1.structured_output.grammar import (Grammar, StructuredOutputKey,
17
if TYPE_CHECKING:
18
import numpy as np
19
import numpy.typing as npt
20
+ import torch
21
import xgrammar as xgr
22
23
from vllm.v1.request import Request
@@ -53,8 +54,7 @@ class StructuredOutputManager:
54
# compilation, so we set it to half the number of CPUs.
55
max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
56
self.executor = ThreadPoolExecutor(max_workers=max_workers)
56
- self._grammar_bitmask = xgr.allocate_token_bitmask(
57
- self.vllm_config.scheduler_config.max_num_seqs, self.vocab_size)
57
+ self._grammar_bitmask: Optional[torch.Tensor] = None
58
59
def __getitem__(self, key: StructuredOutputKey) -> Optional[Grammar]:
60
# We need to pop and re-insert the grammar here for LRU cache
@@ -134,6 +134,11 @@ class StructuredOutputManager:
134
if not structured_output_request_ids:
135
return None
136
137
+ if self._grammar_bitmask is None:
138
+ self._grammar_bitmask = xgr.allocate_token_bitmask(
139
+ self.vllm_config.scheduler_config.max_num_seqs,
140
+ self.vocab_size)
141
+
142
# Fill the bitmask using the index of each request equal to its
143
# position in the batch. Resize the bitmask down to the size of
144
# the batch.