From 517b1f3e1db4a6855a359cb87f88d1e1c17f79f7 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Wed, 23 Sep 2026 23:35:57 -0700 Subject: [PATCH 1/6] [None][perf] Use the prefix-tokenization cache for MiniMax-M3 text-only prompts enable_tokenization_cache only reaches DefaultInputProcessor. MiniMax-M3 checkpoints resolve to MiniMaxM3VLInputProcessor, which tokenizes every prompt through the HF processor, so the flag has no effect on them. create_input_processor now passes enable_tokenization_cache to model-specific input processors that set supports_tokenization_cache; the others would reject the unknown kwarg. MiniMaxM3VLInputProcessor opts in and tokenizes text-only prompts through the cache, built on the HF processor's own tokenizer. For such prompts the HF processor only runs that tokenizer, so the ids are identical as long as the tokenizer adds no special tokens, which the constructor checks. Requests with images or videos are unchanged. Signed-off-by: Zheyu Fu --- .../features/prefix-tokenization-cache.md | 9 +- .../_torch/models/modeling_minimaxm3_vl.py | 24 +++ tensorrt_llm/inputs/registry.py | 6 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../test_minimax_m3_vl_prefix_token_cache.py | 144 ++++++++++++++++++ 5 files changed, 180 insertions(+), 4 deletions(-) create mode 100644 tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index 4855e01901f1..b6011a23b041 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -54,11 +54,14 @@ feature you turned on. shared across tokenizers. - The cache requires a fast (Rust-backed) tokenizer, because it relies on character offsets. With a slow tokenizer the cache is disabled with a warning. -- The cache is used only when the tokenizer would be called exactly as the - cache calls it: `add_special_tokens=False` and no prompt truncation. Chat - completions apply the chat template and tokenize with +- `DefaultInputProcessor` uses the cache only when the tokenizer would be + called exactly as the cache calls it: `add_special_tokens=False` and no prompt + truncation. Chat completions apply the chat template and tokenize with `add_special_tokens=False`, so they benefit. `/v1/completions` defaults to `add_special_tokens=True` and is not accelerated. +- Model-specific (multimodal) input processors do not use the cache, except + MiniMax-M3's, which uses it for every text-only prompt. Requests with images + or videos bypass the cache. - A prompt that extends a cached entry replaces that entry, so a conversation costs one entry regardless of how many turns it has. Lookup is bucketed by a hash of the first `TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS` characters, and each diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index 32475c5fc1b7..54909118e2b8 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py @@ -1735,17 +1735,26 @@ class MiniMaxM3VLInputProcessor: video placeholders are framed by the image start/end tokens (``MINIMAX_M3_VL_VISION_START_TOKEN`` / ``MINIMAX_M3_VL_VISION_END_TOKEN`` above, resolved via the tokenizer). + + With ``enable_tokenization_cache``, text-only prompts are tokenized through + the prefix-tokenization cache instead of the HF processor. """ + # Makes create_input_processor pass enable_tokenization_cache. + supports_tokenization_cache = True + def __init__( self, model_path: str, config: Any, tokenizer: Any = None, trust_remote_code: bool = True, + enable_tokenization_cache: bool = False, **kwargs: Any, ): + from tensorrt_llm.inputs.prefix_token_cache import create_prefix_token_cache from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor + from tensorrt_llm.logger import logger BaseMultimodalInputProcessor.__init__( self, @@ -1770,6 +1779,18 @@ def __init__( use_fast=self._use_fast, trust_remote_code=trust_remote_code, ) + # The HF processor tokenizes a text-only prompt with add_special_tokens=True + # (the tokenizer default) and the cache with add_special_tokens=False, so the + # cache is exact only if the tokenizer adds no special tokens; MiniMax-M3's adds none. + self._prefix_token_cache = None + if enable_tokenization_cache: + if self._processor.tokenizer.num_special_tokens_to_add() == 0: + self._prefix_token_cache = create_prefix_token_cache(self._processor.tokenizer) + else: + logger.warning( + "enable_tokenization_cache is ignored: the MiniMax-M3 tokenizer adds " + "special tokens, so cached ids would differ from the HF processor's." + ) text_cfg = getattr(config, "text_config", None) if isinstance(text_cfg, dict): self._dtype = getattr(text_cfg, "torch_dtype", torch.bfloat16) @@ -2050,6 +2071,9 @@ def call_with_text_prompt( templated_text = "\n".join(explicit) else: templated_text = text_prompt or "" + if self._prefix_token_cache is not None: + ids = self._prefix_token_cache.encode(self._processor.tokenizer, templated_text) + return ids, {"multimodal_data": {}} # Run the HF processor. ``return_tensors='pt'`` yields tensors # in the BatchFeature output. diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 985075db25ff..4bff7111d5aa 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -1129,7 +1129,8 @@ def create_input_processor( model-provided Python code. enable_tokenization_cache: Whether the ``DefaultInputProcessor`` caches the tokenization of recent prompts. Ignored for model-specific - (multimodal) input processors. + (multimodal) input processors unless they set + ``supports_tokenization_cache``. **kwargs: Additional arguments passed to input processor constructors (e.g., video_pruning_rate for multimodal models). @@ -1177,6 +1178,9 @@ def create_input_processor( logger.info("Unregistered model, using DefaultInputProcessor") input_processor_cls = None if input_processor_cls is not None: + if getattr(input_processor_cls, "supports_tokenization_cache", + False): + kwargs["enable_tokenization_cache"] = enable_tokenization_cache # Input processors build an AutoTokenizer/AutoProcessor with # trust_remote_code; doing so copies the checkpoint's .py files # into the shared HF module cache non-atomically, and a rank that diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 643adfd03076..66dc56d6e23a 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -73,6 +73,7 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_piecewise_attention_boundary_preserves_indexer_cache_contract - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip + - unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py - unittest/_torch/models/checkpoints - unittest/_torch/modules - unittest/_torch/moe diff --git a/tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py b/tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py new file mode 100644 index 000000000000..1db1008e9ad2 --- /dev/null +++ b/tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py @@ -0,0 +1,144 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-M3 VL input processor with the prefix-tokenization cache. + +The processor is built through ``create_input_processor``, as the LLM API does, with +a stub tokenizer and HF processor, so these tests cover the wiring rather than the +cache itself, which tests/unittest/inputs/test_prefix_token_cache.py covers. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any +from unittest.mock import patch + +import pytest +import torch +import transformers + +from tensorrt_llm._torch.models.modeling_minimaxm3_vl import get_minimax_m3_vl_input_processor_cls +from tensorrt_llm.inputs.registry import INPUT_PROCESSOR_REGISTRY, create_input_processor + +pytestmark = pytest.mark.cpu_only + + +class _Tokenizer: + """One id per space-separated word, prefixed by ``num_special_tokens`` ids.""" + + is_fast = True + + def __init__(self, num_special_tokens: int = 0) -> None: + self._num_special_tokens = num_special_tokens + + def __call__( + self, text: str, add_special_tokens: bool = True, return_offsets_mapping: bool = False + ) -> dict[str, Any]: + ids, offsets, pos = [], [], 0 + for word in text.split(" "): + ids.append(sum(map(ord, word)) % 50_000 + 1) + offsets.append((pos, pos + len(word))) + pos += len(word) + 1 + if add_special_tokens: + ids = [1] * self._num_special_tokens + ids + return {"input_ids": ids, "offset_mapping": offsets} + + def num_special_tokens_to_add(self) -> int: + return self._num_special_tokens + + +class _HFProcessor: + """Stands in for ``MiniMaxVLProcessor``: tokenizes ``text`` and counts its calls.""" + + def __init__(self, tokenizer: _Tokenizer) -> None: + self.tokenizer = tokenizer + self.calls = 0 + + def __call__( + self, text: list[str], images: Any = None, videos: Any = None, return_tensors: Any = None + ) -> dict[str, Any]: + self.calls += 1 + out = {"input_ids": torch.tensor([self.tokenizer(text[0])["input_ids"]])} + if images: + out["pixel_values"] = torch.zeros(1, 4) + out["image_grid_thw"] = torch.tensor([[1, 2, 2]]) + return out + + +def _create_input_processor(processor_cls: type, tokenizer: Any, enabled: bool) -> Any: + model_cls = object() + config = SimpleNamespace(image_token_index=200_025, video_token_index=200_026) + with ( + patch( + "tensorrt_llm._torch.model_config.ModelConfig.from_pretrained", + return_value=SimpleNamespace(pretrained_config=config), + ), + patch("tensorrt_llm._torch.models.get_model_architecture", return_value=(model_cls, None)), + patch.dict( + INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type, {model_cls: processor_cls} + ), + ): + return create_input_processor("unused", tokenizer, enable_tokenization_cache=enabled) + + +def _input_processor( + monkeypatch: pytest.MonkeyPatch, enabled: bool = True, num_special_tokens: int = 0 +) -> Any: + monkeypatch.setenv("TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS", "16") + # Prompts are tokenized by the HF processor's tokenizer; the LLM tokenizer + # only resolves the vision marker tokens. + llm_tokenizer = SimpleNamespace(convert_tokens_to_ids=lambda token: 200_029) + hf_processor = _HFProcessor(_Tokenizer(num_special_tokens)) + with patch.object(transformers.AutoProcessor, "from_pretrained", return_value=hf_processor): + return _create_input_processor( + get_minimax_m3_vl_input_processor_cls(), llm_tokenizer, enabled + ) + + +def _turns(n: int) -> list[str]: + """Prompts that each extend the previous one.""" + return [" ".join(f"w{i}" for i in range(100 * k)) for k in range(1, n + 1)] + + +def test_text_only_prompts_use_the_cache(monkeypatch: pytest.MonkeyPatch) -> None: + ip = _input_processor(monkeypatch) + for prompt in _turns(4): + ids, extra = ip({"prompt": prompt}, None) + assert ids == _Tokenizer()(prompt)["input_ids"] + assert extra == {"multimodal_data": {}} + assert ip.processor.calls == 0 + assert ip._prefix_token_cache.hits == 3 + + +def test_cache_is_bypassed_when_disabled(monkeypatch: pytest.MonkeyPatch) -> None: + ip = _input_processor(monkeypatch, enabled=False) + assert ip._prefix_token_cache is None + ip({"prompt": _turns(1)[0]}, None) + assert ip.processor.calls == 1 + + +def test_tokenizer_that_adds_special_tokens_disables_the_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + ip = _input_processor(monkeypatch, num_special_tokens=1) + assert ip._prefix_token_cache is None + ids, _ = ip({"prompt": _turns(1)[0]}, None) + assert ip.processor.calls == 1 + assert ids[0] == 1 + + +def test_image_requests_bypass_the_cache(monkeypatch: pytest.MonkeyPatch) -> None: + ip = _input_processor(monkeypatch) + _, extra = ip({"prompt": _turns(1)[0], "multi_modal_data": {"image": [object()]}}, None) + assert ip.processor.calls == 1 + assert "image" in extra["multimodal_data"] + assert ip._prefix_token_cache.hits == ip._prefix_token_cache.misses == 0 + + +def test_other_input_processors_do_not_receive_the_flag() -> None: + class _InputProcessor: + def __init__(self, model_path: str, config: Any, tokenizer: Any, **kwargs: Any) -> None: + self.kwargs = kwargs + + ip = _create_input_processor(_InputProcessor, None, enabled=True) + assert "enable_tokenization_cache" not in ip.kwargs From b321c7136c4bae7452efb395035e03e25f203821 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Mon, 28 Sep 2026 17:28:55 -0700 Subject: [PATCH 2/6] [None][perf] Drop the MiniMax-M3 prefix-token-cache unit test and docs change Signed-off-by: Zheyu Fu --- .../features/prefix-tokenization-cache.md | 9 +- .../integration/test_lists/test-db/l0_cpu.yml | 1 - .../test_minimax_m3_vl_prefix_token_cache.py | 144 ------------------ 3 files changed, 3 insertions(+), 151 deletions(-) delete mode 100644 tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py diff --git a/docs/source/features/prefix-tokenization-cache.md b/docs/source/features/prefix-tokenization-cache.md index b6011a23b041..4855e01901f1 100644 --- a/docs/source/features/prefix-tokenization-cache.md +++ b/docs/source/features/prefix-tokenization-cache.md @@ -54,14 +54,11 @@ feature you turned on. shared across tokenizers. - The cache requires a fast (Rust-backed) tokenizer, because it relies on character offsets. With a slow tokenizer the cache is disabled with a warning. -- `DefaultInputProcessor` uses the cache only when the tokenizer would be - called exactly as the cache calls it: `add_special_tokens=False` and no prompt - truncation. Chat completions apply the chat template and tokenize with +- The cache is used only when the tokenizer would be called exactly as the + cache calls it: `add_special_tokens=False` and no prompt truncation. Chat + completions apply the chat template and tokenize with `add_special_tokens=False`, so they benefit. `/v1/completions` defaults to `add_special_tokens=True` and is not accelerated. -- Model-specific (multimodal) input processors do not use the cache, except - MiniMax-M3's, which uses it for every text-only prompt. Requests with images - or videos bypass the cache. - A prompt that extends a cached entry replaces that entry, so a conversation costs one entry regardless of how many turns it has. Lookup is bucketed by a hash of the first `TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS` characters, and each diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 66dc56d6e23a..643adfd03076 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -73,7 +73,6 @@ l0_cpu: - unittest/_torch/models/test_minimax_m3.py::test_piecewise_attention_boundary_preserves_indexer_cache_contract - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_projection_shard_geometry - unittest/_torch/models/test_minimax_m3.py::test_minimax_m3_five_way_loader_returns_exact_generic_skip - - unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py - unittest/_torch/models/checkpoints - unittest/_torch/modules - unittest/_torch/moe diff --git a/tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py b/tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py deleted file mode 100644 index 1db1008e9ad2..000000000000 --- a/tests/unittest/_torch/models/test_minimax_m3_vl_prefix_token_cache.py +++ /dev/null @@ -1,144 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -"""MiniMax-M3 VL input processor with the prefix-tokenization cache. - -The processor is built through ``create_input_processor``, as the LLM API does, with -a stub tokenizer and HF processor, so these tests cover the wiring rather than the -cache itself, which tests/unittest/inputs/test_prefix_token_cache.py covers. -""" - -from __future__ import annotations - -from types import SimpleNamespace -from typing import Any -from unittest.mock import patch - -import pytest -import torch -import transformers - -from tensorrt_llm._torch.models.modeling_minimaxm3_vl import get_minimax_m3_vl_input_processor_cls -from tensorrt_llm.inputs.registry import INPUT_PROCESSOR_REGISTRY, create_input_processor - -pytestmark = pytest.mark.cpu_only - - -class _Tokenizer: - """One id per space-separated word, prefixed by ``num_special_tokens`` ids.""" - - is_fast = True - - def __init__(self, num_special_tokens: int = 0) -> None: - self._num_special_tokens = num_special_tokens - - def __call__( - self, text: str, add_special_tokens: bool = True, return_offsets_mapping: bool = False - ) -> dict[str, Any]: - ids, offsets, pos = [], [], 0 - for word in text.split(" "): - ids.append(sum(map(ord, word)) % 50_000 + 1) - offsets.append((pos, pos + len(word))) - pos += len(word) + 1 - if add_special_tokens: - ids = [1] * self._num_special_tokens + ids - return {"input_ids": ids, "offset_mapping": offsets} - - def num_special_tokens_to_add(self) -> int: - return self._num_special_tokens - - -class _HFProcessor: - """Stands in for ``MiniMaxVLProcessor``: tokenizes ``text`` and counts its calls.""" - - def __init__(self, tokenizer: _Tokenizer) -> None: - self.tokenizer = tokenizer - self.calls = 0 - - def __call__( - self, text: list[str], images: Any = None, videos: Any = None, return_tensors: Any = None - ) -> dict[str, Any]: - self.calls += 1 - out = {"input_ids": torch.tensor([self.tokenizer(text[0])["input_ids"]])} - if images: - out["pixel_values"] = torch.zeros(1, 4) - out["image_grid_thw"] = torch.tensor([[1, 2, 2]]) - return out - - -def _create_input_processor(processor_cls: type, tokenizer: Any, enabled: bool) -> Any: - model_cls = object() - config = SimpleNamespace(image_token_index=200_025, video_token_index=200_026) - with ( - patch( - "tensorrt_llm._torch.model_config.ModelConfig.from_pretrained", - return_value=SimpleNamespace(pretrained_config=config), - ), - patch("tensorrt_llm._torch.models.get_model_architecture", return_value=(model_cls, None)), - patch.dict( - INPUT_PROCESSOR_REGISTRY._input_processors_cls_by_model_type, {model_cls: processor_cls} - ), - ): - return create_input_processor("unused", tokenizer, enable_tokenization_cache=enabled) - - -def _input_processor( - monkeypatch: pytest.MonkeyPatch, enabled: bool = True, num_special_tokens: int = 0 -) -> Any: - monkeypatch.setenv("TLLM_PREFIX_TOKEN_CACHE_MIN_CHARS", "16") - # Prompts are tokenized by the HF processor's tokenizer; the LLM tokenizer - # only resolves the vision marker tokens. - llm_tokenizer = SimpleNamespace(convert_tokens_to_ids=lambda token: 200_029) - hf_processor = _HFProcessor(_Tokenizer(num_special_tokens)) - with patch.object(transformers.AutoProcessor, "from_pretrained", return_value=hf_processor): - return _create_input_processor( - get_minimax_m3_vl_input_processor_cls(), llm_tokenizer, enabled - ) - - -def _turns(n: int) -> list[str]: - """Prompts that each extend the previous one.""" - return [" ".join(f"w{i}" for i in range(100 * k)) for k in range(1, n + 1)] - - -def test_text_only_prompts_use_the_cache(monkeypatch: pytest.MonkeyPatch) -> None: - ip = _input_processor(monkeypatch) - for prompt in _turns(4): - ids, extra = ip({"prompt": prompt}, None) - assert ids == _Tokenizer()(prompt)["input_ids"] - assert extra == {"multimodal_data": {}} - assert ip.processor.calls == 0 - assert ip._prefix_token_cache.hits == 3 - - -def test_cache_is_bypassed_when_disabled(monkeypatch: pytest.MonkeyPatch) -> None: - ip = _input_processor(monkeypatch, enabled=False) - assert ip._prefix_token_cache is None - ip({"prompt": _turns(1)[0]}, None) - assert ip.processor.calls == 1 - - -def test_tokenizer_that_adds_special_tokens_disables_the_cache( - monkeypatch: pytest.MonkeyPatch, -) -> None: - ip = _input_processor(monkeypatch, num_special_tokens=1) - assert ip._prefix_token_cache is None - ids, _ = ip({"prompt": _turns(1)[0]}, None) - assert ip.processor.calls == 1 - assert ids[0] == 1 - - -def test_image_requests_bypass_the_cache(monkeypatch: pytest.MonkeyPatch) -> None: - ip = _input_processor(monkeypatch) - _, extra = ip({"prompt": _turns(1)[0], "multi_modal_data": {"image": [object()]}}, None) - assert ip.processor.calls == 1 - assert "image" in extra["multimodal_data"] - assert ip._prefix_token_cache.hits == ip._prefix_token_cache.misses == 0 - - -def test_other_input_processors_do_not_receive_the_flag() -> None: - class _InputProcessor: - def __init__(self, model_path: str, config: Any, tokenizer: Any, **kwargs: Any) -> None: - self.kwargs = kwargs - - ip = _create_input_processor(_InputProcessor, None, enabled=True) - assert "enable_tokenization_cache" not in ip.kwargs From 39a71796b1a30e3c6d71821e2fffd14e9b3379ab Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Sat, 3 Oct 2026 15:44:14 -0700 Subject: [PATCH 3/6] [None][perf] Serve MiniMax-M3 prompts from the cache under DefaultInputProcessor's rules Use the cache for a text-only prompt only when add_special_tokens is False and the prompt is not truncated, the same rule DefaultInputProcessor applies, so one documented rule covers every processor and a future truncation fix in the HF-processor path cannot be bypassed. Declare supports_tokenization_cache on BaseMultimodalInputProcessor next to supports_token_id_mm_expansion. Signed-off-by: Zheyu Fu --- tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py | 11 ++++++++--- tensorrt_llm/inputs/registry.py | 5 +++++ 2 files changed, 13 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index 54909118e2b8..7dd6a5f16508 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py @@ -1737,10 +1737,10 @@ class MiniMaxM3VLInputProcessor: ``MINIMAX_M3_VL_VISION_END_TOKEN`` above, resolved via the tokenizer). With ``enable_tokenization_cache``, text-only prompts are tokenized through - the prefix-tokenization cache instead of the HF processor. + the prefix-tokenization cache instead of the HF processor, under the same + rules as ``DefaultInputProcessor``. """ - # Makes create_input_processor pass enable_tokenization_cache. supports_tokenization_cache = True def __init__( @@ -2071,7 +2071,12 @@ def call_with_text_prompt( templated_text = "\n".join(explicit) else: templated_text = text_prompt or "" - if self._prefix_token_cache is not None: + if ( + self._prefix_token_cache is not None + and sampling_params is not None + and not sampling_params.add_special_tokens + and sampling_params.truncate_prompt_tokens is None + ): ids = self._prefix_token_cache.encode(self._processor.tokenizer, templated_text) return ids, {"multimodal_data": {}} diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 4bff7111d5aa..9b3d0e407632 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -251,6 +251,11 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False + # Whether the subclass takes `enable_tokenization_cache` in `__init__` and + # serves text-only prompts from the prefix-tokenization cache. Only such + # subclasses receive the flag from `create_input_processor`. + supports_tokenization_cache: ClassVar[bool] = False + def get_mm_encoder_item_metadata( self, prompt_token_ids: List[int], From 0801b9de47f75371386a062dc21ab62200a2dcbc Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 6 Oct 2026 00:06:09 -0700 Subject: [PATCH 4/6] [None][perf] Share the tokenization-cache hooks across multimodal input processors Move the prefix-tokenization cache setup and the per-request eligibility rule from MiniMaxM3VLInputProcessor into BaseMultimodalInputProcessor, so other multimodal processors can opt in with the same two calls: _init_tokenization_cache builds the cache on the tokenizer the subclass supplies, and _encode_with_tokenization_cache applies DefaultInputProcessor's rule. The check that cached ids match the HF processor's stays in MiniMax-M3, since it depends on how its processor tokenizes. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zheyu Fu --- .../_torch/models/modeling_minimaxm3_vl.py | 26 +++++--------- tensorrt_llm/inputs/registry.py | 35 +++++++++++++++++-- 2 files changed, 42 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index 7dd6a5f16508..a9ce2184d90c 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py @@ -1752,7 +1752,6 @@ def __init__( enable_tokenization_cache: bool = False, **kwargs: Any, ): - from tensorrt_llm.inputs.prefix_token_cache import create_prefix_token_cache from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor from tensorrt_llm.logger import logger @@ -1782,15 +1781,13 @@ def __init__( # The HF processor tokenizes a text-only prompt with add_special_tokens=True # (the tokenizer default) and the cache with add_special_tokens=False, so the # cache is exact only if the tokenizer adds no special tokens; MiniMax-M3's adds none. - self._prefix_token_cache = None - if enable_tokenization_cache: - if self._processor.tokenizer.num_special_tokens_to_add() == 0: - self._prefix_token_cache = create_prefix_token_cache(self._processor.tokenizer) - else: - logger.warning( - "enable_tokenization_cache is ignored: the MiniMax-M3 tokenizer adds " - "special tokens, so cached ids would differ from the HF processor's." - ) + if enable_tokenization_cache and self._processor.tokenizer.num_special_tokens_to_add() != 0: + logger.warning( + "enable_tokenization_cache is ignored: the MiniMax-M3 tokenizer adds " + "special tokens, so cached ids would differ from the HF processor's." + ) + enable_tokenization_cache = False + self._init_tokenization_cache(enable_tokenization_cache, self._processor.tokenizer) text_cfg = getattr(config, "text_config", None) if isinstance(text_cfg, dict): self._dtype = getattr(text_cfg, "torch_dtype", torch.bfloat16) @@ -2071,13 +2068,8 @@ def call_with_text_prompt( templated_text = "\n".join(explicit) else: templated_text = text_prompt or "" - if ( - self._prefix_token_cache is not None - and sampling_params is not None - and not sampling_params.add_special_tokens - and sampling_params.truncate_prompt_tokens is None - ): - ids = self._prefix_token_cache.encode(self._processor.tokenizer, templated_text) + ids = self._encode_with_tokenization_cache(templated_text, sampling_params) + if ids is not None: return ids, {"multimodal_data": {}} # Run the HF processor. ``return_tensors='pt'`` yields tensors diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 9b3d0e407632..8564c5690a6e 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -251,8 +251,10 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False - # Whether the subclass takes `enable_tokenization_cache` in `__init__` and - # serves text-only prompts from the prefix-tokenization cache. Only such + # Whether the subclass serves text-only prompts from the prefix-tokenization + # cache: it takes `enable_tokenization_cache` in `__init__`, passes it to + # `_init_tokenization_cache`, and tries `_encode_with_tokenization_cache` + # before running the HF processor on a text-only prompt. Only such # subclasses receive the flag from `create_input_processor`. supports_tokenization_cache: ClassVar[bool] = False @@ -282,6 +284,35 @@ def __init__(self, self._trust_remote_code = trust_remote_code self._multimodal_hashing_supported: Optional[bool] = None + def _init_tokenization_cache(self, enable_tokenization_cache: bool, + tokenizer: PreTrainedTokenizerBase) -> None: + """Set up the prefix-tokenization cache for text-only prompts. + + `tokenizer` is the tokenizer the HF processor applies to a text-only + prompt. The cache returns `tokenizer(prompt, add_special_tokens=False)` + ids; the subclass passes `enable_tokenization_cache=False` unless those + match its HF processor's ids for text-only prompts. + """ + self._prefix_token_cache_tokenizer = tokenizer + self._prefix_token_cache = (create_prefix_token_cache(tokenizer) + if enable_tokenization_cache else None) + + def _encode_with_tokenization_cache( + self, prompt: str, + sampling_params: Optional[SamplingParams]) -> Optional[List[int]]: + """Cached ids of a text-only `prompt`, or None to run the HF processor. + + Same rule as `DefaultInputProcessor`: only requests with + `add_special_tokens=False` and no prompt truncation, because the cache + never adds special tokens and never truncates. + """ + if (self._prefix_token_cache is None or sampling_params is None + or sampling_params.add_special_tokens + or sampling_params.truncate_prompt_tokens is not None): + return None + return self._prefix_token_cache.encode( + self._prefix_token_cache_tokenizer, prompt) + def attach_multimodal_embeddings( self, inputs: TextPrompt, From f39f7bf8bc2bf5f9071194d1a09dd96c92eed087 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 6 Oct 2026 00:16:54 -0700 Subject: [PATCH 5/6] [None][perf] Trim comments and drop getattr in the tokenization-cache hooks Keep only the comments that say why the MiniMax-M3 special-token check exists and what _init_tokenization_cache requires of its caller. create_input_processor now checks issubclass(..., BaseMultimodalInputProcessor) instead of getattr; the only registered processor outside that base, WhisperInputProcessor, never opts in. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zheyu Fu --- .../_torch/models/modeling_minimaxm3_vl.py | 9 ++----- tensorrt_llm/inputs/registry.py | 24 ++++--------------- 2 files changed, 6 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py index a9ce2184d90c..024e7ae68b45 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3_vl.py @@ -1735,10 +1735,6 @@ class MiniMaxM3VLInputProcessor: video placeholders are framed by the image start/end tokens (``MINIMAX_M3_VL_VISION_START_TOKEN`` / ``MINIMAX_M3_VL_VISION_END_TOKEN`` above, resolved via the tokenizer). - - With ``enable_tokenization_cache``, text-only prompts are tokenized through - the prefix-tokenization cache instead of the HF processor, under the same - rules as ``DefaultInputProcessor``. """ supports_tokenization_cache = True @@ -1778,9 +1774,8 @@ def __init__( use_fast=self._use_fast, trust_remote_code=trust_remote_code, ) - # The HF processor tokenizes a text-only prompt with add_special_tokens=True - # (the tokenizer default) and the cache with add_special_tokens=False, so the - # cache is exact only if the tokenizer adds no special tokens; MiniMax-M3's adds none. + # The HF processor tokenizes with add_special_tokens=True and the cache with False, + # so their ids match only if the tokenizer adds no special tokens. if enable_tokenization_cache and self._processor.tokenizer.num_special_tokens_to_add() != 0: logger.warning( "enable_tokenization_cache is ignored: the MiniMax-M3 tokenizer adds " diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 8564c5690a6e..bbe0d0b46151 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -251,11 +251,7 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False - # Whether the subclass serves text-only prompts from the prefix-tokenization - # cache: it takes `enable_tokenization_cache` in `__init__`, passes it to - # `_init_tokenization_cache`, and tries `_encode_with_tokenization_cache` - # before running the HF processor on a text-only prompt. Only such - # subclasses receive the flag from `create_input_processor`. + # Whether the subclass accepts `enable_tokenization_cache` in `__init__`. supports_tokenization_cache: ClassVar[bool] = False def get_mm_encoder_item_metadata( @@ -286,13 +282,7 @@ def __init__(self, def _init_tokenization_cache(self, enable_tokenization_cache: bool, tokenizer: PreTrainedTokenizerBase) -> None: - """Set up the prefix-tokenization cache for text-only prompts. - - `tokenizer` is the tokenizer the HF processor applies to a text-only - prompt. The cache returns `tokenizer(prompt, add_special_tokens=False)` - ids; the subclass passes `enable_tokenization_cache=False` unless those - match its HF processor's ids for text-only prompts. - """ + """Enable only if cached ids match the HF processor's for text-only prompts.""" self._prefix_token_cache_tokenizer = tokenizer self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None) @@ -300,12 +290,6 @@ def _init_tokenization_cache(self, enable_tokenization_cache: bool, def _encode_with_tokenization_cache( self, prompt: str, sampling_params: Optional[SamplingParams]) -> Optional[List[int]]: - """Cached ids of a text-only `prompt`, or None to run the HF processor. - - Same rule as `DefaultInputProcessor`: only requests with - `add_special_tokens=False` and no prompt truncation, because the cache - never adds special tokens and never truncates. - """ if (self._prefix_token_cache is None or sampling_params is None or sampling_params.add_special_tokens or sampling_params.truncate_prompt_tokens is not None): @@ -1214,8 +1198,8 @@ def create_input_processor( logger.info("Unregistered model, using DefaultInputProcessor") input_processor_cls = None if input_processor_cls is not None: - if getattr(input_processor_cls, "supports_tokenization_cache", - False): + if (issubclass(input_processor_cls, BaseMultimodalInputProcessor) + and input_processor_cls.supports_tokenization_cache): kwargs["enable_tokenization_cache"] = enable_tokenization_cache # Input processors build an AutoTokenizer/AutoProcessor with # trust_remote_code; doing so copies the checkpoint's .py files From 4eec68b98d0f4a04b1e9825ef1eb1702971fce45 Mon Sep 17 00:00:00 2001 From: Zheyu Fu Date: Tue, 6 Oct 2026 01:14:21 -0700 Subject: [PATCH 6/6] [None][perf] Document the tokenization-cache opt-in contract Name the two helpers in the supports_tokenization_cache comment, and state in _init_tokenization_cache's docstring what the caller must guarantee before enabling the cache. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Zheyu Fu --- tensorrt_llm/inputs/registry.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index bbe0d0b46151..7c2c2a1b0e30 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -251,7 +251,8 @@ class BaseMultimodalInputProcessor(ABC): # inputs to `call_with_token_ids` instead of detokenizing upstream. supports_token_id_mm_expansion: ClassVar[bool] = False - # Whether the subclass accepts `enable_tokenization_cache` in `__init__`. + # Whether the subclass takes `enable_tokenization_cache` in `__init__` and + # uses `_init_tokenization_cache` and `_encode_with_tokenization_cache`. supports_tokenization_cache: ClassVar[bool] = False def get_mm_encoder_item_metadata( @@ -282,7 +283,11 @@ def __init__(self, def _init_tokenization_cache(self, enable_tokenization_cache: bool, tokenizer: PreTrainedTokenizerBase) -> None: - """Enable only if cached ids match the HF processor's for text-only prompts.""" + """Build the cache on `tokenizer` if `enable_tokenization_cache`. + + Enable it only if this processor's ids for a text-only prompt equal + `tokenizer(prompt, add_special_tokens=False)`. + """ self._prefix_token_cache_tokenizer = tokenizer self._prefix_token_cache = (create_prefix_token_cache(tokenizer) if enable_tokenization_cache else None)