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 9, 03:45 UTC
[Redo] #33110 with threading limit (#33502)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: YunzhuLu <lucia.yunzhu@gmail.com>
Cyrus Leung committed
Feb 1, 2026 at 17:18 UTC
21997f45b10c17f44276cf3872e5f85c61dc7dfd
2 files changed
+64
-17
vllm/utils/torch_utils.py
+30
-7
@@ -17,6 +17,7 @@ from packaging.version import Version
17
from torch.library import Library, infer_schema
18
19
import vllm.envs as envs
20
+from vllm.logger import init_logger
21
22
if TYPE_CHECKING:
23
from vllm.config import ModelConfig
@@ -25,9 +26,7 @@ else:
26
ModelConfig = object
27
IntermediateTensors = object
28
28
-import logging
29
-
30
-logger = logging.getLogger(__name__)
29
+logger = init_logger(__name__)
30
31
32
STR_DTYPE_TO_TORCH_DTYPE = {
@@ -104,12 +103,36 @@ def set_default_torch_dtype(dtype: torch.dtype):
103
104
105
@contextlib.contextmanager
107
-def set_default_torch_num_threads(num_threads: int):
108
- """Sets the default number of threads for PyTorch to the given value."""
106
+def set_default_torch_num_threads(num_threads: int | None = None):
107
+ """
108
+ Sets the default number of threads for PyTorch to the given value.
109
+
110
+ `None` means using the value of the environment variable `OMP_NUM_THREADS`
111
+ (or `1` if that is not available).
112
+ """
113
+ if num_threads is None:
114
+ num_threads = 1
115
+
116
+ try:
117
+ num_threads = int(os.environ["OMP_NUM_THREADS"])
118
+ except KeyError:
119
+ logger.debug_once(
120
+ "OMP_NUM_THREADS is not set; defaulting Torch threads to %d.",
121
+ num_threads,
122
+ )
123
+ except ValueError:
124
+ logger.warning_once(
125
+ "OMP_NUM_THREADS is invalid; defaulting Torch threads to %d.",
126
+ num_threads,
127
+ )
128
+
129
old_num_threads = torch.get_num_threads()
130
torch.set_num_threads(num_threads)
111
- yield
112
- torch.set_num_threads(old_num_threads)
131
+
132
+ try:
133
+ yield
134
+ finally:
135
+ torch.set_num_threads(old_num_threads)
136
137
138
@contextlib.contextmanager
vllm/v1/engine/input_processor.py
+34
-10
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import os
4
import time
5
from collections.abc import Mapping
6
from typing import Any, Literal, cast
@@ -35,6 +34,7 @@ from vllm.tokenizers import TokenizerLike
34
from vllm.tokenizers.mistral import MistralTokenizer
35
from vllm.utils import length_from_prompt_token_ids_or_embeds, random_uuid
36
from vllm.utils.torch_utils import set_default_torch_num_threads
37
+from vllm.v1.core.encoder_cache_manager import compute_mm_encoder_budget
38
from vllm.v1.engine import EngineCoreRequest
39
from vllm.v1.metrics.stats import MultiModalCacheStats
40
from vllm.v1.structured_output.backend_guidance import (
@@ -68,6 +68,19 @@ class InputProcessor:
68
69
self.mm_registry = mm_registry
70
self.mm_processor_cache = mm_registry.processor_cache_from_config(vllm_config)
71
+ self.mm_encoder_cache_size = None
72
+ if (
73
+ self.mm_registry.supports_multimodal_inputs(self.model_config)
74
+ and not self.model_config.skip_tokenizer_init
75
+ ):
76
+ with set_default_torch_num_threads():
77
+ max_tokens_by_modality = (
78
+ mm_registry.get_max_tokens_per_item_by_modality(self.model_config)
79
+ )
80
+
81
+ _, self.mm_encoder_cache_size = compute_mm_encoder_budget(
82
+ self.vllm_config.scheduler_config, max_tokens_by_modality
83
+ )
84
85
self.input_preprocessor = InputPreprocessor(
86
self.model_config,
@@ -534,15 +547,7 @@ class InputProcessor:
547
# 1. Tokenize text prompt, with LoRA request if one exists.
548
# 2. For multimodal models with a merged preprocessor, preprocess
549
# multimodal data and expand prompt token ids accordingly.
537
- num_threads = int(os.environ.get("OMP_NUM_THREADS", "1"))
538
- if "OMP_NUM_THREADS" not in os.environ:
539
- logger.debug_once(
540
- "OMP_NUM_THREADS is not set; defaulting Torch threads to %d for "
541
- "input preprocessing.",
542
- num_threads,
543
- )
544
-
545
- with set_request_id(request_id), set_default_torch_num_threads(num_threads):
550
+ with set_request_id(request_id), set_default_torch_num_threads():
551
processed_inputs: ProcessorInputs = self.input_preprocessor.preprocess(
552
prompt,
553
tokenization_kwargs=tokenization_kwargs,
@@ -743,6 +748,25 @@ class InputProcessor:
748
f"model length of {max_prompt_len}. {suggestion}"
749
)
750
751
+ if (
752
+ prompt_type == "decoder"
753
+ and prompt_inputs["type"] == "multimodal"
754
+ and self.mm_encoder_cache_size is not None
755
+ ):
756
+ decoder_mm_positions = prompt_inputs["mm_placeholders"]
757
+ for modality, mm_positions in decoder_mm_positions.items():
758
+ for mm_position in mm_positions:
759
+ embed_length = mm_position.get_num_embeds
760
+ if embed_length > self.mm_encoder_cache_size:
761
+ raise ValueError(
762
+ f"The {prompt_type} prompt contains a(n) {modality} item "
763
+ f"with length {embed_length}, which exceeds the "
764
+ f"pre-allocated encoder cache size "
765
+ f"{self.mm_encoder_cache_size}. Please reduce the input "
766
+ f"size or increase the encoder cache size "
767
+ f"by setting --limit-mm-per-prompt at startup."
768
+ )
769
+
770
def stat_mm_cache(self) -> MultiModalCacheStats | None:
771
return self.input_preprocessor.stat_mm_cache()
772