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] Delay all xgrammar usage until needed (#14616)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Mar 11, 2025 at 16:21 UTC
61a01b27a7bd3aa2c6e067d7c4624e020615722b
1 file changed
+20
-13
vllm/v1/structured_output/__init__.py
+20
-13
@@ -14,7 +14,6 @@ from vllm.v1.structured_output.grammar import Grammar, StructuredOutputOptions
14
if TYPE_CHECKING:
15
import numpy as np
16
import numpy.typing as npt
17
- import torch
17
import xgrammar as xgr
18
19
from vllm.v1.request import Request
@@ -27,14 +26,18 @@ logger = init_logger(__name__)
26
class StructuredOutputManager:
27
28
def __init__(self, vllm_config: VllmConfig):
30
- tokenizer_group = init_tokenizer_from_configs(
31
- model_config=vllm_config.model_config,
32
- scheduler_config=vllm_config.scheduler_config,
33
- parallel_config=vllm_config.parallel_config,
34
- lora_config=vllm_config.lora_config) # type: ignore[arg-type]
35
- tokenizer_group.ping()
29
self.vocab_size = vllm_config.model_config.get_vocab_size()
30
self.vllm_config = vllm_config
31
+ self.init_complete = False
32
+
33
+ def _delayed_init(self):
34
+ """Initialization delayed until we know it is needed."""
35
+ tokenizer_group = init_tokenizer_from_configs(
36
+ model_config=self.vllm_config.model_config,
37
+ scheduler_config=self.vllm_config.scheduler_config,
38
+ parallel_config=self.vllm_config.parallel_config,
39
+ lora_config=self.vllm_config.lora_config) # type: ignore[arg-type]
40
+ tokenizer_group.ping()
41
42
tokenizer = tokenizer_group.get_lora_tokenizer(None)
43
tokenizer_info = xgr.TokenizerInfo.from_huggingface(
@@ -47,12 +50,21 @@ class StructuredOutputManager:
50
# compilation, so we set it to half the number of CPUs.
51
max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
52
self.executor = ThreadPoolExecutor(max_workers=max_workers)
50
- self._grammar_bitmask: Optional[torch.Tensor] = None
53
+ self._grammar_bitmask = xgr.allocate_token_bitmask(
54
+ self.vllm_config.scheduler_config.max_num_seqs, self.vocab_size)
55
+
56
+ self.init_complete = True
57
58
def grammar_init(self, request: Request) -> None:
59
if request.structured_output_request is None:
60
return
61
62
+ # The first time this is called, we need to finish initialization
63
+ # of xgrammar. We defer it to avoid the import of xgrammar and
64
+ # initialization cost if it is not going to be used.
65
+ if not self.init_complete:
66
+ self._delayed_init()
67
+
68
grammar: Future[Grammar] = self.executor.submit(
69
self._async_create_grammar, request)
70
request.structured_output_request.grammar = grammar # type: ignore[assignment]
@@ -100,11 +112,6 @@ class StructuredOutputManager:
112
if not structured_output_request_ids:
113
return None
114
103
- if self._grammar_bitmask is None:
104
- self._grammar_bitmask = xgr.allocate_token_bitmask(
105
- self.vllm_config.scheduler_config.max_num_seqs,
106
- self.vocab_size)
107
-
115
# Fill the bitmask using the index of each request equal to its
116
# position in the batch. Resize the bitmask down to the size of
117
# the batch.