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, 14:45 UTC
[Core] Upgrade to xgrammar 0.1.18, add cache size limit (#16283)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Apr 8, 2025 at 22:13 UTC
cb84e45ac75b42ba6795145923e8eb323bb825ad
4 files changed
+22
-3
requirements/common.txt
+1
-1
@@ -22,7 +22,7 @@ lm-format-enforcer >= 0.10.11, < 0.11
22
llguidance >= 0.7.9, < 0.8.0; platform_machine == "x86_64" or platform_machine == "arm64" or platform_machine == "aarch64"
23
outlines == 0.1.11
24
lark == 1.2.2
25
-xgrammar == 0.1.17; platform_machine == "x86_64" or platform_machine == "aarch64"
25
+xgrammar == 0.1.18; platform_machine == "x86_64" or platform_machine == "aarch64"
26
typing_extensions >= 4.10
27
filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/317
28
partial-json-parser # used for parsing partial JSON outputs
vllm/envs.py
+7
@@ -106,6 +106,7 @@ if TYPE_CHECKING:
106
VLLM_TPU_DISABLE_TOPK_TOPP_OPTIMIZATION: bool = False
107
VLLM_TPU_BUCKET_PADDING_GAP: int = 0
108
VLLM_USE_DEEP_GEMM: bool = False
109
+ VLLM_XGRAMMAR_CACHE_MB: int = 0
110
111
112
def get_default_cache_root():
@@ -697,6 +698,12 @@ environment_variables: dict[str, Callable[[], Any]] = {
698
# Allow use of DeepGemm kernels for fused moe ops.
699
"VLLM_USE_DEEP_GEMM":
700
lambda: bool(int(os.getenv("VLLM_USE_DEEP_GEMM", "0"))),
701
+
702
+ # Control the cache sized used by the xgrammar compiler. The default
703
+ # of 512 MB should be enough for roughly 1000 JSON schemas.
704
+ # It can be changed with this variable if needed for some reason.
705
+ "VLLM_XGRAMMAR_CACHE_MB":
706
+ lambda: int(os.getenv("VLLM_XGRAMMAR_CACHE_MB", "512")),
707
}
708
709
# end-env-vars-definition
vllm/model_executor/guided_decoding/xgrammar_decoding.py
+7
-1
@@ -10,6 +10,7 @@ from typing import TYPE_CHECKING, Any, List
10
11
import torch
12
13
+import vllm.envs
14
from vllm.logger import init_logger
15
16
try:
@@ -131,8 +132,13 @@ class GrammarCompilerCache:
132
encoded_vocab=config_data.encoded_vocab,
133
metadata=config_data.metadata,
134
)
135
+ cache_size = vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024
136
cls._cache[cache_key] = xgr.GrammarCompiler(
135
- tokenizer_info, max_threads=config.max_threads)
137
+ tokenizer_info,
138
+ max_threads=config.max_threads,
139
+ cache_enabled=True,
140
+ cache_limit_bytes=cache_size,
141
+ )
142
143
return cls._cache[cache_key]
144
vllm/v1/structured_output/backend_xgrammar.py
+7
-1
@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING
5
6
import torch
7
8
+import vllm.envs
9
from vllm.config import VllmConfig
10
from vllm.logger import init_logger
11
from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
@@ -76,7 +77,12 @@ class XgrammarBackend(StructuredOutputBackend):
77
tokenizer,
78
vocab_size=self.vocab_size,
79
)
79
- self.compiler = xgr.GrammarCompiler(tokenizer_info, max_threads=8)
80
+ self.compiler = xgr.GrammarCompiler(
81
+ tokenizer_info,
82
+ max_threads=8,
83
+ cache_enabled=True,
84
+ cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024,
85
+ )
86
87
def compile_grammar(self, request_type: StructuredOutputOptions,
88
grammar_spec: str) -> StructuredOutputGrammar: