diff --git a/docs/source/commands/trtllm-serve/trtllm-serve.rst b/docs/source/commands/trtllm-serve/trtllm-serve.rst index 7782d4998eef..23ce61c1d853 100644 --- a/docs/source/commands/trtllm-serve/trtllm-serve.rst +++ b/docs/source/commands/trtllm-serve/trtllm-serve.rst @@ -39,6 +39,24 @@ The following abbreviated command syntax shows the commonly used arguments to st For the full syntax and argument descriptions, refer to :ref:`syntax`. +Preprocessing Workers +~~~~~~~~~~~~~~~~~~~~~ + +Set ``--num_input_processor_workers N`` or use the same setting in ``--config``: + +.. code-block:: yaml + + num_input_processor_workers: 8 + +The default remains eight input-processing threads per frontend for chat and +completion requests, including text-only models. Zero runs inline on the request +event loop. Qwen3-VL uses a separate HF processor and tokenizer per thread. +Additional workers increase host memory use. Offline ``LLM`` calls remain synchronous. + +An explicit flag overrides the YAML value, and ``--set +num_input_processor_workers=N`` has final precedence. Media loading +uses its separate ``--num_media_load_workers`` pool. + .. _tool-parser-selection: Tool-call Parser Selection diff --git a/tensorrt_llm/_torch/models/modeling_qwen3vl.py b/tensorrt_llm/_torch/models/modeling_qwen3vl.py index 82c35ee4534e..fcb2f080f70e 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3vl.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3vl.py @@ -5,6 +5,7 @@ import dataclasses import math import re +import threading from functools import lru_cache from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Tuple, Union from typing import Mapping as TypingMapping @@ -258,6 +259,35 @@ def __init__( # Qwen3-VL keeps ``torch_dtype`` only on ``text_config`` under transformers 5.x. self._dtype = self.config.text_config.dtype + def configure_preprocessing(self, num_workers: int) -> None: + if num_workers > 1: + if getattr(self, "_preprocess_thread_state", None) is None: + self._preprocess_thread_state = threading.local() + self._preprocess_copy_lock = threading.Lock() + self._get_preprocessing_state() + else: + self._preprocess_thread_state = None + + def _get_preprocessing_state(self) -> Optional[threading.local]: + state = getattr(self, "_preprocess_thread_state", None) + if state is not None and not hasattr(state, "processor"): + # HF processors and fast tokenizers can change padding/truncation + # state during a call. Each preprocessing thread owns its copies. + with self._preprocess_copy_lock: + state.processor = copy.deepcopy(self._processor) + state.tokenizer = copy.deepcopy(self._tokenizer) + return state + + @property + def processor(self): + state = self._get_preprocessing_state() + return self._processor if state is None else state.processor + + @property + def tokenizer(self): + state = self._get_preprocessing_state() + return self._tokenizer if state is None else state.tokenizer + @classmethod def _build_temporal_block( cls, @@ -457,7 +487,7 @@ def get_num_tokens_per_video( # wrong-formula fallback that would break chunked prefill, so we must # run the full processor here. do_rescale = not (video and isinstance(video[0], torch.Tensor)) - processed = self._processor( + processed = self.processor( text=["<|vision_start|><|video_pad|><|vision_end|>"], videos=[video], padding=True, diff --git a/tensorrt_llm/_torch/shared_tensor/shared_tensor.py b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py index 24666511249d..05529e990f9d 100644 --- a/tensorrt_llm/_torch/shared_tensor/shared_tensor.py +++ b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py @@ -156,6 +156,8 @@ class SharedTensorContainer: """ + _cpu_storage_lock = threading.Lock() + def __init__(self, method_key: int, tensor_handle: Dict[str, Any]): """Initialize the SharedTensorContainer. @@ -452,14 +454,16 @@ def dump_to_dict(self) -> Dict[str, Any]: tensor_dict = SharedTensorContainer.cuda_handle_to_dict( self.tensor_handle) elif self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU: - sharing_strategy = get_sharing_strategy() - # Here we use file_system sharing strategy to make it serializable between two non-python independent processes - set_sharing_strategy("file_system") - storage = self.tensor_handle[1] - meta_data = self.tensor_handle[2] - storage_handle = reduce_storage(storage) - # restore the original sharing strategy - set_sharing_strategy(sharing_strategy) + # The sharing strategy also affects preprocessing threads. + with self._cpu_storage_lock: + sharing_strategy = get_sharing_strategy() + try: + set_sharing_strategy("file_system") + storage = self.tensor_handle[1] + meta_data = self.tensor_handle[2] + storage_handle = reduce_storage(storage) + finally: + set_sharing_strategy(sharing_strategy) # exclude the first element which is the type of the storage storage_metadata = storage_handle[-1][1:] tensor_dict = SharedTensorContainer.cpu_handle_to_dict( diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 0873ddcf5251..61f452ffabf1 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -276,6 +276,7 @@ def get_llm_args( enable_chunked_prefill: bool = False, enable_attention_dp: bool = False, video_pruning_rate: Optional[float] = None, + num_input_processor_workers: Optional[int] = None, telemetry: bool = True, agent_percentage: float = 0.0, agent_types: Optional[str] = None, @@ -363,6 +364,8 @@ def get_llm_args( build_cli_config("multimodal_config", video_pruning_rate=video_pruning_rate) if video_pruning_rate is not None else None, + "num_input_processor_workers": + num_input_processor_workers, "telemetry_config": _telemetry_config.TelemetryConfig( disabled=not telemetry, @@ -530,7 +533,7 @@ def launch_server( multimodal_server_config: Optional[MultimodalServerConfig] = None, served_model_name: Optional[str] = None, allow_request_chat_template: bool = False, - num_input_processor_workers: int = 8, + num_input_processor_workers: Optional[int] = None, num_media_load_workers: int = 8, multi_frontend_enabled: bool = True, internal_disagg_auth_key: Optional[str] = None, @@ -544,6 +547,10 @@ def launch_server( os.environ[ DisaggLauncherEnvs.TLLM_DISAGG_ROLE] = server_role.name.lower() + if num_input_processor_workers is not None: + llm_args = update_llm_args_with_extra_dict( + llm_args, + {"num_input_processor_workers": num_input_processor_workers}) backend = llm_args["backend"] model = served_model_name or llm_args["model"] @@ -642,7 +649,6 @@ def launch_server( multimodal_server_config=multimodal_server_config, chat_template=chat_template, allow_request_chat_template=allow_request_chat_template, - input_processor_workers=num_input_processor_workers, media_load_workers=num_media_load_workers, internal_disagg_auth_key=internal_disagg_auth_key, enable_runtime_control_endpoints= @@ -1047,11 +1053,11 @@ def launch_visual_gen_server( "SO_REUSEPORT (classic IPC executor path only).", status="prototype") @stability_option("--num_input_processor_workers", - type=click.IntRange(min=1), - default=8, - help="Size of the dedicated thread pool that runs the HF " - "input processor (multimodal preprocess) on the chat " - "and completion endpoints.", + type=click.IntRange(min=0), + default=None, + help="Input-processing threads per frontend for chat and " + "completion requests (default: 8). 0 runs inline. Sets " + "num_input_processor_workers.", status="prototype") @stability_option("--num_media_load_workers", type=click.IntRange(min=1), @@ -1284,7 +1290,7 @@ def serve( kv_cache_dtype: str, num_postprocess_workers: int, num_serve_frontends: int, - num_input_processor_workers: int, + num_input_processor_workers: Optional[int], num_media_load_workers: int, trust_remote_code: bool, revision: Optional[str], @@ -1456,6 +1462,7 @@ def _serve_llm(): enable_attention_dp=enable_attention_dp, enable_locality_domains=enable_locality_domains, video_pruning_rate=video_pruning_rate, + num_input_processor_workers=num_input_processor_workers, telemetry=telemetry, agent_percentage=agent_percentage, agent_types=agent_types, @@ -1617,7 +1624,6 @@ def _serve_llm(): multimodal_server_config, served_model_name=served_model_name, allow_request_chat_template=allow_request_chat_template, - num_input_processor_workers=num_input_processor_workers, num_media_load_workers=num_media_load_workers, internal_disagg_auth_key=internal_disagg_auth_key, report_addr=report_addr, diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index 9a0fa952eb18..e22b55be0809 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -244,6 +244,9 @@ def __init__(self, self._trust_remote_code = trust_remote_code self._multimodal_hashing_supported: Optional[bool] = None + def configure_preprocessing(self, num_workers: int) -> None: + """Prepare processor state for the serving frontend's worker count.""" + def _init_tokenization_cache(self, enable_tokenization_cache: bool, tokenizer: PreTrainedTokenizerBase) -> None: """Build the cache on `tokenizer` if `enable_tokenization_cache`. diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index d1bd36c2e982..f77670d85087 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -2120,6 +2120,9 @@ def _build_model(self): trust_remote_code=self.args.trust_remote_code, enable_tokenization_cache=self.args.enable_tokenization_cache, **input_processor_kwargs) + if isinstance(self.input_processor, BaseMultimodalInputProcessor): + self.input_processor.configure_preprocessing( + self.args.num_input_processor_workers) self._tokenizer = self.input_processor.tokenizer # Resolve encode_only mode (opt-in only) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 355c77c0b319..45a3bf6ed76f 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6473,6 +6473,16 @@ def _check_removed_sampler_type(value: Any) -> None: class TorchLlmArgs(BaseLlmArgs): # PyTorch backend specific configurations + num_input_processor_workers: NonNegativeInt = Field( + default=8, + description=( + "Number of input-processing threads per serving frontend for " + "chat and completion requests, including text-only models. " + "0 runs inline on the request event loop. Offline LLM calls " + "remain synchronous."), + status="prototype", + ) + generation_config: Literal["auto", "trtllm"] = Field( default="trtllm", description= diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 3eac7f79bf0c..1513b4da7abb 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -344,6 +344,13 @@ "additionalProperties": false, "description": "Configuration for KV cache block reuse policies.", "properties": { + "enable_branch_snapshot": { + "default": false, + "description": "Keep reusable state where a request's prompt diverges from the prefix cache, so later requests sharing that prefix can reuse up to the fork. For sliding-window attention (SWA), the SWA window at the fork is committed. This helps when only the SWA window at a chunk end is kept: policy 'per_request' or 'per_conversation', or SWA scratch reuse. The fork is aligned down to a block boundary and becomes a forced context chunk boundary. SWA layers skip this with joint target/draft KV cache reuse. For hybrid Mamba models, the recurrent state is snapshotted at the fork. If no periodic or additional snapshot point applies to the prompt, the prompt end is snapshotted as a fallback. Requires KV cache manager v2 and block reuse.", + "status": "prototype", + "title": "Enable Branch Snapshot", + "type": "boolean" + }, "max_num_turns": { "default": 1, "description": "Maximum number of completed conversation turns whose committed SWA-window blocks and Mamba stable-boundary state are retained by KV cache manager v2. Only used when `policy` is 'per_conversation'.", @@ -3782,9 +3789,11 @@ }, "enable_branch_snapshot": { "default": false, - "description": "Snapshot the Mamba recurrent state where a request's content diverges from the prefix cache, so that later requests sharing the same prefix can reuse up to the fork instead of being truncated to an earlier snapshot. If no periodic or additional snapshot point applies to the prompt, the prompt end is snapshotted as a fallback. Requires KV cache manager V2.", - "status": "prototype", - "telemetry": {}, + "description": "Deprecated alias for kv_cache_config.block_reuse_config.enable_branch_snapshot. It takes effect only when block reuse is enabled.", + "status": "deprecated", + "telemetry": { + "exclude": true + }, "title": "Enable Branch Snapshot", "type": "boolean" }, @@ -7237,6 +7246,14 @@ "description": "ModelExpress (MX) P2P checkpoint loading config.", "status": "prototype" }, + "num_input_processor_workers": { + "default": 8, + "description": "Number of input-processing threads per serving frontend for chat and completion requests, including text-only models. 0 runs inline on the request event loop. Offline LLM calls remain synchronous.", + "minimum": 0, + "status": "prototype", + "title": "Num Input Processor Workers", + "type": "integer" + }, "num_postprocess_workers": { "default": 0, "description": "The number of processes used for postprocessing the generated tokens, including detokenization.", diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index 5e9ffbc0a526..929b370085e9 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -344,6 +344,13 @@ "additionalProperties": false, "description": "Configuration for KV cache block reuse policies.", "properties": { + "enable_branch_snapshot": { + "default": false, + "description": "Keep reusable state where a request's prompt diverges from the prefix cache, so later requests sharing that prefix can reuse up to the fork. For sliding-window attention (SWA), the SWA window at the fork is committed. This helps when only the SWA window at a chunk end is kept: policy 'per_request' or 'per_conversation', or SWA scratch reuse. The fork is aligned down to a block boundary and becomes a forced context chunk boundary. SWA layers skip this with joint target/draft KV cache reuse. For hybrid Mamba models, the recurrent state is snapshotted at the fork. If no periodic or additional snapshot point applies to the prompt, the prompt end is snapshotted as a fallback. Requires KV cache manager v2 and block reuse.", + "status": "prototype", + "title": "Enable Branch Snapshot", + "type": "boolean" + }, "max_num_turns": { "default": 1, "description": "Maximum number of completed conversation turns whose committed SWA-window blocks and Mamba stable-boundary state are retained by KV cache manager v2. Only used when `policy` is 'per_conversation'.", @@ -2552,6 +2559,14 @@ "description": "ModelExpress (MX) P2P checkpoint loading config.", "status": "prototype" }, + "num_input_processor_workers": { + "default": 8, + "description": "Number of input-processing threads per serving frontend for chat and completion requests, including text-only models. 0 runs inline on the request event loop. Offline LLM calls remain synchronous.", + "minimum": 0, + "status": "prototype", + "title": "Num Input Processor Workers", + "type": "integer" + }, "num_instances": { "default": 1, "title": "Num Instances", @@ -5346,9 +5361,11 @@ }, "enable_branch_snapshot": { "default": false, - "description": "Snapshot the Mamba recurrent state where a request's content diverges from the prefix cache, so that later requests sharing the same prefix can reuse up to the fork instead of being truncated to an earlier snapshot. If no periodic or additional snapshot point applies to the prompt, the prompt end is snapshotted as a fallback. Requires KV cache manager V2.", - "status": "prototype", - "telemetry": {}, + "description": "Deprecated alias for kv_cache_config.block_reuse_config.enable_branch_snapshot. It takes effect only when block reuse is enabled.", + "status": "deprecated", + "telemetry": { + "exclude": true + }, "title": "Enable Branch Snapshot", "type": "boolean" }, @@ -8934,6 +8951,14 @@ "null" ] }, + "num_input_processor_workers": { + "default": 8, + "description": "Number of input-processing threads per serving frontend for chat and completion requests, including text-only models. 0 runs inline on the request event loop. Offline LLM calls remain synchronous.", + "minimum": 0, + "status": "prototype", + "title": "Num Input Processor Workers", + "type": "integer" + }, "num_postprocess_workers": { "default": 0, "description": "The number of processes used for postprocessing the generated tokens, including detokenization.", diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 4d0a9c22b315..e2cc5bb631bd 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -23,7 +23,7 @@ from http import HTTPStatus from pathlib import Path from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, - AsyncIterator, Dict, List, Optional, Tuple, Union) + AsyncIterator, Callable, Dict, List, Optional, Tuple, Union) import msgspec import uvicorn @@ -32,7 +32,7 @@ from fastapi.responses import (FileResponse, JSONResponse, Response, StreamingResponse) from fastapi.routing import APIRoute -from pydantic import ValidationError +from pydantic import TypeAdapter, ValidationError from starlette.routing import Mount from transformers import AutoProcessor @@ -57,6 +57,7 @@ from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig, MetadataServerConfig, ServerRole) from tensorrt_llm.llmapi.llm import LLM, RequestOutput +from tensorrt_llm.llmapi.llm_args import TorchLlmArgs from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.llmapi.thinking_budget import \ add_thinking_budget_logits_processor @@ -569,6 +570,12 @@ def resolve_spec_decode_num_spec_tokens(args: Any) -> Optional[int]: class OpenAIServer(_VideoRoutesMixin): + async def _run_input_processor(self, call: Callable[[], Any]) -> Any: + if self._input_proc_executor is None: + return call() + return await asyncio.get_running_loop().run_in_executor( + self._input_proc_executor, call) + @staticmethod def _iteration_stats_buffer_maxlen( iter_stats_max_iterations: Optional[int]) -> Optional[int]: @@ -591,7 +598,7 @@ def __init__( allow_request_chat_template: bool = False, embedding_max_queue_delay: float = 0.005, embedding_max_queue_size: int = 2048, - input_processor_workers: int = 8, + input_processor_workers: Optional[int] = None, media_load_workers: int = 8, internal_disagg_auth_key: Optional[str] = None, enable_rl_control_endpoints: bool = False, @@ -674,10 +681,22 @@ def __init__( # multimodal preprocessing and decode work off the asyncio default # executor avoids contention with unrelated `to_thread` callers and # lets the two stages be sized independently. + args = getattr(self.generator, "args", None) + worker_field = TorchLlmArgs.model_fields["num_input_processor_workers"] + workers = getattr(args, "num_input_processor_workers", + worker_field.default) + if input_processor_workers is not None: + workers = TypeAdapter( + worker_field.rebuild_annotation()).validate_python( + input_processor_workers) + if args is not None: + args.num_input_processor_workers = workers + if isinstance(ip, BaseMultimodalInputProcessor): + ip.configure_preprocessing(workers) self._input_proc_executor = ThreadPoolExecutor( - max_workers=input_processor_workers, + max_workers=workers, thread_name_prefix="trtllm_inputproc", - ) + ) if workers else None self._media_load_executor = ThreadPoolExecutor( max_workers=media_load_workers, thread_name_prefix="trtllm_media_load", @@ -846,6 +865,9 @@ async def lifespan(app: FastAPI): if self.resource_governor is not None: self.resource_governor.close() self.generator.shutdown() + if self._input_proc_executor is not None: + self._input_proc_executor.shutdown(wait=False, + cancel_futures=True) # A stats read still in flight fails fast now that the executor is # down; the /metrics handler already treats that as no stats. self._metrics_stats_executor.shutdown(wait=False, @@ -2423,8 +2445,8 @@ async def chat_stream_generator( parser.build_required_structural_tag, request.tools, thinking=thinking) - required_format = await asyncio.get_running_loop( - ).run_in_executor(self._input_proc_executor, build_required) + required_format = await self._run_input_processor(build_required + ) if required_format is None: raise ValueError( "The configured tool parser does not support required tools." @@ -2457,9 +2479,7 @@ async def chat_stream_generator( generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) if preprocess_fn is not None: - loop = asyncio.get_event_loop() - generate_inputs = await loop.run_in_executor( - self._input_proc_executor, + generate_inputs = await self._run_input_processor( functools.partial(preprocess_fn, prompt, sampling_params, disaggregated_params)) @@ -3120,9 +3140,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): prompt = prompt_inputs(prompt) if prompt.get("prompt") is not None: - loop = asyncio.get_event_loop() - prompt_token_ids, extra_processed_inputs = await loop.run_in_executor( - self._input_proc_executor, + prompt_token_ids, extra_processed_inputs = await self._run_input_processor( functools.partial(self.generator.input_processor, prompt, sampling_params)) tokens_prompt = TokensPrompt( diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index c541b31bd9fa..4a89034efe53 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -1109,6 +1109,11 @@ "kind": "value", "path": "mx_config.server_query_timeout_s" }, + { + "capture_policy": "int", + "kind": "value", + "path": "num_input_processor_workers" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py index c430a47081d8..921d8bfcc693 100644 --- a/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py +++ b/tests/unittest/_torch/modeling/test_modeling_qwen3vl.py @@ -3,6 +3,8 @@ import copy import os +import threading +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from types import SimpleNamespace from typing import List, Optional @@ -1379,3 +1381,53 @@ def test_qwen3vl_vision_block_forward_end_to_end(): attn_metadata=attn_metadata, ) assert out.shape == hidden.shape + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("workers", [0, 1, 4]) +@pytest.mark.parametrize("modality", ["image", "video"]) +def test_preprocessing_thread_local_processor_and_tokenizer(workers, modality): + class FakeProcessor(SimpleNamespace): + def __call__(self, *, videos, **kwargs): + self.video_processor["values"].append(videos[0][0]) + return {"video_grid_thw": torch.tensor([[1, 4, 4]])} + + processor = object.__new__(Qwen3VLInputProcessorBase) + processor._processor = FakeProcessor( + tokenizer={"values": []}, image_processor={"values": []}, video_processor={"values": []} + ) + processor._tokenizer = {"values": []} + processor._config = SimpleNamespace(vision_config=SimpleNamespace(spatial_merge_size=2)) + processor.configure_preprocessing(workers) + if workers <= 1: + assert processor.processor is processor._processor + assert processor.tokenizer is processor._tokenizer + if modality == "video": + assert processor.get_num_tokens_per_video(video=[0]) == 4 + return + gate = threading.Barrier(workers) + + def use_processor(value): + hf_processor, tokenizer = processor.processor, processor.tokenizer + hf_processor.tokenizer["values"].append(value) + if modality == "video": + assert processor.get_num_tokens_per_video(video=[value]) == 4 + else: + hf_processor.image_processor["values"].append(value) + tokenizer["values"].append(value) + gate.wait(timeout=5) + assert processor.processor is hf_processor + assert processor.tokenizer is tokenizer + assert hf_processor.tokenizer["values"] == [value] + media_processor = getattr(hf_processor, f"{modality}_processor") + assert media_processor["values"] == [value] + assert tokenizer["values"] == [value] + return hf_processor, tokenizer + + with ThreadPoolExecutor(max_workers=workers) as pool: + copies = list(pool.map(use_processor, range(workers))) + assert len({id(p) for p, _ in copies}) == workers + assert len({id(t) for _, t in copies}) == workers + assert processor._processor.tokenizer["values"] == [] + assert processor._processor.video_processor["values"] == [] + assert processor._tokenizer["values"] == [] diff --git a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml index 8c06a6d0ea57..a54fdea1d86d 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_cli.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_cli.yaml @@ -353,7 +353,7 @@ commands: - "--num_serve_frontends" num_input_processor_workers: type: int - default: 8 + default: null status: prototype required: false multiple: false diff --git a/tests/unittest/inputs/test_multimodal.py b/tests/unittest/inputs/test_multimodal.py index c1aa7bc3befa..a53934e9b28d 100644 --- a/tests/unittest/inputs/test_multimodal.py +++ b/tests/unittest/inputs/test_multimodal.py @@ -2,12 +2,14 @@ # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Tests for MultimodalRuntimeData cumsum math and the flat-mask producer.""" +from concurrent.futures import ThreadPoolExecutor from unittest.mock import patch import numpy as np import pytest import torch +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm.inputs.multimodal import ( MULTIMODAL_ENCODER_ITEM_METADATA_KEY, DisaggPrefillMultimodalInputs, @@ -488,3 +490,17 @@ def test_runtime_data_requires_cumsum(): """embed_mask_cumsum is required.""" with pytest.raises(TypeError): MultimodalRuntimeData(past_seen_token_num=0, chunk_end_pos=5) + + +def test_concurrent_cpu_handle_exports_restore_values_and_strategy(): + strategy = torch.multiprocessing.get_sharing_strategy() + + def roundtrip(value): + tensor = torch.full((3, 8), value, dtype=torch.int32) + payload = SharedTensorContainer.from_tensor(tensor).dump_to_dict() + restored = SharedTensorContainer.from_dict(payload).get_local_view() + assert torch.equal(restored, tensor) + + with ThreadPoolExecutor(max_workers=4) as pool: + list(pool.map(roundtrip, range(40))) + assert torch.multiprocessing.get_sharing_strategy() == strategy diff --git a/tests/unittest/inputs/test_multimodal_input_processor.py b/tests/unittest/inputs/test_multimodal_input_processor.py index 34356979e8cf..572d47afe845 100644 --- a/tests/unittest/inputs/test_multimodal_input_processor.py +++ b/tests/unittest/inputs/test_multimodal_input_processor.py @@ -129,3 +129,9 @@ def test_get_dummy_mm_data_default_raises_not_implemented(): builder = _StubBuilder() with pytest.raises(NotImplementedError): builder.get_dummy_mm_data(max_num_encoder_tokens=1024, mm_counts={"image": 4}) + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("workers", [0, 1, 8]) +def test_preprocessing_worker_modes(workers): + _make_subclass().configure_preprocessing(workers) diff --git a/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py b/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py index 9178486417c5..0e1593b7e752 100644 --- a/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py +++ b/tests/unittest/llmapi/apps/test_openai_server_iteration_stats.py @@ -366,3 +366,56 @@ async def test_metrics_endpoint_reads_visual_gen_stats(): assert _response_content(response) == stats assert server.generator.stats_timeouts == [None] assert server.metrics_collector.logged_stats == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("workers", [0, 1, 4]) +async def test_preprocessing_pool_and_inline_thread_ownership(workers): + server = object.__new__(OpenAIServer) + server._input_proc_executor = ThreadPoolExecutor(max_workers=workers) if workers else None + try: + thread = await server._run_input_processor(threading.get_ident) + assert (thread == threading.get_ident()) == (workers == 0) + + def fail(): + raise ValueError("invalid input") + + with pytest.raises(ValueError, match="invalid input"): + await server._run_input_processor(fail) + finally: + if server._input_proc_executor is not None: + server._input_proc_executor.shutdown(wait=True) + + +@pytest.mark.parametrize("workers", [None, 0, 1, 4, 8]) +@pytest.mark.parametrize("override", [None, 0, 2]) +def test_preprocessing_constructor_uses_typed_config(monkeypatch, workers, override): + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + from tensorrt_llm.serve import openai_server + + for name in ("_init_llm", "register_routes"): + monkeypatch.setattr(OpenAIServer, name, lambda *args: None) + monkeypatch.setattr(openai_server, "load_builtin_extensions", lambda: None) + monkeypatch.setattr(openai_server, "create_metadata_server", lambda _: None) + monkeypatch.setattr(openai_server.BaseMediaIO, "set_executor", lambda _: None) + config = ( + TorchLlmArgs(model="dummy/model") + if workers is None + else TorchLlmArgs(model="dummy/model", num_input_processor_workers=workers) + ) + generator = SimpleNamespace(args=config) + server = OpenAIServer( + generator, "dummy/model", None, None, None, input_processor_workers=override + ) + expected = override if override is not None else config.num_input_processor_workers + try: + assert generator.args.num_input_processor_workers == expected + if expected: + assert server._input_proc_executor._max_workers == expected + else: + assert server._input_proc_executor is None + finally: + if server._input_proc_executor is not None: + server._input_proc_executor.shutdown(wait=True) + server._media_load_executor.shutdown(wait=True) + server._metrics_stats_executor.shutdown(wait=True) diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index cf1869e4d68e..f7f61d12f70e 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -1809,6 +1809,26 @@ def test_rejects_buckets_with_too_few_tokens(self): MultimodalEncoderCudaGraphConfig(buckets=[(1, 2)]) +@pytest.mark.cpu_only +class TestInputProcessorWorkers: + + @pytest.mark.parametrize("workers", [0, 1, 2, 4, 8, 16]) + def test_num_input_processor_workers(self, workers): + assert TorchLlmArgs( + model=llama_model_path).num_input_processor_workers == 8 + args = TorchLlmArgs( + model=llama_model_path, + num_input_processor_workers=workers) + assert args.num_input_processor_workers == workers + assert "num_input_processor_workers" not in MultimodalConfig.model_fields + + @pytest.mark.parametrize("workers", [-1, 1.5, "invalid"]) + def test_invalid_num_input_processor_workers(self, workers): + with pytest.raises(ValidationError, match="num_input_processor_workers"): + TorchLlmArgs(model=llama_model_path, + num_input_processor_workers=workers) + + @pytest.mark.cpu_only class TestMultimodalConfig: @@ -5467,3 +5487,53 @@ def test_deepseek_resolves_auto_to_python_on_nixl(self) -> None: cfg = self._pretrained_config(["DeepseekV3ForCausalLM"], "deepseek_v3") _resolve_transceiver_runtime_auto(args, DeepseekV3ForCausalLM, cfg) assert args.cache_transceiver_config.transceiver_runtime == "PYTHON" + + +@pytest.mark.cpu_only +@pytest.mark.parametrize( + ("extra_args", "expected"), + [ + ([], 3), + (["--num_input_processor_workers", "0"], 0), + (["--num_input_processor_workers", "1"], 1), + (["--num_input_processor_workers", "2"], 2), + (["--num_input_processor_workers", "8"], 8), + ([ + "--num_input_processor_workers", "2", "--set", + "num_input_processor_workers=4" + ], 4), + ([ + "--set", "num_input_processor_workers=4", + "--num_input_processor_workers", "2" + ], 4), + ], +) +def test_serve_preprocess_workers_precedence(tmp_path, extra_args, expected): + config = tmp_path / "config.yaml" + config.write_text( + "num_input_processor_workers: 3\nmultimodal_config:\n" + " encoder_cache_max_bytes: 1024\n", + encoding="utf-8") + with ( + patch("tensorrt_llm.commands.serve.device_count", return_value=1), + patch("tensorrt_llm.commands.serve.get_is_diffusion_only_model", + return_value=False), + patch("tensorrt_llm.commands.serve.launch_server") as launch, + ): + serve_main(args=["dummy/model", "--config", + str(config), *extra_args], + standalone_mode=False) + args = launch.call_args.args[2] + mm_config = args["multimodal_config"] + if isinstance(mm_config, BaseModel): + mm_config = mm_config.model_dump() + assert args["num_input_processor_workers"] == expected + assert mm_config["encoder_cache_max_bytes"] == 1024 + assert "num_input_processor_workers" not in mm_config + + +@pytest.mark.cpu_only +def test_serve_preprocess_workers_rejects_negative_cli(): + with pytest.raises(click.BadParameter, match="x>=0"): + serve_main(args=["dummy/model", "--num_input_processor_workers", "-1"], + standalone_mode=False) diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 251a3223aa9b..073c3a56bdd5 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -80,6 +80,13 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: (generator.SERVE_SCHEMA, {"env_overrides": {"FLAG": 1, "OTHER": True}}, True), (generator.SERVE_SCHEMA, {"telemetry_config": {"disabled": True}}, True), (generator.SERVE_SCHEMA, {"multimodal_config": None}, True), + (generator.SERVE_SCHEMA, {"num_input_processor_workers": 0}, True), + (generator.SERVE_SCHEMA, {"num_input_processor_workers": 1}, True), + (generator.SERVE_SCHEMA, {"num_input_processor_workers": 8}, True), + (generator.SERVE_SCHEMA, {"num_input_processor_workers": -1}, False), + (generator.SERVE_SCHEMA, {"num_input_processor_workers": 1.5}, False), + (generator.SERVE_SCHEMA, {"num_input_processor_workers": "eight"}, False), + (generator.SERVE_SCHEMA, {"multimodal_config": {"num_input_processor_workers": 8}}, False), ( generator.SERVE_SCHEMA, {"multimodal_config": {"encoder_cuda_graph": {"vision": {"max_batch_size": 8}}}},