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
[Chore] Separate out `vllm.utils.collections` (#26990)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Oct 16, 2025 at 16:35 UTC
d2740fafbf2527f4026863854bf1aeaa383c1a8e
29 files changed
+218
-184
tests/conftest.py
+2
-1
@@ -57,7 +57,8 @@ from vllm.multimodal.utils import fetch_image
57
from vllm.outputs import RequestOutput
58
from vllm.sampling_params import BeamSearchParams
59
from vllm.transformers_utils.utils import maybe_model_redirect
60
-from vllm.utils import is_list_of, set_default_torch_num_threads
60
+from vllm.utils import set_default_torch_num_threads
61
+from vllm.utils.collections import is_list_of
62
63
logger = init_logger(__name__)
64
tests/models/multimodal/generation/vlm_utils/model_utils.py
+1
-1
@@ -25,7 +25,7 @@ from transformers import (
25
from transformers.video_utils import VideoMetadata
26
27
from vllm.logprobs import SampleLogprobs
28
-from vllm.utils import is_list_of
28
+from vllm.utils.collections import is_list_of
29
30
from .....conftest import HfRunner, ImageAsset, ImageTestAssets
31
from .types import RunnerOutput
tests/models/multimodal/processing/test_tensor_schema.py
+1
-1
@@ -35,7 +35,7 @@ from vllm.multimodal import MULTIMODAL_REGISTRY, BatchedTensorInputs
35
from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
36
from vllm.multimodal.utils import group_mm_kwargs_by_modality
37
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
38
-from vllm.utils import is_list_of
38
+from vllm.utils.collections import is_list_of
39
40
from ...registry import _MULTIMODAL_EXAMPLE_MODELS, HF_EXAMPLE_MODELS
41
from ...utils import dummy_hf_overrides
tests/utils_/test_collections.py
new
+31
@@ -0,0 +1,31 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+from vllm.utils.collections import swap_dict_values
6
+
7
+
8
+@pytest.mark.parametrize(
9
+ "obj,key1,key2",
10
+ [
11
+ # Tests for both keys exist
12
+ ({1: "a", 2: "b"}, 1, 2),
13
+ # Tests for one key does not exist
14
+ ({1: "a", 2: "b"}, 1, 3),
15
+ # Tests for both keys do not exist
16
+ ({1: "a", 2: "b"}, 3, 4),
17
+ ],
18
+)
19
+def test_swap_dict_values(obj, key1, key2):
20
+ original_obj = obj.copy()
21
+
22
+ swap_dict_values(obj, key1, key2)
23
+
24
+ if key1 in original_obj:
25
+ assert obj[key2] == original_obj[key1]
26
+ else:
27
+ assert key2 not in obj
28
+ if key2 in original_obj:
29
+ assert obj[key1] == original_obj[key2]
30
+ else:
31
+ assert key1 not in obj
tests/utils_/test_utils.py
-25
@@ -38,7 +38,6 @@ from vllm.utils import (
38
sha256,
39
split_host_port,
40
split_zmq_path,
41
- swap_dict_values,
41
unique_filepath,
42
)
43
@@ -516,30 +515,6 @@ def test_placeholder_module_error_handling():
515
_ = placeholder_attr.module
516
517
519
-@pytest.mark.parametrize(
520
- "obj,key1,key2",
521
- [
522
- # Tests for both keys exist
523
- ({1: "a", 2: "b"}, 1, 2),
524
- # Tests for one key does not exist
525
- ({1: "a", 2: "b"}, 1, 3),
526
- # Tests for both keys do not exist
527
- ({1: "a", 2: "b"}, 3, 4),
528
- ],
529
-)
530
-def test_swap_dict_values(obj, key1, key2):
531
- original_obj = obj.copy()
532
- swap_dict_values(obj, key1, key2)
533
- if key1 in original_obj:
534
- assert obj[key2] == original_obj[key1]
535
- else:
536
- assert key2 not in obj
537
- if key2 in original_obj:
538
- assert obj[key1] == original_obj[key2]
539
- else:
540
- assert key1 not in obj
541
-
542
-
518
def test_model_specification(
519
parser_with_config, cli_config_file, cli_config_file_with_model
520
):
vllm/entrypoints/llm.py
+2
-1
@@ -75,7 +75,8 @@ from vllm.transformers_utils.tokenizer import (
75
get_cached_tokenizer,
76
)
77
from vllm.usage.usage_lib import UsageContext
78
-from vllm.utils import Counter, Device, as_iter, is_list_of
78
+from vllm.utils import Counter, Device
79
+from vllm.utils.collections import as_iter, is_list_of
80
from vllm.v1.engine import EngineCoreRequest
81
from vllm.v1.engine.llm_engine import LLMEngine
82
from vllm.v1.sample.logits_processor import LogitsProcessor
vllm/entrypoints/openai/serving_chat.py
+1
-1
@@ -70,7 +70,7 @@ from vllm.transformers_utils.tokenizers import (
70
truncate_tool_call_ids,
71
validate_request_params,
72
)
73
-from vllm.utils import as_list
73
+from vllm.utils.collections import as_list
74
75
logger = init_logger(__name__)
76
vllm/entrypoints/openai/serving_completion.py
+1
-1
@@ -34,8 +34,8 @@ from vllm.logprobs import Logprob
34
from vllm.outputs import RequestOutput
35
from vllm.sampling_params import BeamSearchParams, SamplingParams
36
from vllm.transformers_utils.tokenizer import AnyTokenizer
37
-from vllm.utils import as_list
37
from vllm.utils.asyncio import merge_async_iterators
38
+from vllm.utils.collections import as_list
39
40
logger = init_logger(__name__)
41
vllm/entrypoints/openai/serving_embedding.py
+1
-1
@@ -39,8 +39,8 @@ from vllm.outputs import (
39
RequestOutput,
40
)
41
from vllm.pooling_params import PoolingParams
42
-from vllm.utils import chunk_list
42
from vllm.utils.asyncio import merge_async_iterators
43
+from vllm.utils.collections import chunk_list
44
45
logger = init_logger(__name__)
46
vllm/entrypoints/openai/serving_engine.py
+2
-1
@@ -90,13 +90,14 @@ from vllm.tracing import (
90
log_tracing_disabled_warning,
91
)
92
from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer
93
-from vllm.utils import is_list_of, random_uuid
93
+from vllm.utils import random_uuid
94
from vllm.utils.asyncio import (
95
AsyncMicrobatchTokenizer,
96
collect_from_async_generator,
97
make_async,
98
merge_async_iterators,
99
)
100
+from vllm.utils.collections import is_list_of
101
from vllm.v1.engine import EngineCoreRequest
102
103
logger = init_logger(__name__)
vllm/entrypoints/openai/tool_parsers/abstract_tool_parser.py
+2
-1
@@ -12,7 +12,8 @@ from vllm.entrypoints.openai.protocol import (
12
)
13
from vllm.logger import init_logger
14
from vllm.transformers_utils.tokenizer import AnyTokenizer
15
-from vllm.utils import import_from_path, is_list_of
15
+from vllm.utils import import_from_path
16
+from vllm.utils.collections import is_list_of
17
18
logger = init_logger(__name__)
19
vllm/inputs/parse.py
+1
-1
@@ -5,7 +5,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, TypeAlias, TypedDict, cas
5
6
from typing_extensions import TypeIs
7
8
-from vllm.utils import is_list_of
8
+from vllm.utils.collections import is_list_of
9
10
from .data import (
11
EmbedsPrompt,
vllm/model_executor/layers/activation.py
+1
-1
@@ -17,7 +17,7 @@ from vllm.logger import init_logger
17
from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.utils import set_weight_attrs
19
from vllm.platforms import current_platform
20
-from vllm.utils import LazyDict
20
+from vllm.utils.collections import LazyDict
21
22
logger = init_logger(__name__)
23
vllm/model_executor/layers/quantization/gptq.py
+1
-1
@@ -28,7 +28,7 @@ from vllm.model_executor.parameter import (
28
RowvLLMParameter,
29
)
30
from vllm.transformers_utils.config import get_safetensors_params_metadata
31
-from vllm.utils import is_list_of
31
+from vllm.utils.collections import is_list_of
32
33
if TYPE_CHECKING:
34
from vllm.model_executor.layers.quantization import QuantizationMethods
vllm/model_executor/layers/quantization/gptq_marlin.py
+1
-1
@@ -57,7 +57,7 @@ from vllm.model_executor.parameter import (
57
from vllm.platforms import current_platform
58
from vllm.scalar_type import scalar_types
59
from vllm.transformers_utils.config import get_safetensors_params_metadata
60
-from vllm.utils import is_list_of
60
+from vllm.utils.collections import is_list_of
61
62
logger = init_logger(__name__)
63
vllm/model_executor/models/deepseek_vl2.py
+1
-1
@@ -49,7 +49,7 @@ from vllm.transformers_utils.configs.deepseek_vl2 import (
49
)
50
from vllm.transformers_utils.processors.deepseek_vl2 import DeepseekVLV2Processor
51
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
52
-from vllm.utils import is_list_of
52
+from vllm.utils.collections import is_list_of
53
from vllm.utils.tensor_schema import TensorSchema, TensorShape
54
55
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
vllm/model_executor/models/llava_next_video.py
+1
-1
@@ -33,7 +33,7 @@ from vllm.multimodal.processing import (
33
)
34
from vllm.multimodal.profiling import BaseDummyInputsBuilder
35
from vllm.sequence import IntermediateTensors
36
-from vllm.utils import is_list_of
36
+from vllm.utils.collections import is_list_of
37
from vllm.utils.tensor_schema import TensorSchema, TensorShape
38
39
from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
vllm/model_executor/models/minicpmv.py
+1
-1
@@ -86,7 +86,7 @@ from vllm.multimodal.processing import (
86
from vllm.multimodal.profiling import BaseDummyInputsBuilder
87
from vllm.platforms import current_platform
88
from vllm.sequence import IntermediateTensors
89
-from vllm.utils import flatten_2d_lists
89
+from vllm.utils.collections import flatten_2d_lists
90
from vllm.utils.tensor_schema import TensorSchema, TensorShape
91
92
from .idefics2_vision_model import Idefics2VisionTransformer
vllm/model_executor/models/qwen3_vl.py
+1
-1
@@ -79,7 +79,7 @@ from vllm.multimodal.processing import (
79
)
80
from vllm.multimodal.profiling import BaseDummyInputsBuilder
81
from vllm.sequence import IntermediateTensors
82
-from vllm.utils import is_list_of
82
+from vllm.utils.collections import is_list_of
83
84
from .interfaces import (
85
MultiModalEmbeddings,
vllm/multimodal/inputs.py
+2
-1
@@ -22,7 +22,8 @@ from typing import (
22
import numpy as np
23
from typing_extensions import NotRequired, TypeVar, deprecated
24
25
-from vllm.utils import LazyLoader, full_groupby, is_list_of
25
+from vllm.utils import LazyLoader
26
+from vllm.utils.collections import full_groupby, is_list_of
27
from vllm.utils.jsontree import json_map_leaves
28
29
if TYPE_CHECKING:
vllm/multimodal/parse.py
+7
-4
@@ -19,7 +19,8 @@ import numpy as np
19
import torch
20
from typing_extensions import assert_never
21
22
-from vllm.utils import LazyLoader, is_list_of
22
+from vllm.utils import LazyLoader
23
+from vllm.utils.collections import is_list_of
24
25
from .audio import AudioResampler
26
from .inputs import (
@@ -364,7 +365,7 @@ class MultiModalDataParser:
365
if isinstance(data, torch.Tensor):
366
return data.ndim == 3
367
if is_list_of(data, torch.Tensor):
367
- return data[0].ndim == 2
368
+ return data[0].ndim == 2 # type: ignore[index]
369
370
return False
371
@@ -422,6 +423,7 @@ class MultiModalDataParser:
423
if self._is_embeddings(data):
424
return AudioEmbeddingItems(data)
425
426
+ data_items: list[AudioItem]
427
if (
428
is_list_of(data, float)
429
or isinstance(data, (np.ndarray, torch.Tensor))
@@ -432,7 +434,7 @@ class MultiModalDataParser:
434
elif isinstance(data, (np.ndarray, torch.Tensor)):
435
data_items = [elem for elem in data]
436
else:
435
- data_items = data
437
+ data_items = data # type: ignore[assignment]
438
439
new_audios = list[np.ndarray]()
440
for data_item in data_items:
@@ -485,6 +487,7 @@ class MultiModalDataParser:
487
if self._is_embeddings(data):
488
return VideoEmbeddingItems(data)
489
490
+ data_items: list[VideoItem]
491
if (
492
is_list_of(data, PILImage.Image)
493
or isinstance(data, (np.ndarray, torch.Tensor))
@@ -496,7 +499,7 @@ class MultiModalDataParser:
499
elif isinstance(data, tuple) and len(data) == 2:
500
data_items = [data]
501
else:
499
- data_items = data
502
+ data_items = data # type: ignore[assignment]
503
504
new_videos = list[tuple[np.ndarray, dict[str, Any] | None]]()
505
metadata_lst: list[dict[str, Any] | None] = []
vllm/multimodal/processing.py
+6
-3
@@ -25,7 +25,7 @@ from typing_extensions import TypeVar, assert_never
25
from vllm.logger import init_logger
26
from vllm.transformers_utils.processor import cached_processor_from_config
27
from vllm.transformers_utils.tokenizer import AnyTokenizer, decode_tokens, encode_tokens
28
-from vllm.utils import flatten_2d_lists, full_groupby
28
+from vllm.utils.collections import flatten_2d_lists, full_groupby
29
from vllm.utils.functools import get_allowed_kwarg_only_overrides
30
from vllm.utils.jsontree import JSONTree, json_map_leaves
31
@@ -484,8 +484,11 @@ _M = TypeVar("_M", bound=_HasModalityAttr | _HasModalityProp)
484
485
486
def full_groupby_modality(values: Iterable[_M]) -> ItemsView[str, list[_M]]:
487
- """Convenience function to apply [`full_groupby`][vllm.utils.full_groupby]
488
- based on modality."""
487
+ """
488
+ Convenience function to apply
489
+ [`full_groupby`][vllm.utils.collections.full_groupby]
490
+ based on modality.
491
+ """
492
return full_groupby(values, key=lambda x: x.modality)
493
494
vllm/multimodal/registry.py
+1
-1
@@ -9,7 +9,7 @@ import torch.nn as nn
9
from vllm.config.multimodal import BaseDummyOptions
10
from vllm.logger import init_logger
11
from vllm.transformers_utils.tokenizer import AnyTokenizer, cached_tokenizer_from_config
12
-from vllm.utils import ClassRegistry
12
+from vllm.utils.collections import ClassRegistry
13
14
from .cache import BaseMultiModalProcessorCache
15
from .processing import (
vllm/reasoning/abs_reasoning_parsers.py
+2
-1
@@ -8,7 +8,8 @@ from functools import cached_property
8
from typing import TYPE_CHECKING, Any
9
10
from vllm.logger import init_logger
11
-from vllm.utils import import_from_path, is_list_of
11
+from vllm.utils import import_from_path
12
+from vllm.utils.collections import is_list_of
13
14
if TYPE_CHECKING:
15
from vllm.entrypoints.openai.protocol import (
vllm/utils/__init__.py
+3
-129
@@ -37,29 +37,19 @@ from argparse import (
37
RawDescriptionHelpFormatter,
38
_ArgumentGroup,
39
)
40
-from collections import UserDict, defaultdict
40
+from collections import defaultdict
41
from collections.abc import (
42
Callable,
43
Collection,
44
Generator,
45
- Hashable,
46
- Iterable,
45
Iterator,
48
- Mapping,
46
Sequence,
47
)
48
from concurrent.futures.process import ProcessPoolExecutor
49
from dataclasses import dataclass, field
50
from functools import cache, lru_cache, partial, wraps
51
from pathlib import Path
55
-from typing import (
56
- TYPE_CHECKING,
57
- Any,
58
- Generic,
59
- Literal,
60
- TextIO,
61
- TypeVar,
62
-)
52
+from typing import TYPE_CHECKING, Any, TextIO, TypeVar
53
from urllib.parse import urlparse
54
from uuid import uuid4
55
@@ -78,7 +68,7 @@ import zmq.asyncio
68
from packaging import version
69
from packaging.version import Version
70
from torch.library import Library
81
-from typing_extensions import Never, TypeIs, assert_never
71
+from typing_extensions import Never
72
73
import vllm.envs as envs
74
from vllm.logger import enable_trace_function_call, init_logger
@@ -170,9 +160,6 @@ def set_default_torch_num_threads(num_threads: int):
160
T = TypeVar("T")
161
U = TypeVar("U")
162
173
-_K = TypeVar("_K", bound=Hashable)
174
-_V = TypeVar("_V")
175
-
163
164
class Device(enum.Enum):
165
GPU = enum.auto()
@@ -421,12 +408,6 @@ def update_environment_variables(envs: dict[str, str]):
408
os.environ[k] = v
409
410
424
-def chunk_list(lst: list[T], chunk_size: int):
425
- """Yield successive chunk_size chunks from lst."""
426
- for i in range(0, len(lst), chunk_size):
427
- yield lst[i : i + chunk_size]
428
-
429
-
411
def cdiv(a: int, b: int) -> int:
412
"""Ceiling division."""
413
return -(a // -b)
@@ -743,53 +724,6 @@ def common_broadcastable_dtype(dtypes: Collection[torch.dtype]):
724
)
725
726
746
-def as_list(maybe_list: Iterable[T]) -> list[T]:
747
- """Convert iterable to list, unless it's already a list."""
748
- return maybe_list if isinstance(maybe_list, list) else list(maybe_list)
749
-
750
-
751
-def as_iter(obj: T | Iterable[T]) -> Iterable[T]:
752
- if isinstance(obj, str) or not isinstance(obj, Iterable):
753
- return [obj] # type: ignore[list-item]
754
- return obj
755
-
756
-
757
-# `collections` helpers
758
-def is_list_of(
759
- value: object,
760
- typ: type[T] | tuple[type[T], ...],
761
- *,
762
- check: Literal["first", "all"] = "first",
763
-) -> TypeIs[list[T]]:
764
- if not isinstance(value, list):
765
- return False
766
-
767
- if check == "first":
768
- return len(value) == 0 or isinstance(value[0], typ)
769
- elif check == "all":
770
- return all(isinstance(v, typ) for v in value)
771
-
772
- assert_never(check)
773
-
774
-
775
-def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]:
776
- """Flatten a list of lists to a single list."""
777
- return [item for sublist in lists for item in sublist]
778
-
779
-
780
-def full_groupby(values: Iterable[_V], *, key: Callable[[_V], _K]):
781
- """
782
- Unlike [`itertools.groupby`][], groups are not broken by
783
- non-contiguous data.
784
- """
785
- groups = defaultdict[_K, list[_V]](list)
786
-
787
- for value in values:
788
- groups[key(value)].append(value)
789
-
790
- return groups.items()
791
-
792
-
727
# TODO: This function can be removed if transformer_modules classes are
728
# serialized by value when communicating between processes
729
def init_cached_hf_modules() -> None:
@@ -1578,50 +1512,6 @@ class AtomicCounter:
1512
return self._value
1513
1514
1581
-# Adapted from: https://stackoverflow.com/a/47212782/5082708
1582
-class LazyDict(Mapping[str, T], Generic[T]):
1583
- def __init__(self, factory: dict[str, Callable[[], T]]):
1584
- self._factory = factory
1585
- self._dict: dict[str, T] = {}
1586
-
1587
- def __getitem__(self, key: str) -> T:
1588
- if key not in self._dict:
1589
- if key not in self._factory:
1590
- raise KeyError(key)
1591
- self._dict[key] = self._factory[key]()
1592
- return self._dict[key]
1593
-
1594
- def __setitem__(self, key: str, value: Callable[[], T]):
1595
- self._factory[key] = value
1596
-
1597
- def __iter__(self):
1598
- return iter(self._factory)
1599
-
1600
- def __len__(self):
1601
- return len(self._factory)
1602
-
1603
-
1604
-class ClassRegistry(UserDict[type[T], _V]):
1605
- def __getitem__(self, key: type[T]) -> _V:
1606
- for cls in key.mro():
1607
- if cls in self.data:
1608
- return self.data[cls]
1609
-
1610
- raise KeyError(key)
1611
-
1612
- def __contains__(self, key: object) -> bool:
1613
- return self.contains(key)
1614
-
1615
- def contains(self, key: object, *, strict: bool = False) -> bool:
1616
- if not isinstance(key, type):
1617
- return False
1618
-
1619
- if strict:
1620
- return key in self.data
1621
-
1622
- return any(cls in self.data for cls in key.mro())
1623
-
1624
-
1515
def weak_ref_tensor(tensor: Any) -> Any:
1516
"""
1517
Create a weak reference to a tensor.
@@ -2588,22 +2478,6 @@ class LazyLoader(types.ModuleType):
2478
return dir(self._module)
2479
2480
2591
-def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None:
2592
- """
2593
- Helper function to swap values for two keys
2594
- """
2595
- v1 = obj.get(key1)
2596
- v2 = obj.get(key2)
2597
- if v1 is not None:
2598
- obj[key2] = v1
2599
- else:
2600
- obj.pop(key2, None)
2601
- if v2 is not None:
2602
- obj[key1] = v2
2603
- else:
2604
- obj.pop(key1, None)
2605
-
2606
-
2481
@contextlib.contextmanager
2482
def cprofile_context(save_file: str | None = None):
2483
"""Run a cprofile
vllm/utils/collections.py
new
+139
@@ -0,0 +1,139 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Contains helpers that are applied to collections.
5
+
6
+This is similar in concept to the `collections` module.
7
+"""
8
+
9
+from collections import UserDict, defaultdict
10
+from collections.abc import Callable, Generator, Hashable, Iterable, Mapping
11
+from typing import Generic, Literal, TypeVar
12
+
13
+from typing_extensions import TypeIs, assert_never
14
+
15
+T = TypeVar("T")
16
+U = TypeVar("U")
17
+
18
+_K = TypeVar("_K", bound=Hashable)
19
+_V = TypeVar("_V")
20
+
21
+
22
+class ClassRegistry(UserDict[type[T], _V]):
23
+ """
24
+ A registry that acts like a dictionary but searches for other classes
25
+ in the MRO if the original class is not found.
26
+ """
27
+
28
+ def __getitem__(self, key: type[T]) -> _V:
29
+ for cls in key.mro():
30
+ if cls in self.data:
31
+ return self.data[cls]
32
+
33
+ raise KeyError(key)
34
+
35
+ def __contains__(self, key: object) -> bool:
36
+ return self.contains(key)
37
+
38
+ def contains(self, key: object, *, strict: bool = False) -> bool:
39
+ if not isinstance(key, type):
40
+ return False
41
+
42
+ if strict:
43
+ return key in self.data
44
+
45
+ return any(cls in self.data for cls in key.mro())
46
+
47
+
48
+class LazyDict(Mapping[str, T], Generic[T]):
49
+ """
50
+ Evaluates dictionary items only when they are accessed.
51
+
52
+ Adapted from: https://stackoverflow.com/a/47212782/5082708
53
+ """
54
+
55
+ def __init__(self, factory: dict[str, Callable[[], T]]):
56
+ self._factory = factory
57
+ self._dict: dict[str, T] = {}
58
+
59
+ def __getitem__(self, key: str) -> T:
60
+ if key not in self._dict:
61
+ if key not in self._factory:
62
+ raise KeyError(key)
63
+ self._dict[key] = self._factory[key]()
64
+ return self._dict[key]
65
+
66
+ def __setitem__(self, key: str, value: Callable[[], T]):
67
+ self._factory[key] = value
68
+
69
+ def __iter__(self):
70
+ return iter(self._factory)
71
+
72
+ def __len__(self):
73
+ return len(self._factory)
74
+
75
+
76
+def as_list(maybe_list: Iterable[T]) -> list[T]:
77
+ """Convert iterable to list, unless it's already a list."""
78
+ return maybe_list if isinstance(maybe_list, list) else list(maybe_list)
79
+
80
+
81
+def as_iter(obj: T | Iterable[T]) -> Iterable[T]:
82
+ if isinstance(obj, str) or not isinstance(obj, Iterable):
83
+ return [obj] # type: ignore[list-item]
84
+ return obj
85
+
86
+
87
+def is_list_of(
88
+ value: object,
89
+ typ: type[T] | tuple[type[T], ...],
90
+ *,
91
+ check: Literal["first", "all"] = "first",
92
+) -> TypeIs[list[T]]:
93
+ if not isinstance(value, list):
94
+ return False
95
+
96
+ if check == "first":
97
+ return len(value) == 0 or isinstance(value[0], typ)
98
+ elif check == "all":
99
+ return all(isinstance(v, typ) for v in value)
100
+
101
+ assert_never(check)
102
+
103
+
104
+def chunk_list(lst: list[T], chunk_size: int) -> Generator[list[T]]:
105
+ """Yield successive chunk_size chunks from lst."""
106
+ for i in range(0, len(lst), chunk_size):
107
+ yield lst[i : i + chunk_size]
108
+
109
+
110
+def flatten_2d_lists(lists: Iterable[Iterable[T]]) -> list[T]:
111
+ """Flatten a list of lists to a single list."""
112
+ return [item for sublist in lists for item in sublist]
113
+
114
+
115
+def full_groupby(values: Iterable[_V], *, key: Callable[[_V], _K]):
116
+ """
117
+ Unlike [`itertools.groupby`][], groups are not broken by
118
+ non-contiguous data.
119
+ """
120
+ groups = defaultdict[_K, list[_V]](list)
121
+
122
+ for value in values:
123
+ groups[key(value)].append(value)
124
+
125
+ return groups.items()
126
+
127
+
128
+def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None:
129
+ """Swap values between two keys."""
130
+ v1 = obj.get(key1)
131
+ v2 = obj.get(key2)
132
+ if v1 is not None:
133
+ obj[key2] = v1
134
+ else:
135
+ obj.pop(key2, None)
136
+ if v2 is not None:
137
+ obj[key1] = v2
138
+ else:
139
+ obj.pop(key1, None)
vllm/v1/engine/async_llm.py
+2
-1
@@ -29,8 +29,9 @@ from vllm.tracing import init_tracer
29
from vllm.transformers_utils.config import maybe_register_config_serialize_by_value
30
from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs
31
from vllm.usage.usage_lib import UsageContext
32
-from vllm.utils import Device, as_list, cdiv
32
+from vllm.utils import Device, cdiv
33
from vllm.utils.asyncio import cancel_task_threadsafe
34
+from vllm.utils.collections import as_list
35
from vllm.utils.functools import deprecate_kwargs
36
from vllm.v1.engine import EngineCoreRequest
37
from vllm.v1.engine.core_client import EngineCoreClient
vllm/v1/worker/gpu_input_batch.py
+2
-1
@@ -12,7 +12,8 @@ from vllm.lora.request import LoRARequest
12
from vllm.multimodal.inputs import MultiModalFeatureSpec
13
from vllm.pooling_params import PoolingParams
14
from vllm.sampling_params import SamplingParams, SamplingType
15
-from vllm.utils import length_from_prompt_token_ids_or_embeds, swap_dict_values
15
+from vllm.utils import length_from_prompt_token_ids_or_embeds
16
+from vllm.utils.collections import swap_dict_values
17
from vllm.v1.outputs import LogprobsTensors
18
from vllm.v1.pool.metadata import PoolingMetadata
19
from vllm.v1.sample.logits_processor import (
vllm/v1/worker/tpu_input_batch.py
+2
-1
@@ -9,7 +9,8 @@ import torch
9
10
from vllm.lora.request import LoRARequest
11
from vllm.sampling_params import SamplingType
12
-from vllm.utils import length_from_prompt_token_ids_or_embeds, swap_dict_values
12
+from vllm.utils import length_from_prompt_token_ids_or_embeds
13
+from vllm.utils.collections import swap_dict_values
14
from vllm.v1.outputs import LogprobsTensors
15
from vllm.v1.worker.block_table import MultiGroupBlockTable
16
from vllm.v1.worker.gpu_input_batch import CachedRequestState