A high-throughput and memory-efficient inference and serving engine for LLMs
Register parsed config classes before tokenizer init (#40299)
Signed-off-by: Bortlesboat <bortstheboat@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Andrew Barnes committed
Jun 16, 2026 at 01:33 UTC
a9a8a32dcdb7e74006ca9d85d3bc4e4536d05488
3 files changed
+85
-5
tests/tokenizers_/test_registry.py
+64
@@ -1,15 +1,24 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import json
4
from pathlib import Path
5
+from types import SimpleNamespace
6
+from unittest.mock import patch
7
8
import pytest
9
+from transformers import AutoConfig
10
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
11
12
from vllm.tokenizers import TokenizerLike
13
from vllm.tokenizers.registry import (
14
TokenizerRegistry,
15
+ cached_get_tokenizer,
16
+ cached_resolve_tokenizer_args,
17
+ cached_tokenizer_from_config,
18
get_tokenizer,
19
resolve_tokenizer_args,
20
)
21
+from vllm.transformers_utils.configs.qwen3_5_moe import Qwen3_5MoeConfig
22
23
24
class TestTokenizer(TokenizerLike):
@@ -75,3 +84,58 @@ def test_customized_tokenizer():
84
assert tokenizer.bos_token_id == 0
85
assert tokenizer.eos_token_id == 1
86
assert tokenizer.pad_token_id == 2
87
+
88
+
89
+def test_cached_tokenizer_from_config_registers_local_config(tmp_path: Path):
90
+ (tmp_path / "config.json").write_text(
91
+ json.dumps({"model_type": "qwen3_5_moe"}),
92
+ encoding="utf-8",
93
+ )
94
+
95
+ model_config = SimpleNamespace(
96
+ skip_tokenizer_init=False,
97
+ tokenizer=str(tmp_path),
98
+ runner_type="generate",
99
+ tokenizer_mode="hf",
100
+ tokenizer_revision=None,
101
+ trust_remote_code=True,
102
+ hf_config=Qwen3_5MoeConfig(),
103
+ )
104
+
105
+ registered_config = CONFIG_MAPPING._extra_content.pop("qwen3_5_moe", None)
106
+ cached_get_tokenizer.cache_clear()
107
+ cached_resolve_tokenizer_args.cache_clear()
108
+
109
+ try:
110
+
111
+ def fake_from_pretrained(path_or_repo_id: str, *args, **kwargs):
112
+ loaded_config = AutoConfig.from_pretrained(
113
+ path_or_repo_id,
114
+ trust_remote_code=False,
115
+ )
116
+ assert isinstance(loaded_config, Qwen3_5MoeConfig)
117
+ return SimpleNamespace(is_fast=True)
118
+
119
+ with (
120
+ patch(
121
+ "vllm.tokenizers.registry.logger.debug_once",
122
+ lambda *args, **kwargs: None,
123
+ ),
124
+ patch(
125
+ "vllm.tokenizers.hf.AutoTokenizer.from_pretrained",
126
+ side_effect=fake_from_pretrained,
127
+ ),
128
+ patch(
129
+ "vllm.tokenizers.hf.get_cached_tokenizer",
130
+ side_effect=lambda tokenizer: tokenizer,
131
+ ),
132
+ ):
133
+ tokenizer = cached_tokenizer_from_config(model_config)
134
+
135
+ assert tokenizer.is_fast is True
136
+ finally:
137
+ cached_get_tokenizer.cache_clear()
138
+ cached_resolve_tokenizer_args.cache_clear()
139
+ CONFIG_MAPPING._extra_content.pop("qwen3_5_moe", None)
140
+ if registered_config is not None:
141
+ CONFIG_MAPPING._extra_content["qwen3_5_moe"] = registered_config
vllm/tokenizers/registry.py
+3
-1
@@ -11,7 +11,7 @@ from typing_extensions import TypeVar, assert_never
11
12
import vllm.envs as envs
13
from vllm.logger import init_logger
14
-from vllm.transformers_utils.config import get_config
14
+from vllm.transformers_utils.config import _maybe_register_hf_config, get_config
15
from vllm.transformers_utils.repo_utils import (
16
any_pattern_in_repo_files,
17
is_mistral_model_repo,
@@ -246,6 +246,8 @@ def cached_tokenizer_from_config(model_config: "ModelConfig", **kwargs):
246
if model_config.skip_tokenizer_init:
247
return None
248
249
+ _maybe_register_hf_config(getattr(model_config, "hf_config", None))
250
+
251
return cached_get_tokenizer(
252
model_config.tokenizer,
253
runner_type=model_config.runner_type,
vllm/transformers_utils/config.py
+18
-4
@@ -141,6 +141,22 @@ _AUTO_CONFIG_KWARGS_OVERRIDES: dict[str, dict[str, Any]] = {
141
}
142
143
144
+def _register_config_class(
145
+ model_type: str, config_class: type[PretrainedConfig]
146
+) -> None:
147
+ config_class.model_type = model_type
148
+ AutoConfig.register(model_type, config_class, exist_ok=True)
149
+
150
+
151
+def _maybe_register_hf_config(config: PretrainedConfig | None) -> None:
152
+ if config is None:
153
+ return
154
+
155
+ model_type = getattr(config, "model_type", None)
156
+ if isinstance(model_type, str) and model_type in _CONFIG_REGISTRY:
157
+ _register_config_class(model_type, _CONFIG_REGISTRY[model_type])
158
+
159
+
160
def is_rope_parameters_nested(rope_parameters: dict[str, Any]) -> bool:
161
"""Check if rope_parameters is nested by layer types."""
162
# Cannot be nested if rope_parameters is empty
@@ -244,8 +260,7 @@ class HFConfigParser(ConfigParserBase):
260
# in future calls to `from_pretrained` (e.g. from
261
# AutoTokenizer or AutoProcessor).
262
config_class = _CONFIG_REGISTRY[model_type]
247
- config_class.model_type = model_type
248
- AutoConfig.register(model_type, config_class, exist_ok=True)
263
+ _register_config_class(model_type, config_class)
264
# If the on-disk model_type differs from the overridden
265
# one, register under both so AutoConfig.from_pretrained
266
# returns the correct class regardless of what the
@@ -253,8 +268,7 @@ class HFConfigParser(ConfigParserBase):
268
if (
269
config_model_type := config_dict.get("model_type")
270
) and config_model_type != model_type:
256
- config_class.model_type = config_model_type
257
- AutoConfig.register(config_model_type, config_class, exist_ok=True)
271
+ _register_config_class(config_model_type, config_class)
272
config_class.model_type = model_type
273
# Now that it is registered, it is not considered remote code anymore
274
trust_remote_code = False