[Bugfix] Support MistralCommonBackend tokenizers in structured output (#52720)
Signed-off-by: Thanh Phan <thanhphantuan1110@gmail.com>
Thanh Phan committed
Aug 21, 2026 at 01:08 UTC
d56bbf3995b820864a94971264cd5d2e5265261a
3 files changed
+118
-2
tests/v1/structured_output/test_mistral_common_tokenizer.py
new
+102
@@ -0,0 +1,102 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Structured output with a `MistralCommonBackend` tokenizer."""
4
+
5
+import pytest
6
+from transformers import AutoTokenizer, MistralCommonBackend
7
+
8
+from vllm.config import StructuredOutputsConfig, VllmConfig
9
+from vllm.tokenizers import get_tokenizer
10
+from vllm.tokenizers.mistral import MistralTokenizer
11
+from vllm.v1.structured_output.backend_outlines import OutlinesBackend
12
+from vllm.v1.structured_output.backend_types import StructuredOutputOptions
13
+from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
14
+from vllm.v1.structured_output.utils import (
15
+ _reduced_vocabulary,
16
+ maybe_wrap_mistral_common_tokenizer,
17
+)
18
+
19
+TOKENIZER = "mistralai/Mistral-Nemo-Instruct-2407"
20
+JSON_SCHEMA = (
21
+ '{"type": "object", "properties": {"x": {"type": "integer"}}, '
22
+ '"required": ["x"], "additionalProperties": false}'
23
+)
24
+VALID_DOCUMENT = '{"x": 1}'
25
+
26
+BACKENDS = [("xgrammar", XgrammarBackend), ("outlines", OutlinesBackend)]
27
+
28
+
29
+@pytest.fixture(scope="module")
30
+def loaded_tokenizer():
31
+ tokenizer = get_tokenizer(tokenizer_name=TOKENIZER, tokenizer_mode="hf")
32
+ assert isinstance(tokenizer, MistralCommonBackend)
33
+ return tokenizer
34
+
35
+
36
+@pytest.fixture(scope="module")
37
+def wrapped_tokenizer(loaded_tokenizer):
38
+ return maybe_wrap_mistral_common_tokenizer(loaded_tokenizer)
39
+
40
+
41
+def _encode(tokenizer, text):
42
+ special = set(tokenizer.all_special_ids)
43
+ return [tid for tid in tokenizer.encode(text) if tid not in special]
44
+
45
+
46
+def _backend(cls, name, tokenizer):
47
+ vllm_config = VllmConfig(
48
+ structured_outputs_config=StructuredOutputsConfig(backend=name)
49
+ )
50
+ return cls(vllm_config, tokenizer=tokenizer, vocab_size=tokenizer.vocab_size)
51
+
52
+
53
+def test_wrapping_produces_a_mistral_tokenizer(loaded_tokenizer, wrapped_tokenizer):
54
+ assert isinstance(wrapped_tokenizer, MistralTokenizer)
55
+ assert len(wrapped_tokenizer.vocab) == loaded_tokenizer.vocab_size
56
+
57
+
58
+def test_other_tokenizers_are_left_alone():
59
+ tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
60
+ assert maybe_wrap_mistral_common_tokenizer(tokenizer) is tokenizer
61
+
62
+
63
+def test_wrapping_fixes_the_outlines_vocabulary(loaded_tokenizer, wrapped_tokenizer):
64
+ unwrapped_vocab = _reduced_vocabulary(loaded_tokenizer)
65
+ wrapped_vocab = _reduced_vocabulary(wrapped_tokenizer)
66
+
67
+ assert unwrapped_vocab.keys() == wrapped_vocab.keys()
68
+
69
+ mismapped = {
70
+ token for token, ids in unwrapped_vocab.items() if ids != wrapped_vocab[token]
71
+ }
72
+ assert mismapped
73
+
74
+ assert len({tuple(wrapped_vocab[t]) for t in mismapped}) == len(mismapped)
75
+ assert len({tuple(unwrapped_vocab[t]) for t in mismapped}) < len(mismapped)
76
+
77
+
78
+@pytest.mark.parametrize("name,cls", BACKENDS)
79
+def test_backend_accepts_a_valid_document(wrapped_tokenizer, name, cls):
80
+ backend = _backend(cls, name, wrapped_tokenizer)
81
+ grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA)
82
+
83
+ for token in _encode(wrapped_tokenizer, VALID_DOCUMENT):
84
+ assert grammar.accept_tokens("req", [token])
85
+
86
+
87
+@pytest.mark.parametrize("name,cls", BACKENDS)
88
+def test_backend_rejects_an_invalid_token(wrapped_tokenizer, name, cls):
89
+ backend = _backend(cls, name, wrapped_tokenizer)
90
+ grammar = backend.compile_grammar(StructuredOutputOptions.JSON, JSON_SCHEMA)
91
+
92
+ assert not grammar.accept_tokens("req", [_encode(wrapped_tokenizer, "]")[-1]])
93
+
94
+
95
+def test_xgrammar_vocab_size_matches_the_tokenizer(wrapped_tokenizer):
96
+ backend = _backend(XgrammarBackend, "xgrammar", wrapped_tokenizer)
97
+ assert backend.vocab_size == len(wrapped_tokenizer.vocab)
98
+
99
+
100
+def test_xgrammar_rejects_the_unwrapped_tokenizer(loaded_tokenizer):
101
+ with pytest.raises(ValueError, match="Unsupported tokenizer type"):
102
+ _backend(XgrammarBackend, "xgrammar", loaded_tokenizer)
vllm/v1/structured_output/__init__.py
+3
-2
@@ -17,6 +17,7 @@ from vllm.v1.structured_output.backend_types import (
17
StructuredOutputGrammar,
18
)
19
from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
20
+from vllm.v1.structured_output.utils import maybe_wrap_mistral_common_tokenizer
21
22
if TYPE_CHECKING:
23
import numpy as np
@@ -75,8 +76,8 @@ class StructuredOutputManager:
76
# of CPUs.
77
max_workers = max(1, (multiprocessing.cpu_count() + 1) // 2)
78
self.executor = ThreadPoolExecutor(max_workers=max_workers)
78
- self.tokenizer = cached_tokenizer_from_config(
79
- model_config=self.vllm_config.model_config
79
+ self.tokenizer = maybe_wrap_mistral_common_tokenizer(
80
+ cached_tokenizer_from_config(model_config=self.vllm_config.model_config)
81
)
82
reasoning_parser_plugin = (
83
self.vllm_config.structured_outputs_config.reasoning_parser_plugin
vllm/v1/structured_output/utils.py
+13
@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, TypeVar
14
import regex as re
15
import torch
16
from cachetools import LRUCache
17
+from transformers import MistralCommonBackend
18
19
import vllm.envs as envs
20
from vllm.logger import init_logger
@@ -305,6 +306,18 @@ re_llama_byte_token = re.compile(r"^<0x[0-9A-F]{2}>$")
306
re_replacement_seq = re.compile(r"^.{0,6}�+.{0,6}$")
307
308
309
+def maybe_wrap_mistral_common_tokenizer(tokenizer: TokenizerLike) -> TokenizerLike:
310
+ """The grammar backends cannot consume a `MistralCommonBackend` directly."""
311
+ if not isinstance(tokenizer, MistralCommonBackend):
312
+ return tokenizer
313
+
314
+ # Deferred: `vllm.tokenizers.mistral` pulls in a large dependency tree, and
315
+ # this module is imported by `vllm.v1.worker.gpu_model_runner`.
316
+ from vllm.tokenizers.mistral import MistralTokenizer
317
+
318
+ return MistralTokenizer(tokenizer)
319
+
320
+
321
def _reduced_vocabulary(tokenizer: TokenizerLike) -> dict[bytes, list[int]]:
322
"""Create a map from vocabulary tokens to lists of equivalent token ids.
323