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 14, 05:30 UTC
[Misc] Clean up utils (#27552)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 27, 2025 at 17:05 UTC
7c2bdb83dca6d71f46a5dd123402242d650df553
12 files changed
+41
-258
docs/mkdocs/hooks/generate_argparse.py
+3
-1
@@ -65,7 +65,9 @@ ChatCommand = auto_mock("vllm.entrypoints.cli.openai", "ChatCommand")
65
CompleteCommand = auto_mock("vllm.entrypoints.cli.openai", "CompleteCommand")
66
cli_args = auto_mock("vllm.entrypoints.openai", "cli_args")
67
run_batch = auto_mock("vllm.entrypoints.openai", "run_batch")
68
-FlexibleArgumentParser = auto_mock("vllm.utils", "FlexibleArgumentParser")
68
+FlexibleArgumentParser = auto_mock(
69
+ "vllm.utils.argparse_utils", "FlexibleArgumentParser"
70
+)
71
72
73
class MarkdownFormatter(HelpFormatter):
tests/utils.py
+1
-3
@@ -45,9 +45,7 @@ from vllm.entrypoints.cli.serve import ServeSubcommand
45
from vllm.model_executor.model_loader import get_model_loader
46
from vllm.platforms import current_platform
47
from vllm.transformers_utils.tokenizer import get_tokenizer
48
-from vllm.utils import (
49
- FlexibleArgumentParser,
50
-)
48
+from vllm.utils.argparse_utils import FlexibleArgumentParser
49
from vllm.utils.mem_constants import GB_bytes
50
from vllm.utils.network_utils import get_open_port
51
from vllm.utils.torch_utils import cuda_device_count_stateless
tests/utils_/test_argparse_utils.py
renamed
+2
-91
@@ -4,23 +4,15 @@
4
5
import json
6
import os
7
-import tempfile
8
-from pathlib import Path
9
-from unittest.mock import patch
7
8
import pytest
12
-import torch
9
import yaml
10
from transformers import AutoTokenizer
11
16
-from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
12
from vllm.transformers_utils.detokenizer_utils import convert_ids_list_to_tokens
13
19
-from vllm.utils import (
20
- FlexibleArgumentParser,
21
- bind_kv_cache,
22
-)
23
-from ..utils import create_new_process_for_each_test, flat_product
14
+from vllm.utils.argparse_utils import FlexibleArgumentParser
15
+from ..utils import flat_product
16
17
18
# Tests for FlexibleArgumentParser
@@ -256,87 +248,6 @@ def test_duplicate_dict_args(caplog_vllm, parser):
248
assert "-O.mode" in caplog_vllm.text
249
250
259
-def test_bind_kv_cache():
260
- from vllm.attention import Attention
261
-
262
- ctx = {
263
- "layers.0.self_attn": Attention(32, 128, 0.1),
264
- "layers.1.self_attn": Attention(32, 128, 0.1),
265
- "layers.2.self_attn": Attention(32, 128, 0.1),
266
- "layers.3.self_attn": Attention(32, 128, 0.1),
267
- }
268
- kv_cache = [
269
- torch.zeros((1,)),
270
- torch.zeros((1,)),
271
- torch.zeros((1,)),
272
- torch.zeros((1,)),
273
- ]
274
- bind_kv_cache(ctx, [kv_cache])
275
- assert ctx["layers.0.self_attn"].kv_cache[0] is kv_cache[0]
276
- assert ctx["layers.1.self_attn"].kv_cache[0] is kv_cache[1]
277
- assert ctx["layers.2.self_attn"].kv_cache[0] is kv_cache[2]
278
- assert ctx["layers.3.self_attn"].kv_cache[0] is kv_cache[3]
279
-
280
-
281
-def test_bind_kv_cache_kv_sharing():
282
- from vllm.attention import Attention
283
-
284
- ctx = {
285
- "layers.0.self_attn": Attention(32, 128, 0.1),
286
- "layers.1.self_attn": Attention(32, 128, 0.1),
287
- "layers.2.self_attn": Attention(32, 128, 0.1),
288
- "layers.3.self_attn": Attention(32, 128, 0.1),
289
- }
290
- kv_cache = [
291
- torch.zeros((1,)),
292
- torch.zeros((1,)),
293
- torch.zeros((1,)),
294
- torch.zeros((1,)),
295
- ]
296
- shared_kv_cache_layers = {
297
- "layers.2.self_attn": "layers.1.self_attn",
298
- "layers.3.self_attn": "layers.0.self_attn",
299
- }
300
- bind_kv_cache(ctx, [kv_cache], shared_kv_cache_layers)
301
- assert ctx["layers.0.self_attn"].kv_cache[0] is kv_cache[0]
302
- assert ctx["layers.1.self_attn"].kv_cache[0] is kv_cache[1]
303
- assert ctx["layers.2.self_attn"].kv_cache[0] is kv_cache[1]
304
- assert ctx["layers.3.self_attn"].kv_cache[0] is kv_cache[0]
305
-
306
-
307
-def test_bind_kv_cache_non_attention():
308
- from vllm.attention import Attention
309
-
310
- # example from Jamba PP=2
311
- ctx = {
312
- "model.layers.20.attn": Attention(32, 128, 0.1),
313
- "model.layers.28.attn": Attention(32, 128, 0.1),
314
- }
315
- kv_cache = [
316
- torch.zeros((1,)),
317
- torch.zeros((1,)),
318
- ]
319
- bind_kv_cache(ctx, [kv_cache])
320
- assert ctx["model.layers.20.attn"].kv_cache[0] is kv_cache[0]
321
- assert ctx["model.layers.28.attn"].kv_cache[0] is kv_cache[1]
322
-
323
-
324
-def test_bind_kv_cache_pp():
325
- with patch("vllm.utils.torch_utils.cuda_device_count_stateless", lambda: 2):
326
- # this test runs with 1 GPU, but we simulate 2 GPUs
327
- cfg = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=2))
328
- with set_current_vllm_config(cfg):
329
- from vllm.attention import Attention
330
-
331
- ctx = {
332
- "layers.0.self_attn": Attention(32, 128, 0.1),
333
- }
334
- kv_cache = [[torch.zeros((1,))], [torch.zeros((1,))]]
335
- bind_kv_cache(ctx, kv_cache)
336
- assert ctx["layers.0.self_attn"].kv_cache[0] is kv_cache[0][0]
337
- assert ctx["layers.0.self_attn"].kv_cache[1] is kv_cache[1][0]
338
-
339
-
251
def test_model_specification(
252
parser_with_config, cli_config_file, cli_config_file_with_model
253
):
tests/utils_/test_serial_utils.py
+1
-1
@@ -14,7 +14,7 @@ from vllm.utils.serial_utils import (
14
15
@pytest.mark.parametrize("endianness", ENDIANNESS)
16
@pytest.mark.parametrize("embed_dtype", EMBED_DTYPE_TO_TORCH_DTYPE.keys())
17
-@torch.inference_mode
17
+@torch.inference_mode()
18
def test_encode_and_decode(embed_dtype: str, endianness: str):
19
for i in range(10):
20
tensor = torch.rand(2, 3, 5, 7, 11, 13, device="cpu", dtype=torch.float32)
vllm/distributed/kv_transfer/kv_connector/v1/decode_bench_connector.py
+1
-1
@@ -42,7 +42,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1 import (
42
)
43
from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
44
from vllm.logger import init_logger
45
-from vllm.utils import cdiv
45
+from vllm.utils.math_utils import cdiv
46
47
if TYPE_CHECKING:
48
from vllm.attention.backends.abstract import AttentionMetadata
vllm/distributed/kv_transfer/kv_connector/v1/lmcache_integration/vllm_v1_adapter.py
+2
-1
@@ -44,7 +44,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.lmcache_integration.utils impo
44
)
45
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank, get_tp_group
46
from vllm.sampling_params import SamplingParams
47
-from vllm.utils import cdiv, get_kv_cache_torch_dtype
47
+from vllm.utils import get_kv_cache_torch_dtype
48
+from vllm.utils.math_utils import cdiv
49
from vllm.v1.core.sched.output import SchedulerOutput
50
from vllm.version import __version__ as VLLM_VERSION
51
vllm/entrypoints/anthropic/api_server.py
+2
-1
@@ -51,7 +51,8 @@ from vllm.entrypoints.utils import (
51
with_cancellation,
52
)
53
from vllm.logger import init_logger
54
-from vllm.utils import FlexibleArgumentParser, set_ulimit
54
+from vllm.utils import set_ulimit
55
+from vllm.utils.argparse_utils import FlexibleArgumentParser
56
from vllm.utils.network_utils import is_valid_ipv6_address
57
from vllm.version import __version__ as VLLM_VERSION
58
vllm/entrypoints/cli/serve.py
+1
-1
@@ -18,7 +18,7 @@ from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_se
18
from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
19
from vllm.logger import init_logger
20
from vllm.usage.usage_lib import UsageContext
21
-from vllm.utils import FlexibleArgumentParser
21
+from vllm.utils.argparse_utils import FlexibleArgumentParser
22
from vllm.utils.network_utils import get_tcp_uri
23
from vllm.utils.system_utils import decorate_logs, set_process_title
24
from vllm.v1.engine.core import EngineCoreProc
vllm/entrypoints/openai/api_server.py
+2
-1
@@ -108,7 +108,8 @@ from vllm.entrypoints.utils import (
108
from vllm.logger import init_logger
109
from vllm.reasoning import ReasoningParserManager
110
from vllm.usage.usage_lib import UsageContext
111
-from vllm.utils import Device, FlexibleArgumentParser, set_ulimit
111
+from vllm.utils import Device, set_ulimit
112
+from vllm.utils.argparse_utils import FlexibleArgumentParser
113
from vllm.utils.network_utils import is_valid_ipv6_address
114
from vllm.utils.system_utils import decorate_logs
115
from vllm.v1.engine.exceptions import EngineDeadError
vllm/lora/punica_wrapper/punica_gpu.py
+1
-1
@@ -13,7 +13,7 @@ import torch
13
14
from vllm.lora.layers import LoRAMapping
15
from vllm.triton_utils import HAS_TRITON, triton
16
-from vllm.utils import round_up
16
+from vllm.utils.math_utils import round_up
17
18
if HAS_TRITON:
19
from vllm.lora.ops.triton_ops import (
vllm/model_executor/layers/quantization/mxfp4.py
+1
-1
@@ -48,9 +48,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_s
48
from vllm.model_executor.utils import set_weight_attrs
49
from vllm.platforms import current_platform
50
from vllm.scalar_type import scalar_types
51
-from vllm.utils import round_up
51
from vllm.utils.flashinfer import has_flashinfer
52
from vllm.utils.import_utils import has_triton_kernels
53
+from vllm.utils.math_utils import round_up
54
from vllm.utils.torch_utils import is_torch_equal_or_newer
55
56
logger = init_logger(__name__)
vllm/utils/__init__.py
+24
-155
@@ -12,12 +12,10 @@ import signal
12
import sys
13
import tempfile
14
import threading
15
-import traceback
15
import uuid
16
import warnings
18
-import weakref
17
from collections.abc import Callable
20
-from functools import cache, partial, wraps
18
+from functools import partial, wraps
19
from typing import TYPE_CHECKING, Any, TypeVar
20
21
import cloudpickle
@@ -28,34 +26,6 @@ import vllm.envs as envs
26
from vllm.logger import enable_trace_function_call, init_logger
27
from vllm.ray.lazy_utils import is_in_ray_actor
28
31
-# Import utilities from specialized modules for backward compatibility
32
-from vllm.utils.argparse_utils import (
33
- FlexibleArgumentParser,
34
- SortedHelpFormatter,
35
- StoreBoolean,
36
-)
37
-from vllm.utils.math_utils import (
38
- cdiv,
39
- next_power_of_2,
40
- prev_power_of_2,
41
- round_down,
42
- round_up,
43
-)
44
-from vllm.utils.platform_utils import cuda_is_initialized, xpu_is_initialized
45
-
46
-__all__ = [
47
- # Argparse utilities
48
- "FlexibleArgumentParser",
49
- "SortedHelpFormatter",
50
- "StoreBoolean",
51
- # Math utilities
52
- "cdiv",
53
- "next_power_of_2",
54
- "prev_power_of_2",
55
- "round_down",
56
- "round_up",
57
-]
58
-
29
_DEPRECATED_MAPPINGS = {
30
"cprofile": "profiling",
31
"cprofile_context": "profiling",
@@ -84,12 +54,8 @@ def __dir__() -> list[str]:
54
55
56
if TYPE_CHECKING:
87
- from argparse import Namespace
88
-
57
from vllm.config import ModelConfig, VllmConfig
58
else:
91
- Namespace = object
92
-
59
ModelConfig = object
60
VllmConfig = object
61
@@ -149,35 +115,33 @@ class Counter:
115
self.counter = 0
116
117
152
-def random_uuid() -> str:
153
- return str(uuid.uuid4().hex)
154
-
118
+class AtomicCounter:
119
+ """An atomic, thread-safe counter"""
120
156
-def update_environment_variables(envs: dict[str, str]):
157
- for k, v in envs.items():
158
- if k in os.environ and os.environ[k] != v:
159
- logger.warning(
160
- "Overwriting environment variable %s from '%s' to '%s'",
161
- k,
162
- os.environ[k],
163
- v,
164
- )
165
- os.environ[k] = v
121
+ def __init__(self, initial=0):
122
+ """Initialize a new atomic counter to given initial value"""
123
+ self._value = initial
124
+ self._lock = threading.Lock()
125
126
+ def inc(self, num=1):
127
+ """Atomically increment the counter by num and return the new value"""
128
+ with self._lock:
129
+ self._value += num
130
+ return self._value
131
168
-@cache
169
-def is_pin_memory_available() -> bool:
170
- from vllm.platforms import current_platform
132
+ def dec(self, num=1):
133
+ """Atomically decrement the counter by num and return the new value"""
134
+ with self._lock:
135
+ self._value -= num
136
+ return self._value
137
172
- return current_platform.is_pin_memory_available()
138
+ @property
139
+ def value(self):
140
+ return self._value
141
142
175
-@cache
176
-def is_uva_available() -> bool:
177
- """Check if Unified Virtual Addressing (UVA) is available."""
178
- # UVA requires pinned memory.
179
- # TODO: Add more requirements for UVA if needed.
180
- return is_pin_memory_available()
143
+def random_uuid() -> str:
144
+ return str(uuid.uuid4().hex)
145
146
147
# TODO: This function can be removed if transformer_modules classes are
@@ -212,47 +176,6 @@ def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None:
176
enable_trace_function_call(log_path)
177
178
215
-def weak_bind(
216
- bound_method: Callable[..., Any],
217
-) -> Callable[..., None]:
218
- """Make an instance method that weakly references
219
- its associated instance and no-ops once that
220
- instance is collected."""
221
- ref = weakref.ref(bound_method.__self__) # type: ignore[attr-defined]
222
- unbound = bound_method.__func__ # type: ignore[attr-defined]
223
-
224
- def weak_bound(*args, **kwargs) -> None:
225
- if inst := ref():
226
- unbound(inst, *args, **kwargs)
227
-
228
- return weak_bound
229
-
230
-
231
-class AtomicCounter:
232
- """An atomic, thread-safe counter"""
233
-
234
- def __init__(self, initial=0):
235
- """Initialize a new atomic counter to given initial value"""
236
- self._value = initial
237
- self._lock = threading.Lock()
238
-
239
- def inc(self, num=1):
240
- """Atomically increment the counter by num and return the new value"""
241
- with self._lock:
242
- self._value += num
243
- return self._value
244
-
245
- def dec(self, num=1):
246
- """Atomically decrement the counter by num and return the new value"""
247
- with self._lock:
248
- self._value -= num
249
- return self._value
250
-
251
- @property
252
- def value(self):
253
- return self._value
254
-
255
-
179
def kill_process_tree(pid: int):
180
"""
181
Kills all descendant processes of the given pid by sending SIGKILL.
@@ -303,13 +226,6 @@ def set_ulimit(target_soft_limit=65535):
226
)
227
228
306
-# Adapted from: https://github.com/sgl-project/sglang/blob/v0.4.1/python/sglang/utils.py#L28 # noqa: E501
307
-def get_exception_traceback():
308
- etype, value, tb = sys.exc_info()
309
- err_str = "".join(traceback.format_exception(etype, value, tb))
310
- return err_str
311
-
312
-
229
def _maybe_force_spawn():
230
"""Check if we need to force the use of the `spawn` multiprocessing start
231
method.
@@ -327,6 +243,8 @@ def _maybe_force_spawn():
243
os.environ["RAY_ADDRESS"] = ray.get_runtime_context().gcs_address
244
reasons.append("In a Ray actor and can only be spawned")
245
246
+ from .platform_utils import cuda_is_initialized, xpu_is_initialized
247
+
248
if cuda_is_initialized():
249
reasons.append("CUDA is initialized")
250
elif xpu_is_initialized():
@@ -356,55 +274,6 @@ def get_mp_context():
274
return multiprocessing.get_context(mp_method)
275
276
359
-def bind_kv_cache(
360
- ctx: dict[str, Any],
361
- kv_cache: list[list[torch.Tensor]], # [virtual_engine][layer_index]
362
- shared_kv_cache_layers: dict[str, str] | None = None,
363
-) -> None:
364
- # Bind the kv_cache tensor to Attention modules, similar to
365
- # ctx[layer_name].kv_cache[ve]=kv_cache[ve][extract_layer_index(layer_name)]
366
- # Special things handled here:
367
- # 1. Some models have non-attention layers, e.g., Jamba
368
- # 2. Pipeline parallelism, each rank only has a subset of layers
369
- # 3. Encoder attention has no kv cache
370
- # 4. Encoder-decoder models, encoder-decoder attention and decoder-only
371
- # attention of the same layer (e.g., bart's decoder.layers.1.self_attn
372
- # and decoder.layers.1.encoder_attn) is mapped to the same kv cache
373
- # tensor
374
- # 5. Some models have attention layers that share kv cache with previous
375
- # layers, this is specified through shared_kv_cache_layers
376
- if shared_kv_cache_layers is None:
377
- shared_kv_cache_layers = {}
378
- from vllm.attention import AttentionType
379
- from vllm.model_executor.models.utils import extract_layer_index
380
-
381
- layer_need_kv_cache = [
382
- layer_name
383
- for layer_name in ctx
384
- if (
385
- hasattr(ctx[layer_name], "attn_type")
386
- and ctx[layer_name].attn_type
387
- in (AttentionType.DECODER, AttentionType.ENCODER_DECODER)
388
- )
389
- and ctx[layer_name].kv_sharing_target_layer_name is None
390
- ]
391
- layer_index_sorted = sorted(
392
- set(extract_layer_index(layer_name) for layer_name in layer_need_kv_cache)
393
- )
394
- for layer_name in layer_need_kv_cache:
395
- kv_cache_idx = layer_index_sorted.index(extract_layer_index(layer_name))
396
- forward_ctx = ctx[layer_name]
397
- assert len(forward_ctx.kv_cache) == len(kv_cache)
398
- for ve, ve_kv_cache in enumerate(kv_cache):
399
- forward_ctx.kv_cache[ve] = ve_kv_cache[kv_cache_idx]
400
- if shared_kv_cache_layers is not None:
401
- for layer_name, target_layer_name in shared_kv_cache_layers.items():
402
- assert extract_layer_index(target_layer_name) < extract_layer_index(
403
- layer_name
404
- ), "v0 doesn't support interleaving kv sharing"
405
- ctx[layer_name].kv_cache = ctx[target_layer_name].kv_cache
406
-
407
-
277
def run_method(
278
obj: Any,
279
method: str | bytes | Callable,