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 15, 06:00 UTC
[Bugfix] Fix auto dtype casting for BatchFeature (#19316)
Signed-off-by: Isotr0py <2037008807@qq.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Isotr0py committed
Jun 14, 2025 at 23:13 UTC
2db9044ab6c04a38dbcc9df6756b70dccacc157a
7 files changed
+85
-57
tests/v1/engine/test_async_llm.py
+16
-9
@@ -15,6 +15,7 @@ from vllm.engine.arg_utils import AsyncEngineArgs
15
from vllm.inputs import PromptType
16
from vllm.platforms import current_platform
17
from vllm.sampling_params import RequestOutputKind
18
+from vllm.utils import set_default_torch_num_threads
19
from vllm.v1.engine.async_llm import AsyncLLM
20
from vllm.v1.metrics.loggers import LoggingStatLogger
21
@@ -107,7 +108,8 @@ async def test_load(
108
with monkeypatch.context() as m, ExitStack() as after:
109
m.setenv("VLLM_USE_V1", "1")
110
110
- engine = AsyncLLM.from_engine_args(engine_args)
111
+ with set_default_torch_num_threads(1):
112
+ engine = AsyncLLM.from_engine_args(engine_args)
113
after.callback(engine.shutdown)
114
115
NUM_REQUESTS = 100
@@ -154,7 +156,8 @@ async def test_abort(
156
with monkeypatch.context() as m, ExitStack() as after:
157
m.setenv("VLLM_USE_V1", "1")
158
157
- engine = AsyncLLM.from_engine_args(engine_args)
159
+ with set_default_torch_num_threads(1):
160
+ engine = AsyncLLM.from_engine_args(engine_args)
161
after.callback(engine.shutdown)
162
163
NUM_REQUESTS = 100
@@ -226,7 +229,8 @@ async def test_finished_flag(
229
with monkeypatch.context() as m, ExitStack() as after:
230
m.setenv("VLLM_USE_V1", "1")
231
229
- engine = AsyncLLM.from_engine_args(engine_args)
232
+ with set_default_torch_num_threads(1):
233
+ engine = AsyncLLM.from_engine_args(engine_args)
234
after.callback(engine.shutdown)
235
236
sampling_params = SamplingParams(
@@ -260,7 +264,8 @@ async def test_mid_stream_cancellation(monkeypatch: pytest.MonkeyPatch,
264
with monkeypatch.context() as m, ExitStack() as after:
265
m.setenv("VLLM_USE_V1", "1")
266
263
- engine = AsyncLLM.from_engine_args(engine_args)
267
+ with set_default_torch_num_threads(1):
268
+ engine = AsyncLLM.from_engine_args(engine_args)
269
after.callback(engine.shutdown)
270
271
NUM_REQUESTS = 100
@@ -322,10 +327,11 @@ async def test_customize_loggers(monkeypatch):
327
with monkeypatch.context() as m, ExitStack() as after:
328
m.setenv("VLLM_USE_V1", "1")
329
325
- engine = AsyncLLM.from_engine_args(
326
- TEXT_ENGINE_ARGS,
327
- stat_loggers=[MockLoggingStatLogger],
328
- )
330
+ with set_default_torch_num_threads(1):
331
+ engine = AsyncLLM.from_engine_args(
332
+ TEXT_ENGINE_ARGS,
333
+ stat_loggers=[MockLoggingStatLogger],
334
+ )
335
after.callback(engine.shutdown)
336
337
await engine.do_log_stats()
@@ -340,7 +346,8 @@ async def test_dp_rank_argument(monkeypatch: pytest.MonkeyPatch):
346
with monkeypatch.context() as m, ExitStack() as after:
347
m.setenv("VLLM_USE_V1", "1")
348
343
- engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
349
+ with set_default_torch_num_threads(1):
350
+ engine = AsyncLLM.from_engine_args(TEXT_ENGINE_ARGS)
351
after.callback(engine.shutdown)
352
353
sampling_params = SamplingParams(max_tokens=100,
tests/v1/engine/test_engine_core.py
+13
-9
@@ -12,6 +12,7 @@ from transformers import AutoTokenizer
12
from vllm import SamplingParams
13
from vllm.engine.arg_utils import EngineArgs
14
from vllm.platforms import current_platform
15
+from vllm.utils import set_default_torch_num_threads
16
from vllm.v1.engine import EngineCoreRequest
17
from vllm.v1.engine.core import EngineCore
18
from vllm.v1.executor.abstract import Executor, UniProcExecutor
@@ -56,9 +57,10 @@ def test_engine_core(monkeypatch: pytest.MonkeyPatch):
57
vllm_config = engine_args.create_engine_config()
58
executor_class = Executor.get_class(vllm_config)
59
59
- engine_core = EngineCore(vllm_config=vllm_config,
60
- executor_class=executor_class,
61
- log_stats=True)
60
+ with set_default_torch_num_threads(1):
61
+ engine_core = EngineCore(vllm_config=vllm_config,
62
+ executor_class=executor_class,
63
+ log_stats=True)
64
"""Test basic request lifecycle."""
65
66
# First request.
@@ -190,9 +192,10 @@ def test_engine_core_advanced_sampling(monkeypatch: pytest.MonkeyPatch):
192
vllm_config = engine_args.create_engine_config()
193
executor_class = Executor.get_class(vllm_config)
194
193
- engine_core = EngineCore(vllm_config=vllm_config,
194
- executor_class=executor_class,
195
- log_stats=True)
195
+ with set_default_torch_num_threads(1):
196
+ engine_core = EngineCore(vllm_config=vllm_config,
197
+ executor_class=executor_class,
198
+ log_stats=True)
199
"""Test basic request lifecycle."""
200
# First request.
201
request: EngineCoreRequest = make_request()
@@ -286,9 +289,10 @@ def test_engine_core_concurrent_batches(monkeypatch: pytest.MonkeyPatch):
289
enforce_eager=True,
290
)
291
vllm_config = engine_args.create_engine_config()
289
- engine_core = EngineCore(vllm_config=vllm_config,
290
- log_stats=False,
291
- executor_class=DummyExecutor)
292
+ with set_default_torch_num_threads(1):
293
+ engine_core = EngineCore(vllm_config=vllm_config,
294
+ log_stats=False,
295
+ executor_class=DummyExecutor)
296
assert engine_core.batch_queue is not None
297
298
# Add two requests in a row. Each request have 12 prompt tokens.
tests/v1/engine/test_engine_core_client.py
+35
-28
@@ -19,6 +19,7 @@ from vllm.distributed.kv_events import (BlockStored, KVEventBatch,
19
from vllm.engine.arg_utils import EngineArgs
20
from vllm.platforms import current_platform
21
from vllm.usage.usage_lib import UsageContext
22
+from vllm.utils import set_default_torch_num_threads
23
from vllm.v1.engine import EngineCoreRequest
24
from vllm.v1.engine.core import EngineCore
25
from vllm.v1.engine.core_client import (AsyncMPClient, EngineCoreClient,
@@ -138,13 +139,15 @@ def test_engine_core_client(monkeypatch: pytest.MonkeyPatch,
139
vllm_config = engine_args.create_engine_config(
140
UsageContext.UNKNOWN_CONTEXT)
141
executor_class = Executor.get_class(vllm_config)
141
- client = EngineCoreClient.make_client(
142
- multiprocess_mode=multiprocessing_mode,
143
- asyncio_mode=False,
144
- vllm_config=vllm_config,
145
- executor_class=executor_class,
146
- log_stats=False,
147
- )
142
+
143
+ with set_default_torch_num_threads(1):
144
+ client = EngineCoreClient.make_client(
145
+ multiprocess_mode=multiprocessing_mode,
146
+ asyncio_mode=False,
147
+ vllm_config=vllm_config,
148
+ executor_class=executor_class,
149
+ log_stats=False,
150
+ )
151
152
MAX_TOKENS = 20
153
params = SamplingParams(max_tokens=MAX_TOKENS)
@@ -223,13 +226,15 @@ async def test_engine_core_client_asyncio(monkeypatch: pytest.MonkeyPatch):
226
vllm_config = engine_args.create_engine_config(
227
usage_context=UsageContext.UNKNOWN_CONTEXT)
228
executor_class = Executor.get_class(vllm_config)
226
- client = EngineCoreClient.make_client(
227
- multiprocess_mode=True,
228
- asyncio_mode=True,
229
- vllm_config=vllm_config,
230
- executor_class=executor_class,
231
- log_stats=True,
232
- )
229
+
230
+ with set_default_torch_num_threads(1):
231
+ client = EngineCoreClient.make_client(
232
+ multiprocess_mode=True,
233
+ asyncio_mode=True,
234
+ vllm_config=vllm_config,
235
+ executor_class=executor_class,
236
+ log_stats=True,
237
+ )
238
239
try:
240
MAX_TOKENS = 20
@@ -312,13 +317,14 @@ def test_kv_cache_events(
317
UsageContext.UNKNOWN_CONTEXT)
318
319
executor_class = Executor.get_class(vllm_config)
315
- client = EngineCoreClient.make_client(
316
- multiprocess_mode=multiprocessing_mode,
317
- asyncio_mode=False,
318
- vllm_config=vllm_config,
319
- executor_class=executor_class,
320
- log_stats=False,
321
- )
320
+ with set_default_torch_num_threads(1):
321
+ client = EngineCoreClient.make_client(
322
+ multiprocess_mode=multiprocessing_mode,
323
+ asyncio_mode=False,
324
+ vllm_config=vllm_config,
325
+ executor_class=executor_class,
326
+ log_stats=False,
327
+ )
328
endpoint = publisher_config.endpoint.replace("*", "127.0.0.1")
329
subscriber = MockSubscriber(endpoint,
330
topic=publisher_config.topic,
@@ -394,13 +400,14 @@ async def test_kv_cache_events_dp(
400
UsageContext.UNKNOWN_CONTEXT)
401
402
executor_class = Executor.get_class(vllm_config)
397
- client = EngineCoreClient.make_client(
398
- multiprocess_mode=multiprocessing_mode,
399
- asyncio_mode=True,
400
- vllm_config=vllm_config,
401
- executor_class=executor_class,
402
- log_stats=False,
403
- )
403
+ with set_default_torch_num_threads(1):
404
+ client = EngineCoreClient.make_client(
405
+ multiprocess_mode=multiprocessing_mode,
406
+ asyncio_mode=True,
407
+ vllm_config=vllm_config,
408
+ executor_class=executor_class,
409
+ log_stats=False,
410
+ )
411
await asyncio.sleep(1)
412
413
# Build endpoints for all DP ranks
vllm/inputs/registry.py
+3
-1
@@ -168,10 +168,12 @@ class InputProcessingContext(InputContext):
168
try:
169
output = hf_processor(**data, **merged_kwargs, return_tensors="pt")
170
# this emulates output.to(dtype=self.model_config.dtype)
171
- cast_output = json_map_leaves(maybe_cast_dtype, output)
171
if isinstance(output, BatchFeature):
172
+ cast_output = json_map_leaves(maybe_cast_dtype, output.data)
173
return BatchFeature(cast_output)
174
175
+ cast_output = json_map_leaves(maybe_cast_dtype, output)
176
+
177
logger.warning_once(
178
f"{type(hf_processor).__name__} did not return `BatchFeature`. "
179
"Make sure to match the behaviour of `ProcessorMixin` when "
vllm/model_executor/models/qwen2_5_vl.py
+4
-5
@@ -965,9 +965,9 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
965
grid_thw_list = grid_thw.tolist()
966
967
if image_input["type"] == "image_embeds":
968
- image_embeds = image_input["image_embeds"].type(self.visual.dtype)
968
+ image_embeds = image_input["image_embeds"]
969
else:
970
- pixel_values = image_input["pixel_values"].type(self.visual.dtype)
970
+ pixel_values = image_input["pixel_values"]
971
image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)
972
973
# Split concatenated embeddings for each image item.
@@ -985,10 +985,9 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,
985
grid_thw_list = grid_thw.tolist()
986
987
if video_input["type"] == "video_embeds":
988
- video_embeds = video_input["video_embeds"].type(self.visual.dtype)
988
+ video_embeds = video_input["video_embeds"]
989
else:
990
- pixel_values_videos = video_input["pixel_values_videos"].type(
991
- self.visual.dtype)
990
+ pixel_values_videos = video_input["pixel_values_videos"]
991
video_embeds = self.visual(pixel_values_videos,
992
grid_thw=grid_thw_list)
993
vllm/model_executor/models/qwen2_vl.py
+4
-5
@@ -1208,9 +1208,9 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1208
assert grid_thw.ndim == 2
1209
1210
if image_input["type"] == "image_embeds":
1211
- image_embeds = image_input["image_embeds"].type(self.visual.dtype)
1211
+ image_embeds = image_input["image_embeds"]
1212
else:
1213
- pixel_values = image_input["pixel_values"].type(self.visual.dtype)
1213
+ pixel_values = image_input["pixel_values"]
1214
image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
1215
1216
# Split concatenated embeddings for each image item.
@@ -1226,10 +1226,9 @@ class Qwen2VLForConditionalGeneration(nn.Module, SupportsMultiModal,
1226
assert grid_thw.ndim == 2
1227
1228
if video_input["type"] == "video_embeds":
1229
- video_embeds = video_input["video_embeds"].type(self.visual.dtype)
1229
+ video_embeds = video_input["video_embeds"]
1230
else:
1231
- pixel_values_videos = video_input["pixel_values_videos"].type(
1232
- self.visual.dtype)
1231
+ pixel_values_videos = video_input["pixel_values_videos"]
1232
video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
1233
1234
# Split concatenated embeddings for each video item.
vllm/utils.py
+10
@@ -190,6 +190,16 @@ TORCH_DTYPE_TO_NUMPY_DTYPE = {
190
torch.int64: np.int64,
191
}
192
193
+
194
+@contextlib.contextmanager
195
+def set_default_torch_num_threads(num_threads: int):
196
+ """Sets the default number of threads for PyTorch to the given value."""
197
+ old_num_threads = torch.get_num_threads()
198
+ torch.set_num_threads(num_threads)
199
+ yield
200
+ torch.set_num_threads(old_num_threads)
201
+
202
+
203
P = ParamSpec('P')
204
T = TypeVar("T")
205
U = TypeVar("U")